Steps — Logit Attribution

LogitAttribution step: decompose a logit (difference) into per-component parts.

Direct logit attribution (DLA) answers “how does each component push the answer.” It splits the model’s final residual into the pieces that wrote it (each attention head, each MLP, the embedding), applies the final normalization with the scale frozen from the full residual, and dots each piece with the unembedding direction for the answer. Freezing the norm scale keeps the decomposition linear, so the per-component contributions sum back to the model’s true logit difference; that completeness is checked and reported rather than assumed.

Per-head contributions are read off the attention output projection directly (o_proj(mask_head(z)) - o_proj(0)) so the math is independent of how a given architecture lays out that weight (e.g. gpt2’s transposed Conv1D). The arithmetic runs in float32 so the completeness check is tight even on bf16 models.

LogitAttributionResult

Per-component contributions to a target logit (difference).

Each value is the signed amount that component adds to the target at the answer position, in logit units, averaged over the batch. The named contributions plus embed_contribution and other_contribution sum to total up to completeness_error.

Attributes:

• contributions: {Node: float} mean contribution of each attention head
• (Node(layer, "self_attn", head=h)) and MLP (Node(layer, "mlp")).
• embed_contribution: Mean contribution of the embedding (pre-layer-0
• residual).
• other_contribution: Mean contribution of everything not named above:
• output-projection biases, the final-norm bias, any unembedding bias,
• and any layers excluded from layers.
• target: "logit_diff" when an incorrect answer was given, else
• "logit".
• total: Mean true target value at the answer position (the quantity being
• attributed).
• completeness_error: Maximum absolute difference, over the batch, between
• the summed contributions and the true target. Near zero confirms the
• frozen-norm decomposition reproduces the model.
• per_example: Optional {Node: list[float]} per-example contributions
• for the named components.
• metadata: Free-form metadata (answer positions, norm type, target).

__init__

def __init__(self, contributions: dict[Node, float], embed_contribution: float, other_contribution: float, target: str, total: float, completeness_error: float, per_example: dict[Node, list[float]] | None = None, metadata: dict[str, Any] = dict()) -> None:

LogitAttribution

Decompose a target logit (difference) into per-component contributions.

Runs the prompts, splits the final residual into the embedding, each attention head’s output, and each MLP’s output, applies the final norm with its scale frozen from the full residual, and projects each piece onto the answer direction W_U[correct] - W_U[incorrect] (or W_U[correct] when no incorrect answer is given). The contributions sum to the model’s true logit difference, which is verified and exposed as completeness_error.

Reads from results:

• results['prompts']: PromptBatch.

Writes to results:

• results[output_key]: LogitAttributionResult.

Args:

• model: Model backend to run.
• correct: Correct-answer token spec (token id, string, per-example
• sequence, or tensor), as accepted by the metric steps.
• incorrect: Optional incorrect-answer spec. When given, the target is the
• logit difference; when None, the target is the correct token’s
• logit.
• layers: Layer indices to attribute, or "all". Layers left out fold
• into other_contribution.
• include_mlp: Attribute each layer’s MLP output when True.
• include_embed: Attribute the embedding (pre-layer-0 residual) when True;
• otherwise it folds into other_contribution.
• positions: Optional explicit answer position(s); otherwise the last real
• token is used (via the attention mask).
• per_example: Keep per-example contributions on the result when True.
• output_key: Results key to write the result under.

Raises:

• ValueError: If a target token id is outside the vocabulary.

__init__

def __init__(self, model: ModelBackend, correct: int | str | Sequence[int] | Sequence[str] | Tensor, incorrect: int | str | Sequence[int] | Sequence[str] | Tensor | None = None, layers: list[int] | str = 'all', include_mlp: bool = True, include_embed: bool = True, positions: int | Sequence[int] | Tensor | None = None, per_example: bool = False, output_key: str = keys.LOGIT_ATTRIBUTION):