Direct logit attribution: who wrote the answer?

Note

Generated from notebooks/applications/logit_attribution.ipynb, which you can run yourself. The outputs below are the ones it produced.

The final logits are a linear readout of the residual stream, and the residual stream is just the sum of what every component wrote into it. So the logit difference between two answers decomposes exactly into one number per component: each attention head, each MLP, and the embedding.

That decomposition is direct logit attribution (DLA). It measures only the direct path from a component to the logits, ignoring effects routed through later components.

Key questions

• Which heads push the model toward the correct answer?
• Do any heads push against it?
• Does the decomposition actually add up?

Structure

1. Set up
2. Attribute the logit difference
3. Check that it adds up
4. Plot the contributions
5. Select the top components
6. Save

Model and data. GPT-2 small in float32. It is small enough that bfloat16 rounding visibly degrades its predictions, so the notebooks load it explicitly in float32. Eight indirect-object prompt pairs from murano.tasks.ioi.

Requirements. pip install murano-interp[plot]

1. Set up

import torch

from murano import MuranoModel, Pipeline, keys
from murano.steps import LoadPrompts, LogitAttribution, Plot, Save, SelectComponents

OUTPUT_DIR = "murano_outputs/logit_attribution"

# GPT-2 is small enough that bfloat16 rounding degrades its predictions.
model = MuranoModel("gpt2", dtype=torch.float32)
print(f"{model.model_id}: {model.n_layers} layers, {model.n_heads} heads")

gpt2: 12 layers, 12 heads

This notebook uses the indirect-object-identification task from murano.tasks. Each clean prompt names two people and then has the subject give a drink, so the next token should be the other name, the indirect object. The corrupt prompt swaps who gives, which flips the answer.

clean  : When Mary and John went to the store, John gave a drink to ___  ->  " Mary"
corrupt: When Mary and John went to the store, Mary gave a drink to ___  ->  " John"

The metric throughout is the logit difference: the correct name’s logit minus the incorrect name’s. Positive means the model prefers the right answer.

from murano.tasks import ioi

task = ioi(n=8)
print("clean  :", task.clean[0], "->", repr(task.correct[0]))
print("corrupt:", task.corrupt[0], "->", repr(task.incorrect[0]))
print(f"{len(task.clean)} prompt pairs")

clean  : When Mary and John went to the store, John gave a drink to -> ' Mary'
corrupt: When Mary and John went to the store, Mary gave a drink to -> ' John'
8 prompt pairs

2. Attribute the logit difference

LogitAttribution takes the correct and incorrect answers and returns the signed contribution of every component to logit(correct) - logit(incorrect).

Positive means the component pushes toward the right name. Negative means it actively pushes toward the wrong one.

results = Pipeline(
    [
        LoadPrompts(task.clean),
        LogitAttribution(model, correct=task.correct, incorrect=task.incorrect),
    ]
).run()

attribution = results[keys.LOGIT_ATTRIBUTION]
print("target :", attribution.target)
print("total  :", f"{attribution.total:.4f}")

target : logit_diff
total  : 2.3971

3. Check that it adds up

If the decomposition is exact, the component contributions plus the embedding and the catch-all other term must reconstruct the total logit difference. completeness_error is that residual. It should be a rounding error.

The one approximation is the final layer norm, which is frozen at its observed scale so the readout stays linear.

print(f"total logit difference : {attribution.total:.6f}")
print(f"completeness error     : {attribution.completeness_error:.2e}")

assert abs(attribution.completeness_error) < 1e-3, "decomposition does not close"
print("\nthe decomposition is exact to float precision")

total logit difference : 2.397078
completeness error     : 4.51e-05

the decomposition is exact to float precision

4. Plot the contributions

Plot draws the largest contributors, signed. Blue bars push toward the correct answer, red bars push away.

Plot(output_dir=OUTPUT_DIR)(results);

Some heads have clearly negative contributions. These are not noise. GPT-2 small contains negative name movers, heads that systematically suppress the correct name. They are a real and reproducible part of the circuit.

contributions = sorted(
    attribution.contributions.items(), key=lambda kv: kv[1], reverse=True
)

print("pushes TOWARD the correct name:")
for node, value in contributions[:5]:
    print(f"  {str(node):<24} {value:+.3f}")

print("\npushes AWAY from the correct name:")
for node, value in contributions[-3:]:
    print(f"  {str(node):<24} {value:+.3f}")

pushes TOWARD the correct name:
  L10.self_attn.h0         +1.546
  L9.self_attn.h9          +1.316
  L9.self_attn.h6          +1.066
  L10.self_attn.h10        +0.492
  L11.self_attn.h3         +0.382

pushes AWAY from the correct name:
  L10.self_attn.h2         -0.458
  L11.self_attn.h1         -0.475
  L10.self_attn.h7         -1.491

5. Select the top components

SelectComponents turns an attribution into a ComponentSelection, ranked by absolute contribution so the negative heads survive the cut. Patch, PathPatch and Ablate accept that selection by key, which is how attribution feeds into a causal experiment.

results = SelectComponents(top_k=8, by="abs", modules="self_attn")(results)

selection = results[keys.SELECTION]
print(f"selected {len(selection.nodes)} heads by |contribution|:\n")
for node in selection.nodes:
    print(f"  {str(node):<24} {attribution.contributions[node]:+.3f}")

selected 8 heads by |contribution|:

  L10.self_attn.h0         +1.546
  L10.self_attn.h7         -1.491
  L9.self_attn.h9          +1.316
  L9.self_attn.h6          +1.066
  L10.self_attn.h10        +0.492
  L11.self_attn.h1         -0.475
  L10.self_attn.h2         -0.458
  L11.self_attn.h10        -0.436

The caveat that motivates the next notebook

DLA sees only the direct path to the logits. A head that matters enormously by feeding another head, and never writes to the logits itself, scores near zero here. The indirect-object circuit has exactly such heads. Finding them needs an intervention, not a decomposition.

6. Save

run_dir = Save(
    output_dir="murano_outputs", run_name="logit_attribution", model_id=model.model_id
)(results)[keys.OUTPUT_DIR]

print("saved to:", run_dir)

saved to: murano_outputs/logit_attribution

What next

• activation_patching.ipynb measures each head’s total causal effect, including the indirect paths DLA cannot see.
• circuit_discovery.ipynb traces how the heads found by patching reach the ones found here.