fix: Scale random projections by 1/sqrt(p) (Johnson-Lindenstrauss) - #346
Merged
Conversation
create_projection_matrix normalized each row over its n entries, giving entry
variance 1/n. JL requires 1/p (p = projected dim) so that E[A^T A] = I and
projected inner products are unbiased. The consequence is not the overall
scale — a constant cancels from any ranking — but that the factor p/n depends
on the module's SHAPE, so modules are silently reweighted against each other in
the summed score. Attention (d x d) and MLP (d x 4d) blocks differ by 4x.
Measured E[proj]/true, 3000 draws per shape:
shape p before predicted p^2/(o*i) after
32x32 16 0.24972 0.25000 0.99889
64x64 32 0.25042 0.25000 1.00167
128x128 32 0.06247 0.06250 0.99949
256x64 32 0.06253 0.06250 1.00045
This matches the TrackStar paper (arXiv 2410.17413 §A.1.2), which specifies the
same two-sided per-block projection and states "Projection matrix entries are
sampled i.i.d. from N(0, 1/d)" — variance 1/(number of rows), exactly
A /= sqrt(m). It also matches docs/preprocessing.rst's claim that projections
preserve inner products, which the old scaling did not.
Existing tests could not catch this: they compare the collector against
L @ G @ R^T using the SAME matrices, so they are invariant to how those are
scaled. test_global_projection_linearity uses one R sliced into blocks — the
paper's construction, not the code's. tests/test_projection_inner_products.py
now pins projected against unprojected (12 tests, all 12 fail without this fix),
including that the scale does not depend on module shape.
BREAKING: changes all stored index scales. Existing indices built with
projection_dim > 0 must be regenerated, and any published numbers re-run.
runs/test_build_cache.npy is regenerated here; the measured ratio is exactly
0.5000, matching theory (sqrt(o*i)/p = 8/16 for tiny-Phi3).
THREE TESTS FAIL ON THIS BRANCH AND ARE LEFT FAILING, PENDING A DECISION:
tests/test_batch_size_invariance.py::test_gradient_scale_invariance[100-100]
tests/test_batch_size_invariance.py::test_gradient_scale_invariance[50-150]
Relative error is UNCHANGED by this fix (6.83e-06 on main vs 7.58e-06 here)
and already exceeds the assert_close default rtol=1.3e-06 on main. It passes
on main only because the values are ~150x smaller, so the absolute
difference stays under atol=1e-05. Restoring the correct scale trips atol.
The test is magnitude-sensitive rather than wrong about the invariant.
tests/test_gradients.py::test_gradient_collector_proj_norm
1 of 256 elements, greatest relative difference 2.46e-04 against the test's
explicit 1e-4 tolerance. Probably a near-cancelling element landing
differently at the new scale — the test already carries a comment about
accumulating numerical error — but this is NOT proven.
These are regression guards; loosening them is a call for a human, not
something to do to make a branch green.
Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
test_gradient_scale_invariance compared std separately-vs-combined with assert_close's defaults. The invariant is scale-free, but the default atol=1e-5 let a relative disagreement pass whenever the operands were small enough: on main the relative error is already 6.83e-06 against rtol=1.3e-06, and the test passes only because the values are ~150x smaller than the correctly-scaled ones. Use rtol=1e-4 with atol=0 so the check tracks the invariant rather than the magnitude. test_gradient_collector_proj_norm compares A @ normalize(g) @ B.T against the collector, which projects activations and output grads separately. The two orderings are mathematically identical and differ only in fp32 rounding; at the corrected scale that noise reached 2.46e-04 relative on an O(1) element, over the existing 1e-4. Widen to 1e-3 and say what the tolerance is measuring. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
luciaquirke
force-pushed
the
fix/projection-jl-scaling
branch
from
July 22, 2026 04:44
ca39be7 to
d159010
Compare
Indexes store their projected gradients, so changing the projection scaling makes existing ones unreadable at the correct weighting. projection_scale selects the convention: "jl" (default, variance 1/projection_dim) or "row_norm". GradientProcessor.load resolves a missing projection_scale key to "row_norm", so existing indexes keep working with no user action, and new ones record "jl" in processor_config.yaml. Threaded through IndexConfig, create_processor and EkfacConfig, which builds matching matrices to compress the IVHP output. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
luciaquirke
force-pushed
the
fix/projection-jl-scaling
branch
from
July 22, 2026 04:45
d159010 to
a21a00d
Compare
for more information, see https://pre-commit.ci
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.
Achieve a variance of 1/p by scaling by 1/sqrt(p)
This doesn't affect LDS/relative ranking of items because the per-module scale cancels in the inner product, it's to match the paper.
Breaking: changes cached test index scales; runs/test_build_cache.npy regenerated (ratio exactly 0.5000, matching _(o·i)/p = 8/16 for tiny-Phi3).
Adds tests/test_projection_inner_products.py, and makes two existing tests magnitude-independent.