refactor(evals): vectorize correlation metrics - #145
Open
KenyaOtsuka wants to merge 2 commits into
Open
KenyaOtsuka wants to merge 2 commits into
KenyaOtsuka wants to merge 2 commits into
Conversation
Add tests for behavior that is currently only exercised indirectly through the pickled 2-d fixtures: undefined correlations on zero-variance columns and rows, NaN column removal (and its opt-out), the mean/std arguments of pattern_correlation, and pairwise_identification against a definition-based oracle built from cdist, including a non-correlation metric.
profile_correlation and pattern_correlation called np.corrcoef once per
unit and per sample, and pairwise_identification went through cdist, whose
correlation metric walks every pair in a Python-level C loop. On the
feature sizes produced by recent vision encoders (hundreds of thousands of
units) this dominates the evaluation runtime.
Compute the same quantities with array operations instead:
- profile_correlation and pattern_correlation accumulate centered dot
products over blocks of the axis they do not reduce over, which keeps the
two-pass formula of np.corrcoef, including its clipping to [-1, 1] and
the NaN it yields for a zero-variance column or row.
- pairwise_identification derives 1 - cdist(p, t, 'correlation') from a
single matrix product when the metric is 'correlation'. Other metrics and
the single-trial branch still go through cdist, and the NaN handling that
follows is untouched, so the emitted warnings are unchanged. scipy clips
the cosine it builds the correlation distance from, so the matrix product
is clipped the same way.
Blocking the work also bounds peak memory, which the previous
standardization and NaN scan did not: they materialized full-size copies.
Measured with random arrays of the given (n_samples, n_units), in seconds,
best of three runs:
500x5000 50x1000000 1500x40000
before after before after before after
profile_correlation 0.17 0.01 28.12 0.30 2.75 0.34
pattern_correlation 0.03 0.01 0.39 0.35 0.45 0.38
pairwise_identification 0.45 0.03 3.26 0.64 89.26 1.43
profile_correlation gains as the unit count grows, since it used to make one
np.corrcoef call per unit. pairwise_identification gains as the sample count
grows, since its cost is pairs times units. pattern_correlation gains least,
and at some shapes is only break-even: it already looped over the sample
axis, the small one, so the gain there is in memory. Peak RSS of the
1500x40000 pairwise_identification run drops from 931 MiB to 149 MiB.
Author
|
Benchmark results, seconds, best of three runs on the same machine:
The largest gains are in Peak RSS for |
||||||||||||||||||||||||||||||||||
KenyaOtsuka
requested review from
HirokiYasuda03,
ganow and
micchu
and removed request for
micchu
September 18, 2026 04:02
KenyaOtsuka
marked this pull request as ready for review
September 18, 2026 04:03
This branch has not been deployed
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Why
profile_correlation,pattern_correlation, andpairwise_identificationbecome very slow on modern vision features with hundreds of thousands of
units. A 1,492 × 425,984 evaluation had not finished after 17 hours.
What
profile_correlationandpattern_correlationusing blockedcentered dot products.
pairwise_identification.single_trial=Trueon the existingcdistpath.
For real-valued inputs, public signatures, output shapes, defaults, NaN
handling, and warning behavior are unchanged. The implementation also
preserves the
[-1, 1]clipping used by NumPy/SciPy correlation calculations.Test
Existing golden tests still pass at
rtol=1e-12, atol=1e-12.Added regression tests for zero-variance inputs, NaN handling,
mean/std,the correlation path of
pairwise_identification, non-correlation fallback,and multi-block execution.
Verified with NumPy 1.26.4 / SciPy 1.13.1 and NumPy 2.4.6 / SciPy 1.17.1.