scikit-learn / scikit-learn/scikit-learn

PERF: Avoid unnecessary n_samples x n_samples diagonal sparse matrix in `_rescale_data`

Open
#33,960 3 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Needs Benchmarks Performance
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_array constructs the diagonal object first, then runs a matrix-multiply
    code path
  • For 100k samples: old code allocates a 100k-element data array and shape tuple for the
    dia_array before touching X; new code works directly on the nnz non-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_weight with sparse X
    LinearRegression, Ridge, ElasticNet, Lasso, and their CV variants
  • Tests: the existing test_rescale_data in sklearn/linear_model/tests/test_base.py
    already covers sparse_container x n_targets combinations — no new test infrastructure needed
  • Fix is ready: branch perf/avoid-dia-array-in-rescale-data on fork

Labels

Enhancement, Needs Triage, module:linear_model

Describe alternatives you've considered, if relevant

No response

Additional context

No response

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. 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

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.