Coverage for transformer_lens/tools/analysis/direct_logit_attribution.py: 94%
69 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +0000
1"""Direct Logit Attribution (DLA).
3Direct Logit Attribution decomposes a model's output logit (or a logit
4*difference* between a correct and an incorrect token) into the additive
5contributions of upstream components — the embedding, each attention and MLP
6sublayer, or each individual attention head. Because the unembedding is linear
7and the residual stream is a sum of component outputs, the final logit is
8(up to the final LayerNorm) a sum of per-component dot products with the
9unembedding direction of the token of interest. DLA reads off those dot
10products. See the `logit lens
11<https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens>`_
12and `Interpretability in the Wild <https://arxiv.org/abs/2211.00593>`_ for the
13canonical uses.
15This module exposes a single entry point, :func:`direct_logit_attribution`,
16that wraps the lower-level ``ActivationCache`` primitives
17(:meth:`~transformer_lens.ActivationCache.ActivationCache.decompose_resid`,
18:meth:`~transformer_lens.ActivationCache.ActivationCache.accumulated_resid`,
19:meth:`~transformer_lens.ActivationCache.ActivationCache.stack_head_results`
20and :meth:`~transformer_lens.ActivationCache.ActivationCache.logit_attrs`) into
21one call. It works with any TransformerLens model that shares the cache API.
23Example::
25 from transformer_lens import TransformerBridge
26 from transformer_lens.tools.analysis import direct_logit_attribution
28 model = TransformerBridge.boot_transformers("gpt2", device="cpu")
29 model.enable_compatibility_mode()
30 result = direct_logit_attribution(
31 model,
32 "The Eiffel Tower is in the city of",
33 answer_tokens=" Paris",
34 incorrect_tokens=" London",
35 unit="component",
36 )
37 for label, value in zip(result.labels, result.attribution.squeeze()):
38 print(f"{label:>12}: {value.item():+.3f}")
39"""
41from dataclasses import dataclass
42from typing import List, Optional, Union
44import torch
45from jaxtyping import Float
47from transformer_lens.ActivationCache import ActivationCache
48from transformer_lens.tools.analysis._model_state import (
49 hybrid_block_types,
50 require_eval_mode,
51)
52from transformer_lens.tools.analysis.logit_lens import (
53 _final_norm,
54 _require_direct_readout,
55 _require_readout_identity,
56)
57from transformer_lens.utilities import Slice, SliceInput
59# Token-like inputs accepted for the correct/incorrect answers, mirroring
60# ActivationCache.logit_attrs.
61AnswerTokens = Union[
62 str,
63 int,
64 torch.Tensor,
65]
66# Former name, kept so existing imports keep working.
67TokenInput = AnswerTokens
69# Which structural unit the residual stream is decomposed into.
70Unit = str # one of: "component", "layer", "head"
72_VALID_UNITS = ("component", "layer", "head")
75@dataclass
76class DirectLogitAttribution:
77 """Result of a :func:`direct_logit_attribution` call.
79 Attributes:
80 attribution:
81 Tensor of logit (or logit-difference) attributions with shape
82 ``[component, *batch_and_pos]``. The leading axis is aligned with
83 ``labels``. When ``pos`` selects a single position (the default) the
84 position axis is dropped, leaving ``[component, batch]`` — or
85 ``[component]`` if the cache had its batch dimension removed.
86 labels:
87 Human-readable name for each component, aligned with the leading
88 axis of ``attribution`` (e.g. ``"embed"``, ``"0_attn_out"``,
89 ``"L3H7"``).
90 unit:
91 The decomposition unit used ("component", "layer", or "head").
92 """
94 attribution: Float[torch.Tensor, "component *batch_and_pos"]
95 labels: List[str]
96 unit: Unit
98 def top(self, k: int = 5) -> List[tuple]:
99 """Return the ``k`` highest-attribution ``(label, value)`` pairs.
101 Attribution is reduced to a scalar per component by meaning over any
102 remaining batch/position dimensions, so this is most meaningful when a
103 single position was selected.
104 """
105 flat = self.attribution
106 if flat.ndim > 1: 106 ↛ 108line 106 didn't jump to line 108 because the condition on line 106 was always true
107 flat = flat.flatten(start_dim=1).mean(dim=-1)
108 values, indices = torch.topk(flat, min(k, flat.shape[0]))
109 return [(self.labels[i], values[j].item()) for j, i in enumerate(indices.tolist())]
112def _residual_stack_and_labels(
113 cache: ActivationCache,
114 unit: Unit,
115 pos_slice: SliceInput,
116):
117 """Decompose the residual stream into ``unit`` components plus labels.
119 LayerNorm is intentionally *not* applied here — ``logit_attrs`` applies the
120 final-layer scaling itself, so applying it twice would double-count.
121 """
122 if unit == "component":
123 # embed (+ pos_embed) and each layer's attn_out / mlp_out.
124 return cache.decompose_resid(apply_ln=False, pos_slice=pos_slice, return_labels=True)
125 if unit == "layer":
126 # Cumulative residual stream after each sublayer — logit-lens style.
127 return cache.accumulated_resid(
128 apply_ln=False, incl_mid=True, pos_slice=pos_slice, return_labels=True
129 )
130 if unit == "head": 130 ↛ 135line 130 didn't jump to line 135 because the condition on line 130 was always true
131 # Each attention head's contribution, plus the MLP/embedding remainder.
132 return cache.stack_head_results(
133 apply_ln=False, pos_slice=pos_slice, incl_remainder=True, return_labels=True
134 )
135 raise ValueError(f"unit must be one of {_VALID_UNITS}, got {unit!r}")
138def _validate_bridge_compatibility(model) -> bool:
139 """Reject Bridges DLA can't produce correct numbers for; return whether ``model`` is one.
141 Compatibility mode alone does not guarantee the final norm is folded:
142 ``fold_ln=False`` and ``no_processing=True`` leave its affine parameters
143 active. The hybrid-arch check catches Mamba/SSM blocks before decomposition.
144 """
145 # Lazy import — keeps the module importable without dragging in the bridge.
146 from transformer_lens.model_bridge import TransformerBridge
148 if not isinstance(model, TransformerBridge): 148 ↛ 149line 148 didn't jump to line 149 because the condition on line 148 was never true
149 return False
151 # Ahead of the compatibility check: these bridges stay wrong in compatibility mode.
152 _require_direct_readout(model, operation="DLA")
154 if not getattr(model, "compatibility_mode", False):
155 raise ValueError(
156 "DLA on a TransformerBridge requires compatibility mode and an identity "
157 "final norm. Call `model.enable_compatibility_mode()` with its default "
158 "weight processing after loading the bridge, then re-run DLA."
159 )
161 final_norm = getattr(model, "ln_final", None)
162 if final_norm is not None:
163 weight = getattr(final_norm, "weight", None)
164 if isinstance(weight, torch.Tensor):
165 identity = (
166 torch.zeros_like(weight)
167 if getattr(model.cfg, "rmsnorm_uses_offset", False)
168 else torch.ones_like(weight)
169 )
170 if not torch.equal(weight, identity):
171 raise ValueError(
172 "DLA requires an identity final norm weight: the learned scale is not "
173 "folded into W_U. Load a fresh bridge and call "
174 "`enable_compatibility_mode()` with fold_ln=True and no_processing=False."
175 )
176 bias = getattr(final_norm, "bias", None)
177 if isinstance(bias, torch.Tensor) and not torch.equal(bias, torch.zeros_like(bias)):
178 raise ValueError(
179 "DLA requires a zero final norm bias: the learned bias is not folded "
180 "into b_U. Load a fresh bridge and call `enable_compatibility_mode()` "
181 "with fold_ln=True and no_processing=False."
182 )
184 # Refused up front: decompose_resid would otherwise fail with a confusing KeyError.
185 hybrid = hybrid_block_types(model)
186 if hybrid:
187 raise NotImplementedError(
188 f"DLA does not yet support hybrid architectures (found block types {hybrid}). "
189 f"Only standard attention + MLP transformers (e.g. GPT-2, LLaMA, Pythia) are "
190 f"supported; hybrid support requires extending ActivationCache.decompose_resid."
191 )
192 return True
195def direct_logit_attribution(
196 model,
197 input: Union[str, List[str], torch.Tensor, None] = None,
198 answer_tokens: Optional[AnswerTokens] = None,
199 incorrect_tokens: Optional[AnswerTokens] = None,
200 *,
201 unit: Unit = "component",
202 pos: SliceInput = -1,
203 cache: Optional[ActivationCache] = None,
204) -> DirectLogitAttribution:
205 """Compute Direct Logit Attribution for a prompt.
207 Decomposes the contribution of model components to the logit of
208 ``answer_tokens`` (or, if ``incorrect_tokens`` is given, to the logit
209 *difference* ``answer - incorrect`` along the ``W_U`` direction, which is
210 usually what you want for circuit analysis).
212 The model is run once with caching unless a precomputed ``cache`` is passed.
214 Note that DLA attributes only the part of a logit that comes from the
215 residual stream through the unembedding direction; the unembedding bias
216 ``b_U`` is a per-token constant that no component produces. So a complete
217 decomposition reconstructs ``logit[token] - b_U[token]`` rather than the raw
218 logit.
220 On a ``TransformerBridge``, compatibility mode must be enabled and the final
221 normalization's affine parameters must be folded into ``W_U`` and ``b_U``
222 (or already be identity);
223 otherwise the projection direction is wrong and DLA would return silently
224 incorrect numbers. Hybrid architectures
225 (Mamba/SSM/Mixer/LinearAttention) are not yet supported because
226 ``decompose_resid`` only understands the ``attn_out + mlp_out`` block layout;
227 both conditions raise an explicit error at call time, as does a Bridge whose
228 readout is not ``ln_final -> unembed`` on the last block's output; a prompt run
229 here is also checked against the model's own output, a precomputed ``cache``
230 only structurally.
232 Args:
233 model:
234 A ``TransformerBridge`` with ``enable_compatibility_mode()``
235 already called and an identity final normalization.
236 input:
237 Prompt to run — a string, list of strings, or token tensor. Optional
238 only when a precomputed ``cache`` is supplied.
239 answer_tokens:
240 The correct token(s) to attribute, as a string, id, or tensor. A
241 string is converted with ``model.to_single_token``.
242 incorrect_tokens:
243 Optional baseline token(s). When given, attribution is computed for
244 the ``answer - incorrect`` residual direction. Must broadcast to the
245 same shape as ``answer_tokens``.
246 unit:
247 Decomposition granularity:
249 - ``"component"`` (default): embedding + each layer's attention and
250 MLP output (via ``decompose_resid``).
251 - ``"layer"``: cumulative residual stream after each sublayer, i.e.
252 logit-lens trajectory (via ``accumulated_resid``).
253 - ``"head"``: each attention head individually, plus a remainder
254 term for everything else (via ``stack_head_results``).
255 pos:
256 Sequence position(s) to attribute. Defaults to ``-1`` (the final
257 token, the usual choice for next-token DLA). Pass ``None`` to keep
258 every position (the result then has a trailing position axis).
259 cache:
260 Optional precomputed ``ActivationCache`` to reuse instead of running
261 the model again.
263 Returns:
264 A :class:`DirectLogitAttribution` with ``attribution`` (shape
265 ``[component, *batch_and_pos]``) and aligned ``labels``.
267 Raises:
268 ValueError: If ``unit`` is invalid, ``answer_tokens`` is ``None``,
269 neither ``input`` nor ``cache`` is provided, or a
270 ``TransformerBridge`` is passed without compatibility mode enabled,
271 with an active learned final-norm weight or bias, or with a readout
272 that is not ``ln_final -> unembed``; also if the model must run while
273 any module is in training mode (a supplied ``cache`` is not checked).
274 NotImplementedError: If a ``TransformerBridge`` reports a hybrid block
275 layout (Mamba/SSM/Mixer/LinearAttention).
276 """
277 if unit not in _VALID_UNITS:
278 raise ValueError(f"unit must be one of {_VALID_UNITS}, got {unit!r}")
279 if answer_tokens is None:
280 raise ValueError("answer_tokens is required")
282 is_bridge = _validate_bridge_compatibility(model)
284 if cache is None:
285 if input is None:
286 raise ValueError("provide either `input` to run the model, or a precomputed `cache`")
287 require_eval_mode(model, operation="direct_logit_attribution()")
288 logits, cache = model.run_with_cache(input)
289 if is_bridge: 289 ↛ 295line 289 didn't jump to line 295 because the condition on line 289 was always true
290 positions = pos if isinstance(pos, Slice) else Slice(pos)
291 _require_readout_identity(
292 model, cache, logits, _final_norm(model), positions, operation="DLA"
293 )
295 residual_stack, labels = _residual_stack_and_labels(cache, unit, pos)
297 # logit_attrs applies the final LayerNorm scaling (with the same pos slice)
298 # and dots each component against the (correct - incorrect) unembed direction.
299 attribution = cache.logit_attrs(
300 residual_stack,
301 tokens=answer_tokens,
302 incorrect_tokens=incorrect_tokens,
303 pos_slice=pos,
304 has_batch_dim=cache.has_batch_dim,
305 )
307 return DirectLogitAttribution(attribution=attribution, labels=labels, unit=unit)