SVD Circuits Demo¶
This notebook decomposes one attention head’s QK (\(W_Q W_K^\top\)) and OV (\(W_V W_O\)) weight maps into orthogonal singular directions with FactoredMatrix.svd, reads the OV directions out through the final norm and unembedding with logit_readout, projects the head’s cached output onto them, and measures the effect of retaining or ablating single directions on one IOI prompt’s logit difference with a hook on hook_result. Mathematically distinct directions need not be
semantically or causally distinct subfunctions.
Based on Areeb Ahmad, Abhinav Joshi, Ashutosh Modi, “Beyond Components: Singular Vector-Based Interpretability of Transformer Circuits”, arXiv 2511.20273.
Setup (Ignore)¶
[1]:
# NBVAL_IGNORE_OUTPUT
# Janky code to do different setup when run in a Colab notebook vs VSCode
import os
DEVELOPMENT_MODE = False
IN_GITHUB = os.getenv("GITHUB_ACTIONS") == "true"
try:
import google.colab
IN_COLAB = True
print("Running as a Colab notebook")
except ImportError:
IN_COLAB = False
print("Running as a Jupyter notebook - intended for development only!")
DEVELOPMENT_MODE = True
from IPython import get_ipython
ipython = get_ipython()
ipython.run_line_magic("load_ext", "autoreload")
ipython.run_line_magic("autoreload", "2")
if IN_COLAB or IN_GITHUB:
%pip install transformer_lens
Running as a Jupyter notebook - intended for development only!
[2]:
# NBVAL_IGNORE_OUTPUT
import torch
from transformers.utils import logging as hf_logging
from transformer_lens import FactoredMatrix
from transformer_lens.model_bridge import TransformerBridge
from transformer_lens.tools.analysis import logit_readout
hf_logging.disable_progress_bar()
# Every step here is a forward pass or a weight read; nothing needs a graph.
torch.set_grad_enabled(False)
MODEL_ID = "gpt2"
device = "cuda" if torch.cuda.is_available() else "cpu"
model = TransformerBridge.boot_transformers(MODEL_ID, dtype=torch.float32, device=device)
# Folding ln1 into W_Q/W_K/W_V puts each map in the normalized basis the head actually reads,
# and centering W_O leaves every write direction mean-free.
model.enable_compatibility_mode()
# Splits the attention output by head into hook_result, which sections 3 and 4 read and patch.
model.set_use_attn_result(True)
model.eval()
print(f"loaded {MODEL_ID}: {model.cfg.n_layers} layers, d_model={model.cfg.d_model}, {device=}")
loaded gpt2: 12 layers, d_model=768, device='cpu'
1. Decompose a name-mover head¶
Layer 9 head 9 is the canonical name-mover head in GPT-2-small’s indirect-object-identification (IOI) circuit. Two weight maps describe what it does:
QK \(= W_Q W_K^\top\) decides which residual-stream directions the head compares when forming its attention pattern.
OV \(= W_V W_O\) decides what the head writes back into the residual stream.
Both are \(d_{\text{model}} \times d_{\text{model}}\) but rank at most \(d_{\text{head}} = 64\). FactoredMatrix keeps each map as its two factors and computes the SVD \(M = U \Sigma V^\top\) without ever materialising the \(768 \times 768\) product. The convention matters, because both factors are \(d_{\text{model}}\)-wide and a swap raises no shape error:
For OV, the columns of \(V\) are the residual-stream write directions, the ones read out through the unembedding below; the columns of \(U\) span the input the values read.
For QK, the columns of \(U\) are the destination (query-side) read directions and the columns of \(V\) the source (key-side) read directions. QK writes nothing, so it has no vocab readout.
On grouped-query models W_K and W_V can carry one row per key/value head, so query head \(h\) reads key/value head \(h\) // (n_heads / n_key_value_heads). head_maps resolves that from the weight shapes, which makes it a no-op on GPT-2.
[3]:
# NBVAL_IGNORE_OUTPUT
# sigma and sigma/sigma_max are float32 reductions, so the fourth decimal can differ between CPU
# kernels. The report shows which directions are degenerate, not a pinned numeric value.
PROMPT = "When Mary and John went to the store, John gave a drink to"
LAYER, HEAD = 9, 9
EPS = 1e-2 # Neighbouring singular values closer than this, relatively, form one block.
def head_maps(model, layer, head):
"""The head's QK and OV maps as FactoredMatrix objects."""
attn = model.blocks[layer].attn
kv_head = head // (attn.W_Q.shape[0] // attn.W_K.shape[0])
W_Q, W_O = attn.W_Q[head].detach(), attn.W_O[head].detach()
W_K, W_V = attn.W_K[kv_head].detach(), attn.W_V[kv_head].detach()
return FactoredMatrix(W_Q, W_K.T), FactoredMatrix(W_V, W_O)
def rank_report(S, eps=EPS):
"""One row per direction; degenerate means inside a near-equal block or numerically null."""
near = (1 - S[1:] / S[:-1]) < eps
blocked = torch.cat([near, near.new_zeros(1)]) | torch.cat([near.new_zeros(1), near])
# The relative tolerance torch.linalg.matrix_rank uses for a d_model x d_model map.
null = S / S[0] < model.cfg.d_model * torch.finfo(S.dtype).eps
return [
{
"idx": i,
"sigma": float(S[i]),
"ratio": float(S[i] / S[0]),
"degenerate": bool(blocked[i] or null[i]),
"null": bool(null[i]),
}
for i in range(len(S))
]
qk, ov = head_maps(model, LAYER, HEAD)
U, S, V = ov.svd() # U and V are [d_model, d_head]; S is sorted descending.
report = rank_report(S)
eligible_ids = [row["idx"] for row in report if not row["degenerate"]]
display_ids = eligible_ids[:5]
qk_isolated = sum(not row["degenerate"] for row in rank_report(qk.S))
print(f"QK: sigma_max {float(qk.S[0]):.4f}, {qk_isolated} of {len(qk.S)} directions isolated")
print(f"OV: sigma_max {float(S[0]):.4f}, {len(eligible_ids)} of {len(S)} directions isolated")
print()
print(f"{'idx':>4} {'sigma':>10} {'sigma/sigma_max':>16} {'degenerate':>11} {'null':>6}")
for row in report[:12]:
print(
f"{row['idx']:>4} {row['sigma']:>10.4f} {row['ratio']:>16.4f} "
f"{str(row['degenerate']):>11} {str(row['null']):>6}"
)
QK: sigma_max 2.4274, 18 of 64 directions isolated
OV: sigma_max 8.4096, 13 of 64 directions isolated
idx sigma sigma/sigma_max degenerate null
0 8.4096 1.0000 False False
1 8.2600 0.9822 False False
2 8.1019 0.9634 False False
3 8.0099 0.9525 True False
4 7.9312 0.9431 True False
5 7.8743 0.9364 True False
6 7.8141 0.9292 True False
7 7.7938 0.9268 True False
8 7.7175 0.9177 True False
9 7.6490 0.9096 True False
10 7.5680 0.8999 True False
11 7.5578 0.8987 True False
Reading the rank report¶
Singular directions are unique only when the singular values are distinct. Near-equal consecutive singular values define a subspace whose basis is only defined up to a rotation, so individual directions inside such a block are not attributable: “direction 3 is the surname subfunction” would be a statement about an arbitrary basis choice rather than about the model. A block is a run of neighbours each within EPS of the previous one, so a block can span more than EPS end to end.
Numerically null directions are arbitrary vectors from the map’s null space and are equally unusable.
The rest of the notebook therefore reads, projects and patches only isolated, non-null directions. A complete degenerate block can still be retained or removed as a subspace; only splitting it is meaningless.
2. Vocab projections of eligible directions¶
logit_readout reads any \(d_{\text{model}}\) vector through the model’s final norm and unembedding. Two adjustments make that a readout of a direction:
Singular vectors are defined only up to sign: flipping a column of \(U\) and the matching column of \(V\) leaves the map unchanged. Both ends of each direction are shown.
The readout of the zero vector is the part of the final norm and unembedding that does not depend on the input (the folded final-norm bias). Subtracting it leaves logits that flip sign with the direction.
[4]:
# NBVAL_IGNORE_OUTPUT
# Token rankings can change across kernels when float32 projection magnitudes nearly tie.
def read_directions(directions):
"""Vocab logits of [k, d_model] directions, less the input-independent offset."""
zero = torch.zeros_like(directions[:1])
return logit_readout(model, directions).values - logit_readout(model, zero).values
def decode(ids):
return [model.tokenizer.decode([int(token)]) for token in ids]
readout = read_directions(V[:, display_ids].T)
for idx, logits in zip(display_ids, readout):
print(f"direction {idx:>2} +: {decode(logits.topk(6).indices)}")
print(f" -: {decode((-logits).topk(6).indices)}")
direction 0 +: ['Sav', 'AV', 'av', 'adh', 'avour', 'aw']
-: ['tein', ' Kling', ' Pes', ' Monk', ' Schmidt', 'enegger']
direction 1 +: ['NI', 'ANI', 'irin', 'Ir', ' NI', ' Irwin']
-: ['tro', 'Tro', 'ero', 'bos', ' Merc', ' tom']
direction 2 +: [' McKenna', 'etr', ' copper', ' Ack', 'eter', 'EP']
-: [' Lindsay', ' 339', 'Lind', ' seiz', 'gae', ' Instr']
direction 23 +: [' Mahjong', ' Buk', 'abal', 'Mp', 'jp', 'hap']
-: [' Sieg', 'tery', 'YN', ' IG', ' Ez', ' Ys']
direction 48 +: [' Pupp', ' DN', 'Doc', ' Blog', ' Ning', ' doct']
-: ['û', 'Rand', 'iris', ' Wade', ' Ashe', 'rain']
Vocab projections can motivate hypotheses but do not establish semantic labels. The next section describes activation coefficients; the intervention section measures changes in the selected prompt’s logit difference, without confirming a named subfunction.
3. Activation projections on a prompt¶
With use_attn_result on, blocks.9.attn.hook_result holds each head’s output separately, shaped [batch, pos, head, d_model]. Caching only that hook and projecting head 9’s slice onto \(V\) gives the coefficient of every write direction at every position. The sign of a coefficient follows the arbitrary sign of its direction, so magnitudes are shown.
[5]:
# NBVAL_IGNORE_OUTPUT
# The coefficients are float32 reductions over the head's output, so the second decimal can
# differ between CPU kernels; the table shows which directions fire at which positions.
hook_name = f"blocks.{LAYER}.attn.hook_result"
_, cache = model.run_with_cache(PROMPT, names_filter=hook_name)
coefficients = cache[hook_name][0, :, HEAD] @ V # [pos, d_head]
print(f"{'token':>10} " + " ".join(f"{f'dir{idx}':>8}" for idx in display_ids))
for token, row in zip(model.to_str_tokens(PROMPT), coefficients):
print(f"{token:>10} " + " ".join(f"{abs(float(row[idx])):>8.2f}" for idx in display_ids))
token dir0 dir1 dir2 dir23 dir48
<|endoftext|> 0.34 0.20 0.17 0.14 0.27
When 0.34 0.21 0.19 0.12 0.29
Mary 0.31 0.21 0.19 0.15 0.27
and 0.20 0.37 0.51 0.35 0.20
John 0.26 0.19 0.18 0.14 0.30
went 0.34 0.20 0.17 0.13 0.28
to 0.27 0.18 0.16 0.15 0.31
the 0.34 0.22 0.20 0.13 0.27
store 0.31 0.24 0.23 0.11 0.31
, 3.46 0.04 0.65 1.49 0.77
John 0.27 0.17 0.16 0.10 0.36
gave 5.59 1.41 2.72 2.42 0.07
a 0.25 0.27 0.22 0.11 0.29
drink 0.29 0.19 0.24 0.07 0.32
to 7.31 1.74 3.41 3.24 0.24
4. Intervention comparisons¶
A hook on hook_result replaces head 9’s output \(r\) at every position with its image under a projector and leaves every other head untouched:
keep direction \(i\): \(r \mapsto (r \cdot v_i)\, v_i\), retaining only that direction;
ablate direction \(i\): \(r \mapsto r - (r \cdot v_i)\, v_i\), removing it.
Each run reports delta, the change in the logit difference " Mary" minus " John" at the final position, which includes all downstream computation. As a reference, the same intervention is applied along 32 random unit directions drawn inside the head’s own OV span, so a control is an arbitrary direction of this head’s output rather than of the whole residual stream. share is the fraction of those controls whose |delta| is smaller than the direction’s own:
In keep mode, a low share means retaining the direction moves the metric less than most random in-span directions do.
In ablate mode, a high share means removing it moves the metric more than most random in-span directions do.
These are comparisons with a small sample of controls on one prompt, not significance tests. The whole-head ablation first gives context: it zeros the head’s output, which a same-width control cannot distinguish from itself.
[6]:
# NBVAL_IGNORE_OUTPUT
# The deltas are float32 reductions, and a value sitting on a rounding boundary can print one
# digit apart across CPU kernels.
mary, john = model.to_single_token(" Mary"), model.to_single_token(" John")
identity = torch.eye(model.cfg.d_model, device=V.device)
def logit_diff(logits):
return float(logits[0, -1, mary] - logits[0, -1, john])
def patched_logit_diff(projector):
"""Logit difference with the head's output replaced by its image under projector."""
def hook(result, hook):
result = result.clone()
result[:, :, HEAD] = result[:, :, HEAD] @ projector
return result
return logit_diff(model.run_with_hooks(PROMPT, fwd_hooks=[(hook_name, hook)]))
def keep(v):
return torch.outer(v, v)
def ablate(v):
return identity - torch.outer(v, v)
original = logit_diff(model(PROMPT))
print(f"original logit difference: {original:.3f}")
print(f"whole-head ablation delta: {patched_logit_diff(0 * identity) - original:+.3f}")
N_CONTROLS = 32
rank = sum(not row["null"] for row in report)
draws = torch.randn(rank, N_CONTROLS, generator=torch.Generator().manual_seed(0)).to(V.device)
controls = V[:, :rank] @ torch.nn.functional.normalize(draws, dim=0)
original logit difference: 3.362
whole-head ablation delta: +0.129
[7]:
# NBVAL_IGNORE_OUTPUT
# Same reason as above: the deltas are float32 reductions that can differ in the last printed
# digit across CPU kernels, which can also move a share by one control.
for mode, projector in (("keep", keep), ("ablate", ablate)):
control_deltas = torch.tensor(
[abs(patched_logit_diff(projector(v)) - original) for v in controls.T]
)
mean = float(control_deltas.mean())
print(f"\n{mode}: mean |delta| over {N_CONTROLS} random in-span directions {mean:.3f}")
print(f"{'dir':>4} {'sigma':>9} {'delta':>10} {'share':>6}")
for idx in eligible_ids:
delta = patched_logit_diff(projector(V[:, idx])) - original
share = float((control_deltas < abs(delta)).float().mean())
print(f"{idx:>4} {float(S[idx]):>9.4f} {delta:>10.3f} {share:>6.2f}")
keep: mean |delta| over 32 random in-span directions 0.132
dir sigma delta share
0 8.4096 0.154 0.66
1 8.2600 0.141 0.62
2 8.1019 0.132 0.53
23 6.8420 0.193 0.84
48 5.5535 0.131 0.53
49 5.4845 0.113 0.31
50 5.3609 0.131 0.53
51 5.2535 0.163 0.78
55 5.0464 0.078 0.16
56 4.9738 0.114 0.31
59 4.7041 0.131 0.53
60 4.6103 0.136 0.62
61 4.4351 0.172 0.78
ablate: mean |delta| over 32 random in-span directions 0.041
dir sigma delta share
0 8.4096 -0.012 0.31
1 8.2600 -0.015 0.34
2 8.1019 -0.014 0.34
23 6.8420 -0.058 0.72
48 5.5535 -0.001 0.03
49 5.4845 0.016 0.34
50 5.3609 -0.005 0.16
51 5.2535 -0.036 0.56
55 5.0464 0.040 0.59
56 4.9738 0.025 0.44
59 4.7041 0.002 0.12
60 4.6103 -0.016 0.34
61 4.4351 -0.032 0.53
5. What this does and does not show¶
This example decomposes one GPT-2-small head and measures prompt-specific effects of retaining or ablating eligible OV directions, next to random in-span controls. It does not establish named subfunctions, statistical separation from arbitrary directions, or replication of the paper’s causal subfunction taxonomy.
Single head, single model, single prompt. Everything above concerns layer 9 head 9 of
gpt2-smallon one IOI prompt. Nothing here assembles a multi-head circuit.Descriptive projections. Orthogonal weight-space directions are mathematically distinct; vocab readouts and activation coefficients do not supply semantic labels.
Comparisons, not significance. The tables report metric changes and how they sit among 32 random controls, not confidence bounds. They do not establish that any direction is distinguishable from an arbitrary one.
Downstream responses. Every delta measures the final logit difference after downstream computation, not the head’s direct write contribution. A small final-logit effect need not imply a small direct contribution, and these measurements do not identify a compensating component.
Repeatability, not confirmation. Shares near the middle can change with the seed or the number of controls. Fixing the seed makes the comparison repeatable, not scientifically conclusive.
Using this on other models¶
The layer must hold an attention block with
W_Q,W_K,W_VandW_O. SSM, mixer, linear-attention and AltUp layers, and encoder-decoder bridges without ablocksstack, have no such head to decompose.Sections 3 and 4 need per-head outputs:
set_use_attn_result(True)raisesNotImplementedErroron attention bridges that cannot split their output by head. Check thathook_resultactually appears in the cache before trusting a patch, since a patch on a hook that never fires changes nothing.Decompose, read out and patch under one weight-processing state. The weights, the final norm and
hook_resultall change when compatibility mode folds and centers them, so a decomposition taken beforeenable_compatibility_mode()does not describe the model afterwards.