Repository navigation
refactor(evals): vectorize correlation metrics - #145
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.
|
Benchmark results, seconds, best of three runs on the same machine:
The largest gains are in Peak RSS for |
||||||||||||||||||||||||||||||||||
HirokiYasuda03
left a comment
HirokiYasuda03
left a comment
There was a problem hiding this comment.
@KenyaOtsuka
Thank you for your work!
I commented one minor point. If that point is intentional, ignore it and merge.
| n_feat = _y.shape[1] | ||
|
|
||
| r = np.array( | ||
| [ | ||
| np.corrcoef( | ||
| _x[:, j].ravel(), | ||
| _y[:, j].ravel() | ||
| )[0, 1] | ||
| for j in range(n_feat) | ||
| ] | ||
| ) | ||
| r = _column_correlation(_x, _y) |
There was a problem hiding this comment.
np.corrcoef raises a RuntimeWarning for division by zero, but currently that warning is ignored by np.errstate(invalid=‘ignore’, divide=‘ignore’):. Since this has little impact on profile correlation, this change may be intentional; however, the PR description states that warning behavior doesn't change, so please verify this just to be sure.
There was a problem hiding this comment.
Thanks for pointing this out. You're right — the result is unchanged (NaN for zero-variance inputs), but the RuntimeWarning is now suppressed. I'll keep this behavior and update the PR description.
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, and NaN handling are unchanged. The implementation also preserves the
[-1, 1]clipping used by NumPy/SciPy correlation calculations. Runtime warnings fromnp.corrcoefon zero-variance inputs are now suppressed, while the resulting correlations remainNaNas before.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.