feat(svd_circuits): singular-direction readout, projection, and causal patch gate - #1775
janmenjayap wants to merge 14 commits into
Conversation
…a cross-check test HeadSVD's docstring had U and V swapped for the OV map: for a factored map A @ B, U's columns live in A's input space and V's columns live in B's output space, so for OV (A = W_V_h, B = W_O_h) it is V, not U, whose columns are the residual-stream directions this head writes into and should be projected through W_U for a vocab or logit readout. The swap never raised a shape error because both spaces are d_model-dimensional, so it would have silently used the wrong basis in every downstream readout, projection, or patch built on top of it. Adds a regression test that cross-checks HeadSVD.OV.V against the already-shipped SVDInterpreter (which projects its OV singular vectors through the unembedding the other way), on a tiny no-download TransformerBridge, so the assertion cannot pass under either labeling by construction.
vocab_readout projects a head's top-k OV output directions (HeadSVD.V, per the corrected convention) through the unembedding, gated on TransformerBridge compatibility mode so LayerNorm stays folded into W_U. logit_signature reconstructs one direction's rank-1 OV output and reads its signed logit effect on given tokens; it requires the direction to be isolated first, since a rotation-ambiguous or null direction's signature is not attributable to it alone.
Add ActivationProjection and project_activations to recover per-position firing coefficients by projecting a head's actual OV output (hook_result) onto its V basis. Summing the coefficients against V reconstructs the head's output because V's columns are orthonormal and define the basis the head writes in. Restore the model's prior use_attn_result setting after the forward pass so read-only analysis does not leave a configuration side effect.
patch_along_directions causally validates a claimed OV subfunction by reconstructing a head's hook_result onto a chosen span of HeadSVD.V directions (keep) or its complement (ablate), then comparing a caller's metric before and after against an equally-sized random-subspace baseline. gated is true only when the requested subspace moves the metric by more than the random baseline does, so a claim cannot be called causal merely because patching moved the metric at all. Before touching the model, the retained set is checked against every degenerate block reported by HeadSVD.degenerate_blocks(): a block must be kept whole or dropped whole, since attributing an effect to part of a rotation-ambiguous or null block would let a caller hand-pick around the same guard require_isolated already enforces per direction. The forward passes run under torch.no_grad(), matching the other analysis tools that repeat a forward pass per call, and the model's prior use_attn_result setting is restored in a finally block so this read-only check leaves no configuration side effect.
…ty mode Adds decompose_head, project_activations, patch_along_directions, vocab_readout, logit_signature, and their supporting types (HeadSVD, HeadDecomposition, RankReportRow, ActivationProjection, LogitSignature, PatchResult, DegenerateDirectionError) to transformer_lens.tools.analysis's imports, __all__, and tool-listing docstring, in the same ASCII-sorted order the module already uses for its other tools. No compatibility-mode logic changes here: vocab_readout and logit_signature already gate on TransformerBridge compatibility mode, and patch_along_directions already restores use_attn_result. This commit only widens the public surface, so the causal gate ships alongside the readouts rather than either landing without the other.
Exercise vocab_readout and patch_along_directions against GPT-2 small's layer 9 head 9 (name-mover head) on an IOI-style prompt with a Mary/John logit-diff metric, checking that the readout and causal gate agree on a real Bridge with compatibility mode enabled.
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for the U/V correction and for cross-checking it against SVDInterpreter rather than the module's own arithmetic, that is the correct test point. Let me know if you have any questions on the below comments!
| def _validate_bridge_compatibility(model) -> None: | ||
| """Reject a ``TransformerBridge`` whose ``W_U`` would give a silently wrong projection. | ||
|
|
||
| ``HookedTransformer`` always has the final LayerNorm folded into ``W_U``, so this |
There was a problem hiding this comment.
HookedTransformer is to be deprecated, please remove any references to it, and any additional code that supports only HookedTransformer.
| raise ValueError( | ||
| "Projecting an OV direction through W_U on a TransformerBridge requires " | ||
| "compatibility mode, so that LayerNorm weights are folded into W_U. Call " | ||
| "`model.enable_compatibility_mode()` after loading the bridge, then retry." |
There was a problem hiding this comment.
enable_compatibility_mode() folds ln1 into W_V and centres W_O and W_U, so a HeadSVD decomposed before the call describes a different OV map (on gpt2 L9H9 the top singular value goes 23.3 → 8.4 and direction 1's top-10 tokens share nothing with a fresh decomposition), yet the guard inspects only the model, so retrying with the stale object returns silently, and project_activations and patch_along_directions have no guard at all. Record the compatibility-mode state on HeadSVD at decomposition time, refuse a mismatch in all four consumers, and point this message at re-running decompose_head.
| if not 0 < k <= rank: | ||
| raise ValueError(f"k must be in (0, {rank}], got {k!r}") | ||
| _validate_bridge_compatibility(model) | ||
| return model.W_U.T @ head_svd.V[:, :k].float() |
There was a problem hiding this comment.
V is cast to fp32 while W_U keeps the model dtype and matmul does not promote, so vocab_readout and logit_signature (line 531) raise expected m1 and m2 to have the same dtype on every bf16, fp16 and fp64 model, while the projection and patch paths cast to the activation's dtype and work at all four. Cast both operands to a common dtype in both functions and add a dtype-parametrised test.
| raise ValueError("patch_along_directions requires exactly one of keep or ablate") | ||
| rank = head_svd.V.shape[1] | ||
| if keep is not None: | ||
| retained = sorted(set(keep)) |
There was a problem hiding this comment.
keep, ablate and logit_signature's direction (line 528) go straight into V[:, ...] and the range(rank) complement, so ablate=[-1] or ablate=[999] subtracts nothing and reports delta ≈ 0, gated False for an intervention that never ran (gpt2 L9H9: −9.5e-07 against −3.2e-02 for the last isolated direction), keep=[-1] passes the degenerate-block guard that keep=[rank-1] trips, keep=[999] raises a raw IndexError, and keep=[] ties the baseline by construction. Validate every index against [0, rank) the way decompose_head checks layer and head, coerce to int, and refuse an empty retained set unless an explicit threshold is given.
|
|
||
| def hook_fn(activation: torch.Tensor, hook) -> torch.Tensor: | ||
| activation = activation.clone() | ||
| activation[:, :, head, :] = activation[:, :, head, :] @ projector.to(activation.dtype) |
There was a problem hiding this comment.
baseline_projector comes from torch.randn(..., generator=rng) with no device, so it is a CPU tensor while kept_projector sits on the model's device, and since this line reconciles dtype only the third forward pass raises Tensor for argument #2 'mat2' is on CPU, but expected it to be on GPU on any MPS or CUDA model (the draw has to stay on CPU: QR is unimplemented on MPS and a CUDA generator cannot feed a CPU randn). Move the finished projector to activation.device here alongside the dtype cast and add a device-skipif test that the call completes.
|
|
||
| def test_project_activations_reconstructs_head_output(tiny_bridge): | ||
| """Coefficients recovered against V reconstruct the actual cached hook_result slice.""" | ||
| layer, head = 0, 0 |
There was a problem hiding this comment.
Every projection and patch test decomposes head 0 (the stub has one head), so hardcoding the head index in project_activations (line 591) or the patch hook (line 655), or writing the projector to every head, passes the whole suite, and the hook's "leaves every other head untouched" promise is asserted nowhere. Run the reconstruction and patch tests on head 1 of the tiny bridge as well and assert the sibling head's hook_result slice is unchanged.
| patch_along_directions(SimpleNamespace(), ov, "prompt", metric) | ||
|
|
||
|
|
||
| def test_patch_along_directions_rejects_partial_degenerate_block(): |
There was a problem hiding this comment.
Both calls here go through keep and no ablate call anywhere asserts retained, so guarding _validate_retained_blocks with if keep is not None passes the suite, as does treating ablate as keep, since the discriminates test only orders |delta| for strong against weak. Add ablate=[1] on the [5, 3, 3, 1] spectrum expecting DegenerateDirectionError and assert retained on an ablate call.
| ) | ||
| assert result.delta_metric == pytest.approx(result.patched_metric - result.original_metric) | ||
| assert math.isfinite(result.baseline_delta_metric) | ||
| expected_gated = abs(result.delta_metric) > abs(result.baseline_delta_metric) |
There was a problem hiding this comment.
expected_gated is recomputed from the same result's fields, so it holds for any self-consistent implementation: gated = True hardcoded, the baseline comparison dropped, threshold ignored, rng ignored and the QR removed all pass, and no test passes threshold= or compares two seeded runs. Assert outcomes with a seeded generator instead, ablate=[weakest isolated] gating False, a threshold above abs(delta_metric) gating False, and the same seed reproducing baseline_delta_metric, updated to whatever the gate's semantics become.
|
|
||
| result = patch_along_directions(gpt2_bridge, ov, prompt, metric, keep=[top_direction]) | ||
|
|
||
| assert result.retained == [top_direction] |
There was a problem hiding this comment.
Every assertion is shape, finiteness, isinstance, or retained echoing the keep= argument, so the only real-model test passes with the hook never installed, the gate hardcoded, vocab_readout returning zeros, or the wrong basis projected. Assert that readout[:, i] matches SVDInterpreter(model).get_singular_vectors("OV", 9, head_index=9) up to sign, that keep=range(rank) leaves the logit diff within 1e-4, and that ablating every direction moves it by more than 1e-2.
| previous = getattr(model.cfg, "use_attn_result", False) | ||
| model.set_use_attn_result(True) | ||
| try: | ||
| _, cache = model.run_with_cache(prompt) |
There was a problem hiding this comment.
run_with_cache(prompt) retains every hook point to read one, about 1 GB against 19 MB on a 512-token gpt2-small prompt, and the batch guard at line 592 only runs after that forward. Pass names_filter for the one hook_result name and check a tensor prompt's batch dimension before the forward, keeping the post-cache check that catches a list of strings.
HookedTransformer is being removed, so the four consumer docstrings (vocab_readout, logit_signature, project_activations, patch_along_directions) and _validate_bridge_compatibility no longer list it as an accepted model or cite its LayerNorm folding as the reason the compatibility-mode guard fires only for TransformerBridge. The contract is now stated in terms of TransformerBridge alone: vocab_readout and logit_signature require compatibility mode so that W_U carries the folded final LayerNorm weights. Docstring-only, no behavior change. The guard still returns early for a non-TransformerBridge model; tightening that early-return to fail loudly is a behavior change, not a documentation one, so it is left out of this commit.
…state enable_compatibility_mode folds ln1 into W_V and centres W_O and W_U, so a HeadSVD taken before the call describes a different OV map than the model then computes. The existing guard inspected only the live model, so re-running a consumer with a decomposition captured under the other state returned a silently wrong result, and project_activations and patch_along_directions had no guard at all. Record the model's compatibility-mode state on HeadSVD at decomposition time and add an orthogonal guard that refuses a mismatch in all four consumers, pointing the caller at re-running decompose_head.
…eject a batched tensor pre-forward
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for working through all those comments. The compatibility-state binding and the dtype and device fixes are excellent. An additional batch of comments added below:
| model.set_use_attn_result(previous) | ||
|
|
||
| delta_metric = patched_metric - original_metric | ||
| baseline_delta_metric = sum(baseline_deltas) / len(baseline_deltas) |
There was a problem hiding this comment.
Averaging the control draws was my ask last round, but this averages them signed and then takes the absolute value, which is never larger than the mean of their magnitudes and is far smaller wherever the controls mix sign. At those widths the threshold understates the typical control effect and shrinks further as n_baseline grows, so ablate passes too easily and keep too rarely; on gpt2 L9H9 the default ablate=[0] gates True on a delta several times smaller than an arbitrary same-width subspace produces. Threshold on the mean of the per-draw magnitudes, and say whether baseline_delta_metric stays signed or becomes the magnitude, since the abs() below and the docstring at line 775 depend on that.
| # _resolve_retained has confirmed exactly one of keep/ablate is set, so the mode the | ||
| # caller expressed is unambiguous and selects the gate's success condition (see below). | ||
| mode = "keep" if keep is not None else "ablate" | ||
| if not retained and threshold is None: |
There was a problem hiding this comment.
width == rank has the tie this guard refuses for the empty set: the control projector is then V Vᵀ, the kept projector, so both runs apply the same projection and gated resolves on floating-point residue. It is reachable, since a full-rank keep passes the block guard, and it is the call the integration test makes at line 82. Extend this condition to a retained set that spans the full rank, give that line-82 call a threshold, and reword the message, which says the set is empty.
| def _validate_bridge_compatibility(model) -> None: | ||
| """Reject a ``TransformerBridge`` whose ``W_U`` would give a silently wrong projection. | ||
|
|
||
| Projecting an OV direction through ``W_U`` requires the final LayerNorm folded |
There was a problem hiding this comment.
The replacement sentence is also untrue: the bridge sets compatibility_mode = True before it processes weights and skips the processing entirely under no_processing, so the flag can be True with W_U unfolded and this guard passes. Check the recorded folding state, _weights_processed and the adapter's _fold_ln_requested, which the sibling guards already read, instead of the flag; while here, name direct_logit_attribution as the one tool this mirrors and fix the staleness guard's docstring below, which says this function checks the folding.
| Returns: | ||
| An :class:`ActivationProjection` with the per-position coefficients. | ||
|
|
||
| Raises: |
There was a problem hiding this comment.
set_use_attn_result(True) raises NotImplementedError on adapters without per-head results, and this function and patch_along_directions (line 852) both call it, but neither Raises block says so. Add it to both.
| slice reconstructed onto ``span(head_svd.V[:, retained])``, and once per random | ||
| control subspace. Each control subspace is drawn *inside the head's own OV span* | ||
| ``span(V)`` (not from the full residual stream, where a width-``w`` random subspace | ||
| would keep only ``w/d_model`` of a head output that lives entirely in ``w/rank`` of |
There was a problem hiding this comment.
The head output lives in span(V), which is rank/d_model of the stream; w/rank is what an in-span control of width w retains of it. Reword.
| rank_report: List[RankReportRow] | ||
| eps: float | ||
| null_rtol: float | ||
| compatibility_mode: bool |
There was a problem hiding this comment.
Can we default this to False? decompose_head passes it explicitly and the test helper already defaults it.
| # forward instead of running the model on every row only to discard the result. A list of | ||
| # prompt strings only reveals its batch size once tokenized, so the post-cache check below | ||
| # still guards that path. | ||
| if isinstance(prompt, torch.Tensor) and (prompt.ndim != 2 or prompt.shape[0] != 1): |
There was a problem hiding this comment.
Requiring exactly two dimensions also rejects a 1-D [pos] tensor, which the previous revision accepted with identical results and which patch_along_directions still accepts. Reject only a tensor whose leading dimension is larger than one.
| """A patch stub whose single head writes ``ov.V @ weights`` at every position and whose | ||
| readout is that same in-span vector. | ||
|
|
||
| Because the head output lies entirely in ``span(ov.V)`` and the readout equals it, removing |
There was a problem hiding this comment.
Because the readout equals the head output here, every control delta has the same sign, so the absolute signed mean equals the mean of magnitudes and no gate test on this stub can tell which one the gate uses; the same geometry lets a full-residual-stream control and a single draw pass every test. Add an assertion against the stub's analytic Haar expectation, E[baseline_delta] = (width/rank − 1)·‖h‖²·pos, with enough draws to separate in-span from full-stream, and one that threading a generator through single-draw calls reproduces the multi-draw average bit for bit.
| def test_patch_along_directions_discriminates_causal_direction(tiny_bridge): | ||
| """Ablating the head's strongest OV direction should move a metric aligned with that | ||
| direction materially more than ablating its weakest, least load-bearing direction.""" | ||
| layer, head = 0, 0 |
There was a problem hiding this comment.
The sibling-head test calls _make_subspace_hook directly and every patch_along_directions test in this file uses head 0, so the control-draw call inside patch_along_directions can be hardcoded to head 0 without failing anything; the kept-projector call is caught by the integration test. Run this test on head 1 as well, with a check on baseline_delta_metric that a control drawn for the wrong head cannot pass.
2 participants
Description
Adds the causal-validation layer on top of PR1's per-head QK/OV SVD (#1768): a direction can now be read out to vocab/logit space, projected against a real forward pass, and, critically, causally patched before any subfunction claim is accepted. Second 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).vocab_readout(model, head_svd, k=10)— projects the top-k OV output directions (HeadSVD.V) through the unembedding, gated onTransformerBridgecompatibility mode.logit_signature(model, head_svd, direction, tokens)— signed logit effect of one OV direction's rank-1 reconstruction; requires the direction to passrequire_isolatedfirst, since a rotation-ambiguous or null direction's signature is not attributable to it alone.project_activations(model, head_svd, prompt)— per-position firing coefficients of a head's actualhook_resultoutput against itsVbasis; summing the coefficients againstVreconstructs the real output to numerical precision, since both read the same basis the head actually writes in.patch_along_directions(model, head_svd, prompt, metric, keep=..., ablate=...)— the mandatory causal gate. Reconstructs or ablates a head's output onto a chosen span ofVdirections via a directhook_resulthook, and reportsdelta_metricagainst an equally-sized random-subspace baseline (gatedis true only when the requested subspace beats that baseline, not merely when the metric moves at all). Degenerate blocks reported byHeadSVD.degenerate_blocks()must be kept or dropped whole; a caller cannot route aroundrequire_isolatedby hand-picking part of a rotation-ambiguous block throughkeep/ablate.decompose_head,project_activations,patch_along_directions,vocab_readout,logit_signature, and their supporting types now ship fromtransformer_lens.tools.analysis— the first time this module is part of the public API, by design, so a readout was never importable without the causal gate attached to it.tests/integration/test_svd_circuits.py(new): exercisesvocab_readoutandpatch_along_directionsend to end against GPT-2 small's layer 9 head 9 (name-mover head) on an IOI-style prompt.Correctness fix folded in first:
HeadSVD's docstring had the OV output direction backwards —U's columns are the value-computation input space,V's are the residual-stream output space this head writes into and the one to project throughW_U. The swap never raised a shape error because both spaces ared_model-dimensional, so every downstream readout, projection, or patch would have silently used the wrong basis. Fixed first, with a regression test that cross-checks against the already-shippedSVDInterpreter(which projects the OV map the other way already), so the assertion cannot pass under either labeling.Two implementation notes for anyone diffing this against the original proposal:
vocab_readoutdoes not wrapSVDInterpreter.get_singular_vectors. That method indexesW_V/W_Kwith a raw query-head index, which is wrong on grouped-query models (the same bug PR1 already fixed for this module's own weight reads).vocab_readoutreuses PR1's already-kv-mappedHeadSVD.Vand projects it throughW_Udirectly instead of introducing a second, buggy SVD path.patch_along_directionscallsmodel.run_with_hooksdirectly rather thangeneric_activation_patch. That helper is built for a clean-vs-corrupted two-run sweep over an index grid; this tool is a single-prompt reconstruct/ablate on a live projection, which does not fit the sweep shape.Unit-tested with a tiny no-download
TransformerBridge(real forward passes, no Hub access) for projection/patch behavior, plus model-free tests for the pure weight-space math. Integration test downloadsgpt2-small.PR3 (docs, demo notebook, slow multi-head oracle-parity check) follows this one.
Part of #1767
Type of change
Checklist: