Contribution decomposition¶
Direct logit attribution over the standard contributions.
layer_input + attention_output + mlp_output == layer_output on every family, so direct logit attribution reads attention_output and mlp_output from each block, pushes each term through the final norm frozen at the real stream and the unembedding, and the terms sum to the model's logit. attention_head_outputs times the output projection's weight splits a block's attention term per head. The method is from Elhage et al. (2021), A Mathematical Framework for Transformer Circuits.
import torch
from nnterp import StandardizedTransformer
model = StandardizedTransformer("openai-community/gpt2", device="cpu", attn_implementation="eager")
prompt = "The Eiffel Tower is in the city of"
paris = model.tokenizer(" Paris", add_special_tokens=False).input_ids[0]
parts = {} # bound outside the trace
with model.trace(prompt):
base = model.layers[0].layer_input.save() # the stream entering block 0
for i, layer in enumerate(model.layers): # block i's attention, then its MLP: forward order
parts["attn", i] = layer.self_attn.attention_output.save()
parts["mlp", i] = layer.mlp.mlp_output.save()
final = model.layers[-1].layer_output.save()
logit = model.lm_head.output[0, -1, paris].save() # the head's output, before any step after it
print(torch.allclose(base + sum(parts.values()), final, atol=1e-3)) # the stream is the sum of its contributions
norm, head = model.norm._module, model.lm_head._module
x = final[0, -1] # the stream the final norm reads at this position
with torch.no_grad():
shift = norm(torch.zeros_like(x)) # a LayerNorm's bias; zeros on an RMSNorm
centers = torch.allclose(norm(torch.ones_like(x)), shift) # a LayerNorm subtracts the mean first
center = (lambda t: t - t.mean(-1, keepdim=True)) if centers else (lambda t: t)
eps = getattr(norm, "variance_epsilon", getattr(norm, "eps", 0.0))
probe = torch.ones_like(x); probe[1::2] = -1 # mean 0, mean square 1
gain = (norm(probe) - shift) / probe * (1 + eps) ** 0.5 # the norm's weight, or 1 + weight on Gemma
scale = (center(x).pow(2).mean() + eps).rsqrt() # the one nonlinear step, frozen at the real stream
frozen = lambda t: gain * center(t) * scale # the final norm as a linear map at this position
row = head.weight[paris]
dla = {key: float(frozen(part[0, -1]) @ row) for key, part in parts.items()}
dla["base"] = float(frozen(base[0, -1]) @ row)
constant = float(shift @ row) + (float(head.bias[paris]) if head.bias is not None else 0.0)
print(f"{sum(dla.values()) + constant:.4f} {float(logit):.4f}")
for i in range(model.num_layers):
print(f"block {i:2d} attn {dla['attn', i]:+.3f} mlp {dla['mlp', i]:+.3f}")
# True
# -92.0873 -92.0873
# block 0 attn -0.264 mlp -0.320
# block 1 attn -0.291 mlp +0.049
# block 2 attn +0.190 mlp -0.232
# block 3 attn +0.051 mlp -0.049
# block 4 attn -0.230 mlp -0.080
# block 5 attn +0.303 mlp -0.476
# block 6 attn -0.480 mlp -1.090
# block 7 attn +0.303 mlp -0.239
# block 8 attn -0.346 mlp -1.012
# block 9 attn +1.006 mlp -1.404
# block 10 attn -1.263 mlp -17.002
# block 11 attn -36.726 mlp -24.635
The attributions plus the constant sum to the model's own Paris logit, -92.0873, to four decimals. The last two blocks dominate the logit's scale (block 11's attention -36.7 and MLP -24.6, block 10's MLP -17.0), and block 9's attention is the largest positive term, +1.006. An attribution to one logit includes what a term adds to every token at once. row = head.weight[paris] - head.weight[rome] attributes the contrast between two tokens instead, with the same code.
The norm has to be frozen because it rescales each position by that position's own statistics, so it does not distribute over the sum. Holding the scale at the real stream's value makes it a linear map (a centering, a gain, a scale), and the LayerNorm's bias plus any lm_head bias enters once, as constant. The recipe reads the centering, the gain and the bias off the norm module itself (its output on zeros is its bias, on ones says whether it centers, on a ±1 vector gives its gain), so an RMSNorm (Llama), a LayerNorm (GPT-2) and Gemma's 1 + weight gain need no separate cases. model.project_on_vocab(contribution) is the other reading, the full norm applied to one term alone. It answers what the term would predict by itself and does not sum.
Per head¶
attention_head_outputs is each head's output before concatenation and the output projection, [batch, seq, num_heads, head_dim]. The projection is linear, so a head's share of attention_output is its slice times the matching rows of the projection's weight. The projection is o_proj on every family (GPT-2's c_proj is aliased), but its weight layout differs: torch.nn.Linear stores [out, in], GPT-2's Conv1D stores [in, out].
LAYER = 9
attention = model.layers[LAYER].self_attn
with model.trace(prompt):
heads = attention.attention_head_outputs.save() # [batch, seq, num_heads, head_dim]
contribution = attention.attention_output.save() # [batch, seq, hidden_size]
proj = attention.o_proj._module
W = proj.weight.t() if isinstance(proj, torch.nn.Linear) else proj.weight # [num_heads * head_dim, hidden_size]
H, D = heads.shape[2], heads.shape[3]
per_head = torch.einsum("bshd,hdo->bsho", heads, W.reshape(H, D, -1)) # [batch, seq, num_heads, hidden_size]
print(torch.allclose(per_head.sum(2) + proj.bias, contribution, atol=1e-3)) # the heads sum to the contribution
with torch.no_grad():
head_dla = [float(frozen(per_head[0, -1, h]) @ row) for h in range(H)]
bias_dla = float(frozen(proj.bias) @ row)
print([round(v, 3) for v in head_dla])
print(round(sum(head_dla) + bias_dla, 3), round(dla["attn", LAYER], 3))
# True
# [-0.084, -0.132, -0.066, 0.791, 0.116, -0.153, -0.0, -0.072, 0.872, 0.048, 0.008, 0.064]
# 1.006 1.006
Block 9's +1.006 is two heads: head 8 (+0.872) and head 3 (+0.791), with the other heads near zero and the projection's bias taking 0.39 back off. Head 8 of block 9 is also the head that Ablation finds by removing it.
The per-head sum equals attention_output where the projection's output is what the block adds, as on GPT-2 and Llama. On a post-norm family (Gemma-2, Gemma-3, OLMo-2) a norm sits between the projection and the residual add, so the heads sum to o_proj.output and not to attention_output, and the split is of the pre-norm tensor. On Granite the contribution is the projection's output times residual_multiplier. Under grouped-query attention the head axis is still the query heads, so the slicing does not change.
On another family¶
The identity does the per-family work. attention_output and mlp_output are what each sublayer adds whether the block is sequential, parallel or post-norm, so the recipe runs unchanged there. Not every family is a plain sum. On Gemma-4 each block multiplies its sum by layer_scalar, so a term enters the final stream weighted by the product of the scalars of its block and every later block. Doge and ZAYA gate the stream per channel, and DeepSeek-V4 keeps parallel streams. Weight each term accordingly before the unembedding. The frozen norm then applies to the weighted terms unchanged.