Direct Logit Attribution

The problem, from first principles

The final residual stream is a sum: the embedding plus the output of every attention head and MLP layer, each added in turn. The final logits are a linear function of that sum, apart from the normalization scale. That means the logit difference between a correct and an incorrect answer token splits exactly into one contribution per component. Direct logit attribution reads those contributions off: positive for components pushing toward the correct answer, negative for those pushing away. It is the cheapest circuit-finding tool because it needs a single forward pass and no interventions, though it only sees direct effects on the output, not effects mediated by later layers.

From theory to code

Implement direct_logit_attribution and attribution_total.

Constraints

• components has shape (C, d): the vectors each component wrote into the final-position residual stream. W_U has shape (d, V). correct and wrong are token ids. scale is the scalar final-normalization divisor (the final stream is sum(components) / scale).
• direct_logit_attribution(components, W_U, correct, wrong, scale) returns a length-C array: for each component, (component @ (W_U[:, correct] - W_U[:, wrong])) / scale.
• attribution_total(contributions) is the sum, which equals the logit difference of the full stream.
• Ignore any normalization gain or bias: the unembedding is exactly W_U.

Hints