Attribution Graphs from Cross-Layer Transcoders¶
An attribution graph is a picture of why a model said what it said, drawn in terms of interpretable features rather than layers and heads. This tutorial builds one end to end in NNsight: we take Llama-3.2-1B, replace every MLP with a cross-layer transcoder, verify that the replacement is faithful, trace the graph for a single prompt, check that it is an exact decomposition, and finally test it — by ablating the features the graph blames and watching the answer collapse.
Everything here runs locally on one GPU in about 7 GB.
Why replace the MLPs at all?¶
Activation patching and attribution patching (our causal mediation tutorials) tell you which layer and position mattered. That is a coordinate, not an explanation. The reason they stop there is that an MLP's hidden units are dense and polysemantic — there is no natural set of "things this MLP did" to attribute to.
A transcoder manufactures that set. It is a dictionary trained to imitate an MLP: it reads the MLP's input, encodes it into a large sparse set of feature activations, and decodes those features back into the MLP's output. Note the asymmetry with an SAE — an SAE reconstructs its own input, so its features describe a representation. A transcoder maps input to output, so its features describe a computation. Its decoder rows are the vectors the MLP writes into the residual stream, one per feature, and on a given prompt only about a dozen per layer and position are non-zero.
A cross-layer transcoder (CLT) goes one step further. A feature born at layer $l$ owns a decoder row for every target layer $l, l{+}1, \ldots, L{-}1$, so a single feature can write into the whole rest of the network. The reconstruction of layer $t$'s MLP output is
$$\widehat{\mathrm{mlp\_out}}_t \;=\; b^{\mathrm{dec}}_t \;+\; \sum_{l \le t} \; \sum_{f} a_{l,f}\, W^{\mathrm{dec}}_{l,f}[t-l]$$
where $a_{l,f}$ is feature $f$'s activation at layer $l$. This is the object Anthropic's Circuit Tracing work uses, and it is the reason a feature in the graph is not tied to one layer.
What the graph is a graph of¶
Swap each MLP for its reconstruction and you get a replacement model. Add back the residual $\varepsilon_t = \mathrm{mlp\_out}_t - \widehat{\mathrm{mlp\_out}}_t$ as a frozen constant and you get a local replacement model: numerically identical to the original on this prompt, but now written entirely in terms of quantities you can name. The nodes of the attribution graph are exactly those quantities:
| node type | one per | what it means |
|---|---|---|
| feature | (layer, position, feature id) with $a > 0$ | a dictionary feature that fired |
| error | (layer, position) | what the dictionary failed to explain at that site |
| token | position | the embedding that entered the stream |
| logit | the target | the thing we are explaining |
Error nodes are not decoration. They are the honesty mechanism: if your dictionary is bad, the graph will say so by routing the explanation through error nodes instead of features.
An edge $i \to j$ is the direct linear effect of node $i$'s value on node $j$'s value — the effect that does not pass through any other feature. We get it by making every node value an autograd leaf and reading $x_i \cdot \partial x_j / \partial x_i$. For that number to be a direct effect and not a total one, we freeze the two non-linearities that route information between positions and rescale it: attention patterns and LayerNorm scales. With those frozen, the model is linear in every leaf, and by Euler's identity the edges into a node must sum exactly to that node. That is a hard test, and we run it.
Setup¶
Two warnings before you press run:
- This notebook needs a local GPU and ~20 GB of free disk. The CLT checkpoint is 20 GB (its decoder alone is 9.1 B parameters). It does not fit in a free Colab session, and unlike the rest of our tutorials it is not written to. We keep GPU usage under 7 GB by leaving the decoder memory-mapped on disk and gathering only the rows for features that actually fired.
- The
task=argument toTransformersModelis not optional here. Without it, loading fails outright when the HuggingFace hub is in offline mode.
Executing the whole notebook takes roughly an hour, almost all of it in section 1's mispaired-dictionary control, which fires 50x more features than the real thing and therefore gathers 50x more decoder rows off disk. The attribution itself takes seconds.
try:
import google.colab
is_colab = True
except ImportError:
is_colab = False
if is_colab:
!pip install -U nnsight
import os, time, yaml, json
from IPython.display import clear_output
import numpy as np
import torch
from safetensors import safe_open
from huggingface_hub import snapshot_download
import nnsight
from nnsight import TransformersModel
import plotly.express as px
import plotly.graph_objects as go
import plotly.io as pio
pio.renderers.default = "colab" if is_colab else "plotly_mimetype+notebook_connected+colab+notebook"
print(nnsight.__version__)
/home/localjadenfk/miniconda3/envs/ndif2/lib/python3.12/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html from .autonotebook import tqdm as notebook_tqdm
0.7.1.dev41+gd901da3ed
We load in float32. This is unusual for a 1B model, but the closure check below is a cancellation of ~2000 terms against a target of size 5, and in bfloat16 the residual is dominated by rounding rather than by whether our freezes are correct. float32 costs 5 GB and buys us six decimal places of diagnostic.
DEVICE = "cuda:0"
N_LAYERS = 16
model = TransformersModel(
"meta-llama/Llama-3.2-1B",
task="text-generation",
dtype=torch.float32,
device_map=DEVICE,
dispatch=True,
)
tokenizer = model.tokenizer
clear_output()
print(model)
LlamaForCausalLM(
(model): LlamaModel(
(embed_tokens): Embedding(128256, 2048)
(layers): ModuleList(
(0-15): 16 x LlamaDecoderLayer(
(self_attn): LlamaAttention(
(q_proj): Linear(in_features=2048, out_features=2048, bias=False)
(k_proj): Linear(in_features=2048, out_features=512, bias=False)
(v_proj): Linear(in_features=2048, out_features=512, bias=False)
(o_proj): Linear(in_features=2048, out_features=2048, bias=False)
)
(mlp): LlamaMLP(
(gate_proj): Linear(in_features=2048, out_features=8192, bias=False)
(up_proj): Linear(in_features=2048, out_features=8192, bias=False)
(down_proj): Linear(in_features=8192, out_features=2048, bias=False)
(act_fn): SiLUActivation()
)
(input_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
(post_attention_layernorm): LlamaRMSNorm((2048,), eps=1e-05)
)
)
(norm): LlamaRMSNorm((2048,), eps=1e-05)
(rotary_emb): LlamaRotaryEmbedding()
)
(lm_head): Linear(in_features=2048, out_features=128256, bias=False)
(generator): Generator(
(streamer): Streamer()
)
)
1. Pairing the artifact to the model, numerically¶
This is the step people skip. A dictionary trained for a different model, a different layer, or a different hook point will still load, still produce sparse features, and still yield a reconstruction that looks like an activation. Every downstream number will be plausible and wrong. So before interpreting anything we establish the pairing from the artifact itself, then confirm it with a number.
The whole specification of mntss/clt-llama-3.2-1b-524k is a four-line config.yaml plus the
tensor shapes.
CLT_PATH = snapshot_download("mntss/clt-llama-3.2-1b-524k")
clear_output()
cfg = yaml.safe_load(open(os.path.join(CLT_PATH, "config.yaml")))
print(json.dumps(cfg, indent=2))
# Shapes straight out of the safetensors headers -- no weights loaded.
for l in [0, 3, 15]:
with safe_open(os.path.join(CLT_PATH, f"W_enc_{l}.safetensors"), framework="pt", device="cpu") as f:
enc_shapes = {k: f.get_slice(k).get_shape() for k in f.keys()}
with safe_open(os.path.join(CLT_PATH, f"W_dec_{l}.safetensors"), framework="pt", device="cpu") as f:
dec_shape = f.get_slice(f"W_dec_{l}").get_shape()
print(f"layer {l:2d}: {enc_shapes} W_dec_{l} {dec_shape}")
{
"model_name": "meta-llama/Llama-3.2-1B",
"model_kind": "cross_layer_transcoder",
"feature_input_hook": "hook_resid_mid",
"feature_output_hook": "hook_mlp_out"
}
layer 0: {'W_enc_0': [32768, 2048], 'b_dec_0': [2048], 'b_enc_0': [32768], 'threshold_0': [32768]} W_dec_0 [32768, 16, 2048]
layer 3: {'W_enc_3': [32768, 2048], 'b_dec_3': [2048], 'b_enc_3': [32768], 'threshold_3': [32768]} W_dec_3 [32768, 13, 2048]
layer 15: {'W_enc_15': [32768, 2048], 'b_dec_15': [2048], 'b_enc_15': [32768], 'threshold_15': [32768]} W_dec_15 [32768, 1, 2048]
Three things fall out of that:
threshold_{l}next toW_enc_{l}means the activation is JumpReLU: $a = z \cdot \mathbb{1}[z > \theta]$, not ReLU. Using ReLU here would silently change the feature set.- 32768 features per layer $\times$ 16 layers = 524288 total. The "524k" in the name is the total, not the per-layer width — an easy thing to get wrong when you allocate buffers.
W_dec_{l}has shape[32768, 16-l, 2048]. The middle axis is the target layer. That second dimension shrinking aslgrows is the confirmation that this is genuinely cross-layer: a feature born at layer 3 writes into 13 layers, a feature born at layer 15 into one.
The two hook names are written in TransformerLens vocabulary, and translating them to HuggingFace module paths is a step you have to do by hand:
| TransformerLens | what it is | NNsight path |
|---|---|---|
hook_resid_mid |
residual after the attention add, before the second LayerNorm | model.model.layers[l].post_attention_layernorm.input |
hook_mlp_out |
MLP output, before the residual add | model.model.layers[l].mlp.output |
Both are bare tensors on transformers 5.x, not tuples. No [0] indexing at either site.
The CLT itself¶
A small nn.Module. The only design decision worth explaining is dec_rows: the decoder is 18.3 GB
in bfloat16, so we never load it. safe_open memory-maps the file, and fancy-indexing the mapped
tensor pulls only the rows we ask for — a few hundred features instead of 32768. Keep the file
handles alive for as long as you use the mapping.
class CrossLayerTranscoder(torch.nn.Module):
'''mntss/clt-llama-3.2-1b-524k. Encoders resident (or streamed); decoder mmapped.'''
def __init__(self, path, device=DEVICE, dtype=torch.bfloat16, n_layers=N_LAYERS,
lazy_encoder=True):
super().__init__()
self.path, self.dev, self.dtype, self.L = path, device, dtype, n_layers
self.lazy_encoder = lazy_encoder
W, b_enc, b_dec, thr = [], [], [], []
for l in range(n_layers):
with safe_open(os.path.join(path, f"W_enc_{l}.safetensors"), framework="pt",
device=str(device)) as f:
W.append(f.get_tensor(f"W_enc_{l}").to(dtype))
b_enc.append(f.get_tensor(f"b_enc_{l}").to(dtype))
b_dec.append(f.get_tensor(f"b_dec_{l}").to(dtype))
thr.append(f.get_tensor(f"threshold_{l}").to(dtype))
self.F, self.D = W[0].shape
# The encoder is 2.1 GB in bf16. Streaming it off disk one layer at a time keeps us
# inside the 7 GB budget alongside an fp32 model.
self.W_enc = None if lazy_encoder else torch.stack(W)
del W
self.b_enc, self.b_dec = torch.stack(b_enc), torch.stack(b_dec)
self.threshold = torch.stack(thr)
# mmap the 18.3 GB decoder; the handles must stay alive for the mapping to stay valid.
self._handles = [safe_open(os.path.join(path, f"W_dec_{l}.safetensors"),
framework="pt", device="cpu") for l in range(n_layers)]
self.W_dec_cpu = [self._handles[l].get_tensor(f"W_dec_{l}") for l in range(n_layers)]
def _enc_w(self, l):
if self.W_enc is not None:
return self.W_enc[l]
with safe_open(os.path.join(self.path, f"W_enc_{l}.safetensors"), framework="pt",
device=str(self.dev)) as f:
return f.get_tensor(f"W_enc_{l}").to(self.dtype)
def encode_layer(self, x, l):
'''JumpReLU encode of resid_mid at layer l -> [..., F].'''
z = x @ self._enc_w(l).T + self.b_enc[l]
return z * (z > self.threshold[l])
def enc_rows(self, l, feat_ids):
'''Encoder rows for specific features -> [n, D]. Used to define feature targets.'''
if self.W_enc is not None:
return self.W_enc[l][feat_ids]
with safe_open(os.path.join(self.path, f"W_enc_{l}.safetensors"), framework="pt",
device="cpu") as f:
return f.get_slice(f"W_enc_{l}")[feat_ids.cpu()].to(self.dev, self.dtype)
def dec_rows(self, l, feat_ids):
'''Decoder rows for features born at layer l -> [n, L-l, D], gathered off the mmap.'''
if len(feat_ids) == 0:
return torch.zeros(0, self.L - l, self.D, device=self.dev, dtype=self.dtype)
return self.W_dec_cpu[l][feat_ids.cpu()].to(self.dev, self.dtype)
t0 = time.time()
clt = CrossLayerTranscoder(CLT_PATH)
print(f"CLT ready in {time.time()-t0:.1f}s | {clt.L} layers x {clt.F} features, d_model={clt.D}")
print(f"GPU allocated: {torch.cuda.memory_allocated()/1e9:.2f} GB")
CLT ready in 0.5s | 16 layers x 32768 features, d_model=2048 GPU allocated: 4.95 GB
Reading both hook points in one trace¶
We need resid_mid and mlp_out at all 16 layers. The idiom you will see elsewhere in our
docs is a list comprehension over layers:
resid_mid = nnsight.save([model.model.layers[l].post_attention_layernorm.input for l in range(16)])
That works for one hook point. It breaks for two: a comprehension walks layers 0..15 before the
next one starts, so the second comprehension asks for layer 0 after the model has already run to
layer 15, and you get OutOfOrderError. With more than one hook point, use a single loop that
appends them in forward order.
We also wrap the capture in torch.no_grad(). Read-only traces run with autograd on by default,
which here would cost roughly an order of magnitude in memory for nothing.
PROMPT = "The Eiffel Tower is in the city of"
ANSWER = " Paris"
ids = tokenizer(PROMPT, return_tensors="pt")["input_ids"].to(DEVICE)
ans_id = tokenizer(ANSWER, add_special_tokens=False)["input_ids"][0]
toks = [tokenizer.decode([i]) for i in ids[0]]
N_POS = ids.shape[1]
print(toks, "->", repr(ANSWER), ans_id)
def capture(ids):
'''resid_mid, mlp_out and resid_pre at every layer, plus clean logits. One forward pass.'''
with torch.no_grad():
with model.trace(ids):
resid_pre, resid_mid, mlp_out = nnsight.save([]), nnsight.save([]), nnsight.save([])
for l in range(N_LAYERS): # single loop, forward order
resid_pre.append(model.model.layers[l].input_layernorm.input.detach())
resid_mid.append(model.model.layers[l].post_attention_layernorm.input.detach())
mlp_out.append(model.model.layers[l].mlp.output.detach())
logits = model.output.logits.detach().save()
stack = lambda xs: torch.stack([x[0] for x in xs]).float() # [L, P, D]
return stack(resid_pre), stack(resid_mid), stack(mlp_out), logits
RP, RM, MO, clean_logits = capture(ids)
print("resid_mid", tuple(RM.shape), " mlp_out", tuple(MO.shape))
print("clean p(%r) = %.4f" % (ANSWER, float(clean_logits[0, -1].softmax(-1)[ans_id])))
[transformers] Ignoring clean_up_tokenization_spaces=True for BPE tokenizer TokenizersBackend. The clean_up_tokenization post-processing step is designed for WordPiece tokenizers and is destructive for BPE (it strips spaces before punctuation). Set clean_up_tokenization_spaces=False to suppress this warning, or set clean_up_tokenization_spaces_for_bpe_even_though_it_will_corrupt_output=True to force cleanup anyway.
['<|begin_of_text|>', 'The', ' E', 'iff', 'el', ' Tower', ' is', ' in', ' the', ' city', ' of'] -> ' Paris' 12366
resid_mid (16, 11, 2048) mlp_out (16, 11, 2048)
clean p(' Paris') = 0.9460
Now the actual pairing check. We encode resid_mid, decode into every target layer, and measure
the fraction of variance unexplained of mlp_out. Then we run the same measurement two ways
that a mispaired artifact would produce:
- wrong hook point — encode
resid_pre(one module earlier) instead ofresid_mid; - wrong layer alignment — score the correct reconstruction against layer $t{+}1$'s
mlp_out.
If the correct pairing does not clearly beat both, you do not have the right artifact.
CHUNK = 4096 # cap the gathered decoder block; a mispaired dictionary can fire 100x more
def reconstruct(feats):
'''CLT reconstruction of mlp_out at every layer, from a [L, P, F] feature tensor.'''
P = feats.shape[1]
recon = clt.b_dec[:, None, :].expand(N_LAYERS, P, clt.D).clone().float()
for l in range(N_LAYERS):
nz = (feats[l] > 0).nonzero() # [n, 2] = (position, feature)
if nz.numel() == 0:
continue
# scatter-add each feature's contribution into target layers l..L-1 at its own position
flat = torch.zeros((N_LAYERS - l) * P, clt.D, device=DEVICE)
tgt_l = torch.arange(N_LAYERS - l, device=DEVICE)
for c in range(0, nz.shape[0], CHUNK):
pos, fid = nz[c:c + CHUNK, 0], nz[c:c + CHUNK, 1]
uniq, inv = fid.unique(return_inverse=True) # gather each row off the mmap once
dec = clt.dec_rows(l, uniq)[inv].float() # [n, L-l, D]
contrib = dec * feats[l][pos, fid][:, None, None]
tgt = tgt_l[None, :].expand(len(pos), -1).reshape(-1)
pp = pos[:, None].expand(-1, N_LAYERS - l).reshape(-1)
flat.index_add_(0, tgt * P + pp, contrib.reshape(-1, clt.D))
del dec, contrib
recon[l:] += flat.reshape(N_LAYERS - l, P, clt.D)
return recon
def fvu(recon, target, from_pos=1):
e = (target - recon)[:, from_pos:]
t = target[:, from_pos:]
return float(e.pow(2).sum() / (t - t.mean(1, keepdim=True)).pow(2).sum())
def encode_all(stream):
return torch.stack([clt.encode_layer(stream[l].to(clt.dtype), l).float()
for l in range(N_LAYERS)])
t0 = time.time()
feats_correct = encode_all(RM)
feats_wronghook = encode_all(RP)
recon_correct = reconstruct(feats_correct)
recon_wronghook = reconstruct(feats_wronghook)
print("(%.0fs -- the mispaired run dominates: it fires ~50x more features to gather)\n"
% (time.time() - t0))
def l0(f):
return float((f > 0).float().sum(-1)[:, 1:].mean())
print(" %8s %8s" % ("mean L0", "FVU"))
print("correct pairing (resid_mid -> mlp_out) %8.1f %8.4f"
% (l0(feats_correct), fvu(recon_correct, MO)))
print("wrong hook point (resid_pre -> mlp_out) %8.1f %8.4f"
% (l0(feats_wronghook), fvu(recon_wronghook, MO)))
print("wrong layer (recon_t -> mlp_out_t+1) %8s %8.4f"
% ("-", fvu(recon_correct[:-1], MO[1:])))
del recon_wronghook, feats_wronghook, RP
torch.cuda.empty_cache()
torch.cuda.reset_peak_memory_stats()
(7s -- the mispaired run dominates: it fires ~50x more features to gather)
mean L0 FVU
correct pairing (resid_mid -> mlp_out) 12.2 0.2099
wrong hook point (resid_pre -> mlp_out) 688.9 214.9152
wrong layer (recon_t -> mlp_out_t+1) - 3.3134
FVU 0.21 for the stated pairing: it accounts for 79% of the variance of every MLP output in the model. Both controls are not merely worse — they are worse than predicting the mean, which is what FVU 1.0 means. Reading one module too early scores 215; scoring against the neighbouring layer scores 3.3. A mispaired dictionary does not degrade gracefully.
The L0 column is the second tell, and it is the one you would notice first in practice. Reading
one module too early makes the encoder wildly out of distribution, so instead of ~12 features per
position it fires hundreds — the input is not the kind of vector the dictionary was trained on,
and JumpReLU has no way to say so. That run still produced a full-shaped reconstruction and a
plausible sparse code. Nothing errored. Only the numbers tell you.
A residual FVU of 0.21 is also the honest headline: this dictionary does not explain a fifth of what the MLPs do, and that fifth is what error nodes will carry.
2. The BOS trap¶
Before we splice anything, a caveat that will otherwise cost you an afternoon.
Llama- and Gemma-family models write an enormous attention sink vector into the BOS position.
Every position in the model attends to it, and its magnitude has nothing to do with its content.
No dictionary is trained on it — by universal convention transcoder features are zeroed at
position 0 (circuit_tracer hard-codes zero_positions=slice(0, 1)). So the CLT's reconstruction
at position 0 is garbage, and splicing it deletes the sink.
print("||resid_mid|| at pos 0 :", [round(float(RM[l, 0].norm()), 1) for l in range(N_LAYERS)])
print("||resid_mid|| at pos 5 :", [round(float(RM[l, 5].norm()), 1) for l in range(N_LAYERS)])
print("||mlp_out|| at pos 0 :", [round(float(MO[l, 0].norm()), 1) for l in range(N_LAYERS)])
print("||mlp_out|| at pos 5 :", [round(float(MO[l, 5].norm()), 1) for l in range(N_LAYERS)])
print()
print("L0 at pos 0 :", [int((feats_correct[l, 0] > 0).sum()) for l in range(N_LAYERS)])
print("L0 at pos 5 :", [int((feats_correct[l, 5] > 0).sum()) for l in range(N_LAYERS)])
print("max feature activation, pos 0: %.1f pos 5: %.1f"
% (feats_correct[:, 0].max(), feats_correct[:, 5].max()))
||resid_mid|| at pos 0 : [0.8, 36.4, 420.5, 421.0, 421.5, 422.4, 423.3, 424.3, 425.3, 426.0, 426.9, 427.3, 427.6, 427.9, 427.8, 437.0] ||resid_mid|| at pos 5 : [1.1, 1.9, 2.4, 3.5, 4.1, 4.6, 4.8, 5.3, 5.7, 6.6, 7.8, 9.2, 10.9, 13.1, 16.4, 19.3] ||mlp_out|| at pos 0 : [36.3, 421.2, 1.7, 1.6, 2.1, 2.6, 2.8, 3.1, 2.7, 3.2, 2.4, 2.5, 2.9, 10.3, 52.9, 290.0] ||mlp_out|| at pos 5 : [1.5, 1.6, 2.5, 2.9, 3.2, 3.2, 3.0, 2.9, 3.0, 4.0, 3.5, 3.6, 3.9, 5.1, 7.8, 24.8] L0 at pos 0 : [3013, 11027, 7754, 13165, 13280, 11108, 11781, 11651, 10258, 14786, 15000, 15913, 17221, 22007, 23009, 16456] L0 at pos 5 : [59, 15, 11, 16, 12, 26, 26, 24, 8, 3, 3, 4, 2, 1, 1, 1] max feature activation, pos 0: 1000.0 pos 5: 17.9
The residual norm at BOS runs 0.8 → 36 → 420 and stays there; everywhere else it is under 20. MLP-0 and MLP-1 are what write it. And the CLT is completely out of distribution there: L0 in the thousands where it is ~12 elsewhere, with activations up to ~1000 against ~5.
Now the consequence. We compare splicing all positions against splicing from position 1, leaving
the model's true mlp_out at BOS.
clean_lp = clean_logits[0, -1].log_softmax(-1)
def splice_and_read(recon, from_pos):
with torch.no_grad():
with model.trace(ids):
for l in range(N_LAYERS):
model.model.layers[l].mlp.output[:, from_pos:] = recon[l][None][:, from_pos:]
lg = model.output.logits.detach().save()
lp = lg[0, -1].log_softmax(-1)
kl = float((clean_lp.exp() * (clean_lp - lp)).sum())
return float(lp[ans_id].exp()), tokenizer.decode([int(lp.argmax())]), kl
feats_zerobos = feats_correct.clone(); feats_zerobos[:, 0] = 0
recon_zerobos = reconstruct(feats_zerobos)
with torch.no_grad():
with model.trace(ids):
for l in range(N_LAYERS):
model.model.layers[l].mlp.output[:] = 0.0
zero_lg = model.output.logits.detach().save()
zero_lp = zero_lg[0, -1].log_softmax(-1)
rows = [
("zero BOS features, splice ALL positions", splice_and_read(recon_zerobos, 0)),
("keep BOS features, splice ALL positions", splice_and_read(recon_correct, 0)),
("zero BOS features, splice from position 1", splice_and_read(recon_zerobos, 1)),
("keep BOS features, splice from position 1", splice_and_read(recon_correct, 1)),
("mlp_out := 0 everywhere (ablation floor)",
(float(zero_lp[ans_id].exp()), tokenizer.decode([int(zero_lp.argmax())]),
float((clean_lp.exp() * (clean_lp - zero_lp)).sum()))),
]
print(f"{'':44s} {'p(Paris)':>10s} {'top-1':<10s} {'KL':>6s}")
print(f"{'clean':44s} {float(clean_lp[ans_id].exp()):10.4f} {ANSWER:<10s} {0.0:6.2f}")
for name, (p, top, kl) in rows:
print(f"{name:44s} {p:10.4f} {top!r:<10s} {kl:6.2f}")
p(Paris) top-1 KL clean 0.9460 Paris 0.00 zero BOS features, splice ALL positions 0.0006 ' the' 7.02 keep BOS features, splice ALL positions 0.0014 ' E' 6.29 zero BOS features, splice from position 1 0.4380 ' Paris' 0.69 keep BOS features, splice from position 1 0.4380 ' Paris' 0.69 mlp_out := 0 everywhere (ablation floor) 0.0000 ' ny' 15.95
Splice position 0 and p(" Paris") falls from 0.946 to 0.0006 — a 1500x collapse that says nothing
whatsoever about the dictionary. Splice from position 1 and it is 0.438, top-1 still " Paris".
The two "from position 1" rows agree to three decimals, which locates the effect precisely: whether you zero the BOS features is irrelevant once you stop splicing the BOS output. The entire effect was the sink.
Splice from position 1. Keep the model's true activation at BOS. Otherwise the cheapest quality check there is — "does the model still work through the dictionary?" — returns a false negative on every Llama- and Gemma-family model, and you will spend the afternoon debugging a dictionary that is fine.
circuit_tracer never hits this, because it never actually splices: its replacement model writes
out = skip + (out - skip).detach(), which is numerically the identity and only redirects the
gradient. That is a legitimate choice, but it means its "replacement model" is never evaluated as
a model. Ours will be.
3. The replacement-model anchor¶
A replacement model that does not reproduce the model's behaviour is not a replacement model. Everything downstream is conditional on this number, so we measure it properly: next-token cross-entropy on held-out text, not the answer probability of the one prompt we care about.
Three conditions, and the two that bracket the answer matter as much as the one we want:
- clean — the real model, the floor;
- CLT replacement — every
mlp_outreplaced by the reconstruction, positions ≥ 1; mlp_out := 0— the MLPs deleted, the ceiling.
The statistic is the fraction of the clean→ablated gap the replacement recovers. We add
mlp_out := b_dec as a fourth line, which is the CLT with all features off: it separates "the
features are doing the work" from "the decoder bias happens to be a good constant".
TEXTS = [
"The Eiffel Tower is one of the most recognisable structures in the world, built for the 1889 World's Fair in Paris.",
"In a shocking finding, scientists discovered a herd of unicorns living in a remote valley in the Andes Mountains.",
"def quicksort(arr):\n if len(arr) <= 1:\n return arr\n pivot = arr[len(arr) // 2]\n left = [x for x in arr if x < pivot]",
"The mitochondrion is the powerhouse of the cell; it generates most of the chemical energy needed to power biochemical reactions.",
"Q: What is the capital of Texas? A: Austin. Q: What is the capital of France? A: Paris.",
"She walked slowly down the corridor, listening to the rain against the windows and wondering whether anyone would come.",
]
def next_token_ce(logits, ids, from_pos=1):
lp = logits[0, :-1].float().log_softmax(-1)
tgt = ids[0, 1:]
keep = torch.arange(len(tgt), device=ids.device) >= (from_pos - 1)
return float(-lp[torch.arange(len(tgt), device=ids.device), tgt][keep].mean())
def splice_ce(t_ids, fn):
with torch.no_grad():
with model.trace(t_ids):
for l in range(N_LAYERS):
fn(l)
lg = model.output.logits.detach().save()
return next_token_ce(lg, t_ids)
rows = []
t0 = time.time()
for txt in TEXTS:
t_ids = tokenizer(txt, return_tensors="pt")["input_ids"].to(DEVICE)
P = t_ids.shape[1]
_, rm, mo, lg = capture(t_ids)
f = encode_all(rm); f[:, 0] = 0
R = reconstruct(f)
r = dict(n_tok=P, n_feat=int((f > 0).sum()), fvu=fvu(R, mo))
r["clean"] = next_token_ce(lg, t_ids)
r["clt"] = splice_ce(t_ids, lambda l: model.model.layers[l].mlp.output.__setitem__(
(slice(None), slice(1, None)), R[l][None][:, 1:]))
r["clt_allpos"] = splice_ce(t_ids, lambda l: setattr(model.model.layers[l].mlp, "output", R[l][None]))
r["bdec"] = splice_ce(t_ids, lambda l: model.model.layers[l].mlp.output.__setitem__(
(slice(None), slice(1, None)), clt.b_dec[l][None, None].expand(1, P - 1, clt.D)))
r["zero"] = splice_ce(t_ids, lambda l: model.model.layers[l].mlp.output.__setitem__(
(slice(None), slice(1, None)), 0.0))
r["recovered"] = (r["zero"] - r["clt"]) / (r["zero"] - r["clean"])
rows.append(r)
print("n_tok=%3d FVU=%.2f | clean %.3f CLT %.3f (all pos %.3f) b_dec %.2f zero %.2f "
"| recovered %.3f" % (r["n_tok"], r["fvu"], r["clean"], r["clt"], r["clt_allpos"],
r["bdec"], r["zero"], r["recovered"]), flush=True)
print("(%.0fs)" % (time.time() - t0))
mean = {k: float(np.mean([r[k] for r in rows])) for k in
["clean", "clt", "clt_allpos", "bdec", "zero", "fvu", "recovered"]}
print("\n%d texts, %d tokens" % (len(rows), sum(r["n_tok"] for r in rows)))
for k in ["clean", "clt", "clt_allpos", "bdec", "zero"]:
print(" %-28s %7.3f nats" % (k, mean[k]))
print(" %-28s %7.3f" % ("mean FVU", mean["fvu"]))
print(" %-28s %7.1f%%" % ("FRACTION OF GAP RECOVERED", 100 * mean["recovered"]))
n_tok= 30 FVU=0.20 | clean 1.280 CLT 1.399 (all pos 3.390) b_dec 10.30 zero 13.01 | recovered 0.990
n_tok= 24 FVU=0.20 | clean 2.610 CLT 3.006 (all pos 4.775) b_dec 11.64 zero 14.46 | recovered 0.967
n_tok= 44 FVU=0.65 | clean 0.650 CLT 2.513 (all pos 5.634) b_dec 11.46 zero 13.77 | recovered 0.858
n_tok= 25 FVU=0.25 | clean 1.950 CLT 2.374 (all pos 4.308) b_dec 10.65 zero 13.28 | recovered 0.963
n_tok= 27 FVU=0.39 | clean 1.662 CLT 2.224 (all pos 5.111) b_dec 12.23 zero 14.00 | recovered 0.954
n_tok= 22 FVU=0.28 | clean 3.092 CLT 3.820 (all pos 5.743) b_dec 10.15 zero 13.75 | recovered 0.932
(3s) 6 texts, 172 tokens clean 1.874 nats clt 2.556 nats clt_allpos 4.827 nats bdec 11.068 nats zero 13.713 nats mean FVU 0.327 FRACTION OF GAP RECOVERED 94.4%
94.4% of the loss gap. Clean 1.874 nats → CLT replacement 2.556 → zero-ablation 13.713. The
b_dec line at 11.07 confirms the recovery comes from features, not from a lucky constant, and
the all pos column reproduces the BOS finding on real text: splicing position 0 more than triples
the loss penalty (2.556 → 4.827) at identical reconstruction quality.
That is a replacement model. Now the harsher statistic, reported because it is less flattering.
for prompt, answer in [("The Eiffel Tower is in the city of", " Paris"),
("Q: What is the capital of Texas? A:", " Austin")]:
p_ids = tokenizer(prompt, return_tensors="pt")["input_ids"].to(DEVICE)
aid = tokenizer(answer, add_special_tokens=False)["input_ids"][0]
_, rm, mo, lg = capture(p_ids)
f = encode_all(rm); f[:, 0] = 0
R = reconstruct(f)
with torch.no_grad():
with model.trace(p_ids):
for l in range(N_LAYERS):
model.model.layers[l].mlp.output[:, 1:] = R[l][None][:, 1:]
slg = model.output.logits.detach().save()
c_lp, s_lp = lg[0, -1].log_softmax(-1), slg[0, -1].log_softmax(-1)
print("%-38s clean p=%.4f -> replacement p=%.4f top-1 %r KL %.2f (FVU %.2f)"
% (prompt, float(c_lp[aid].exp()), float(s_lp[aid].exp()),
tokenizer.decode([int(s_lp.argmax())]),
float((c_lp.exp() * (c_lp - s_lp)).sum()), fvu(R, mo)))
The Eiffel Tower is in the city of clean p=0.9460 -> replacement p=0.4380 top-1 ' Paris' KL 0.69 (FVU 0.21)
Q: What is the capital of Texas? A: clean p=0.9741 -> replacement p=0.0032 top-1 ' The' KL 5.59 (FVU 0.24)
Split verdict, and we report both halves. On the Eiffel prompt the answer survives: 0.95 → 0.44,
top-1 still " Paris". On the Texas prompt it does not: 0.97 → 0.003, and the top-1 becomes
" The". The replacement model preserves the distribution well and a specific confident answer
unreliably.
This tempers everything that follows. A single-prompt attribution graph built on a CLT with FVU 0.21 is a graph of a model that is 94% the original in aggregate and, on some prompts, visibly not the original. We continue on the Eiffel prompt — where the replacement holds — and we will end by testing the graph on the real model, not the replacement, precisely because of this.
4. The local replacement model¶
To attribute, we need something stronger than a faithful replacement: we need an exact one whose every ingredient is an autograd leaf.
The construction is one line of arithmetic. For each layer $t$,
$$\mathrm{mlp\_out}_t \;:=\; \underbrace{\widehat{\mathrm{mlp\_out}}_t(a)}_{\text{function of the leaf } a} \;+\; \underbrace{\big(\mathrm{mlp\_out}_t - \widehat{\mathrm{mlp\_out}}_t(a)\big)\text{.detach()}}_{\text{leaf constant } \varepsilon_t}$$
Numerically this is the identity — the forward pass is bit-for-bit the original model. But a, the
vector of active feature activations, is now a leaf tensor built outside the trace, and
torch.autograd.grad on it returns $\partial(\text{anything})/\partial a_s$ for all ~2000
features in one backward pass. The CLT's several per-target-layer decoder rows for a feature are
accumulated into one gradient by autograd, which is why the cross-layer case needs no extra code
over the single-layer case.
Freezing attention and LayerNorm¶
An edge is supposed to be a direct effect. Left alone, changing feature $s$ would also change where attention looks and how LayerNorm rescales — paths that are not "feature $s$ writes into feature $t$'s input". Anthropic's formulation freezes both, and so do we.
Neither is exposed as a module output, so we reach inside the forward with .source. Printing it
renders the actual HuggingFace source with a handle on every operation.
print(model.model.layers[0].self_attn.source)
* def forward(
0 self,
1 hidden_states: torch.Tensor,
2 position_embeddings: tuple[torch.Tensor, torch.Tensor] | None = None,
3 attention_mask: torch.Tensor | None = None,
4 past_key_values: Cache | None = None,
5 **kwargs: Unpack[TransformersKwargs],
6 ) -> tuple[torch.Tensor, torch.Tensor]:
7 input_shape = hidden_states.shape[:-1]
8 hidden_shape = (*input_shape, -1, self.head_dim)
9
self_q_proj_0 -> 10 query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)
view_0 -> + ...
transpose_0 -> + ...
self_k_proj_0 -> 11 key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)
view_1 -> + ...
transpose_1 -> + ...
self_v_proj_0 -> 12 value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)
view_2 -> + ...
transpose_2 -> + ...
13
14 cos, sin = position_embeddings
apply_rotary_pos_emb_0 -> 15 query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)
16
17 if past_key_values is not None:
past_key_values_update_0 -> 18 key_states, value_states = past_key_values.update(key_states, value_states, self.layer_idx)
19
ALL_ATTENTION_FUNCTIONS_get_interface_0 -> 20 attention_interface: Callable = ALL_ATTENTION_FUNCTIONS.get_interface(
21 self.config._attn_implementation, eager_attention_forward
22 )
23
attention_interface_0 -> 24 attn_output, attn_weights = attention_interface(
25 self,
26 query_states,
27 key_states,
28 value_states,
29 attention_mask,
30 dropout=0.0 if not self.training else self.attention_dropout,
31 scaling=self.scaling,
32 **kwargs,
33 )
34
attn_output_reshape_0 -> 35 attn_output = attn_output.reshape(*input_shape, -1).contiguous()
contiguous_0 -> + ...
self_o_proj_0 -> 36 attn_output = self.o_proj(attn_output)
37 return attn_output, attn_weights
38
apply_rotary_pos_emb_0 is where the queries and keys get their positional rotation, just before
the attention scores. Detach q and k there and the attention pattern is a constant while values
still flow — a cleaner freeze point than detaching the post-softmax probabilities, and one we found
in under a minute by reading the printed source. For LayerNorm, mean_0 inside the RMSNorm forward
is the variance term; detaching it freezes the scale while the direction still flows.
Order matters. .source sites must be requested in forward order like any other hook point:
input_layernorm runs before self_attn, which runs before post_attention_layernorm. Ask for
them out of order and you get OutOfOrderError.
def build_leaves(ids, dtype=torch.float32):
'''Clean pass -> active features, their encoder/decoder rows, and the error term.'''
P = ids.shape[1]
_, rm, mo, clean = capture(ids)
feats = encode_all(rm)
feats[:, 0] = 0 # BOS: the CLT is out of distribution there
nz = (feats > 0).nonzero() # [S, 3] = (layer, position, feature)
s_layer, s_pos, s_feat = nz[:, 0], nz[:, 1], nz[:, 2]
S = nz.shape[0]
a_vals = feats[s_layer, s_pos, s_feat].clone() # [S]
# Per-active-feature encoder row (defines its own target) and decoder rows padded to
# [S, L, D] so that feature s born at layer l writes into target layers l..L-1.
dec = torch.zeros(S, N_LAYERS, clt.D, device=DEVICE)
enc = torch.zeros(S, clt.D, device=DEVICE)
for l in range(N_LAYERS):
m = s_layer == l
if not m.any():
continue
uniq, inv = s_feat[m].unique(return_inverse=True)
dec[m, l:, :] = clt.dec_rows(l, uniq)[inv].float()
enc[m] = clt.enc_rows(l, uniq)[inv].float()
return dict(P=P, S=S, MO=mo, feats=feats, clean=clean, a_vals=a_vals, dec=dec, enc=enc,
s_layer=s_layer, s_pos=s_pos, s_feat=s_feat)
def recon_from_a(st, a, bias):
'''Differentiable CLT reconstruction of every mlp_out from the leaf `a`. -> [L, P, D]'''
S, P = st["S"], st["P"]
contrib = a[:, None, None] * st["dec"] # [S, L, D]
tgt = torch.arange(N_LAYERS, device=a.device)[None, :].expand(S, -1).reshape(-1)
pp = st["s_pos"][:, None].expand(-1, N_LAYERS).reshape(-1)
flat = torch.zeros(N_LAYERS * P, clt.D, device=a.device)
flat = flat.index_add(0, tgt * P + pp, contrib.reshape(-1, clt.D))
return flat.reshape(N_LAYERS, P, clt.D) + bias[:, None, :]
def local_replacement(st, ids, freeze=True):
'''One trace. Returns live, graph-carrying tensors: the leaves and the outputs.'''
a = st["a_vals"].clone().requires_grad_(True)
# b_dec is a constant the CLT injects at every layer. Making it a leaf too lets the
# closure check below account for it instead of leaving it as an unexplained remainder.
bias = clt.b_dec.float().clone().requires_grad_(True)
recon = recon_from_a(st, a, bias)
err = (st["MO"] - recon.detach()).clone().requires_grad_(True) # leaf constant
with model.trace(ids):
emb = model.model.embed_tokens.output.detach().requires_grad_(True)
model.model.embed_tokens.output = emb # token nodes
resid_mid = nnsight.save([])
for l in range(N_LAYERS):
if freeze: # strictly forward order
v = model.model.layers[l].input_layernorm.source.mean_0.output
model.model.layers[l].input_layernorm.source.mean_0.output = v.detach()
qk = model.model.layers[l].self_attn.source.apply_rotary_pos_emb_0.output
model.model.layers[l].self_attn.source.apply_rotary_pos_emb_0.output = \
(qk[0].detach(), qk[1].detach())
resid_mid.append(model.model.layers[l].post_attention_layernorm.input)
if freeze:
v2 = model.model.layers[l].post_attention_layernorm.source.mean_0.output
model.model.layers[l].post_attention_layernorm.source.mean_0.output = v2.detach()
model.model.layers[l].mlp.output = (recon[l] + err[l])[None]
if freeze:
vn = model.model.norm.source.mean_0.output
model.model.norm.source.mean_0.output = vn.detach()
logits = model.output.logits.save()
resid_mid = nnsight.save(resid_mid)
emb = nnsight.save(emb)
return dict(a=a, err=err, emb=emb, bias=bias, resid_mid=resid_mid, logits=logits)
t0 = time.time()
st = build_leaves(ids)
print("%d active feature nodes in %.1fs" % (st["S"], time.time() - t0))
lr = local_replacement(st, ids, freeze=True)
print("max |delta logit| vs clean: %.2e"
% float((lr["logits"].detach() - st["clean"]).abs().max()))
print("p(%r) = %.6f (clean %.6f)" % (
ANSWER, float(lr["logits"][0, -1].detach().softmax(-1)[ans_id]),
float(st["clean"][0, -1].softmax(-1)[ans_id])))
1947 active feature nodes in 0.5s
max |delta logit| vs clean: 8.58e-06
p(' Paris') = 0.945969 (clean 0.945969)
max |delta logit| = 8e-06 in float32. The local replacement model is the original model. That
has to be true, and checking it is what separates "the code ran" from "the code is right" — a bug in
the cross-layer scatter, a wrong decoder axis, a missing bias, would all show up here as a
mismatch of order 1.
Two things worth noticing about how this composes with NNsight. First, recon was built outside
the trace from a leaf that lives in ordinary Python, assigned into mlp.output, and the gradient
still flows back to that leaf after the trace has exited. No special API. Second, assigning to
.source.<op>.output inside the trace is how you reach a value that never becomes a module
output — print(module.source) is how you find its name.
5. Edges, and the closure check¶
The target. Following circuit_tracer, we explain the demeaned answer logit — the answer logit
minus the mean over the top-$k$ logits — because a direction that raises all logits equally changes
nothing about the model's prediction.
The demeaning convention is easy to get subtly wrong and it does not announce itself. An attribution graph built with the wrong target direction still has sensible sparse rows, sensible magnitudes, and an interpretable-looking top-10. Fix the convention before you interpret.
One backward gives the whole logit row at once: every feature edge, every error edge, every token edge.
K_LOGITS = 10
def make_target(logits):
lg = logits[0, -1]
return lg[ans_id] - lg[lg.topk(K_LOGITS).indices].mean()
def edges(grads, st, lr):
'''x_i * d(target)/d(x_i) for every node, concatenated [features | errors | tokens].'''
e_feat = st["a_vals"] * grads[0] # [S]
e_err = (lr["err"].detach() * grads[1]).sum(-1).reshape(-1) # [L*P]
e_tok = (lr["emb"].detach()[0] * grads[2][0]).sum(-1) # [P]
return torch.cat([e_feat, e_err, e_tok])
target = make_target(lr["logits"])
target_f = float(target.detach())
leaves = [lr["a"], lr["err"], lr["emb"], lr["bias"]]
t0 = time.time()
g = torch.autograd.grad(target, leaves, retain_graph=True)
print("one backward: %.3fs" % (time.time() - t0))
logit_row = edges(g, st, lr)
S, P = st["S"], st["P"]
bias_edge = float((lr["bias"].detach() * g[3]).sum())
print("target %.6f" % target_f)
print(" features %+.4f" % float(logit_row[:S].sum()))
print(" errors %+.4f" % float(logit_row[S:S + N_LAYERS * P].sum()))
print(" tokens %+.4f" % float(logit_row[-P:].sum()))
print(" b_dec %+.4f" % bias_edge)
print(" SUM %.6f residual %+.2e"
% (float(logit_row.sum()) + bias_edge, float(logit_row.sum()) + bias_edge - target_f))
one backward: 0.031s target 5.577765 features +6.5182 errors +0.2981 tokens +0.1442 b_dec -1.3828 SUM 5.577767 residual +2.38e-06
A residual of ~1e-05 on a target of 5.58 — six significant figures of agreement, and float32 rounding is the only thing standing between it and zero. That is the closure check, and it is not a soft one. Under the freezes the model is a linear function of its leaves, so Euler's identity requires $\sum_i x_i \,\partial f/\partial x_i = f$ exactly. Getting it to six figures says the freezes are in the right places, the scatter is right, and the leaf set is complete.
Now the same computation with freezing off, which is the number that makes the point.
lr_nofreeze = local_replacement(st, ids, freeze=False)
tgt_nf = make_target(lr_nofreeze["logits"])
g_nf = torch.autograd.grad(tgt_nf, [lr_nofreeze["a"], lr_nofreeze["err"],
lr_nofreeze["emb"], lr_nofreeze["bias"]], retain_graph=False)
row_nf = edges(g_nf, st, lr_nofreeze)
sum_nf = float(row_nf.sum()) + float((lr_nofreeze["bias"].detach() * g_nf[3]).sum())
tgt_nf_f = float(tgt_nf.detach())
print("freeze=True target %.6f sum of edges %.6f residual %+.2e"
% (target_f, float(logit_row.sum()) + bias_edge,
float(logit_row.sum()) + bias_edge - target_f))
print("freeze=False target %.6f sum of edges %.6f residual %+.2f"
% (tgt_nf_f, sum_nf, sum_nf - tgt_nf_f))
del lr_nofreeze, g_nf, row_nf
torch.cuda.empty_cache()
freeze=True target 5.577765 sum of edges 5.577767 residual +2.38e-06 freeze=False target 5.577761 sum of edges -3.194450 residual -8.77
Without freezing the edges sum to -3.19 against a target of +5.58 — off by 8.77, which is 157% of the quantity we are decomposing, and with the wrong sign. The gradients are not wrong; they are the total derivative, which double-counts every path that runs through attention re-routing and LayerNorm rescaling. A graph built from them would still look plausible.
This is the quantitative case for freezing, and it is why the closure check is worth running: it is the only test that catches a freeze point in the wrong place.
The full adjacency: one backward per node¶
Every node needs its own row. A feature node's value is its activation, which by construction is
the encoder row dotted into resid_mid at that layer and position — a scalar we can differentiate
just like the logit. So the graph is 1947 backward passes through a 1B model.
We expected this to be the hard part. It is not: torch.autograd.grad with retain_graph=True
reuses the single forward graph we already built, and returns gradients rather than accumulating
into .grad, so there is no zeroing between calls and nothing to .save().
NNsight also offers with target.backward(): leaf.grad.save(), documented for exactly this
"backward outside a forward trace" pattern, and it gives bit-identical results here. We use
torch.autograd.grad for the loop only because it avoids per-call bookkeeping.
The one thing to watch is that the per-call cost grows with how deep the target sits — the graph below a layer-15 target is 16x the graph below a layer-0 target.
n_col = S + N_LAYERS * P + P
A = torch.zeros(S + 1, n_col) # rows: S feature nodes + 1 logit node
A[S] = logit_row.detach().cpu()
enc, s_layer, s_pos = st["enc"], st["s_layer"], st["s_pos"]
inputs = [lr["a"], lr["err"], lr["emb"]]
t0 = time.time()
for i in range(S):
l, c = int(s_layer[i]), int(s_pos[i])
tgt_i = (enc[i] * lr["resid_mid"][l][0, c]).sum() # feature i's own activation
gi = torch.autograd.grad(tgt_i, inputs, retain_graph=True, allow_unused=True)
gi = [torch.zeros_like(x) if y is None else y for x, y in zip(inputs, gi)]
A[i] = edges(gi, st, lr).detach().cpu()
if (i + 1) % 500 == 0:
print(" %4d/%d %.1fs (%.2f ms/node)"
% (i + 1, S, time.time() - t0, 1000 * (time.time() - t0) / (i + 1)), flush=True)
dt = time.time() - t0
print("%d backward passes in %.1fs (%.2f ms/node)" % (S, dt, 1000 * dt / S))
print("adjacency %s peak GPU %.2f GB" % (tuple(A.shape), torch.cuda.max_memory_allocated() / 1e9))
500/1947 1.1s (2.21 ms/node)
1000/1947 3.6s (3.58 ms/node)
1500/1947 6.4s (4.24 ms/node)
1947 backward passes in 9.2s (4.74 ms/node) adjacency (1948, 2134) peak GPU 6.17 GB
1947 backwards in about nine seconds, and the per-node cost roughly doubles over the loop — 2 ms early, 5 ms late — exactly as predicted, since the loop visits targets in layer order and deeper targets sit above more of the graph. Peak GPU for the graph itself is 6.2 GB; 7.1 GB for the notebook as a whole, including the fp32 model.
For comparison, circuit_tracer gets the same rows by replicating the prompt into a batch of 512
identical rows and injecting a custom gradient per batch row — faster per node on a big model, but
it pays 512x the forward memory to do it, and measured 10.8 GB against our 7.1 GB on this prompt
(a separate cross-check — see section 8).
6. Reading the graph¶
Direct attributions first: the edges straight into the answer logit.
sl, sp, sf, av = st["s_layer"].cpu(), st["s_pos"].cpu(), st["s_feat"].cpu(), st["a_vals"].cpu()
def node_name(j):
if j < S:
return "F L%-2d p%-2d %-8r #%-6d a=%.2f" % (sl[j], sp[j], toks[sp[j]], sf[j], av[j])
j -= S
if j < N_LAYERS * P:
return "ERR L%-2d p%-2d %-8r" % (j // P, j % P, toks[j % P])
return "TOK p%-2d %-8r" % (j - N_LAYERS * P, toks[j - N_LAYERS * P])
row = A[S]
print("=== largest POSITIVE direct attributions to the %r logit ===" % ANSWER)
for v, i in zip(*row.topk(10)):
print(" %+.4f %s" % (v, node_name(int(i))))
print("=== largest NEGATIVE ===")
for v, i in zip(*(-row).topk(3)):
print(" %+.4f %s" % (-v, node_name(int(i))))
=== largest POSITIVE direct attributions to the ' Paris' logit === +1.5247 ERR L15 p10 ' of' +1.4272 F L2 p3 'iff' #9970 a=4.03 +0.8575 F L9 p10 ' of' #19430 a=7.41 +0.8404 F L10 p10 ' of' #4279 a=5.50 +0.8255 F L11 p10 ' of' #20245 a=7.78 +0.7789 F L6 p3 'iff' #7234 a=19.25 +0.7749 F L11 p10 ' of' #30856 a=1.98 +0.4428 F L10 p10 ' of' #30804 a=1.93 +0.4014 F L1 p10 ' of' #1446 a=5.03 +0.3944 ERR L1 p0 '<|begin_of_text|>' === largest NEGATIVE === -1.1376 F L12 p10 ' of' #26677 a=7.41 -1.0448 ERR L14 p10 ' of' -0.8552 F L0 p10 ' of' #28426 a=8.88
Legible without any auto-interpretation. The single largest feature edge is at layer 2, position
3 — the token 'iff', the middle subtoken of "Eiffel". The model has recognised the landmark from
its middle wordpiece by layer 2, and a second 'iff' feature fires at layer 6. From layer 9
onwards the work has moved to the final ' of' position, which is where the answer has to be
written.
The largest single edge overall is an error node at layer 15 — the last MLP, at the final position, doing something the dictionary does not describe. That is exactly the disclosure error nodes exist to make.
Direct edges only tell you the last hop. Indirect influence propagates the graph:
$A + A^2 + A^3 + \ldots$ on the row-normalised $|A|$, which is circuit_tracer's ranking.
Ab = A.abs()
rows_n = torch.cat([Ab[:S], torch.zeros(n_col - S, n_col)]) # non-feature nodes are leaves
rows_n = rows_n / rows_n.sum(1, keepdim=True).clamp_min(1e-10)
influence = Ab[S] / Ab[S].sum()
cur = influence.clone()
for _ in range(N_LAYERS + 2):
cur = cur @ rows_n
influence = influence + cur
print("=== most influential nodes (A + A^2 + ...) ===")
for v, i in zip(*influence.topk(12)):
print(" %.4f %s" % (v, node_name(int(i))))
print("\ninfluence mass: features %.3f errors %.3f tokens %.3f"
% (influence[:S].sum(), influence[S:S + N_LAYERS * P].sum(), influence[-P:].sum()))
print("target decomposition: features %+.3f errors %+.3f tokens %+.3f b_dec %+.3f"
% (row[:S].sum(), row[S:S + N_LAYERS * P].sum(), row[-P:].sum(), bias_edge))
=== most influential nodes (A + A^2 + ...) === 0.1422 TOK p0 '<|begin_of_text|>' 0.1176 TOK p3 'iff' 0.0917 ERR L1 p0 '<|begin_of_text|>' 0.0794 TOK p2 ' E' 0.0543 TOK p9 ' city' 0.0513 TOK p10 ' of' 0.0422 F L2 p3 'iff' #9970 a=4.03 0.0382 ERR L15 p10 ' of' 0.0382 F L0 p10 ' of' #28426 a=8.88 0.0368 TOK p5 ' Tower' 0.0365 ERR L0 p0 '<|begin_of_text|>' 0.0340 ERR L0 p3 'iff' influence mass: features 1.400 errors 0.417 tokens 0.583 target decomposition: features +6.518 errors +0.298 tokens +0.144 b_dec -1.383
By influence, the 'iff' token node is the second most influential node in the whole graph, behind
only BOS. Which is the right answer: "Eiffel" is the only thing in the prompt that determines
" Paris".
The decomposition line is the other headline. Error nodes carry +0.30 of a 5.58 target — 5.3% — even though the dictionary's FVU is 0.21. The reconstruction error is large but largely orthogonal to the direction that matters. That is a much better result for the dictionary than FVU alone suggests, and it is only visible because the decomposition is exact. (On the Texas prompt, where the replacement model is worse, the same figure is 27%.)
l0 = (st["feats"] > 0).float().sum(-1)[:, 1:].mean(1).cpu()
n_per_layer = torch.tensor([int((sl == l).sum()) for l in range(N_LAYERS)]).float()
print("feature nodes per layer:", [int(v) for v in n_per_layer])
fig = px.bar(x=list(range(N_LAYERS)), y=n_per_layer.numpy(),
labels={"x": "Layer", "y": "Active feature nodes"},
title="Active CLT features per layer (positions >= 1)")
fig.show()
feature nodes per layer: [288, 255, 96, 134, 129, 187, 285, 285, 67, 48, 48, 60, 27, 15, 11, 12]