feat(svd_circuits): per-head QK/OV SVD with degeneracy guard - #1768
Merged
jlarson4 merged 8 commits intoSep 11, 2026
Conversation
… tests Add model-free unit tests for a per-head QK/OV singular-vector decomposition. Synthetic weight tensors only, so no model is loaded and no pretrained weights are downloaded. Coverage: - Oracle parity of the factored SVD against torch.linalg.svd for the OV (W_V @ W_O) and QK (W_Q @ W_K.T) maps, plus reconstruction, descending order, the d_head rank bound, and the sigma-ratio rank report. - The .V-not-.Vh convention, asserted by treating the FactoredMatrix.Vh DeprecationWarning as an error. - The degeneracy guard: equal and near-equal singular values group into a block, require_isolated refuses per-direction attribution inside a block and points at the subspace, well-separated spectra flag nothing, eps is honored, near-null runs group, and the blocks partition every direction exactly once. Written test-first: the suite fails at collection until the decomposition module lands.
Decompose a single attention head's query-key (W_Q W_K^T) and output-value (W_V W_O) maps via FactoredMatrix.svd(). Weight-space only: no forward pass, activation cache, or compatibility mode. Each map stays factored, so the d_model x d_model product is never materialized and the rank is bounded by d_head. Right singular vectors are read from FactoredMatrix.V; the deprecated .Vh alias is never touched. Each direction is summarized in a rank report with its singular value, its ratio to the top value, and the contiguous block it groups into, so near-equal singular values surface as rotation-ambiguous subspaces. The module is importable by full path only and is intentionally absent from the analysis package __all__.
…bution Add HeadSVD.is_degenerate, block_of, degenerate_blocks, and require_isolated so callers can detect rotation-ambiguous singular blocks and refuse per-direction attribution inside them. Near-equal singular values leave the singular directions defined only up to a rotation within their block, so require_isolated raises DegenerateDirectionError and points at the block as a subspace instead of returning a rotation-dependent direction. The methods read the existing per-direction report (block_id and is_degenerate are already populated by _degeneracy_blocks and _build_rank_report), so no grouping logic is recomputed.
6 tasks
jlarson4
reviewed
Sep 11, 2026
jlarson4
left a comment
jlarson4
left a comment
Collaborator
There was a problem hiding this comment.
Hi @janmenjayap! Thanks for building out this tool! it is very cool. Just a couple notes on areas where we might be able to improve our efficiency or close a couple testing gaps.
Additionally, I noticed this includes references to HookedTransformer. That class will be official removed soon, any new features do not require supporting HookedTransformer going forward.
Sorry for the delay on this, I wrote up all my comments this morning and forgot to post them. Let me know if you have any questions.
…y-value heads - Read the attention block's own per-head W_Q/W_K/W_V/W_O instead of the full-model W_Q/W_K/W_V/W_O stacks, so decompose_head materializes one layer instead of the whole model's weights. - On grouped-query attention W_K/W_V carry one row per key-value head, so a query head index was out of range or wrong once key-value heads outnumbered query heads; map the query head to its key-value head (h // (n_heads // n_kv_heads)) before indexing, mirroring the bridge's own kv-head expansion rule. - A no-op for multi-head attention, where the head counts already match. - Add model-free unit tests exercising decompose_head directly through a lightweight stub model (no from_pretrained, no download): head selection under multi-head attention, the query-to-kv-head mapping under grouped-query attention, the QK transpose wiring, the which filter, and the layer/head/which input guards.
…nchor - Require a new direction to fall within eps of both its immediate predecessor and the block's anchor (its first, largest member) before it joins a near-equal block, capping how far a block can spread. - Without the anchor bound, a spectrum decaying by just under eps at every step chained every direction into one block whose extremes differed by far more than eps, so require_isolated wrongly refused the top direction over the whole spectrum. - Add a regression test for a slowly decaying spectrum: it no longer collapses into one block, and the well-separated top direction stays isolated. - Add a test pinning the exact block degenerate_blocks() returns for a repeated singular value, rather than only checking non-degeneracy.
- Add a dedicated null_rtol parameter to decompose_head, threaded through _factored_head_svd into _degeneracy_blocks, so the null-run grouping no longer reuses the near-equal eps threshold - Default null_rtol to d_model * torch.finfo(S.dtype).eps, matching torch.linalg.matrix_rank's default relative tolerance, so directions group as null only when they reflect the map's numerical rank - Store the resolved null_rtol on HeadSVD for reproducibility - Update _degeneracy_blocks and the module docstring to describe the null cutoff as distinct from the near-equal gap threshold eps - Rewrite test_null_run_grouped to use a spectrum the relative-gap rule cannot group, grouped only via an explicit null_rtol - Add test_null_default_does_not_overgroup and test_null_run_groups_true_null_tail to pin the default cutoff's behavior on both a spectrum it should not over-group and a genuinely null tail it should still catch
… HookedTransformer - Rewrite the module docstring Example:: to boot a TransformerBridge instead of HookedTransformer, matching the per-block accessor path decompose_head already reads. - Drop the "Works with both HookedTransformer and TransformerBridge" sentence and the HookedTransformer mention in the Args: block; state the per-block bridge accessors decompose_head reads instead. - Docstring-only: no HookedTransformer-only code path existed to remove.
6 tasks
…tion as degenerate, read per-head weights detached, widen which to Sequence[str], adjust tests to match above
Collaborator
|
@janmenjayap After reviewing your changes, I realized I provided some incorrect feedback in my latest review. I requested you cap how wide a degenerate block can get, but that was bad advice. Under the capped rule, a long run of near-equal singular values gets cut at arbitrary points. So a direction can be reported as isolated while its neighbor is within I went ahead and cleaned up that error and pushed the fix to your branch. Assuming this passes CI, I am going to merge it so that we can guarantee inclusion in 3.9.0. Thanks you for putting this together, and apologies for my mistake. |
7 tasks
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
2 participants
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.
Description
Adds
svd_circuits.decompose_head, the first of three PRs implementing per-head singular-vector decomposition of a head's QK (W_Q W_K^T) and OV (W_V W_O) maps (tracking issue: #1767 ).This PR is weight-space only: it reads
W_Q/W_K/W_V/W_Oand needs no forward pass, no activation cache, and no compatibility mode.decompose_head(model, layer, head, which=("QK", "OV"))— factored SVD viaFactoredMatrix, never materializing thed_model x d_modelproduct.HeadSVD/HeadDecompositiondataclasses with a per-directionrank_report.require_isolatedraisesDegenerateDirectionErrorfor per-direction attribution inside a block instead of silently returning a rotation-dependent direction.FactoredMatrix.V; the deprecated.Vhalias is never touched (enforced by a test that turns itsDeprecationWarninginto an error).By design, this PR does not export a public API.
svd_circuitsis intentionally absent fromtransformer_lens/tools/analysis/__init__.py__all__— the causal patch-along-direction gate (PR2) is what turns areadout into a validated claim, so the public surface should never ship without it attached.
Unit-tested model-free: synthetic weight tensors only, no
from_pretrained, no HF download. SVD parity againsttorch.linalg.svdwithinatol=1e-5; reconstruction withinatol=1e-4.Follow-ups (not in this PR): OV vocab/logit readout,
project_activations, the mandatorypatch_along_directionscausal gate, and public exports (PR2); demo notebook, slow paper-sanity check, and docs (PR3).Part of #1767
Type of change
Checklist: