Plot Elastic Net Precomputed Gram Matrix With Weighted SamplesΒΆ

========================================================================== Fitting an Elastic Net with a precomputed Gram Matrix and Weighted SamplesΒΆ

The following example shows how to precompute the gram matrix while using weighted samples with an :class:~sklearn.linear_model.ElasticNet.

If weighted samples are used, the design matrix must be centered and then rescaled by the square root of the weight vector before the gram matrix is computed.

… note:: sample_weight vector is also rescaled to sum to n_samples, see the documentation for the sample_weight parameter to

meth:

~sklearn.linear_model.ElasticNet.fit.

Imports for Elastic Net with Precomputed Gram Matrix and Sample WeightsΒΆ

When fitting Elastic Net on very large datasets, precomputing the Gram matrix (X^T X) can significantly speed up the coordinate descent optimization because the inner products needed at each iteration are already available. However, combining precomputed Gram matrices with sample weights requires careful preprocessing: the design matrix must first be centered using weighted means, then rescaled by the square root of the normalized weights before computing the Gram matrix.

Why this matters: Precomputation trades memory (storing an n_features x n_features matrix) for speed, which is beneficial when n_features is moderate but n_samples is very large. The weighted centering and rescaling ensure that the Gram matrix correctly represents the weighted inner products, maintaining mathematical equivalence with the non-precomputed solution. This technique is commonly used in production ML pipelines where the same feature matrix is reused across multiple regularization strengths during hyperparameter search.

# Authors: The scikit-learn developers
# SPDX-License-Identifier: BSD-3-Clause

# %%
# Let's start by loading the dataset and creating some sample weights.
import numpy as np

from sklearn.datasets import make_regression

rng = np.random.RandomState(0)

n_samples = int(1e5)
X, y = make_regression(n_samples=n_samples, noise=0.5, random_state=rng)

sample_weight = rng.lognormal(size=n_samples)
# normalize the sample weights
normalized_weights = sample_weight * (n_samples / (sample_weight.sum()))

# %%
# To fit the elastic net using the `precompute` option together with the sample
# weights, we must first center the design matrix,  and rescale it by the
# normalized weights prior to computing the gram matrix.
X_offset = np.average(X, axis=0, weights=normalized_weights)
X_centered = X - np.average(X, axis=0, weights=normalized_weights)
X_scaled = X_centered * np.sqrt(normalized_weights)[:, np.newaxis]
gram = np.dot(X_scaled.T, X_scaled)

# %%
# We can now proceed with fitting. We must passed the centered design matrix to
# `fit` otherwise the elastic net estimator will detect that it is uncentered
# and discard the gram matrix we passed. However, if we pass the scaled design
# matrix, the preprocessing code will incorrectly rescale it a second time.
from sklearn.linear_model import ElasticNet

lm = ElasticNet(alpha=0.01, precompute=gram)
lm.fit(X_centered, y, sample_weight=normalized_weights)