From e6cc413bc98029fc5bf9180b700f701ae60385c1 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 12 Sep 2026 10:40:03 +0530 Subject: [PATCH 01/14] fix(svd_circuits): correct the OV output-direction docstring and add 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. --- tests/unit/tools/test_svd_circuits.py | 63 ++++++++++++++++++- .../tools/analysis/svd_circuits.py | 28 ++++++--- 2 files changed, 82 insertions(+), 9 deletions(-) diff --git a/tests/unit/tools/test_svd_circuits.py b/tests/unit/tools/test_svd_circuits.py index 913e8dca5..855087f5e 100644 --- a/tests/unit/tools/test_svd_circuits.py +++ b/tests/unit/tools/test_svd_circuits.py @@ -1,11 +1,14 @@ """Unit tests for per-head QK/OV singular-vector decomposition. -Model-free: they build synthetic weight tensors and exercise the factored SVD and the -degeneracy guard directly, so no model is loaded and no pretrained weights are downloaded. +Mostly model-free: most tests build synthetic weight tensors and exercise the factored SVD +and the degeneracy guard directly, so no model is loaded and no pretrained weights are +downloaded. A few tests instantiate a tiny, randomly-initialized TransformerBridge (CPU-only, +no Hub access) where a real per-head decomposition is needed for a cross-check. """ import warnings from types import SimpleNamespace +from unittest.mock import MagicMock import pytest import torch @@ -470,3 +473,59 @@ def test_decompose_head_rejects_bare_string_which(): model = _StubModel(n_heads=2, n_kv_heads=2) with pytest.raises(ValueError, match="sequence"): decompose_head(model, layer=0, head=0, which="QK") + + +# --------------------------------------------------------------------------- # +# OV output-direction convention (cross-check against SVDInterpreter) +# --------------------------------------------------------------------------- # +@pytest.fixture(scope="module") +def tiny_bridge(): + """Tiny GPT-2 bridge: MHA (n_heads == n_kv_heads), CPU-only, no Hub access.""" + from transformers import GPT2Config, GPT2LMHeadModel + + from transformer_lens.config.transformer_bridge_config import ( + TransformerBridgeConfig, + ) + from transformer_lens.model_bridge import TransformerBridge + from transformer_lens.model_bridge.supported_architectures.gpt2 import ( + GPT2ArchitectureAdapter, + ) + + torch.manual_seed(0) + hf_model = GPT2LMHeadModel( + GPT2Config(n_embd=32, n_layer=1, n_head=2, vocab_size=64, n_positions=32) + ).eval() + cfg = TransformerBridgeConfig( + d_model=32, + d_head=16, + n_heads=2, + n_layers=1, + n_ctx=32, + d_vocab=64, + architecture="GPT2LMHeadModel", + ) + return TransformerBridge(hf_model, GPT2ArchitectureAdapter(cfg), tokenizer=MagicMock()) + + +def test_ov_output_direction_matches_svd_interpreter(tiny_bridge): + """OV's write/vocab-readout direction is ``.V``, not ``.U`` - checked against the + already-shipped ``SVDInterpreter``, not just internal self-consistency, so the + assertion cannot pass under either U/V labeling by construction.""" + from transformer_lens.SVDInterpreter import SVDInterpreter + + layer, head = 0, 0 + decomposition = decompose_head(tiny_bridge, layer, head, which=("OV",)) + ov = decomposition.OV + W_U = tiny_bridge.W_U + + interpreter = SVDInterpreter(tiny_bridge) + reference = interpreter.get_singular_vectors("OV", layer, head_index=head, num_vectors=1) + + projected_via_V = ov.V[:, 0] @ W_U + assert torch.allclose(projected_via_V, reference[:, 0, 0], atol=1e-4) or torch.allclose( + projected_via_V, -reference[:, 0, 0], atol=1e-4 + ) + + projected_via_U = ov.U[:, 0] @ W_U + assert not torch.allclose(projected_via_U, reference[:, 0, 0], atol=1e-2) + assert not torch.allclose(projected_via_U, -reference[:, 0, 0], atol=1e-2) diff --git a/transformer_lens/tools/analysis/svd_circuits.py b/transformer_lens/tools/analysis/svd_circuits.py index fea61efae..1ce524e67 100644 --- a/transformer_lens/tools/analysis/svd_circuits.py +++ b/transformer_lens/tools/analysis/svd_circuits.py @@ -5,7 +5,7 @@ ``W_V W_O`` that writes the attended value back into the residual stream. This tool takes the singular value decomposition of each map for one head and exposes the singular values (how much each direction matters) together with the left and -right singular vectors (the output and input directions they act on). +right singular vectors that span the map's input and output spaces. The decomposition is weight-space only: it reads ``W_Q``/``W_K``/``W_V``/``W_O`` and needs no forward pass, no activation cache, and no compatibility mode. Each @@ -14,9 +14,18 @@ ``d_model x d_model`` product is never materialized and the returned rank is bounded by ``d_head``. -Right singular vectors are read from :attr:`~transformer_lens.FactoredMatrix.FactoredMatrix.V` -(its columns are the right singular vectors). The historical ``.Vh`` alias is -deprecated and returns the same tensor, so it is never used here. +For a factored map ``A @ B`` (``A: [ldim, mdim]``, ``B: [mdim, rdim]``), the SVD's +``U`` columns live in ``A``'s input space (``ldim``) and ``V`` columns live in +``B``'s output space (``rdim``): feeding ``x = U[:, i]`` through the map gives +``x @ (A @ B) == S[i] * V[:, i]``, never the reverse. For OV (``A = W_V_h``, +``B = W_O_h``), ``U``'s columns are therefore the value-computation *input* +directions this head reads from the residual stream, and ``V``'s columns are the +*output* directions it writes back into the residual stream - the ones to +project through ``W_U`` for a vocab or logit readout. For QK (``A = W_Q_h``, +``B = W_K_h.transpose(-1, -2)``), both ``U`` (destination/query-read) and ``V`` +(source/key-read) are read directions; QK only ever produces a scalar attention +score, so neither is a write direction. The historical ``.Vh`` alias returns the +same tensor as ``.V`` and is never used here. Adjacent singular values closer than a relative gap ``eps`` leave their singular directions defined only up to a rotation, so the result carries a per-direction @@ -104,10 +113,15 @@ class HeadSVD: which: Which map this decomposes, ``"QK"`` or ``"OV"``. layer: Layer of the decomposed head. head: Head index within the layer. - U: Left singular vectors, ``[d_model, rank]`` (column i is output direction i). + U: Left singular vectors, ``[d_model, rank]``: column i is the map's input + direction i (for OV, the residual-stream direction this head's value + computation reads from; for QK, the destination/query-read direction). S: Singular values, ``[rank]``, sorted descending. - V: Right singular vectors, ``[d_model, rank]`` (column i is input direction i). - The reconstruction is ``U @ S.diag() @ V.transpose(-2, -1)``. + V: Right singular vectors, ``[d_model, rank]``: column i is the map's output + direction i for OV (the residual-stream direction this head writes into, + the one to project through ``W_U``), or the source/key-read direction for + QK (QK produces no write direction). The reconstruction is + ``U @ S.diag() @ V.transpose(-2, -1)``. rank_report: Per-direction :class:`RankReportRow` list, aligned with the columns of ``U``/``V``. eps: Relative gap below which adjacent directions share a block; every block From 7c4e283db6222e39690a15c3921da19e3387bd1f Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 12 Sep 2026 15:05:10 +0530 Subject: [PATCH 02/14] feat(svd_circuits): add OV vocab readout and logit signature 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. --- tests/unit/tools/test_svd_circuits.py | 59 +++++++++ .../tools/analysis/svd_circuits.py | 115 +++++++++++++++++- 2 files changed, 173 insertions(+), 1 deletion(-) diff --git a/tests/unit/tools/test_svd_circuits.py b/tests/unit/tools/test_svd_circuits.py index 855087f5e..a4db78883 100644 --- a/tests/unit/tools/test_svd_circuits.py +++ b/tests/unit/tools/test_svd_circuits.py @@ -16,10 +16,13 @@ from transformer_lens.tools.analysis.svd_circuits import ( DegenerateDirectionError, HeadSVD, + LogitSignature, RankReportRow, _degeneracy_blocks, _factored_head_svd, decompose_head, + logit_signature, + vocab_readout, ) D_MODEL, D_HEAD = 12, 4 @@ -529,3 +532,59 @@ def test_ov_output_direction_matches_svd_interpreter(tiny_bridge): projected_via_U = ov.U[:, 0] @ W_U assert not torch.allclose(projected_via_U, reference[:, 0, 0], atol=1e-2) assert not torch.allclose(projected_via_U, -reference[:, 0, 0], atol=1e-2) + + +# --------------------------------------------------------------------------- # +# vocab_readout / logit_signature (OV output-direction readout) +# --------------------------------------------------------------------------- # +def test_vocab_readout_shape_and_which_guard(): + """vocab_readout returns [d_vocab, k] for OV, and rejects QK input and out-of-range k.""" + ov = _factored_head_svd(*_random_ov(), which="OV", layer=0, head=0, eps=1e-2) + vocab_size = 20 + model = SimpleNamespace(W_U=torch.randn(D_MODEL, vocab_size)) + + result = vocab_readout(model, ov, k=3) + assert result.shape == (vocab_size, 3) + + qk = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="QK", layer=0, head=0, eps=1e-2 + ) + with pytest.raises(ValueError, match="OV"): + vocab_readout(model, qk, k=1) + + with pytest.raises(ValueError, match="k must be"): + vocab_readout(model, ov, k=D_HEAD + 1) + + +def test_vocab_readout_raises_without_compatibility_mode(tiny_bridge): + """A TransformerBridge without compatibility mode enabled refuses to project through W_U.""" + decomposition = decompose_head(tiny_bridge, layer=0, head=0, which=("OV",)) + with pytest.raises(ValueError, match="enable_compatibility_mode"): + vocab_readout(tiny_bridge, decomposition.OV) + + +def test_logit_signature_matches_manual_projection(): + """logit_signature's reconstruction matches a hand-computed S[i] * V[:, i] projection.""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + vocab_size = 16 + W_U = torch.randn(D_MODEL, vocab_size) + model = SimpleNamespace(W_U=W_U) + token_ids = torch.tensor([1, 5, 9]) + + result = logit_signature(model, ov, direction=0, tokens=token_ids) + assert isinstance(result, LogitSignature) + expected = (ov.S[0] * ov.V[:, 0]) @ W_U[:, token_ids] + assert torch.allclose(result.values, expected, atol=1e-5) + assert result.direction == 0 + + +def test_logit_signature_raises_on_degenerate_direction(): + """A rotation-ambiguous direction's 'signature' is not attributable, so this must raise.""" + ov = _factored_head_svd( + *_factored_with_spectrum([5.0, 3.0, 3.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + model = SimpleNamespace(W_U=torch.randn(D_MODEL, 8)) + with pytest.raises(DegenerateDirectionError): + logit_signature(model, ov, direction=1, tokens=torch.tensor([0])) diff --git a/transformer_lens/tools/analysis/svd_circuits.py b/transformer_lens/tools/analysis/svd_circuits.py index 1ce524e67..6642eb55b 100644 --- a/transformer_lens/tools/analysis/svd_circuits.py +++ b/transformer_lens/tools/analysis/svd_circuits.py @@ -52,7 +52,7 @@ """ from dataclasses import dataclass -from typing import List, Literal, Optional, Sequence, Tuple +from typing import List, Literal, Optional, Sequence, Tuple, Union import torch from jaxtyping import Float @@ -420,3 +420,116 @@ def decompose_head( W_V_h, W_O_h, which="OV", layer=layer, head=head, eps=eps, null_rtol=null_rtol ) return HeadDecomposition(layer=layer, head=head, QK=qk, OV=ov) + + +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 + only fires for ``TransformerBridge``. Mirrors the compatibility-mode check other + unembedding-touching analysis tools already run, without any hybrid-architecture + restriction: projecting a rank-1 OV direction through ``W_U`` does not depend on + the block-layout assumptions that check exists for elsewhere. + """ + # Lazy import - keeps the module importable without the bridge as a hard dependency. + from transformer_lens.model_bridge import TransformerBridge + + if not isinstance(model, TransformerBridge): + return + if not getattr(model, "compatibility_mode", False): + 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." + ) + + +def vocab_readout( + model, head_svd: HeadSVD, *, k: int = 10 +) -> Float[torch.Tensor, "d_vocab k"]: + """Project the top-k OV output directions through the unembedding. + + Requires ``head_svd.which == "OV"``: QK produces no write direction to project + (see the module docstring). On a ``TransformerBridge``, compatibility mode must + be enabled so ``W_U`` carries the folded final LayerNorm weights; + ``HookedTransformer`` always has this folding applied. + + Does not call ``head_svd.require_isolated``: a degenerate direction's vocab + readout is still a well-defined projection, unlike a per-direction causal claim, + so it is not gated here. The contract that no direction is reported without a + passing causal patch is enforced by :func:`patch_along_directions`. + + Args: + model: A ``TransformerBridge`` (with compatibility mode enabled) or a + ``HookedTransformer``; only its ``W_U`` is read. + head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. + k: Number of top singular directions to project. + + Returns: + ``W_U.T @ head_svd.V[:, :k]``, shape ``[d_vocab, k]``: column i is + direction i's projection through the unembedding. + + Raises: + ValueError: If ``head_svd.which != "OV"``, if ``k`` is not in + ``(0, rank]``, or if ``model`` is a ``TransformerBridge`` without + compatibility mode enabled. + """ + if head_svd.which != "OV": + raise ValueError(f"vocab_readout requires an OV HeadSVD, got which={head_svd.which!r}") + rank = head_svd.V.shape[1] + 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() + + +@dataclass +class LogitSignature: + """Rank-1-reconstruction logit effect for one OV direction, per requested token. + + Attributes: + direction: Which ``HeadSVD`` column this reconstructs. + values: Signed logit contribution, aligned with the requested tokens. + """ + + direction: int + values: Float[torch.Tensor, "token"] + + +def logit_signature( + model, + head_svd: HeadSVD, + direction: int, + tokens: Union[int, Sequence[int], torch.Tensor], +) -> LogitSignature: + """Signed logit effect of one OV direction's rank-1 reconstruction on the given tokens. + + Pure weight-space computation: reconstructs the head's OV output along a single + singular direction (``S[direction] * V[:, direction]``, never ``U`` - see the + module docstring) and projects it through ``W_U`` restricted to ``tokens``. Runs + no forward pass and builds no cache. + + Args: + model: A ``TransformerBridge`` (with compatibility mode enabled) or a + ``HookedTransformer``; only its ``W_U`` is read. + head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. + direction: Column index of the singular direction to reconstruct. + tokens: Token id(s) to read the logit effect for. + + Returns: + A :class:`LogitSignature` with one value per requested token. + + Raises: + ValueError: If ``head_svd.which != "OV"``, or if ``model`` is a + ``TransformerBridge`` without compatibility mode enabled. + DegenerateDirectionError: If ``direction`` is not attributable alone (see + :meth:`HeadSVD.require_isolated`). + """ + if head_svd.which != "OV": + raise ValueError(f"logit_signature requires an OV HeadSVD, got which={head_svd.which!r}") + head_svd.require_isolated(direction) + _validate_bridge_compatibility(model) + token_ids = torch.as_tensor(tokens, dtype=torch.long).reshape(-1) + reconstruction = (head_svd.S[direction] * head_svd.V[:, direction]).float() + values = reconstruction @ model.W_U[:, token_ids] + return LogitSignature(direction=direction, values=values) From 7a4c8eb89d22eeb3379736038f9b9c61e2db8dd1 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 12 Sep 2026 15:25:17 +0530 Subject: [PATCH 03/14] feat(svd_circuits): add project_activations firing coefficients 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. --- tests/unit/tools/test_svd_circuits.py | 58 ++++++++++++++- .../tools/analysis/svd_circuits.py | 70 ++++++++++++++++++- 2 files changed, 124 insertions(+), 4 deletions(-) diff --git a/tests/unit/tools/test_svd_circuits.py b/tests/unit/tools/test_svd_circuits.py index a4db78883..5e5737c93 100644 --- a/tests/unit/tools/test_svd_circuits.py +++ b/tests/unit/tools/test_svd_circuits.py @@ -14,6 +14,7 @@ import torch from transformer_lens.tools.analysis.svd_circuits import ( + ActivationProjection, DegenerateDirectionError, HeadSVD, LogitSignature, @@ -22,6 +23,7 @@ _factored_head_svd, decompose_head, logit_signature, + project_activations, vocab_readout, ) @@ -507,7 +509,11 @@ def tiny_bridge(): d_vocab=64, architecture="GPT2LMHeadModel", ) - return TransformerBridge(hf_model, GPT2ArchitectureAdapter(cfg), tokenizer=MagicMock()) + tokenizer = MagicMock() + # A real tokenizer isn't available (no-download fixture), but to_str_tokens's return + # value is jaxtyped-checked as List[str], so batch_decode must return actual strings. + tokenizer.batch_decode = lambda tokens_list, **kwargs: [str(ids[0]) for ids in tokens_list] + return TransformerBridge(hf_model, GPT2ArchitectureAdapter(cfg), tokenizer=tokenizer) def test_ov_output_direction_matches_svd_interpreter(tiny_bridge): @@ -588,3 +594,53 @@ def test_logit_signature_raises_on_degenerate_direction(): model = SimpleNamespace(W_U=torch.randn(D_MODEL, 8)) with pytest.raises(DegenerateDirectionError): logit_signature(model, ov, direction=1, tokens=torch.tensor([0])) + + +# --------------------------------------------------------------------------- # +# project_activations (per-position firing coefficients in the OV output basis) +# --------------------------------------------------------------------------- # +def test_project_activations_which_guard(): + """QK has no write direction to project onto, so a QK HeadSVD is refused.""" + qk = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="QK", layer=0, head=0, eps=1e-2 + ) + with pytest.raises(ValueError, match="OV"): + project_activations(SimpleNamespace(), qk, "hello world") + + +def test_project_activations_reconstructs_head_output(tiny_bridge): + """Coefficients recovered against V reconstruct the actual cached hook_result slice.""" + layer, head = 0, 0 + decomposition = decompose_head(tiny_bridge, layer, head, which=("OV",)) + # A raw token tensor, not a string: tiny_bridge's tokenizer is a MagicMock and cannot + # tokenize text, but a token-id tensor bypasses that path entirely. + prompt = torch.tensor([[5, 63, 7, 9]]) + + projection = project_activations(tiny_bridge, decomposition.OV, prompt) + assert isinstance(projection, ActivationProjection) + assert projection.head_svd is decomposition.OV + + previous = tiny_bridge.cfg.use_attn_result + tiny_bridge.set_use_attn_result(True) + try: + _, cache = tiny_bridge.run_with_cache(prompt) + finally: + tiny_bridge.set_use_attn_result(previous) + expected = cache[("result", layer, "attn")][0, :, head, :] + + reconstructed = projection.coefficients @ decomposition.OV.V.transpose(-2, -1) + assert torch.allclose(reconstructed, expected, atol=1e-4) + assert projection.str_tokens == tiny_bridge.to_str_tokens(prompt) + + +def test_project_activations_restores_use_attn_result(tiny_bridge): + """use_attn_result is restored to its prior value regardless of what it started as.""" + decomposition = decompose_head(tiny_bridge, 0, 0, which=("OV",)) + original = tiny_bridge.cfg.use_attn_result + try: + for initial in (False, True): + tiny_bridge.set_use_attn_result(initial) + project_activations(tiny_bridge, decomposition.OV, torch.tensor([[5, 63, 7, 9]])) + assert tiny_bridge.cfg.use_attn_result == initial + finally: + tiny_bridge.set_use_attn_result(original) diff --git a/transformer_lens/tools/analysis/svd_circuits.py b/transformer_lens/tools/analysis/svd_circuits.py index 6642eb55b..f42b12309 100644 --- a/transformer_lens/tools/analysis/svd_circuits.py +++ b/transformer_lens/tools/analysis/svd_circuits.py @@ -444,9 +444,7 @@ def _validate_bridge_compatibility(model) -> None: ) -def vocab_readout( - model, head_svd: HeadSVD, *, k: int = 10 -) -> Float[torch.Tensor, "d_vocab k"]: +def vocab_readout(model, head_svd: HeadSVD, *, k: int = 10) -> Float[torch.Tensor, "d_vocab k"]: """Project the top-k OV output directions through the unembedding. Requires ``head_svd.which == "OV"``: QK produces no write direction to project @@ -533,3 +531,69 @@ def logit_signature( reconstruction = (head_svd.S[direction] * head_svd.V[:, direction]).float() values = reconstruction @ model.W_U[:, token_ids] return LogitSignature(direction=direction, values=values) + + +@dataclass +class ActivationProjection: + """Per-position coefficients of a head's actual output in its OV output basis. + + Attributes: + head_svd: The OV decomposition this was projected against. + coefficients: ``[pos, rank]``; ``coefficients[:, i]`` is the signed amount of + singular direction ``i`` (``head_svd.V[:, i]``) present in the head's actual + output at each position. Summing ``coefficients[:, i] * head_svd.V[:, i]`` + over ``i`` reconstructs the head's real per-position output to numerical + precision, since ``V``'s columns are orthonormal and this projects onto the + exact basis the head writes in. + str_tokens: Tokenized prompt, aligned with the position axis, for display. + """ + + head_svd: HeadSVD + coefficients: Float[torch.Tensor, "pos rank"] + str_tokens: List[str] + + +def project_activations( + model, head_svd: HeadSVD, prompt: Union[str, torch.Tensor] +) -> ActivationProjection: + """Project a head's actual per-position output onto its OV singular directions. + + Requires ``head_svd.which == "OV"``: this projects onto the write/output basis + ``V``, and QK has no such vector (see the module docstring). Runs a real forward + pass with ``use_attn_result`` enabled to read the per-head output + (``hook_result``), then projects it onto ``head_svd.V``. Restores the model's + prior ``use_attn_result`` setting afterward, since flipping that config flag as a + side effect of a read-only analysis call would surprise a caller who already had + hooks or a cache built around its prior state. + + Args: + model: A ``TransformerBridge`` or ``HookedTransformer``. + head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. + prompt: A single prompt (not a batch): a string or a ``[1, pos]`` token tensor. + + Returns: + An :class:`ActivationProjection` with the per-position coefficients. + + Raises: + ValueError: If ``head_svd.which != "OV"``, or if ``prompt`` is not a single + (batch-size-1) prompt. + """ + if head_svd.which != "OV": + raise ValueError( + f"project_activations requires an OV HeadSVD, got which={head_svd.which!r}" + ) + previous = getattr(model.cfg, "use_attn_result", False) + model.set_use_attn_result(True) + try: + _, cache = model.run_with_cache(prompt) + finally: + model.set_use_attn_result(previous) + result = cache[("result", head_svd.layer, "attn")][..., head_svd.head, :] + if result.shape[0] != 1: + raise ValueError( + f"project_activations requires a single prompt, got batch={result.shape[0]}" + ) + result = result.squeeze(0).to(head_svd.V.dtype) + coefficients = result @ head_svd.V + str_tokens = model.to_str_tokens(prompt) + return ActivationProjection(head_svd=head_svd, coefficients=coefficients, str_tokens=str_tokens) From 5b8118d8009b0356fe18fb2b09486f5d8962c9b9 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 12 Sep 2026 15:42:24 +0530 Subject: [PATCH 04/14] feat(svd_circuits): add patch_along_directions causal gate 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. --- tests/unit/tools/test_svd_circuits.py | 143 ++++++++++++++ .../tools/analysis/svd_circuits.py | 181 +++++++++++++++++- 2 files changed, 323 insertions(+), 1 deletion(-) diff --git a/tests/unit/tools/test_svd_circuits.py b/tests/unit/tools/test_svd_circuits.py index 5e5737c93..6dcc98a70 100644 --- a/tests/unit/tools/test_svd_circuits.py +++ b/tests/unit/tools/test_svd_circuits.py @@ -6,6 +6,7 @@ no Hub access) where a real per-head decomposition is needed for a cross-check. """ +import math import warnings from types import SimpleNamespace from unittest.mock import MagicMock @@ -18,11 +19,13 @@ DegenerateDirectionError, HeadSVD, LogitSignature, + PatchResult, RankReportRow, _degeneracy_blocks, _factored_head_svd, decompose_head, logit_signature, + patch_along_directions, project_activations, vocab_readout, ) @@ -644,3 +647,143 @@ def test_project_activations_restores_use_attn_result(tiny_bridge): assert tiny_bridge.cfg.use_attn_result == initial finally: tiny_bridge.set_use_attn_result(original) + + +# --------------------------------------------------------------------------- # +# patch_along_directions (mandatory causal gate) +# --------------------------------------------------------------------------- # +class _PatchStubModel: + """Model-free stand-in for patch_along_directions' run_with_hooks/metric protocol. + + Exposes a single fixed per-head "hook_result" activation that a hook can intercept + exactly as ``run_with_hooks`` would present it (the hook receives the activation and + a hook object, and returns the replacement), then reduces it to a scalar the same + way for the unmodified, patched, and baseline calls, so the block guard and gating + arithmetic can be tested without a real forward pass. + """ + + def __init__(self, d_model, n_heads, pos=3, seed=0): + g = torch.Generator().manual_seed(seed) + self.cfg = SimpleNamespace(use_attn_result=False) + self._result = torch.randn(1, pos, n_heads, d_model, generator=g) + self._readout = torch.randn(d_model, generator=g) + + def set_use_attn_result(self, value): + self.cfg.use_attn_result = value + + def _logits(self, result): + return result.sum(dim=2) @ self._readout # [batch, pos] + + def __call__(self, prompt): + return self._logits(self._result) + + def run_with_hooks(self, prompt, fwd_hooks): + activation = self._result + for name, hook_fn in fwd_hooks: + activation = hook_fn(activation, hook=SimpleNamespace(name=name)) + return self._logits(activation) + + +def test_patch_along_directions_which_guard(): + """QK has no write direction to reconstruct onto, so a QK HeadSVD is refused.""" + qk = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="QK", layer=0, head=0, eps=1e-2 + ) + with pytest.raises(ValueError, match="OV"): + patch_along_directions(SimpleNamespace(), qk, "prompt", lambda logits: 0.0, keep=[0]) + + +def test_patch_along_directions_requires_keep_xor_ablate(): + """Passing both keep and ablate, or neither, is refused before any model access.""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + metric = lambda logits: 0.0 + with pytest.raises(ValueError, match="exactly one"): + patch_along_directions(SimpleNamespace(), ov, "prompt", metric, keep=[0], ablate=[1]) + with pytest.raises(ValueError, match="exactly one"): + patch_along_directions(SimpleNamespace(), ov, "prompt", metric) + + +def test_patch_along_directions_rejects_partial_degenerate_block(): + """keep must take a degenerate block whole or not at all; a partial slice is refused, + and the whole block plus an isolated direction is accepted.""" + ov = _factored_head_svd( + *_factored_with_spectrum([5.0, 3.0, 3.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + metric = lambda logits: 0.0 + with pytest.raises(DegenerateDirectionError): + patch_along_directions(SimpleNamespace(), ov, "prompt", metric, keep=[0, 1]) + + stub = _PatchStubModel(d_model=D_MODEL, n_heads=1) + result = patch_along_directions( + stub, ov, "prompt", lambda logits: float(logits.sum()), keep=[0, 1, 2] + ) + assert isinstance(result, PatchResult) + assert result.retained == [0, 1, 2] + + +def test_patch_along_directions_reports_delta_and_restores_use_attn_result(): + """delta_metric is patched minus original metric, gated follows the documented + comparison against baseline_delta_metric, and use_attn_result is restored after.""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + stub = _PatchStubModel(d_model=D_MODEL, n_heads=1) + metric = lambda logits: float(logits.sum()) + for initial in (False, True): + stub.set_use_attn_result(initial) + result = patch_along_directions( + stub, ov, "prompt", metric, keep=[0], rng=torch.Generator().manual_seed(0) + ) + 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) + assert result.gated == expected_gated + assert stub.cfg.use_attn_result == initial + + +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 + decomposition = decompose_head(tiny_bridge, layer, head, which=("OV",)) + ov = decomposition.OV + strong = ov.rank_report[0].idx + weak = ov.rank_report[-1].idx + assert not ov.is_degenerate(strong) + assert not ov.is_degenerate(weak) + + target_token = int((ov.V[:, strong] @ tiny_bridge.W_U).argmax()) + prompt = torch.tensor([[5, 63, 7, 9]]) + + def metric(logits): + return logits[0, -1, target_token].item() + + strong_result = patch_along_directions( + tiny_bridge, ov, prompt, metric, ablate=[strong], rng=torch.Generator().manual_seed(0) + ) + weak_result = patch_along_directions( + tiny_bridge, ov, prompt, metric, ablate=[weak], rng=torch.Generator().manual_seed(0) + ) + assert isinstance(strong_result, PatchResult) + assert math.isfinite(strong_result.baseline_delta_metric) + assert math.isfinite(weak_result.baseline_delta_metric) + assert abs(strong_result.delta_metric) > abs(weak_result.delta_metric) + + +def test_patch_along_directions_restores_use_attn_result(tiny_bridge): + """use_attn_result is restored to its prior value regardless of what it started as.""" + decomposition = decompose_head(tiny_bridge, 0, 0, which=("OV",)) + ov = decomposition.OV + metric = lambda logits: float(logits.sum()) + original = tiny_bridge.cfg.use_attn_result + try: + for initial in (False, True): + tiny_bridge.set_use_attn_result(initial) + patch_along_directions( + tiny_bridge, ov, torch.tensor([[5, 63, 7, 9]]), metric, keep=[ov.rank_report[0].idx] + ) + assert tiny_bridge.cfg.use_attn_result == initial + finally: + tiny_bridge.set_use_attn_result(original) diff --git a/transformer_lens/tools/analysis/svd_circuits.py b/transformer_lens/tools/analysis/svd_circuits.py index f42b12309..2a5f48606 100644 --- a/transformer_lens/tools/analysis/svd_circuits.py +++ b/transformer_lens/tools/analysis/svd_circuits.py @@ -52,7 +52,7 @@ """ from dataclasses import dataclass -from typing import List, Literal, Optional, Sequence, Tuple, Union +from typing import Callable, List, Literal, Optional, Sequence, Tuple, Union import torch from jaxtyping import Float @@ -597,3 +597,182 @@ def project_activations( coefficients = result @ head_svd.V str_tokens = model.to_str_tokens(prompt) return ActivationProjection(head_svd=head_svd, coefficients=coefficients, str_tokens=str_tokens) + + +def _validate_retained_blocks(head_svd: HeadSVD, retained: Sequence[int]) -> None: + """Raise if ``retained`` splits a degenerate block instead of keeping it whole or empty. + + A degenerate block's members are defined only up to a rotation (or, for a null + block, arbitrary null-space vectors), so attributing a causal effect to part of the + block while dropping the rest would let a caller route around the guard + :meth:`HeadSVD.require_isolated` already enforces per direction. + """ + retained_set = set(retained) + for block in head_svd.degenerate_blocks(): + block_set = set(block) + overlap = retained_set & block_set + if overlap and overlap != block_set: + raise DegenerateDirectionError( + f"retained directions {sorted(overlap)} split block {sorted(block_set)} of " + f"the {head_svd.which} SVD of head L{head_svd.layer}H{head_svd.head}, whose " + f"members are not separated by a relative gap of eps={head_svd.eps:g} or are " + f"jointly null. Keep or ablate the whole block, not part of it." + ) + + +def _resolve_retained( + head_svd: HeadSVD, keep: Optional[Sequence[int]], ablate: Optional[Sequence[int]] +) -> List[int]: + """Resolve ``keep``/``ablate`` to the sorted list of retained direction indices. + + Exactly one of ``keep``/``ablate`` must be given; ``ablate``'s complement over the + map's full rank becomes the retained set. Raises :class:`DegenerateDirectionError` + if the result would split a degenerate block (see :func:`_validate_retained_blocks`). + """ + if (keep is None) == (ablate is None): + 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)) + else: + assert ablate is not None + ablate_set = set(ablate) + retained = [i for i in range(rank) if i not in ablate_set] + _validate_retained_blocks(head_svd, retained) + return retained + + +def _make_subspace_hook(head: int, projector: Float[torch.Tensor, "d_model d_model"]): + """Build a ``hook_result`` hook that reconstructs one head's output onto ``span(projector)``. + + Leaves every other head's slice of the ``[batch, pos, head_index, d_model]`` tensor + untouched. Clones before mutating so the hook never writes into the activation the + forward pass itself is still using. + """ + + def hook_fn(activation: torch.Tensor, hook) -> torch.Tensor: + activation = activation.clone() + activation[:, :, head, :] = activation[:, :, head, :] @ projector.to(activation.dtype) + return activation + + return hook_fn + + +@dataclass +class PatchResult: + """Result of causally patching a head's output onto a chosen OV singular subspace. + + Attributes: + head_svd: The OV decomposition patched against. + retained: Direction indices whose span the head's output was reconstructed + onto; the complement was zeroed. + original_metric: Metric value on the unmodified prompt. + patched_metric: Metric value after the subspace reconstruction. + delta_metric: ``patched_metric - original_metric``. + baseline_delta_metric: ``delta_metric`` from reconstructing onto a random + subspace of the same rank as ``retained``, instead of the requested one. + gated: True only if ``abs(delta_metric)`` exceeds ``abs(baseline_delta_metric)`` + (or an explicit threshold, if one was passed): a subfunction is causally + load-bearing only if it beats an equally-sized random subspace, not merely + "moves the metric at all". + """ + + head_svd: HeadSVD + retained: List[int] + original_metric: float + patched_metric: float + delta_metric: float + baseline_delta_metric: float + gated: bool + + +@torch.no_grad() +def patch_along_directions( + model, + head_svd: HeadSVD, + prompt: Union[str, torch.Tensor], + metric: Callable[[torch.Tensor], float], + *, + keep: Optional[Sequence[int]] = None, + ablate: Optional[Sequence[int]] = None, + threshold: Optional[float] = None, + rng: Optional[torch.Generator] = None, +) -> PatchResult: + """Causally validate a claimed OV subfunction by reconstructing the head's output onto it. + + Requires ``head_svd.which == "OV"``: this reconstructs the write/output basis + ``V``, and QK has no such vector (see the module docstring). Runs the prompt three + times with ``use_attn_result`` enabled: once unmodified, once with the head's + ``hook_result`` slice reconstructed onto ``span(head_svd.V[:, retained])``, and once + onto a random orthonormal subspace of the same width, so a moved metric can be + compared against the effect of an equally-sized but arbitrary subspace instead of + being read as significant on its own. Restores the model's prior ``use_attn_result`` + setting afterward. + + Args: + model: A ``TransformerBridge`` or ``HookedTransformer``. + head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. + prompt: A single prompt: a string or a ``[1, pos]`` token tensor. + metric: A function from the model's logits to a scalar. + keep: Direction indices to retain; the rest are zeroed. Exactly one of + ``keep``/``ablate`` must be given. + ablate: Direction indices to zero; the rest are retained. + threshold: Explicit gate threshold. Defaults to ``None``, which uses + ``abs(baseline_delta_metric)`` instead. + rng: Optional generator for the random baseline subspace, for reproducibility. + + Returns: + A :class:`PatchResult` describing the patched, baseline, and original metrics. + + Raises: + ValueError: If ``head_svd.which != "OV"``, or if ``keep``/``ablate`` are both + given or both omitted. + DegenerateDirectionError: If the retained directions split a degenerate block + (see :func:`_validate_retained_blocks`). + """ + if head_svd.which != "OV": + raise ValueError( + f"patch_along_directions requires an OV HeadSVD, got which={head_svd.which!r}" + ) + retained = _resolve_retained(head_svd, keep, ablate) + + V = head_svd.V + kept_projector = V[:, retained] @ V[:, retained].transpose(-2, -1) + + d_model = V.shape[0] + width = len(retained) + random_input = torch.randn(d_model, width, generator=rng, dtype=V.dtype) + random_basis, _ = torch.linalg.qr(random_input) + baseline_projector = random_basis @ random_basis.transpose(-2, -1) + + hook_name = f"blocks.{head_svd.layer}.attn.hook_result" + previous = getattr(model.cfg, "use_attn_result", False) + model.set_use_attn_result(True) + try: + original_metric = float(metric(model(prompt))) + patched_logits = model.run_with_hooks( + prompt, fwd_hooks=[(hook_name, _make_subspace_hook(head_svd.head, kept_projector))] + ) + patched_metric = float(metric(patched_logits)) + baseline_logits = model.run_with_hooks( + prompt, + fwd_hooks=[(hook_name, _make_subspace_hook(head_svd.head, baseline_projector))], + ) + baseline_metric = float(metric(baseline_logits)) + finally: + model.set_use_attn_result(previous) + + delta_metric = patched_metric - original_metric + baseline_delta_metric = baseline_metric - original_metric + gate_threshold = abs(baseline_delta_metric) if threshold is None else threshold + gated = abs(delta_metric) > gate_threshold + + return PatchResult( + head_svd=head_svd, + retained=retained, + original_metric=original_metric, + patched_metric=patched_metric, + delta_metric=delta_metric, + baseline_delta_metric=baseline_delta_metric, + gated=gated, + ) From 5a2945481c64a621c4222b2c61db22e0afa25e99 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 12 Sep 2026 17:04:58 +0530 Subject: [PATCH 05/14] feat(svd_circuits): export public API and validate Bridge compatibility 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. --- transformer_lens/tools/analysis/__init__.py | 30 +++++++++++++++++++++ 1 file changed, 30 insertions(+) diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index 5d93a557f..be84d770c 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -22,6 +22,10 @@ anchored coordinate patching (offline and dynamic/hooked). - projection_kernel: Basis-invariant subspace overlap and TransformerBridge attention-head OQ/OK/OV affinity. + - svd_circuits: Per-head QK/OV singular-vector decomposition with a + degeneracy guard, OV vocab and logit readout, per-position activation + projection onto singular directions, and a mandatory causal patch gate + that reconstructs a head's output onto a chosen singular subspace. """ from transformer_lens.tools.analysis.attribution_patching import ( @@ -76,8 +80,23 @@ projection_kernel, random_projection_kernel_moments, ) +from transformer_lens.tools.analysis.svd_circuits import ( + ActivationProjection, + DegenerateDirectionError, + HeadDecomposition, + HeadSVD, + LogitSignature, + PatchResult, + RankReportRow, + decompose_head, + logit_signature, + patch_along_directions, + project_activations, + vocab_readout, +) __all__ = [ + "ActivationProjection", "AttentionHeadRef", "AttributionResult", "BackwardLens", @@ -85,33 +104,44 @@ "BackwardLensMatrixResult", "BackwardLensResult", "CoordinatePatch", + "DegenerateDirectionError", "DirectLogitAttribution", "EdgeAttributionConfig", "HeadAffinityPair", "HeadAffinityResult", + "HeadDecomposition", + "HeadSVD", "JSpaceDecomposition", "JSpaceOccupancy", "JSpaceVarianceProfile", "JacobianLens", "JacobianLensReadout", "LinearGradientFactors", + "LogitSignature", "Node", + "PatchResult", "ProjectedFactor", "ProjectionKernelResult", "RandomSubspaceReference", + "RankReportRow", "SubspaceBasis", "VocabularyRanking", "WeightLayout", "attention_head_subspace_affinity", "attribution_patch", + "decompose_head", "direct_logit_attribution", "estimate_occupancy", "get_act_patch_direct_path", "get_act_patch_direct_path_all_sources", "get_sparse_decomposition", + "logit_signature", "orthonormal_subspace", + "patch_along_directions", + "project_activations", "projection_kernel", "random_projection_kernel_moments", "solve_coordinate_patch", "solve_coordinate_patch_positions", + "vocab_readout", ] From eb2bb3a6d1b9573da766dff9e7d0c463f25d916c Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Sat, 12 Sep 2026 22:09:26 +0530 Subject: [PATCH 06/14] test(svd_circuits): add GPT-2 small IOI patch integration 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. --- tests/integration/test_svd_circuits.py | 66 ++++++++++++++++++++++++++ 1 file changed, 66 insertions(+) create mode 100644 tests/integration/test_svd_circuits.py diff --git a/tests/integration/test_svd_circuits.py b/tests/integration/test_svd_circuits.py new file mode 100644 index 000000000..256f29743 --- /dev/null +++ b/tests/integration/test_svd_circuits.py @@ -0,0 +1,66 @@ +"""Integration test: OV vocab readout and the causal patch gate on a real GPT-2 head. + +The only file in the ``svd_circuits`` suite that downloads a pretrained model. Drives the +whole read-then-patch path against layer 9 head 9 (the paper's canonical name-mover head) on +a minimal IOI-style prompt, to check that the readout and the causal gate agree on a real +Bridge rather than only on the unit suite's synthetic and tiny-model fixtures. +""" + +from __future__ import annotations + +import math + +import pytest +import torch + +from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.tools.analysis.svd_circuits import ( + decompose_head, + patch_along_directions, + vocab_readout, +) + +CLEAN_PROMPT = "When Mary and John went to the store, John gave a drink to" +LAYER, HEAD = 9, 9 + + +@pytest.fixture(scope="module") +def gpt2_bridge(): + model = TransformerBridge.boot_transformers("gpt2", device="cpu", dtype=torch.float32) + model.enable_compatibility_mode() + return model + + +def _logit_diff_metric(model): + mary_token = model.to_single_token(" Mary") + john_token = model.to_single_token(" John") + + def metric(logits: torch.Tensor) -> float: + return float(logits[0, -1, mary_token] - logits[0, -1, john_token]) + + return metric + + +def test_svd_circuits_readout_and_patch_gate_on_name_mover_head(gpt2_bridge) -> None: + decomposition = decompose_head(gpt2_bridge, layer=LAYER, head=HEAD, which=("OV",)) + ov = decomposition.OV + assert ov is not None + + readout = vocab_readout(gpt2_bridge, ov, k=10) + assert readout.shape == (gpt2_bridge.cfg.d_vocab, 10) + assert torch.isfinite(readout).all() + + # Pick the first non-degenerate direction rather than assuming index 0 is isolated. + top_direction = next(row.idx for row in ov.rank_report if not row.is_degenerate) + + prompt = gpt2_bridge.to_tokens(CLEAN_PROMPT) + metric = _logit_diff_metric(gpt2_bridge) + + result = patch_along_directions(gpt2_bridge, ov, prompt, metric, keep=[top_direction]) + + assert result.retained == [top_direction] + assert math.isfinite(result.original_metric) + assert math.isfinite(result.patched_metric) + assert math.isfinite(result.delta_metric) + assert math.isfinite(result.baseline_delta_metric) + assert isinstance(result.gated, bool) From e6407962073bf3fa016c94f25bf44b9521a6db1d Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 00:37:50 +0530 Subject: [PATCH 07/14] docs(svd_circuits): drop HookedTransformer references from the module 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. --- .../tools/analysis/svd_circuits.py | 26 +++++++++---------- 1 file changed, 13 insertions(+), 13 deletions(-) diff --git a/transformer_lens/tools/analysis/svd_circuits.py b/transformer_lens/tools/analysis/svd_circuits.py index 2a5f48606..f527c2bca 100644 --- a/transformer_lens/tools/analysis/svd_circuits.py +++ b/transformer_lens/tools/analysis/svd_circuits.py @@ -425,11 +425,12 @@ def decompose_head( 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 - only fires for ``TransformerBridge``. Mirrors the compatibility-mode check other - unembedding-touching analysis tools already run, without any hybrid-architecture - restriction: projecting a rank-1 OV direction through ``W_U`` does not depend on - the block-layout assumptions that check exists for elsewhere. + Projecting an OV direction through ``W_U`` requires the final LayerNorm folded + into ``W_U``; on a ``TransformerBridge`` that folding is only present in + compatibility mode, so this check requires it. Mirrors the compatibility-mode + check other unembedding-touching analysis tools already run, without any + hybrid-architecture restriction: projecting a rank-1 OV direction through ``W_U`` + does not depend on the block-layout assumptions that check exists for elsewhere. """ # Lazy import - keeps the module importable without the bridge as a hard dependency. from transformer_lens.model_bridge import TransformerBridge @@ -449,8 +450,7 @@ def vocab_readout(model, head_svd: HeadSVD, *, k: int = 10) -> Float[torch.Tenso Requires ``head_svd.which == "OV"``: QK produces no write direction to project (see the module docstring). On a ``TransformerBridge``, compatibility mode must - be enabled so ``W_U`` carries the folded final LayerNorm weights; - ``HookedTransformer`` always has this folding applied. + be enabled so ``W_U`` carries the folded final LayerNorm weights. Does not call ``head_svd.require_isolated``: a degenerate direction's vocab readout is still a well-defined projection, unlike a per-direction causal claim, @@ -458,8 +458,8 @@ def vocab_readout(model, head_svd: HeadSVD, *, k: int = 10) -> Float[torch.Tenso passing causal patch is enforced by :func:`patch_along_directions`. Args: - model: A ``TransformerBridge`` (with compatibility mode enabled) or a - ``HookedTransformer``; only its ``W_U`` is read. + model: A ``TransformerBridge`` with compatibility mode enabled; only its + ``W_U`` is read. head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. k: Number of top singular directions to project. @@ -508,8 +508,8 @@ def logit_signature( no forward pass and builds no cache. Args: - model: A ``TransformerBridge`` (with compatibility mode enabled) or a - ``HookedTransformer``; only its ``W_U`` is read. + model: A ``TransformerBridge`` with compatibility mode enabled; only its + ``W_U`` is read. head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. direction: Column index of the singular direction to reconstruct. tokens: Token id(s) to read the logit effect for. @@ -567,7 +567,7 @@ def project_activations( hooks or a cache built around its prior state. Args: - model: A ``TransformerBridge`` or ``HookedTransformer``. + model: A ``TransformerBridge``. head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. prompt: A single prompt (not a batch): a string or a ``[1, pos]`` token tensor. @@ -710,7 +710,7 @@ def patch_along_directions( setting afterward. Args: - model: A ``TransformerBridge`` or ``HookedTransformer``. + model: A ``TransformerBridge``. head_svd: An OV :class:`HeadSVD` from :func:`decompose_head`. prompt: A single prompt: a string or a ``[1, pos]`` token tensor. metric: A function from the model's logits to a scalar. From b66092d6f830950dc1da05b46fdfc224bfdaa50c Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 00:47:27 +0530 Subject: [PATCH 08/14] fix(svd_circuits): bind each decomposition to its compatibility-mode 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. --- tests/unit/tools/test_svd_circuits.py | 53 +++++++++++++- .../tools/analysis/svd_circuits.py | 71 ++++++++++++++++--- 2 files changed, 111 insertions(+), 13 deletions(-) diff --git a/tests/unit/tools/test_svd_circuits.py b/tests/unit/tools/test_svd_circuits.py index 6dcc98a70..8899a1e5b 100644 --- a/tests/unit/tools/test_svd_circuits.py +++ b/tests/unit/tools/test_svd_circuits.py @@ -486,9 +486,8 @@ def test_decompose_head_rejects_bare_string_which(): # --------------------------------------------------------------------------- # # OV output-direction convention (cross-check against SVDInterpreter) # --------------------------------------------------------------------------- # -@pytest.fixture(scope="module") -def tiny_bridge(): - """Tiny GPT-2 bridge: MHA (n_heads == n_kv_heads), CPU-only, no Hub access.""" +def _make_tiny_bridge(): + """Build a tiny GPT-2 bridge: MHA (n_heads == n_kv_heads), CPU-only, no Hub access.""" from transformers import GPT2Config, GPT2LMHeadModel from transformer_lens.config.transformer_bridge_config import ( @@ -519,6 +518,14 @@ def tiny_bridge(): return TransformerBridge(hf_model, GPT2ArchitectureAdapter(cfg), tokenizer=tokenizer) +@pytest.fixture(scope="module") +def tiny_bridge(): + """Module-scoped tiny bridge. Compatibility mode folds weights in place and has no + inverse, so tests that enable it must build their own bridge via ``_make_tiny_bridge`` + rather than mutate this shared instance.""" + return _make_tiny_bridge() + + def test_ov_output_direction_matches_svd_interpreter(tiny_bridge): """OV's write/vocab-readout direction is ``.V``, not ``.U`` - checked against the already-shipped ``SVDInterpreter``, not just internal self-consistency, so the @@ -543,6 +550,46 @@ def test_ov_output_direction_matches_svd_interpreter(tiny_bridge): assert not torch.allclose(projected_via_U, -reference[:, 0, 0], atol=1e-2) +# --------------------------------------------------------------------------- # +# Compatibility-mode binding: a decomposition is tied to the state it was built under +# --------------------------------------------------------------------------- # +def test_decompose_head_records_compatibility_mode(): + """Each HeadSVD records the model's compatibility-mode state at decomposition time, + for both maps and both toggle values.""" + off = decompose_head(_make_tiny_bridge(), 0, 0, which=("QK", "OV")) + assert off.OV.compatibility_mode is False + assert off.QK.compatibility_mode is False + + on_model = _make_tiny_bridge() + on_model.enable_compatibility_mode() + on = decompose_head(on_model, 0, 0, which=("QK", "OV")) + assert on.OV.compatibility_mode is True + assert on.QK.compatibility_mode is True + + +@pytest.mark.parametrize( + "call_consumer", + [ + lambda model, ov: vocab_readout(model, ov), + lambda model, ov: logit_signature(model, ov, direction=0, tokens=torch.tensor([0])), + lambda model, ov: project_activations(model, ov, torch.tensor([[5, 63, 7, 9]])), + lambda model, ov: patch_along_directions( + model, ov, torch.tensor([[5, 63, 7, 9]]), lambda logits: float(logits.sum()), keep=[0] + ), + ], + ids=["vocab_readout", "logit_signature", "project_activations", "patch_along_directions"], +) +def test_consumers_reject_stale_compatibility_state(call_consumer): + """A HeadSVD decomposed before enable_compatibility_mode() describes the pre-folding OV + map, so every consumer refuses it once the model's state has changed and points the + caller back at decompose_head. The refusal fires before any forward pass runs.""" + bridge = _make_tiny_bridge() + stale = decompose_head(bridge, 0, 0, which=("OV",)).OV + bridge.enable_compatibility_mode() + with pytest.raises(ValueError, match="decompose_head"): + call_consumer(bridge, stale) + + # --------------------------------------------------------------------------- # # vocab_readout / logit_signature (OV output-direction readout) # --------------------------------------------------------------------------- # diff --git a/transformer_lens/tools/analysis/svd_circuits.py b/transformer_lens/tools/analysis/svd_circuits.py index f527c2bca..87b526311 100644 --- a/transformer_lens/tools/analysis/svd_circuits.py +++ b/transformer_lens/tools/analysis/svd_circuits.py @@ -128,6 +128,11 @@ class HeadSVD: boundary sits at a gap of at least ``eps``. null_rtol: Relative-to-top-singular-value tolerance below which a direction is numerically null. + compatibility_mode: The model's compatibility-mode state at the time this was + decomposed. ``enable_compatibility_mode`` folds ``ln1`` into ``W_V`` and + centres ``W_O``/``W_U``, so a decomposition describes the OV map only under + the state it was built in; the readout and patch consumers refuse a + decomposition whose state no longer matches the model. """ which: Which @@ -139,6 +144,7 @@ class HeadSVD: rank_report: List[RankReportRow] eps: float null_rtol: float + compatibility_mode: bool def is_degenerate(self, i: int) -> bool: """Whether direction ``i`` is refused: rotation-ambiguous inside a block, or null.""" @@ -320,6 +326,7 @@ def _factored_head_svd( head: int, eps: float, null_rtol: Optional[float] = None, + compatibility_mode: bool = False, ) -> HeadSVD: """Decompose the factored map ``A @ B`` for one head into a :class:`HeadSVD`. @@ -332,6 +339,9 @@ def _factored_head_svd( the same relative tolerance :func:`torch.linalg.matrix_rank` uses by default for a square ``d_model x d_model`` map, so the null cutoff tracks the decomposition's own numerical rank rather than a hand-tuned constant. + + ``compatibility_mode`` is the model's compatibility-mode state, recorded on the + result so the consumers can refuse a decomposition taken under a different state. """ U, S, V = FactoredMatrix(A, B).svd() d_model = U.shape[0] @@ -348,6 +358,7 @@ def _factored_head_svd( rank_report=rank_report, eps=eps, null_rtol=resolved_null_rtol, + compatibility_mode=compatibility_mode, ) @@ -364,7 +375,9 @@ def decompose_head( Weight-space only: this reads the head's per-block weights via the bridge's ``model.blocks[layer].attn`` accessors and needs no forward pass and no - compatibility mode. The returned factors are detached from the model. + compatibility mode. The returned factors are detached from the model. The model's + compatibility-mode state is recorded on each returned :class:`HeadSVD` so the readout + and patch consumers can refuse a decomposition taken under a different state. Args: model: A ``TransformerBridge``. @@ -403,6 +416,7 @@ def decompose_head( raise ValueError(f"which entries must be in {_VALID_WHICH}, got {invalid!r}") W_Q_h, W_K_h, W_V_h, W_O_h = _head_weights(model, layer, head) + compatibility_mode = getattr(model, "compatibility_mode", False) qk = None ov = None if "QK" in requested: @@ -414,10 +428,18 @@ def decompose_head( head=head, eps=eps, null_rtol=null_rtol, + compatibility_mode=compatibility_mode, ) if "OV" in requested: ov = _factored_head_svd( - W_V_h, W_O_h, which="OV", layer=layer, head=head, eps=eps, null_rtol=null_rtol + W_V_h, + W_O_h, + which="OV", + layer=layer, + head=head, + eps=eps, + null_rtol=null_rtol, + compatibility_mode=compatibility_mode, ) return HeadDecomposition(layer=layer, head=head, QK=qk, OV=ov) @@ -445,6 +467,27 @@ def _validate_bridge_compatibility(model) -> None: ) +def _validate_decomposition_matches_model(model, head_svd: HeadSVD) -> None: + """Refuse a decomposition built under a different compatibility-mode state than the model. + + ``enable_compatibility_mode`` folds ``ln1`` into ``W_V`` and centres ``W_O``/``W_U``, + so a :class:`HeadSVD` decomposed before the call describes a different OV map than the + model now computes: the cached ``U``/``S``/``V`` are stale, and the returned readout or + patch would be silently wrong rather than raise a shape error. This guard is orthogonal + to :func:`_validate_bridge_compatibility` (which only checks that ``W_U`` carries the + folded LayerNorm): here the model may be in either state, only mismatched from the + decomposition's. + """ + current = getattr(model, "compatibility_mode", False) + if current != head_svd.compatibility_mode: + raise ValueError( + f"This HeadSVD was decomposed with compatibility_mode=" + f"{head_svd.compatibility_mode} but the model now has compatibility_mode=" + f"{current}; the cached singular vectors describe a different OV map. " + f"Re-run decompose_head under the current state, then retry." + ) + + def vocab_readout(model, head_svd: HeadSVD, *, k: int = 10) -> Float[torch.Tensor, "d_vocab k"]: """Project the top-k OV output directions through the unembedding. @@ -469,11 +512,13 @@ def vocab_readout(model, head_svd: HeadSVD, *, k: int = 10) -> Float[torch.Tenso Raises: ValueError: If ``head_svd.which != "OV"``, if ``k`` is not in - ``(0, rank]``, or if ``model`` is a ``TransformerBridge`` without - compatibility mode enabled. + ``(0, rank]``, if ``head_svd`` was decomposed under a different + compatibility-mode state than ``model`` now has, or if ``model`` is a + ``TransformerBridge`` without compatibility mode enabled. """ if head_svd.which != "OV": raise ValueError(f"vocab_readout requires an OV HeadSVD, got which={head_svd.which!r}") + _validate_decomposition_matches_model(model, head_svd) rank = head_svd.V.shape[1] if not 0 < k <= rank: raise ValueError(f"k must be in (0, {rank}], got {k!r}") @@ -518,13 +563,15 @@ def logit_signature( A :class:`LogitSignature` with one value per requested token. Raises: - ValueError: If ``head_svd.which != "OV"``, or if ``model`` is a - ``TransformerBridge`` without compatibility mode enabled. + ValueError: If ``head_svd.which != "OV"``, if ``head_svd`` was decomposed under + a different compatibility-mode state than ``model`` now has, or if ``model`` + is a ``TransformerBridge`` without compatibility mode enabled. DegenerateDirectionError: If ``direction`` is not attributable alone (see :meth:`HeadSVD.require_isolated`). """ if head_svd.which != "OV": raise ValueError(f"logit_signature requires an OV HeadSVD, got which={head_svd.which!r}") + _validate_decomposition_matches_model(model, head_svd) head_svd.require_isolated(direction) _validate_bridge_compatibility(model) token_ids = torch.as_tensor(tokens, dtype=torch.long).reshape(-1) @@ -575,13 +622,15 @@ def project_activations( An :class:`ActivationProjection` with the per-position coefficients. Raises: - ValueError: If ``head_svd.which != "OV"``, or if ``prompt`` is not a single - (batch-size-1) prompt. + ValueError: If ``head_svd.which != "OV"``, if ``head_svd`` was decomposed under + a different compatibility-mode state than ``model`` now has, or if ``prompt`` + is not a single (batch-size-1) prompt. """ if head_svd.which != "OV": raise ValueError( f"project_activations requires an OV HeadSVD, got which={head_svd.which!r}" ) + _validate_decomposition_matches_model(model, head_svd) previous = getattr(model.cfg, "use_attn_result", False) model.set_use_attn_result(True) try: @@ -725,8 +774,9 @@ def patch_along_directions( A :class:`PatchResult` describing the patched, baseline, and original metrics. Raises: - ValueError: If ``head_svd.which != "OV"``, or if ``keep``/``ablate`` are both - given or both omitted. + ValueError: If ``head_svd.which != "OV"``, if ``head_svd`` was decomposed under a + different compatibility-mode state than ``model`` now has, or if + ``keep``/``ablate`` are both given or both omitted. DegenerateDirectionError: If the retained directions split a degenerate block (see :func:`_validate_retained_blocks`). """ @@ -734,6 +784,7 @@ def patch_along_directions( raise ValueError( f"patch_along_directions requires an OV HeadSVD, got which={head_svd.which!r}" ) + _validate_decomposition_matches_model(model, head_svd) retained = _resolve_retained(head_svd, keep, ablate) V = head_svd.V From 06b4f6c143cb1165616695fdd8884846f14904f3 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 01:00:35 +0530 Subject: [PATCH 09/14] fix(svd_circuits): project W_U and V in a common dtype --- tests/unit/tools/test_svd_circuits.py | 51 +++++++++++++++++++ .../tools/analysis/svd_circuits.py | 12 +++-- 2 files changed, 60 insertions(+), 3 deletions(-) diff --git a/tests/unit/tools/test_svd_circuits.py b/tests/unit/tools/test_svd_circuits.py index 8899a1e5b..8752307b3 100644 --- a/tests/unit/tools/test_svd_circuits.py +++ b/tests/unit/tools/test_svd_circuits.py @@ -646,6 +646,57 @@ def test_logit_signature_raises_on_degenerate_direction(): logit_signature(model, ov, direction=1, tokens=torch.tensor([0])) +def _ov_headsvd_in_dtype(dtype): + """A minimal, well-separated OV HeadSVD with U/S/V cast to ``dtype``. + + Built by casting a float32 decomposition rather than decomposing in ``dtype`` directly, + since torch.linalg.svd does not run in half precision on CPU; the readout paths under + test only read U/S/V and W_U, so the cast stands in for a genuine reduced-precision model. + """ + base = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + return HeadSVD( + which="OV", + layer=base.layer, + head=base.head, + U=base.U.to(dtype), + S=base.S.to(dtype), + V=base.V.to(dtype), + rank_report=base.rank_report, + eps=base.eps, + null_rtol=base.null_rtol, + compatibility_mode=False, + ) + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +def test_vocab_readout_dtype_parametrised(dtype): + """vocab_readout promotes V to W_U's dtype, so the projection matmul runs on every model + precision instead of raising a dtype mismatch on bf16, fp16, or fp64.""" + ov = _ov_headsvd_in_dtype(dtype) + vocab_size = 20 + model = SimpleNamespace(W_U=torch.randn(D_MODEL, vocab_size, dtype=dtype)) + result = vocab_readout(model, ov, k=3) + assert result.shape == (vocab_size, 3) + assert result.dtype == dtype + assert torch.isfinite(result.float()).all() + + +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +def test_logit_signature_dtype_parametrised(dtype): + """logit_signature casts its rank-1 reconstruction to W_U's dtype, so it returns one + finite value per token on every model precision instead of raising a dtype mismatch.""" + ov = _ov_headsvd_in_dtype(dtype) + vocab_size = 16 + model = SimpleNamespace(W_U=torch.randn(D_MODEL, vocab_size, dtype=dtype)) + token_ids = torch.tensor([1, 5, 9]) + result = logit_signature(model, ov, direction=0, tokens=token_ids) + assert result.values.shape == (token_ids.numel(),) + assert result.values.dtype == dtype + assert torch.isfinite(result.values.float()).all() + + # --------------------------------------------------------------------------- # # project_activations (per-position firing coefficients in the OV output basis) # --------------------------------------------------------------------------- # diff --git a/transformer_lens/tools/analysis/svd_circuits.py b/transformer_lens/tools/analysis/svd_circuits.py index 87b526311..c396f9367 100644 --- a/transformer_lens/tools/analysis/svd_circuits.py +++ b/transformer_lens/tools/analysis/svd_circuits.py @@ -523,7 +523,10 @@ def vocab_readout(model, head_svd: HeadSVD, *, k: int = 10) -> Float[torch.Tenso 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() + # Promote V to W_U's dtype rather than forcing float32: matmul does not promote its + # operands, so a bare .float() raises a dtype mismatch on any bf16, fp16, or fp64 model. + W_U = model.W_U + return W_U.T @ head_svd.V[:, :k].to(W_U.dtype) @dataclass @@ -575,8 +578,11 @@ def logit_signature( head_svd.require_isolated(direction) _validate_bridge_compatibility(model) token_ids = torch.as_tensor(tokens, dtype=torch.long).reshape(-1) - reconstruction = (head_svd.S[direction] * head_svd.V[:, direction]).float() - values = reconstruction @ model.W_U[:, token_ids] + # Match the rank-1 reconstruction to W_U's dtype for the projection: matmul does not + # promote its operands, so a float32 reconstruction breaks on bf16, fp16, or fp64 models. + W_U = model.W_U + reconstruction = (head_svd.S[direction] * head_svd.V[:, direction]).to(W_U.dtype) + values = reconstruction @ W_U[:, token_ids] return LogitSignature(direction=direction, values=values) From 9b94861b180238667b6f769ee0539cd9cbea4ef7 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 01:11:39 +0530 Subject: [PATCH 10/14] fix(svd_circuits): validate direction indices against the map rank --- tests/unit/tools/test_svd_circuits.py | 83 +++++++++++++++++++ .../tools/analysis/svd_circuits.py | 52 ++++++++++-- 2 files changed, 126 insertions(+), 9 deletions(-) diff --git a/tests/unit/tools/test_svd_circuits.py b/tests/unit/tools/test_svd_circuits.py index 8752307b3..14d75e729 100644 --- a/tests/unit/tools/test_svd_circuits.py +++ b/tests/unit/tools/test_svd_circuits.py @@ -23,6 +23,7 @@ RankReportRow, _degeneracy_blocks, _factored_head_svd, + _resolve_retained, decompose_head, logit_signature, patch_along_directions, @@ -646,6 +647,19 @@ def test_logit_signature_raises_on_degenerate_direction(): logit_signature(model, ov, direction=1, tokens=torch.tensor([0])) +def test_logit_signature_rejects_out_of_range_direction(): + """A direction outside [0, rank) is a clear ValueError, not a wrap-around (direction=-1 + reading the last column) or a bare IndexError (direction=rank).""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + model = SimpleNamespace(W_U=torch.randn(D_MODEL, 8)) + rank = ov.V.shape[1] + for bad in (rank, -1): + with pytest.raises(ValueError, match="out of range"): + logit_signature(model, ov, direction=bad, tokens=torch.tensor([0])) + + def _ov_headsvd_in_dtype(dtype): """A minimal, well-separated OV HeadSVD with U/S/V cast to ``dtype``. @@ -821,6 +835,75 @@ def test_patch_along_directions_rejects_partial_degenerate_block(): assert result.retained == [0, 1, 2] +def test_resolve_retained_rejects_out_of_range(): + """keep/ablate indices outside [0, rank) raise ValueError, not a silent wrap-around for a + negative index or a silent no-op for an out-of-range ablate.""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + for bad in ([999], [-1]): + with pytest.raises(ValueError, match="out of range"): + _resolve_retained(ov, keep=bad, ablate=None) + with pytest.raises(ValueError, match="out of range"): + _resolve_retained(ov, keep=None, ablate=bad) + + +def test_patch_rejects_empty_retained_without_threshold(): + """An empty retained set (keep=[] or ablate over the full rank) reconstructs onto the zero + subspace and ties the baseline by construction, so it is refused unless the caller passes + an explicit threshold, and accepted when one is.""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + rank = ov.V.shape[1] + stub = _PatchStubModel(d_model=D_MODEL, n_heads=1) + metric = lambda logits: float(logits.sum()) + for empty in (dict(keep=[]), dict(ablate=list(range(rank)))): + with pytest.raises(ValueError, match="threshold"): + patch_along_directions(stub, ov, "prompt", metric, **empty) + result = patch_along_directions( + stub, + ov, + "prompt", + metric, + threshold=0.0, + rng=torch.Generator().manual_seed(0), + **empty, + ) + assert isinstance(result, PatchResult) + assert result.retained == [] + + +def test_patch_ablate_rejects_partial_degenerate_block(): + """ablate must take a degenerate block whole: ablating one member of block [1, 2] leaves + the other in the retained complement, splitting the block, and is refused before any model + access the same way the keep path is.""" + ov = _factored_head_svd( + *_factored_with_spectrum([5.0, 3.0, 3.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + with pytest.raises(DegenerateDirectionError): + patch_along_directions(SimpleNamespace(), ov, "prompt", lambda logits: 0.0, ablate=[1]) + + +def test_patch_ablate_reports_retained(): + """An ablate call retains the complement of the ablated set and reports it, so an ablate + path is asserted end to end rather than only the keep path.""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + stub = _PatchStubModel(d_model=D_MODEL, n_heads=1) + result = patch_along_directions( + stub, + ov, + "prompt", + lambda logits: float(logits.sum()), + ablate=[1, 3], + rng=torch.Generator().manual_seed(0), + ) + assert isinstance(result, PatchResult) + assert result.retained == [0, 2] + + def test_patch_along_directions_reports_delta_and_restores_use_attn_result(): """delta_metric is patched minus original metric, gated follows the documented comparison against baseline_delta_metric, and use_attn_result is restored after.""" diff --git a/transformer_lens/tools/analysis/svd_circuits.py b/transformer_lens/tools/analysis/svd_circuits.py index c396f9367..d25e98b25 100644 --- a/transformer_lens/tools/analysis/svd_circuits.py +++ b/transformer_lens/tools/analysis/svd_circuits.py @@ -566,15 +566,23 @@ def logit_signature( A :class:`LogitSignature` with one value per requested token. Raises: - ValueError: If ``head_svd.which != "OV"``, if ``head_svd`` was decomposed under - a different compatibility-mode state than ``model`` now has, or if ``model`` - is a ``TransformerBridge`` without compatibility mode enabled. + ValueError: If ``head_svd.which != "OV"``, if ``direction`` is not in + ``[0, rank)``, if ``head_svd`` was decomposed under a different + compatibility-mode state than ``model`` now has, or if ``model`` is a + ``TransformerBridge`` without compatibility mode enabled. DegenerateDirectionError: If ``direction`` is not attributable alone (see :meth:`HeadSVD.require_isolated`). """ if head_svd.which != "OV": raise ValueError(f"logit_signature requires an OV HeadSVD, got which={head_svd.which!r}") _validate_decomposition_matches_model(model, head_svd) + # Bounds-check before indexing: a negative direction would wrap into V/rank_report and a + # too-large one would raise a bare IndexError, both hiding a caller mistake as a wrong or + # cryptic result rather than a clear refusal. + rank = head_svd.V.shape[1] + direction = int(direction) + if not 0 <= direction < rank: + raise ValueError(f"direction index {direction} out of range [0, {rank})") head_svd.require_isolated(direction) _validate_bridge_compatibility(model) token_ids = torch.as_tensor(tokens, dtype=torch.long).reshape(-1) @@ -681,17 +689,31 @@ def _resolve_retained( """Resolve ``keep``/``ablate`` to the sorted list of retained direction indices. Exactly one of ``keep``/``ablate`` must be given; ``ablate``'s complement over the - map's full rank becomes the retained set. Raises :class:`DegenerateDirectionError` - if the result would split a degenerate block (see :func:`_validate_retained_blocks`). + map's full rank becomes the retained set. Every supplied index is coerced to ``int`` + and bounds-checked against ``[0, rank)`` before use, so a negative index raises rather + than wrapping into ``V[:, ...]`` and an out-of-range ``ablate`` raises rather than + silently subtracting nothing from the complement. Raises + :class:`DegenerateDirectionError` if the result would split a degenerate block (see + :func:`_validate_retained_blocks`). """ if (keep is None) == (ablate is None): raise ValueError("patch_along_directions requires exactly one of keep or ablate") rank = head_svd.V.shape[1] + + def _checked(indices: Sequence[int], name: str) -> List[int]: + resolved: List[int] = [] + for raw in indices: + index = int(raw) + if not 0 <= index < rank: + raise ValueError(f"{name} index {index} out of range [0, {rank})") + resolved.append(index) + return resolved + if keep is not None: - retained = sorted(set(keep)) + retained = sorted(set(_checked(keep, "keep"))) else: assert ablate is not None - ablate_set = set(ablate) + ablate_set = set(_checked(ablate, "ablate")) retained = [i for i in range(rank) if i not in ablate_set] _validate_retained_blocks(head_svd, retained) return retained @@ -781,8 +803,10 @@ def patch_along_directions( Raises: ValueError: If ``head_svd.which != "OV"``, if ``head_svd`` was decomposed under a - different compatibility-mode state than ``model`` now has, or if - ``keep``/``ablate`` are both given or both omitted. + different compatibility-mode state than ``model`` now has, if ``keep``/``ablate`` + are both given or both omitted, if any index is out of ``[0, rank)``, or if the + retained set is empty (``keep=[]`` or ``ablate`` over the full rank) and no + explicit ``threshold`` is supplied. DegenerateDirectionError: If the retained directions split a degenerate block (see :func:`_validate_retained_blocks`). """ @@ -792,6 +816,16 @@ def patch_along_directions( ) _validate_decomposition_matches_model(model, head_svd) retained = _resolve_retained(head_svd, keep, ablate) + if not retained and threshold is None: + # An empty retained set reconstructs the head onto the zero subspace, so its delta and + # the equal-width random baseline's delta are both the full-ablation effect: the gate + # compares a quantity against itself and is meaningless. Require an explicit threshold + # to gate an empty set on purpose. + raise ValueError( + "keep/ablate retain no directions, so the reconstruction is the zero subspace " + "and its delta ties the random baseline by construction; pass an explicit " + "threshold to gate an empty retained set." + ) V = head_svd.V kept_projector = V[:, retained] @ V[:, retained].transpose(-2, -1) From 3b6d7dfe95ed121582ba0e14888a28b5fcae0082 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 01:19:19 +0530 Subject: [PATCH 11/14] fix(svd_circuits): place the subspace projector on the activation's device --- tests/unit/tools/test_svd_circuits.py | 28 +++++++++++++++++-- .../tools/analysis/svd_circuits.py | 6 +++- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/tests/unit/tools/test_svd_circuits.py b/tests/unit/tools/test_svd_circuits.py index 14d75e729..099ac0653 100644 --- a/tests/unit/tools/test_svd_circuits.py +++ b/tests/unit/tools/test_svd_circuits.py @@ -774,11 +774,11 @@ class _PatchStubModel: arithmetic can be tested without a real forward pass. """ - def __init__(self, d_model, n_heads, pos=3, seed=0): + def __init__(self, d_model, n_heads, pos=3, seed=0, device="cpu"): g = torch.Generator().manual_seed(seed) self.cfg = SimpleNamespace(use_attn_result=False) - self._result = torch.randn(1, pos, n_heads, d_model, generator=g) - self._readout = torch.randn(d_model, generator=g) + self._result = torch.randn(1, pos, n_heads, d_model, generator=g).to(device) + self._readout = torch.randn(d_model, generator=g).to(device) def set_use_attn_result(self, value): self.cfg.use_attn_result = value @@ -924,6 +924,28 @@ def test_patch_along_directions_reports_delta_and_restores_use_attn_result(): assert stub.cfg.use_attn_result == initial +@pytest.mark.skipif( + not (torch.backends.mps.is_available() or torch.cuda.is_available()), + reason="needs an MPS or CUDA device to exercise cross-device projector placement", +) +def test_patch_completes_on_accelerator_device(): + """The kept and baseline projectors are drawn on CPU, but the activation lives on the + model's device, so the hook must move the projector before the matmul. On an + accelerator the patch completes without a device-mismatch RuntimeError.""" + device = "mps" if torch.backends.mps.is_available() else "cuda" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + stub = _PatchStubModel(d_model=D_MODEL, n_heads=1, device=device) + metric = lambda logits: float(logits.sum()) + result = patch_along_directions( + stub, ov, "prompt", metric, keep=[0], rng=torch.Generator().manual_seed(0) + ) + assert isinstance(result, PatchResult) + assert math.isfinite(result.delta_metric) + assert math.isfinite(result.baseline_delta_metric) + + 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.""" diff --git a/transformer_lens/tools/analysis/svd_circuits.py b/transformer_lens/tools/analysis/svd_circuits.py index d25e98b25..e4dbd91e3 100644 --- a/transformer_lens/tools/analysis/svd_circuits.py +++ b/transformer_lens/tools/analysis/svd_circuits.py @@ -729,7 +729,11 @@ def _make_subspace_hook(head: int, projector: Float[torch.Tensor, "d_model d_mod def hook_fn(activation: torch.Tensor, hook) -> torch.Tensor: activation = activation.clone() - activation[:, :, head, :] = activation[:, :, head, :] @ projector.to(activation.dtype) + # The projector is drawn on CPU (QR is unimplemented on MPS and a CUDA generator + # cannot feed a CPU randn), so move it to the activation's device and dtype here. + activation[:, :, head, :] = activation[:, :, head, :] @ projector.to( + device=activation.device, dtype=activation.dtype + ) return activation return hook_fn From 85a1f8f5e675d8bae3281b9ad5ff107f9259ff29 Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 01:32:14 +0530 Subject: [PATCH 12/14] fix(svd_circuits): draw the causal baseline in-span, average it, and gate by mode --- tests/unit/tools/test_svd_circuits.py | 138 +++++++++++++++++- .../tools/analysis/svd_circuits.py | 111 ++++++++++---- 2 files changed, 216 insertions(+), 33 deletions(-) diff --git a/tests/unit/tools/test_svd_circuits.py b/tests/unit/tools/test_svd_circuits.py index 099ac0653..79c6f93df 100644 --- a/tests/unit/tools/test_svd_circuits.py +++ b/tests/unit/tools/test_svd_circuits.py @@ -905,8 +905,13 @@ def test_patch_ablate_reports_retained(): def test_patch_along_directions_reports_delta_and_restores_use_attn_result(): - """delta_metric is patched minus original metric, gated follows the documented - comparison against baseline_delta_metric, and use_attn_result is restored after.""" + """delta_metric is patched minus original metric, the averaged baseline is finite, and + use_attn_result is restored to its prior value regardless of what it started as. + + The gate's outcome is not recomputed from the same result's fields here (that would hold + for any self-consistent implementation); the mode-specific outcomes are pinned under a + fixed seed by the dedicated gate-semantics tests below. + """ ov = _factored_head_svd( *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 ) @@ -919,11 +924,136 @@ def test_patch_along_directions_reports_delta_and_restores_use_attn_result(): ) 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) - assert result.gated == expected_gated + assert isinstance(result.gated, bool) assert stub.cfg.use_attn_result == initial +def _span_aligned_stub(ov, weights): + """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 + any in-span component can only reduce the metric, so the per-draw baseline deltas share a + sign and their average does not cancel toward zero. That makes the gate outcomes analytic: + the requested subspace's delta and the averaged random-subspace delta are both exact + functions of ``weights``, so a chosen weight vector fixes which side of the gate a mode + lands on with a wide margin rather than by numerical luck. + """ + V = ov.V + head_out = V @ torch.tensor(weights, dtype=V.dtype) + stub = _PatchStubModel(d_model=V.shape[0], n_heads=1) + stub._result = torch.zeros_like(stub._result) + stub._result[:, :, 0, :] = head_out + stub._readout = head_out + return stub + + +def test_patch_ablate_weak_direction_gates_false(): + """Ablating the weakest isolated direction moves the metric less than removing an + equally-sized random in-span subspace does, so in ablate mode the gate is False.""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + weak = ov.rank_report[-1].idx + assert not ov.is_degenerate(weak) + # Starve the weakest direction of energy while the strong ones carry it, so removing the + # weak one barely moves the metric relative to removing an arbitrary in-span direction. + weights = [1.0, 1.0, 1.0, 1.0] + weights[weak] = 1e-3 + stub = _span_aligned_stub(ov, weights) + metric = lambda logits: float(logits.sum()) + result = patch_along_directions( + stub, ov, "prompt", metric, ablate=[weak], rng=torch.Generator().manual_seed(0) + ) + assert result.gated is False + assert abs(result.delta_metric) < abs(result.baseline_delta_metric) + + +def test_patch_threshold_above_delta_gates_false(): + """An explicit threshold above abs(delta_metric) overrides the averaged baseline and gates + False in ablate mode, even for a direction that gates True against the baseline alone.""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + strong = ov.rank_report[0].idx + assert not ov.is_degenerate(strong) + # Concentrate the head output in the direction to be ablated so removing it beats the + # random in-span baseline; the gate is then True until an explicit threshold overrides it. + weights = [1e-3, 1e-3, 1e-3, 1e-3] + weights[strong] = 1.0 + stub = _span_aligned_stub(ov, weights) + metric = lambda logits: float(logits.sum()) + default = patch_along_directions( + stub, ov, "prompt", metric, ablate=[strong], rng=torch.Generator().manual_seed(0) + ) + assert default.gated is True + raised = patch_along_directions( + stub, + ov, + "prompt", + metric, + ablate=[strong], + threshold=abs(default.delta_metric) + 1.0, + rng=torch.Generator().manual_seed(0), + ) + assert raised.gated is False + + +def test_patch_baseline_is_reproducible(): + """The averaged baseline is drawn from the passed generator, so the same seed reproduces it + bit-for-bit and a different seed gives a different average.""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + stub = _span_aligned_stub(ov, [1.0, 1.0, 1.0, 1.0]) + metric = lambda logits: float(logits.sum()) + first = patch_along_directions( + stub, ov, "prompt", metric, ablate=[0], rng=torch.Generator().manual_seed(0) + ) + same_seed = patch_along_directions( + stub, ov, "prompt", metric, ablate=[0], rng=torch.Generator().manual_seed(0) + ) + other_seed = patch_along_directions( + stub, ov, "prompt", metric, ablate=[0], rng=torch.Generator().manual_seed(1) + ) + assert same_seed.baseline_delta_metric == first.baseline_delta_metric + assert other_seed.baseline_delta_metric != first.baseline_delta_metric + + +def test_patch_keep_mode_gate_semantics(): + """In keep mode the gate asks whether the retained subspace reconstructs the head: keeping + the single direction the head output lies along preserves the metric better than keeping an + equally-sized random in-span subspace, so it gates True, where the ablate-style single-sided + "moved more than baseline" test would have gated the same retained set False.""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + strong = ov.rank_report[0].idx + assert not ov.is_degenerate(strong) + weights = [0.0, 0.0, 0.0, 0.0] + weights[strong] = 1.0 + stub = _span_aligned_stub(ov, weights) + metric = lambda logits: float(logits.sum()) + result = patch_along_directions( + stub, ov, "prompt", metric, keep=[strong], rng=torch.Generator().manual_seed(0) + ) + assert result.gated is True + assert abs(result.delta_metric) < abs(result.baseline_delta_metric) + assert not abs(result.delta_metric) > abs(result.baseline_delta_metric) + + +def test_patch_rejects_non_positive_n_baseline(): + """The averaged baseline needs at least one draw, so n_baseline below 1 is refused before + any model access.""" + ov = _factored_head_svd( + *_factored_with_spectrum([8.0, 4.0, 2.0, 1.0]), which="OV", layer=0, head=0, eps=1e-2 + ) + with pytest.raises(ValueError, match="n_baseline"): + patch_along_directions( + SimpleNamespace(), ov, "prompt", lambda logits: 0.0, keep=[0], n_baseline=0 + ) + + @pytest.mark.skipif( not (torch.backends.mps.is_available() or torch.cuda.is_available()), reason="needs an MPS or CUDA device to exercise cross-device projector placement", diff --git a/transformer_lens/tools/analysis/svd_circuits.py b/transformer_lens/tools/analysis/svd_circuits.py index e4dbd91e3..31e8f8742 100644 --- a/transformer_lens/tools/analysis/svd_circuits.py +++ b/transformer_lens/tools/analysis/svd_circuits.py @@ -71,6 +71,15 @@ # Absolute floor so relative-gap and ratio computations never divide by ~0. _SIGMA_FLOOR = 1e-12 +# Default seed for the causal gate's random baseline, so a bare patch_along_directions +# call is reproducible instead of drawing from the global RNG. Matches estimate_occupancy, +# which seeds its random controls by default. +_DEFAULT_BASELINE_SEED = 0 + +# Number of random in-span control subspaces the gate averages its baseline delta over, +# so one lucky or unlucky draw does not decide the gate. +_DEFAULT_N_BASELINE = 8 + class DegenerateDirectionError(ValueError): """Raised when per-direction attribution is requested for a degenerate direction. @@ -750,12 +759,21 @@ class PatchResult: original_metric: Metric value on the unmodified prompt. patched_metric: Metric value after the subspace reconstruction. delta_metric: ``patched_metric - original_metric``. - baseline_delta_metric: ``delta_metric`` from reconstructing onto a random - subspace of the same rank as ``retained``, instead of the requested one. - gated: True only if ``abs(delta_metric)`` exceeds ``abs(baseline_delta_metric)`` - (or an explicit threshold, if one was passed): a subfunction is causally - load-bearing only if it beats an equally-sized random subspace, not merely - "moves the metric at all". + baseline_delta_metric: the mean ``delta_metric`` over several random control + subspaces of the same width as ``retained``, each drawn inside the head's + own OV span ``span(V)`` rather than from the full residual stream, so the + control is the effect of an arbitrary same-size subspace of this head's + output rather than of an unrelated residual-stream direction. + gated: whether the retained subspace passed the causal test for the mode it was + expressed in, against ``abs(baseline_delta_metric)`` (or an explicit + threshold, if one was passed). For ``ablate`` (retain the complement), + removing a load-bearing subspace should move the metric more than removing an + arbitrary same-size one, so ``gated`` is ``abs(delta_metric) > threshold``. + For ``keep`` (retain only the given subspace), a subspace that reconstructs + the head's behavior should move the metric less than keeping an arbitrary + same-size one, so ``gated`` is ``abs(delta_metric) < threshold``. A single + "moved more than baseline" test cannot answer both, since ``keep=S`` and + ``ablate=complement(S)`` resolve to the same retained set. """ head_svd: HeadSVD @@ -778,17 +796,26 @@ def patch_along_directions( ablate: Optional[Sequence[int]] = None, threshold: Optional[float] = None, rng: Optional[torch.Generator] = None, + n_baseline: int = _DEFAULT_N_BASELINE, ) -> PatchResult: """Causally validate a claimed OV subfunction by reconstructing the head's output onto it. Requires ``head_svd.which == "OV"``: this reconstructs the write/output basis - ``V``, and QK has no such vector (see the module docstring). Runs the prompt three - times with ``use_attn_result`` enabled: once unmodified, once with the head's - ``hook_result`` slice reconstructed onto ``span(head_svd.V[:, retained])``, and once - onto a random orthonormal subspace of the same width, so a moved metric can be - compared against the effect of an equally-sized but arbitrary subspace instead of - being read as significant on its own. Restores the model's prior ``use_attn_result`` - setting afterward. + ``V``, and QK has no such vector (see the module docstring). Runs the prompt with + ``use_attn_result`` enabled once unmodified, once with the head's ``hook_result`` + 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 + the stream), so a moved metric is compared against the effect of an arbitrary + subspace of this head's output of the same width. The control delta is averaged over + ``n_baseline`` draws so one lucky or unlucky draw does not decide the gate. Restores + the model's prior ``use_attn_result`` setting afterward. + + The gate's success condition depends on the mode the caller expressed, because + ``keep=S`` and ``ablate=complement(S)`` resolve to the same retained set and a single + "moved more than the control" test would answer only the ``ablate`` question. See + :attr:`PatchResult.gated`. Args: model: A ``TransformerBridge``. @@ -800,7 +827,11 @@ def patch_along_directions( ablate: Direction indices to zero; the rest are retained. threshold: Explicit gate threshold. Defaults to ``None``, which uses ``abs(baseline_delta_metric)`` instead. - rng: Optional generator for the random baseline subspace, for reproducibility. + rng: Optional generator for the random control subspaces, for reproducibility. + Defaults to a generator seeded with ``_DEFAULT_BASELINE_SEED`` so a bare + call is reproducible rather than drawing from the global RNG. + n_baseline: Number of random in-span control subspaces to average the baseline + delta over. Must be at least 1. Returns: A :class:`PatchResult` describing the patched, baseline, and original metrics. @@ -808,9 +839,9 @@ def patch_along_directions( Raises: ValueError: If ``head_svd.which != "OV"``, if ``head_svd`` was decomposed under a different compatibility-mode state than ``model`` now has, if ``keep``/``ablate`` - are both given or both omitted, if any index is out of ``[0, rank)``, or if the + are both given or both omitted, if any index is out of ``[0, rank)``, if the retained set is empty (``keep=[]`` or ``ablate`` over the full rank) and no - explicit ``threshold`` is supplied. + explicit ``threshold`` is supplied, or if ``n_baseline < 1``. DegenerateDirectionError: If the retained directions split a degenerate block (see :func:`_validate_retained_blocks`). """ @@ -818,8 +849,13 @@ def patch_along_directions( raise ValueError( f"patch_along_directions requires an OV HeadSVD, got which={head_svd.which!r}" ) + if n_baseline < 1: + raise ValueError(f"n_baseline must be at least 1, got {n_baseline}") _validate_decomposition_matches_model(model, head_svd) retained = _resolve_retained(head_svd, keep, ablate) + # _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: # An empty retained set reconstructs the head onto the zero subspace, so its delta and # the equal-width random baseline's delta are both the full-ablation effect: the gate @@ -831,14 +867,13 @@ def patch_along_directions( "threshold to gate an empty retained set." ) - V = head_svd.V - kept_projector = V[:, retained] @ V[:, retained].transpose(-2, -1) + if rng is None: + rng = torch.Generator().manual_seed(_DEFAULT_BASELINE_SEED) - d_model = V.shape[0] + V = head_svd.V + rank = V.shape[1] width = len(retained) - random_input = torch.randn(d_model, width, generator=rng, dtype=V.dtype) - random_basis, _ = torch.linalg.qr(random_input) - baseline_projector = random_basis @ random_basis.transpose(-2, -1) + kept_projector = V[:, retained] @ V[:, retained].transpose(-2, -1) hook_name = f"blocks.{head_svd.layer}.attn.hook_result" previous = getattr(model.cfg, "use_attn_result", False) @@ -849,18 +884,36 @@ def patch_along_directions( prompt, fwd_hooks=[(hook_name, _make_subspace_hook(head_svd.head, kept_projector))] ) patched_metric = float(metric(patched_logits)) - baseline_logits = model.run_with_hooks( - prompt, - fwd_hooks=[(hook_name, _make_subspace_hook(head_svd.head, baseline_projector))], - ) - baseline_metric = float(metric(baseline_logits)) + + baseline_deltas: List[float] = [] + for _ in range(n_baseline): + # Draw a random width-of-rank subspace inside the head's own OV span. The QR + # is drawn on CPU (unimplemented on MPS, and a CUDA generator cannot feed a + # CPU randn); the columns are mapped into span(V) on V's device, and the hook + # moves the finished projector to the activation's device. + random_rank = torch.randn(rank, rank, generator=rng, dtype=V.dtype) + random_basis, _ = torch.linalg.qr(random_rank) + baseline_directions = V @ random_basis[:, :width].to(V.device) + baseline_projector = baseline_directions @ baseline_directions.transpose(-2, -1) + baseline_logits = model.run_with_hooks( + prompt, + fwd_hooks=[(hook_name, _make_subspace_hook(head_svd.head, baseline_projector))], + ) + baseline_deltas.append(float(metric(baseline_logits)) - original_metric) finally: model.set_use_attn_result(previous) delta_metric = patched_metric - original_metric - baseline_delta_metric = baseline_metric - original_metric + baseline_delta_metric = sum(baseline_deltas) / len(baseline_deltas) gate_threshold = abs(baseline_delta_metric) if threshold is None else threshold - gated = abs(delta_metric) > gate_threshold + # keep retains only the claimed subspace, so it passes when it reconstructs the head's + # behavior better than an arbitrary same-width one (moves the metric less); ablate + # removes it, so it passes when removing it matters more than removing an arbitrary + # same-width one (moves the metric more). + if mode == "keep": + gated = abs(delta_metric) < gate_threshold + else: + gated = abs(delta_metric) > gate_threshold return PatchResult( head_svd=head_svd, From fb1c6c690f705a40ff0f8d1a33e63a7e8eceedea Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 01:39:09 +0530 Subject: [PATCH 13/14] test(svd_circuits): pin the OV write basis and per-head isolation against mutation --- tests/integration/test_svd_circuits.py | 25 +++++++++ tests/unit/tools/test_svd_circuits.py | 74 ++++++++++++++++++++++++-- 2 files changed, 95 insertions(+), 4 deletions(-) diff --git a/tests/integration/test_svd_circuits.py b/tests/integration/test_svd_circuits.py index 256f29743..68bb2c372 100644 --- a/tests/integration/test_svd_circuits.py +++ b/tests/integration/test_svd_circuits.py @@ -14,6 +14,7 @@ import torch from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.SVDInterpreter import SVDInterpreter from transformer_lens.tools.analysis.svd_circuits import ( decompose_head, patch_along_directions, @@ -53,6 +54,15 @@ def test_svd_circuits_readout_and_patch_gate_on_name_mover_head(gpt2_bridge) -> # Pick the first non-degenerate direction rather than assuming index 0 is isolated. top_direction = next(row.idx for row in ov.rank_report if not row.is_degenerate) + # readout reproduces the shipped SVDInterpreter reference up to sign, so a U-for-V swap + # or an all-zeros return in vocab_readout fails here rather than passing a finiteness check. + reference = SVDInterpreter(gpt2_bridge).get_singular_vectors( + "OV", LAYER, head_index=HEAD, num_vectors=top_direction + 1 + )[:, 0, top_direction] + assert torch.allclose(readout[:, top_direction], reference, atol=1e-4) or torch.allclose( + readout[:, top_direction], -reference, atol=1e-4 + ) + prompt = gpt2_bridge.to_tokens(CLEAN_PROMPT) metric = _logit_diff_metric(gpt2_bridge) @@ -64,3 +74,18 @@ def test_svd_circuits_readout_and_patch_gate_on_name_mover_head(gpt2_bridge) -> assert math.isfinite(result.delta_metric) assert math.isfinite(result.baseline_delta_metric) assert isinstance(result.gated, bool) + + rank = ov.V.shape[1] + + # Keeping every direction reconstructs the head onto its own full OV span: a no-op, so the + # metric barely moves. A hook never installed, or one projecting the wrong basis, breaks this. + full_keep = patch_along_directions(gpt2_bridge, ov, prompt, metric, keep=list(range(rank))) + assert abs(full_keep.delta_metric) < 1e-4 + + # Ablating every direction zeroes the head's whole output, which must move the logit diff. + # An empty retained set ties the baseline by construction, so gate it with an explicit + # threshold; only delta_metric's magnitude is asserted here. + full_ablate = patch_along_directions( + gpt2_bridge, ov, prompt, metric, ablate=list(range(rank)), threshold=0.0 + ) + assert abs(full_ablate.delta_metric) > 1e-2 diff --git a/tests/unit/tools/test_svd_circuits.py b/tests/unit/tools/test_svd_circuits.py index 79c6f93df..84621bdeb 100644 --- a/tests/unit/tools/test_svd_circuits.py +++ b/tests/unit/tools/test_svd_circuits.py @@ -23,6 +23,7 @@ RankReportRow, _degeneracy_blocks, _factored_head_svd, + _make_subspace_hook, _resolve_retained, decompose_head, logit_signature, @@ -551,6 +552,33 @@ def test_ov_output_direction_matches_svd_interpreter(tiny_bridge): assert not torch.allclose(projected_via_U, -reference[:, 0, 0], atol=1e-2) +def test_vocab_readout_matches_svd_interpreter_through_public_api(): + """The public vocab_readout, not just a hand-rolled V @ W_U, reproduces SVDInterpreter's + OV readout up to sign. + + The shape-only check elsewhere passes a U-for-V swap or an all-zeros return inside + vocab_readout; matching the shipped reference through the public function fails both. + Needs compatibility mode for the W_U-folding contract, so this builds its own bridge + rather than using the shared no-compat fixture. + """ + from transformer_lens.SVDInterpreter import SVDInterpreter + + model = _make_tiny_bridge() + model.enable_compatibility_mode() + layer, head = 0, 0 + ov = decompose_head(model, layer, head, which=("OV",)).OV + top = next(row.idx for row in ov.rank_report if not row.is_degenerate) + + readout = vocab_readout(model, ov, k=top + 1) + reference = SVDInterpreter(model).get_singular_vectors( + "OV", layer, head_index=head, num_vectors=top + 1 + )[:, 0, top] + + assert torch.allclose(readout[:, top], reference, atol=1e-4) or torch.allclose( + readout[:, top], -reference, atol=1e-4 + ) + + # --------------------------------------------------------------------------- # # Compatibility-mode binding: a decomposition is tied to the state it was built under # --------------------------------------------------------------------------- # @@ -723,9 +751,14 @@ def test_project_activations_which_guard(): project_activations(SimpleNamespace(), qk, "hello world") -def test_project_activations_reconstructs_head_output(tiny_bridge): - """Coefficients recovered against V reconstruct the actual cached hook_result slice.""" - layer, head = 0, 0 +@pytest.mark.parametrize("head", [0, 1]) +def test_project_activations_reconstructs_head_output(tiny_bridge, head): + """Coefficients recovered against V reconstruct the actual cached hook_result slice. + + Runs on both heads of the two-head bridge, so a hardcoded head-0 read cannot pass by + only ever being asked about head 0. + """ + layer = 0 decomposition = decompose_head(tiny_bridge, layer, head, which=("OV",)) # A raw token tensor, not a string: tiny_bridge's tokenizer is a MagicMock and cannot # tokenize text, but a token-id tensor bypasses that path entirely. @@ -761,6 +794,36 @@ def test_project_activations_restores_use_attn_result(tiny_bridge): tiny_bridge.set_use_attn_result(original) +@pytest.mark.parametrize("head", [0, 1]) +def test_subspace_hook_leaves_sibling_head_untouched(tiny_bridge, head): + """The patch hook rewrites only its own head's slice of the [batch, pos, head, d_model] + activation: the sibling head's block is bit-for-bit unchanged while the patched head's + block moves. + + Pins _make_subspace_hook's per-head scope on the two-head bridge, which a hook that + wrote every head, or hardcoded head 0, would break. The clean hook_result comes from a + real forward pass; applying the hook directly makes the per-head write the sole variable. + """ + ov = decompose_head(tiny_bridge, 0, head, which=("OV",)).OV + other = 1 - head + prompt = torch.tensor([[5, 63, 7, 9]]) + + previous = tiny_bridge.cfg.use_attn_result + tiny_bridge.set_use_attn_result(True) + try: + _, cache = tiny_bridge.run_with_cache(prompt) + finally: + tiny_bridge.set_use_attn_result(previous) + clean = cache[("result", 0, "attn")] + + keep = [ov.rank_report[0].idx] + projector = ov.V[:, keep] @ ov.V[:, keep].transpose(-2, -1) + patched = _make_subspace_hook(head, projector)(clean, hook=SimpleNamespace(name="hook_result")) + + assert torch.equal(patched[:, :, other, :], clean[:, :, other, :]) + assert not torch.allclose(patched[:, :, head, :], clean[:, :, head, :]) + + # --------------------------------------------------------------------------- # # patch_along_directions (mandatory causal gate) # --------------------------------------------------------------------------- # @@ -1102,7 +1165,10 @@ def metric(logits): assert isinstance(strong_result, PatchResult) assert math.isfinite(strong_result.baseline_delta_metric) assert math.isfinite(weak_result.baseline_delta_metric) - assert abs(strong_result.delta_metric) > abs(weak_result.delta_metric) + # A wide margin, not a bare ordering: swapping the projected basis inside the patch hook + # collapses the strong-vs-weak ratio toward 1 while still ordering the two, so a bare `>` + # would pass the swap this decomposition's OV convention exists to prevent. + assert abs(strong_result.delta_metric) > 5 * abs(weak_result.delta_metric) def test_patch_along_directions_restores_use_attn_result(tiny_bridge): From 095893aa9e20f4b4ae3817a7a14bc0bba04393df Mon Sep 17 00:00:00 2001 From: janmenjayap Date: Wed, 16 Sep 2026 02:05:08 +0530 Subject: [PATCH 14/14] fix(svd_circuits): filter project_activations cache to one hook and reject a batched tensor pre-forward --- tests/unit/tools/test_svd_circuits.py | 38 +++++++++++++++++++ .../tools/analysis/svd_circuits.py | 17 ++++++++- 2 files changed, 53 insertions(+), 2 deletions(-) diff --git a/tests/unit/tools/test_svd_circuits.py b/tests/unit/tools/test_svd_circuits.py index 84621bdeb..48e039afc 100644 --- a/tests/unit/tools/test_svd_circuits.py +++ b/tests/unit/tools/test_svd_circuits.py @@ -794,6 +794,44 @@ def test_project_activations_restores_use_attn_result(tiny_bridge): tiny_bridge.set_use_attn_result(original) +def test_project_activations_caches_only_the_read_hook(tiny_bridge, monkeypatch): + """Only the one hook_result being read is cached, not every hook point of the forward. + + run_with_cache retains every activation it is not told to filter, so reading one head's + output without a names_filter would keep the whole forward pass. Pin that a filter is + passed and that it selects exactly the read hook and nothing else. + """ + layer, head = 0, 0 + ov = decompose_head(tiny_bridge, layer, head, which=("OV",)).OV + captured = {} + original = tiny_bridge.run_with_cache + + def spy(*args, **kwargs): + captured["names_filter"] = kwargs.get("names_filter") + return original(*args, **kwargs) + + monkeypatch.setattr(tiny_bridge, "run_with_cache", spy) + project_activations(tiny_bridge, ov, torch.tensor([[5, 63, 7, 9]])) + + names_filter = captured["names_filter"] + assert names_filter is not None + assert names_filter(f"blocks.{layer}.attn.hook_result") + assert not names_filter(f"blocks.{layer}.attn.hook_z") + + +def test_project_activations_rejects_batched_tensor_before_forward(tiny_bridge): + """A batched token tensor is refused before the forward pass, not after caching it. + + The ids are out of vocab range, so a check that ran only after the forward would surface + the embedding's IndexError instead of this refusal; catching the batch dimension first both + saves the forward and gives the caller the actionable message. + """ + ov = decompose_head(tiny_bridge, 0, 0, which=("OV",)).OV + batched = torch.full((2, 4), 10**6, dtype=torch.long) + with pytest.raises(ValueError, match="single prompt"): + project_activations(tiny_bridge, ov, batched) + + @pytest.mark.parametrize("head", [0, 1]) def test_subspace_hook_leaves_sibling_head_untouched(tiny_bridge, head): """The patch hook rewrites only its own head's slice of the [batch, pos, head, d_model] diff --git a/transformer_lens/tools/analysis/svd_circuits.py b/transformer_lens/tools/analysis/svd_circuits.py index 31e8f8742..826b05b7e 100644 --- a/transformer_lens/tools/analysis/svd_circuits.py +++ b/transformer_lens/tools/analysis/svd_circuits.py @@ -654,13 +654,26 @@ def project_activations( f"project_activations requires an OV HeadSVD, got which={head_svd.which!r}" ) _validate_decomposition_matches_model(model, head_svd) + # A batched token tensor carries its batch dimension up front, so reject it before the + # 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): + raise ValueError( + "project_activations requires a single prompt, got a token tensor of shape " + f"{tuple(prompt.shape)}; pass a [1, pos] tensor or a single string." + ) + # Cache only the one hook this reads. run_with_cache otherwise retains every hook point of + # the forward pass (about 1 GB against 19 MB on a 512-token gpt2-small prompt) to read a + # single head's output. + hook_name = f"blocks.{head_svd.layer}.attn.hook_result" previous = getattr(model.cfg, "use_attn_result", False) model.set_use_attn_result(True) try: - _, cache = model.run_with_cache(prompt) + _, cache = model.run_with_cache(prompt, names_filter=lambda name: name == hook_name) finally: model.set_use_attn_result(previous) - result = cache[("result", head_svd.layer, "attn")][..., head_svd.head, :] + result = cache[hook_name][..., head_svd.head, :] if result.shape[0] != 1: raise ValueError( f"project_activations requires a single prompt, got batch={result.shape[0]}"