Ablation¶
Zero a component and measure what the answer loses. Every ablation is an in-place write at the
component's own location — an MLP's output, one head's slice of z — and a sweep over
components is one trace with one invoke per component, batched by the scheduler. A fixed
steering vocabulary that writes only at residual points cannot express either write.
Every MLP, one at a time¶
from nnsight.modeling.vllm import VLLM
model = VLLM("Qwen/Qwen3-8B", dispatch=True)
tok = model.tokenizer
PROMPT = "The capital of France is"
paris = tok.encode(" Paris")[0]
with model.trace(temperature=0.0, max_tokens=1) as tracer:
with tracer.invoke(PROMPT):
base = model.logits.save()
for i in range(36):
with tracer.invoke(PROMPT):
model.model.layers[i].mlp.output[:] = 0
logits = model.logits.save()
p = lambda l: round(l.float().softmax(-1)[0, paris].item(), 3)
print(p(base), [p(l) for l in logits])
# 0.537 [0.0, 0.48, 0.478, 0.52, 0.785, 0.508, 0.0, 0.776, 0.611, 0.749, 0.146, 0.6, 0.603, 0.246,
# 0.624, 0.709, 0.385, 0.515, 0.597, 0.747, 0.522, 0.415, 0.468, 0.686, 0.678, 0.638, 0.564,
# 0.304, 0.457, 0.351, 0.637, 0.431, 0.417, 0.241, 0.29, 0.983]
P(Paris) is 0.537 unablated. Zeroing the MLP of layer 0 or 6 takes it to zero — those two are
load-bearing for the whole forward, not for this fact — while layers 10, 13, 33 and 34 each cost
about half the probability, and removing the last MLP raises it to 0.98 (it spreads mass over
continuations like " the"). Thirty-seven requests, one trace.
Every head at one layer¶
o_proj.input is z, [pos, n_heads * head_dim]; head h is its slice.
L, d = 24, 128
with model.trace(temperature=0.0, max_tokens=1) as tracer:
for h in range(32):
with tracer.invoke(PROMPT):
model.model.layers[L].self_attn.o_proj.input[:, h * d:(h + 1) * d] = 0
logits = model.logits.save()
print([p(l) for l in logits])
# [0.563, 0.524, 0.536, 0.531, 0.538, 0.539, 0.526, 0.527, 0.518, 0.534, 0.505, 0.532, 0.543, 0.508,
# 0.523, 0.52, 0.539, 0.521, 0.52, 0.505, 0.538, 0.532, 0.52, 0.549, 0.464, 0.478, 0.54, 0.457,
# 0.536, 0.529, 0.546, 0.537]
No single head at layer 24 is necessary on the clean prompt — including head 26, which patching shows is sufficient to carry the answer from the subject to the last position. Necessity and sufficiency are different questions; the two sweeps answer both in two traces.
Over generation¶
Keep the ablation on for every step by putting it under tracer.iter; a bare write fires on the
prefill only.
with model.trace(PROMPT, temperature=0.0, max_tokens=8) as tracer:
for _ in tracer.iter[:8]:
model.model.layers[10].mlp.output[:] = 0
out = tracer.result.save()
Variations¶
- Mean ablation instead of zero: replace with the component's mean output over a reference
set, captured once with
tracer.cache()and shipped with the block. - One position: index the rows (
mlp.output[3] = 0) — on the prefill the rows are the prompt positions, on a decode step there is one. - One neuron:
mlp.act_fn.output[:, j] = 0zeroes neuronjbefore the down-projection. - An expert in a mixture-of-experts model: mask its router logit (Steering).
- Under tensor parallelism the same writes work unchanged:
o_proj.inputis gathered before the block sees it and re-sharded after (Tensor parallelism).