scikit-learn / scikit-learn/scikit-learn
PERF: Avoid unnecessary n_samples x n_samples diagonal sparse matrix in `_rescale_data`
Nobody has claimed this yet.
- Dominant language
- Python
- Stars
- 67.3k
- Forks
- 27.4k
- Avg merge
- 1d 15h
- Merged PRs (30d)
- 58
Description
[!WARNING]
This issue is not yet ready for a PR. If you are interested in contributing to scikit-learn, please have a look at our contributing guidelines, and in particular the sections for new contributors and the "Needs triage" label.
Describe the workflow you want to enable
When fitting any linear model (e.g. LinearRegression, Ridge, ElasticNet) on sparse data
with sample_weight, the internal function _rescale_data is called to bake sample weights
into the data before solving. Currently it constructs an intermediate (n_samples x n_samples)
diagonal sparse matrix just to perform a simple row scaling, then runs a full sparse matrix
multiply through safe_sparse_dot.
For large sparse datasets this creates unnecessary memory allocation and computation on every
fit() call. The issue affects all linear models that support sample_weight with sparse input.
Describe your proposed solution
Replace the dia_array + safe_sparse_dot approach with a direct elementwise row-scaling via
csr_array.multiply(), which is built for this exact operation and needs no intermediate matrix.
Current code in sklearn/linear_model/_base.py:252:
if sp.issparse(X) or sp.issparse(y):
sw_matrix = sparse.dia_array(
(sample_weight_sqrt, 0), shape=(n_samples, n_samples)
)
if sp.issparse(X):
X = safe_sparse_dot(sw_matrix, X)
if sp.issparse(y):
y = safe_sparse_dot(sw_matrix, y)
Proposed code:
if sp.issparse(X):
X = X.multiply(sample_weight_sqrt[:, None])
if sp.issparse(y):
sw = sample_weight_sqrt[:, None] if y.ndim > 1 else sample_weight_sqrt
y = y.multiply(sw)
Both produce identical results — verified against the existing test_rescale_data test and
directly against the old implementation across all input combinations.
Why it is better:
- Eliminates the
(n_samples x n_samples)diagonal sparse matrix allocation entirely csr_array.multiply(vector)is O(nnz) with no intermediate object;
dia_array @ csr_arrayconstructs the diagonal object first, then runs a matrix-multiply
code path- For 100k samples: old code allocates a 100k-element
dataarray and shape tuple for the
dia_arraybefore touchingX; new code works directly on thennznon-zero entries
Describe alternatives you've considered
scipy.sparse.diags(sample_weight_sqrt) @ X— same problem, still builds an intermediate
diagonal sparse matrix- Keeping
safe_sparse_dot— unnecessary here since.multiply()handles elementwise
scaling natively and always returns a sparse result
Additional context
- Affected function:
sklearn.linear_model._base._rescale_data(line 220) - Affected models: all linear models that accept
sample_weightwith sparseX—
LinearRegression,Ridge,ElasticNet,Lasso, and their CV variants - Tests: the existing
test_rescale_datainsklearn/linear_model/tests/test_base.py
already coverssparse_container x n_targetscombinations — no new test infrastructure needed - Fix is ready: branch
perf/avoid-dia-array-in-rescale-dataon fork
Labels
Enhancement, Needs Triage, module:linear_model
Describe alternatives you've considered, if relevant
No response
Additional context
No response
Contributor guide
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- Open a pull request that references the issue number.
Research direction
Start in sklearn/linear_model/_base.py at _rescale_data and read test_rescale_data in sklearn/linear_model/tests/test_base.py. Review the sparse X and y target-shape cases covered there, then verify that results remain identical without constructing the n_samples × n_samples diagonal matrix. Done means the existing tests pass and the sparse sample_weight path avoids that allocation.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- python
- Domain
- machine-learning, performance
- Issue type
- Refactor
- Difficulty
- 3/5
- Estimated time
- 1-2 days
- Activity status
- Active
- Clarity
- Clearly specified
- Newbie friendliness
- 35/100