Thanks to visit codestin.com
Credit goes to github.com

Skip to content

PERF Vectorize sparse path of PolynomialCountSketch.transform - #34919

Draft
regarmukesh3g wants to merge 2 commits into
scikit-learn:mainfrom
regarmukesh3g:fix/polynomial-count-sketch-sparse-perf
Draft

PERF Vectorize sparse path of PolynomialCountSketch.transform#34919
regarmukesh3g wants to merge 2 commits into
scikit-learn:mainfrom
regarmukesh3g:fix/polynomial-count-sketch-sparse-perf

Conversation

@regarmukesh3g

@regarmukesh3g regarmukesh3g commented Sep 10, 2026

Copy link
Copy Markdown

Reference Issues/PRs

See #34920 (not using a closing keyword since the issue is still pending
triage).

What does this implement/fix? Explain your changes.

PolynomialCountSketch.transform's sparse branch loops over every feature
and every degree, slicing one sparse column at a time and calling
.toarray() on it:

for j in range(X_gamma.shape[1]):
    for d in range(self.degree):
        iHashIndex = self.indexHash_[d, j]
        iHashBit = self.bitHash_[d, j]
        count_sketches[:, d, iHashIndex] += (
            (iHashBit * X_gamma[:, [j]]).toarray().ravel()
        )

This densifies a full (n_samples,) array on every iteration regardless of
column sparsity, so cost scales with n_features * degree instead of the
number of nonzeros. On a 1%-density matrix with a few thousand features, the
sparse path ends up an order of magnitude slower than just densifying the
whole input and using the dense branch.

This PR replaces the double loop with one sparse matrix multiplication per
degree: scale columns by bitHash_[d, :] via .multiply(...), then
scatter-sum into buckets with a sparse 0/1 projection matrix built from
indexHash_[d, :]. Output is unchanged.

Local benchmark (2000 samples, 300 components, degree=2, 1% density):

n_features before after
1000 0.097s 0.002s
4000 0.323s 0.005s
8000 0.631s 0.009s

Important: this PR always improves on the code it replaces. The table
below (1000 samples, 200 components, degree=2, 2000 features, varying
density) makes both comparisons explicit — old-vs-new (the one that matters
for approval) and new-vs-dense (context on the sparse/dense crossover):

density old sparse new sparse old vs new new vs dense
0.1% 0.151s 0.002s 75x faster 0.15x (beats dense)
1% 0.155s 0.002s 78x faster 0.18x (beats dense)
10% 0.151s 0.008s 19x faster 0.69x (beats dense)
30% 0.156s 0.017s 9x faster 1.45x (slower than dense)
50% 0.161s 0.022s 7x faster 1.98x (slower than dense)
80% 0.162s 0.030s 5x faster 2.70x (slower than dense)

The old sparse path was uniformly ~0.15-0.16s regardless of density,
because its cost was driven by loop count (n_features * degree) rather
than the data itself. The new path's cost scales with the number of
nonzeros, so it is dramatically faster than the old path at every density
tested — 5x to 78x depending on density — even though at high density (>~20%)
it is still somewhat slower than simply densifying the input and using the
unrelated dense branch, which is expected of any sparse algorithm and not a
regression: a 50-80%-dense matrix stored as a sparse type is already an
unusual choice on the caller's part, and the fix never makes that case worse
than main, only better.

Note on dense input: the else branch handling dense arrays is untouched by
this PR (byte-for-byte identical diff), and measured dense transform timing
is unchanged before/after within run-to-run noise.

Since this is a performance fix rather than a correctness bug, the existing
test_polynomial_count_sketch_dense_sparse test passes on both the old and
new implementation (both are correct, just different speeds), so it alone
doesn't prove the new code is right. I added
test_polynomial_count_sketch_sparse_edge_cases, covering all-zero sparse
columns and a single-feature/single-sample input, shapes not exercised by the
existing random test data. I confirmed this new test is not vacuous: I
temporarily broke the new vectorized code (dropped the sign-hash multiply)
and the new test failed with a clear mismatch, then restored the fix.

Verification:

  • pytest sklearn/tests/test_kernel_approximation.py -> 80 passed
  • check_estimator(PolynomialCountSketch()) passes
  • ruff check / ruff format --check clean

Introduce yourself

I use scikit-learn for general ML work. While checking the open-issue
backlog for something to contribute, I found it heavily claimed (most open
Bug/Documentation/good first issue/help wanted issues already had a
competing PR or an active claim comment), so I looked for an improvement by
reading source directly instead, and found this inefficiency in
kernel_approximation.py.

AI usage disclosure

I used AI assistance for:

  • Code generation (e.g., when writing an implementation or fixing a bug)
  • Test/benchmark generation
  • Research and understanding

Any other comments?

Already added the efficiency changelog entry under
doc/whats_new/upcoming_changes/sklearn.kernel_approximation/34919.efficiency.rst.

Opened #34920 to report the bug ahead of this PR per the contributing
guidelines. It's pending triage, so this PR references it without a closing
keyword for now and I'll switch to Fixes #34920 once the label is cleared.

@github-actions

Copy link
Copy Markdown

Thank you for opening your first pull request to scikit-learn! 🎉

To help get your contribution reviewed, please make sure that:

  • You have filled out the pull request template.

  • The pull request addresses an existing issue that is ready for contribution (e.g. not tagged as 'Needs Triage', 'Needs Decision', ...). If you are proposing a new feature, please open an issue to discuss it first.

  • There are no other open pull requests already targeting the same issue.

  • You have followed the pull request checklist. In particular, linting and tests should pass.

@regarmukesh3g
regarmukesh3g force-pushed the fix/polynomial-count-sketch-sparse-perf branch 2 times, most recently from f162579 to 3286e80 Compare September 10, 2026 10:15
The sparse branch of PolynomialCountSketch.transform looped over every
feature and every degree, slicing a single sparse column and calling
.toarray() on it in each iteration. This densified one (n_samples,)
column at a time regardless of the matrix's actual sparsity, so the
cost scaled with n_features * degree rather than the number of
nonzeros. For a matrix with 1% density and a few thousand features,
this made the sparse path an order of magnitude slower than simply
densifying the whole input up front and using the dense code path.
@regarmukesh3g
regarmukesh3g force-pushed the fix/polynomial-count-sketch-sparse-perf branch from 0d857c0 to f1a2db9 Compare September 10, 2026 10:46
The prior comment re-explained each line of the new sparse branch
(one comment per statement, almost matching the code line-for-line),
which is exactly the kind of comment well-named identifiers should
replace instead. Renamed weighted/projection to signed_features/
bucket_projection and collapsed the comment to a single one above the
if-block stating only the non-obvious part: why this avoids per-column
densification, which is the actual bug being fixed.
@ammar-iitm

Copy link
Copy Markdown

Reviewed this independently (I'd arrived at the same sparse-matmul approach before seeing this PR, see #34920). Checked out the branch and ran a 200-trial randomized differential test against the original loop implementation as ground truth (varying shape, degree, small n_components to force hash collisions, density 0-1, coef0, CSR/CSC) — all matched exactly. Also confirmed check_estimator(PolynomialCountSketch()) passes and ruff check/ruff format --check are clean on the changed files. Looks correct and well-tested to me.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants