Skip to content

[Proposal] SVD Circuits: singular-vector decomposition of a head's QK/ OV into causally-validated subfunctions #1767

Description

@janmenjayap

Proposal

Component-level circuit analysis (direct_logit_attribution, path patching, head detectors) treats an attention head as an indivisible unit. But a single canonical head — e.g. a GPT-2-small IOI name-mover — empirically superposes multiple subfunctions on distinct low-rank directions of its QK and OV maps. This proposes svd_circuits, a self-contained tool in transformer_lens/tools/analysis/ that takes the SVD of a head's QK (W_Q W_Kᵀ) and OV (W_V W_O) matrices, exposes each orthogonal singular direction as an interpretable sub-component, and — crucially — causally gates every claimed subfunction by patching activations along that individual singular direction. This is a finer-than-component decomposition that resolves within-head superposition component-level tools miss.

Capabilities:

  1. Per-head QK/OV SVD via FactoredMatrix — get singular values/vectors of W_Q W_Kᵀ and W_V W_O without materializing the d_model × d_model matrix.
  2. Direction → vocab readout — project each OV singular direction into the unembedding basis (reusing SVDInterpreter) and each singular direction into logit space via direct_logit_attribution, so a direction gets a human-readable "what tokens does this subfunction move" signature.
  3. Activation projection / attribution — project cached per-head activations onto the singular directions to measure how much each subfunction fires on a given prompt (per token/position).
  4. Causal validation (patch-along-a-direction) — reconstruct the head output using only a chosen singular subspace (or ablate a single direction) via generic_activation_patch, and report the behavior change (e.g. IOI logit-diff), so a subfunction claim is accepted only if its direction is causally load-bearing.
  5. Degeneracy guard — detect near-equal singular values (rotation/degeneracy ambiguity) and refuse to attribute individual directions inside a degenerate block, reporting the block as a subspace instead.

**Status:**PR1 (weight-space core: decompose_head, HeadSVD, degeneracy guard) implemented and verified on branch feat/svd-circuits-core — unit tests, mypy, and make unit-test all green (#1768 ). PR2/PR3 not started.
Suggested labels: enhancement, tooling, TransformerBridge, complexity: moderate


Motivation

The existing TL baseline for "what does this head do" is direct_logit_attribution (transformer_lens/tools/analysis/direct_logit_attribution.py) and head_detector — both operate at head granularity. The IOI work established name-mover / S-inhibition / duplicate-token heads as the atomic units of the circuit. The claim in Beyond Components is that this granularity is too coarse: an IOI name-mover head is not one function but several, packed onto near-orthogonal low-rank directions of its QK/OV weights, and only an SVD-based decomposition surfaces them. This tool lets a TL user go inside a canonical head and attribute/patch its subfunctions — the natural next rung below component-level circuit analysis.

Why no maintained TL implementation exists. TL already ships the two hardest primitives — FactoredMatrix (low-rank OV/QK SVD/eigendecomposition without materializing the full matrix) and SVDInterpreter (projects OV/w_in/w_out singular vectors into vocab) — but neither closes the loop from "singular direction" to "causally-validated subfunction." SVDInterpreter is a static weight-inspection utility (its own docstring warns SVD directions are not reliably interpretable); it never touches activations or patching. The missing piece is the activation-projection + causal-patch layer, which is exactly what makes the decomposition trustworthy.

Implementation status elsewhere. The authors' research repo exists — github.com/Exploration-Lab/Beyond-Components (~7 commits; IOI / GP / GT examples) — but it is not maintained or TL-integrated. This would be the first TL-native implementation. This is not SAE/dictionary-learning territory (no learned features, no overcomplete basis) — it is a closed-form weight decomposition, so it sits inside TL rather than at the SAELens boundary.


Pitch

Add svd_circuits as a self-contained analysis tool under transformer_lens/tools/analysis/. It builds on FactoredMatrix for the weight-space SVD, SVDInterpreter for vocab readout, ActivationCache/direct_logit_attribution for activation projection and logit signatures, and generic_activation_patch for the mandatory causal gate.

Proposed API

Names adjustable to maintainer preference.

from transformer_lens.model_bridge import TransformerBridge
from transformer_lens.tools.analysis import svd_circuits

# TransformerBridge is the supported TL 3.x path; HookedTransformer is deprecated (removed next major).
model = TransformerBridge.boot_transformers("gpt2-small")
model.enable_compatibility_mode()  # folds final LN into W_U — required for the OV vocab/logit readout

# 1. Decompose a head's QK and OV into singular directions (FactoredMatrix under the hood).
decomp = svd_circuits.decompose_head(model, layer=9, head=9, which=("QK", "OV"))
decomp.OV.S            # singular values, sorted desc
decomp.OV.rank_report  # (idx, sigma, sigma/sigma_max, is_degenerate) per direction

# 2. Readout: each OV singular direction -> top tokens in the unembedding basis.
decomp.OV.vocab_readout(k=10)     # wraps SVDInterpreter.get_singular_vectors
decomp.OV.logit_signature(model, prompt=ioi_prompt)  # via direct_logit_attribution

# 3. Attribution: how much does each direction fire on this prompt?
proj = svd_circuits.project_activations(model, decomp, prompt=ioi_prompt)  # [pos, direction]

# 4. Causal gate: keep only directions {0,3}, measure IOI logit-diff change.
result = svd_circuits.patch_along_directions(
    model, decomp.OV, keep=[0, 3], prompt=ioi_prompt,
    metric=ioi_logit_diff,      # any callable(logits) -> scalar
)
result.delta_metric   # behavior change attributable to that subspace
result.gated          # True only if |delta| exceeds the causal threshold

Design (algorithm)

For a head (layer ℓ, head h):

  1. Build factored maps. OV = FactoredMatrix(W_V[h], W_O[h]) (shape d_model × d_model, rank ≤ d_head); QK = FactoredMatrix(W_Q[h], W_K[h].T). Both stay factored — never materialize d_model².
  2. SVD. U, S, V = M.svd() (M == U @ S.diag() @ V.transpose(-2, -1); U and V are both [d_model, rank], with the i-th singular direction in column iU[:, i], V[:, i]). Right singular vectors V[:, i] are the input directions (which residual-stream directions the subfunction reads); left singular vectors U[:, i] are the output directions. Access them via the .V property — .Vh is a deprecated alias that returns the same tensor as .V (it emits a DeprecationWarning), not its Hermitian transpose; do not introduce a new call site that relies on it.
  3. Degeneracy guard. Flag any run of singular values with σᵢ/σᵢ₊₁ − 1 < ε (default ε = 1e-2) as a degenerate block: individual directions inside it are rotation-ambiguous, so attribute the block as a subspace, not per-direction.
  4. Vocab readout (OV). For each output direction U[:, i], project through the (LN-folded) unembedding — reuse SVDInterpreter.get_singular_vectors("OV", ℓ, head_index=h) — to get its top-token signature.
  5. Activation projection. Cache the head's per-position input (residual stream into the head; for OV attribution use stack_head_results / the head's value stream). Project onto V[:, i] to get a [pos] firing coefficient per direction i.
  6. Logit signature. Feed the rank-1 reconstruction σᵢ · U[:, i] V[:, i]ᵀ of the head output through direct_logit_attribution to get each direction's signed logit effect on the task metric.
  7. Causal patch-along-direction. Using generic_activation_patch, replace the head's output (at blocks.ℓ.attn.hook_z / hook_result) with its projection onto the chosen singular subspace span(U[:, keep]) (or zero out a single direction for ablation), rerun, and record Δmetric. A subfunction is accepted only if its direction/subspace produces a causal Δmetric above threshold — the decomposition alone never suffices.

Reuse map:

Sub-step TL primitive file:symbol
Factored QK/OV, low-rank SVD/eigen FactoredMatrix (.svd(), .U/.S/.V, .eigenvalues) transformer_lens/FactoredMatrix.py:FactoredMatrix
OV direction → vocab readout SVDInterpreter transformer_lens/SVDInterpreter.py:SVDInterpreter.get_singular_vectors
Per-head activations for projection ActivationCache transformer_lens/ActivationCache.py:stack_head_results (:957)
Direction → logit effect direct_logit_attribution transformer_lens/tools/analysis/direct_logit_attribution.py:direct_logit_attribution
Bridge LN-folding guard (reuse pattern) _validate_bridge_compatibility transformer_lens/tools/analysis/direct_logit_attribution.py:_validate_bridge_compatibility
Characterize resulting subfunction head detector transformer_lens/head_detector.py
Causal patch-along-direction generic_activation_patch + setter transformer_lens/patching.py:generic_activation_patch
Weights / unembed / final LN W_Q/W_K/W_V/W_O, model.W_U Bridge tl_parameters() / legacy transformer_lens/HookedTransformer.py

Main correctness risk (named): interpreting individual singular directions under rotation/degeneracy ambiguity

SVD directions are unique only when singular values are distinct; equal (or near-equal) singular values leave the corresponding singular subspace defined only up to an arbitrary rotation, so any "direction 3 = the surname subfunction" claim inside a degenerate block is meaningless. Compounding this, "the SVD direction is interpretable" is not guaranteed even for well-separated directions (SVDInterpreter's own docstring flags numerical instability and cross-device inconsistency). Mitigation, not a warning: (a) the degeneracy guard (step 3) raises/refuses per-direction attribution inside a near-degenerate block and downgrades it to a subspace; (b) every claimed subfunction must pass the causal patch-along-direction gate (step 7) — a direction that is not causally load-bearing is never reported as a subfunction, regardless of how clean its vocab readout looks. The tool's contract is "causal evidence gates interpretive claims," not "SVD is interpretable."

Validation plan (falsifiable)

  1. Analytic SVD oracle. On a tiny synthetic head with hand-set W_V/W_O (and W_Q/W_K), the tool's S, U, V match torch.linalg.svd of the explicitly materialized matrix to atol=1e-5 (sign/degeneracy handled).
  2. Reconstruction fidelity. Summing all rank-1 direction reconstructions reproduces the full head output (via the cache) to atol=1e-4; keeping top-k directions monotonically increases recovered Δmetric.
  3. Degeneracy guard fires. A synthetic head with two exactly-equal singular values → the guard flags the block and per-direction attribution raises; subspace attribution still works.
  4. Causal gate discriminates. Ablating a high-σ, causally-relevant direction changes the IOI logit-diff materially; ablating a null-space / random direction does not (baseline reported alongside).
  5. Paper sanity check (slow). On gpt2-small head L9H9 (name-mover), the tool surfaces ≥2 causally-gated OV subfunctions whose vocab readouts qualitatively match the paper's name-mover split; reported as a qualitative sanity check, not a pinned numeric threshold.

Scope of the vertical slice

Per-head QK/OV SVD + direction→vocab readout + direction→logit signature + one causal patch-along-direction validation, on gpt2-small IOI. Defer full multi-head circuit assembly and automated subfunction labeling. Ships as 3 sequential PRs, not one — base off dev:

  1. feat(svd_circuits): per-head QK/OV SVD with degeneracy guarddecompose_head, HeadSVD, degeneracy guard. Pure weight-space, no public export yet.
  2. feat(svd_circuits): singular-direction readout, projection, and causal patch gate — vocab/logit readout, project_activations, patch_along_directions, public exports, Bridge compatibility-mode validation, gpt2-small IOI integration test.
  3. docs(svd_circuits): add demo notebook, slow name-mover parity, and tool docs — slow oracle-parity test, SVD_Circuits_Demo.ipynb, docs section.

Files:

  • New: transformer_lens/tools/analysis/svd_circuits.py
  • Export: transformer_lens/tools/analysis/__init__.py
  • Unit: tests/unit/tools/test_svd_circuits.py (synthetic head; checks 1–4, no HF download)
  • Integration: tests/integration/test_svd_circuits.py (gpt2-small, IOI patch-along-direction end-to-end)
  • Oracle-parity (slow): tests/integration/test_svd_circuits_oracle_parity.py (check 5, @pytest.mark.slow)
  • Demo: demos/SVD_Circuits_Demo.ipynb (nbval)

Follow-up work

Deferred to a separate tiered issue, filed after PR3 merges: multi-head circuit assembly, automated subfunction labeling, QK-side pattern attribution, more models.

Model coverage

  • CI reference / oracle: tiny synthetic head (no download) for unit correctness; gpt2-small for the integration + slow paper sanity check (CI-cacheable, canonical IOI).
  • Demo model: gpt2-small.
  • Honesty caveat: the paper's headline subfunction taxonomy is scale- and model-dependent; gpt2-small reproduces the mechanism (a name-mover splits into causally-distinct directions) but not necessarily the exact subfunction inventory of larger models. Baselines (random-direction ablation) are reported next to every headline Δmetric.

Alternatives

  • Stay at head granularity (direct_logit_attribution / head detectors). Rejected: cannot resolve within-head superposition — the whole point.
  • Use SVDInterpreter alone. Rejected: static weight readout with no activation projection and no causal gate; its own docstring disclaims interpretability of the directions.
  • Sparse autoencoders on head outputs (SAELens). Rejected here: needs training, an overcomplete learned basis, and lives outside TL; svd_circuits is closed-form, training-free, and reuses TL primitives.
  • Full path-patching over a rank-1 direction grid. Rejected as the first slice: combinatorially expensive and premature before the single-direction gate is validated; folded into Tier 2 of the follow-up issue.

Correctness oracle

Brute-force / analytic reference (no external frozen impl exists). Primary oracle: torch.linalg.svd on an explicitly materialized synthetic head, compared to the FactoredMatrix-based path within threshold (atol=1e-5), plus a reconstruction-fidelity check. Secondary (slow, @pytest.mark.slow): the gpt2-small name-mover paper sanity check, treated as qualitative because the paper ships no numeric reference to pin against.


Additional context

  • Paper: Areeb Ahmad, Abhinav Joshi, Ashutosh Modi (IIT Kanpur), "Beyond Components: Singular Vector-Based Interpretability of Transformer Circuits," arXiv 2511.20273, NeurIPS 2025 poster (poster 119702, OpenReview 7UbXEQNny7).
  • External implementation status: authors' research repo Exploration-Lab/Beyond-Components exists (research-grade, ~7 commits, 4 stars, no transformer_lens dependency); this would be the first TL-native implementation. Paper nuance: it augments weight matrices with biases folded in and applies SVD uniformly to attention and MLP — this slice covers attention QK/OV (OV logit-facing) only.
  • Artifact sources: none required — the decomposition is computed from model weights at call time; no downloaded lens/SAE artifacts, so no registry entry is needed for the vertical slice.
  • Algorithm note for reviewers: OV singular directions are attributed through the unembedding (output side); QK singular directions describe which residual directions the head compares (attention-pattern side) and are attributed via the pattern, not the unembedding — the vertical slice focuses on OV for logit-facing claims and includes QK SVD as a readout only. The causal gate is what separates this from a pretty-picture SVD tool; keep it mandatory.
  • Substrate history: SVD convention fix [Proposal] Update and discuss behavior of SVD #341 (closed) → Issue resolution for #341, #644, and #210 #1300 (merged) (torch.linalg.svd; .Vh deprecated in favour of .V); [Bug Report] SVD tests fail on GPU #328 (closed, 2023, "SVD tests fail on GPU") is direct evidence for the cross-device instability the degeneracy guard + causal gate backstop.
  • Provisional-citation flags: arXiv id 2511.20273 and NeurIPS-2025 venue are confirmed (no open citation flags).

Checklist

  • I have checked that there is no similar issue in the repo (required).
  • Checked no similar issue / tool exists (SVDInterpreter is static weight readout only; no causal singular-direction tool in tools/analysis/).
  • Scope limited to a vertical slice (single head, OV logit-facing, one causal gate); follow-up tiered separately.
  • Correctness oracle identified (analytic torch.linalg.svd + reconstruction; slow qualitative paper check).
  • Targets TransformerBridge (the supported TL 3.x path) via shared ActivationCache; also runs on the deprecated HookedTransformer compat layer.
  • Claim-in-comments before starting a tier item.

Activity

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

Metadata

Metadata

Assignees

Labels

TransformerBridgeBug specific to the new TransformerBridge systemcomplexity-highVery complicated changes for people to address who are quite familiar with the codeenhancementNew feature or request

Type

No type

Projects

No projects

    Milestone

    No milestone

    Relationships

    None yet

    Development

    No branches or pull requests

    Issue actions