Coverage for transformer_lens/tools/analysis/direct_logit_attribution.py: 95%
62 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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.utilities import SliceInput
50# Token-like inputs accepted for the correct/incorrect answers, mirroring
51# ActivationCache.logit_attrs.
52TokenInput = Union[
53 str,
54 int,
55 torch.Tensor,
56]
58# Which structural unit the residual stream is decomposed into.
59Unit = str # one of: "component", "layer", "head"
61_VALID_UNITS = ("component", "layer", "head")
63# Block variants that lack the attn_out + mlp_out structure decompose_resid expects.
64# When TransformerBridge.layer_types() reports any of these we refuse early — the
65# downstream decompose_resid would otherwise raise a confusing KeyError.
66_HYBRID_VARIANT_NAMES = ("mamba", "ssm", "mixer", "linear_attn")
69@dataclass
70class DirectLogitAttribution:
71 """Result of a :func:`direct_logit_attribution` call.
73 Attributes:
74 attribution:
75 Tensor of logit (or logit-difference) attributions with shape
76 ``[component, *batch_and_pos]``. The leading axis is aligned with
77 ``labels``. When ``pos`` selects a single position (the default) the
78 position axis is dropped, leaving ``[component, batch]`` — or
79 ``[component]`` if the cache had its batch dimension removed.
80 labels:
81 Human-readable name for each component, aligned with the leading
82 axis of ``attribution`` (e.g. ``"embed"``, ``"0_attn_out"``,
83 ``"L3H7"``).
84 unit:
85 The decomposition unit used ("component", "layer", or "head").
86 """
88 attribution: Float[torch.Tensor, "component *batch_and_pos"]
89 labels: List[str]
90 unit: Unit
92 def top(self, k: int = 5) -> List[tuple]:
93 """Return the ``k`` highest-attribution ``(label, value)`` pairs.
95 Attribution is reduced to a scalar per component by meaning over any
96 remaining batch/position dimensions, so this is most meaningful when a
97 single position was selected.
98 """
99 flat = self.attribution
100 if flat.ndim > 1: 100 ↛ 102line 100 didn't jump to line 102 because the condition on line 100 was always true
101 flat = flat.flatten(start_dim=1).mean(dim=-1)
102 values, indices = torch.topk(flat, min(k, flat.shape[0]))
103 return [(self.labels[i], values[j].item()) for j, i in enumerate(indices.tolist())]
106def _residual_stack_and_labels(
107 cache: ActivationCache,
108 unit: Unit,
109 pos_slice: SliceInput,
110):
111 """Decompose the residual stream into ``unit`` components plus labels.
113 LayerNorm is intentionally *not* applied here — ``logit_attrs`` applies the
114 final-layer scaling itself, so applying it twice would double-count.
115 """
116 if unit == "component":
117 # embed (+ pos_embed) and each layer's attn_out / mlp_out.
118 return cache.decompose_resid(apply_ln=False, pos_slice=pos_slice, return_labels=True)
119 if unit == "layer":
120 # Cumulative residual stream after each sublayer — logit-lens style.
121 return cache.accumulated_resid(
122 apply_ln=False, incl_mid=True, pos_slice=pos_slice, return_labels=True
123 )
124 if unit == "head": 124 ↛ 129line 124 didn't jump to line 129 because the condition on line 124 was always true
125 # Each attention head's contribution, plus the MLP/embedding remainder.
126 return cache.stack_head_results(
127 apply_ln=False, pos_slice=pos_slice, incl_remainder=True, return_labels=True
128 )
129 raise ValueError(f"unit must be one of {_VALID_UNITS}, got {unit!r}")
132def _validate_bridge_compatibility(model) -> None:
133 """Reject Bridge inputs that DLA can't produce correct numbers for.
135 Compatibility mode alone does not guarantee the final norm is folded:
136 ``fold_ln=False`` and ``no_processing=True`` leave its affine parameters
137 active. The hybrid-arch check catches Mamba/SSM blocks before decomposition.
138 """
139 # Lazy import — keeps the module importable without dragging in the bridge.
140 from transformer_lens.model_bridge import TransformerBridge
142 if not isinstance(model, TransformerBridge): 142 ↛ 143line 142 didn't jump to line 143 because the condition on line 142 was never true
143 return
145 if not getattr(model, "compatibility_mode", False):
146 raise ValueError(
147 "DLA on a TransformerBridge requires compatibility mode and an identity "
148 "final norm. Call `model.enable_compatibility_mode()` with its default "
149 "weight processing after loading the bridge, then re-run DLA."
150 )
152 final_norm = getattr(model, "ln_final", None)
153 if final_norm is not None:
154 weight = getattr(final_norm, "weight", None)
155 if isinstance(weight, torch.Tensor):
156 identity = (
157 torch.zeros_like(weight)
158 if getattr(model.cfg, "rmsnorm_uses_offset", False)
159 else torch.ones_like(weight)
160 )
161 if not torch.equal(weight, identity):
162 raise ValueError(
163 "DLA requires an identity final norm weight: the learned scale is not "
164 "folded into W_U. Load a fresh bridge and call "
165 "`enable_compatibility_mode()` with fold_ln=True and no_processing=False."
166 )
167 bias = getattr(final_norm, "bias", None)
168 if isinstance(bias, torch.Tensor) and not torch.equal(bias, torch.zeros_like(bias)):
169 raise ValueError(
170 "DLA requires a zero final norm bias: the learned bias is not folded "
171 "into b_U. Load a fresh bridge and call `enable_compatibility_mode()` "
172 "with fold_ln=True and no_processing=False."
173 )
175 layer_types = model.layer_types()
176 hybrid = [lt for lt in layer_types if any(p in _HYBRID_VARIANT_NAMES for p in lt.split("+"))]
177 if hybrid:
178 raise NotImplementedError(
179 f"DLA does not yet support hybrid architectures (found block types {hybrid}). "
180 f"Only standard attention + MLP transformers (e.g. GPT-2, LLaMA, Pythia) are "
181 f"supported; hybrid support requires extending ActivationCache.decompose_resid."
182 )
185def direct_logit_attribution(
186 model,
187 input: Union[str, List[str], torch.Tensor, None] = None,
188 answer_tokens: Optional[TokenInput] = None,
189 incorrect_tokens: Optional[TokenInput] = None,
190 *,
191 unit: Unit = "component",
192 pos: SliceInput = -1,
193 cache: Optional[ActivationCache] = None,
194) -> DirectLogitAttribution:
195 """Compute Direct Logit Attribution for a prompt.
197 Decomposes the contribution of model components to the logit of
198 ``answer_tokens`` (or, if ``incorrect_tokens`` is given, to the logit
199 *difference* ``answer - incorrect`` along the ``W_U`` direction, which is
200 usually what you want for circuit analysis).
202 The model is run once with caching unless a precomputed ``cache`` is passed.
204 Note that DLA attributes only the part of a logit that comes from the
205 residual stream through the unembedding direction; the unembedding bias
206 ``b_U`` is a per-token constant that no component produces. So a complete
207 decomposition reconstructs ``logit[token] - b_U[token]`` rather than the raw
208 logit.
210 On a ``TransformerBridge``, compatibility mode must be enabled and the final
211 normalization's affine parameters must be folded into ``W_U`` and ``b_U``
212 (or already be identity);
213 otherwise the projection direction is wrong and DLA would return silently
214 incorrect numbers. Hybrid architectures
215 (Mamba/SSM/Mixer/LinearAttention) are not yet supported because
216 ``decompose_resid`` only understands the ``attn_out + mlp_out`` block layout;
217 both conditions raise an explicit error at call time.
219 Args:
220 model:
221 A ``TransformerBridge`` with ``enable_compatibility_mode()``
222 already called and an identity final normalization.
223 input:
224 Prompt to run — a string, list of strings, or token tensor. Optional
225 only when a precomputed ``cache`` is supplied.
226 answer_tokens:
227 The correct token(s) to attribute, as a string, id, or tensor. A
228 string is converted with ``model.to_single_token``.
229 incorrect_tokens:
230 Optional baseline token(s). When given, attribution is computed for
231 the ``answer - incorrect`` residual direction. Must broadcast to the
232 same shape as ``answer_tokens``.
233 unit:
234 Decomposition granularity:
236 - ``"component"`` (default): embedding + each layer's attention and
237 MLP output (via ``decompose_resid``).
238 - ``"layer"``: cumulative residual stream after each sublayer, i.e.
239 logit-lens trajectory (via ``accumulated_resid``).
240 - ``"head"``: each attention head individually, plus a remainder
241 term for everything else (via ``stack_head_results``).
242 pos:
243 Sequence position(s) to attribute. Defaults to ``-1`` (the final
244 token, the usual choice for next-token DLA). Pass ``None`` to keep
245 every position (the result then has a trailing position axis).
246 cache:
247 Optional precomputed ``ActivationCache`` to reuse instead of running
248 the model again.
250 Returns:
251 A :class:`DirectLogitAttribution` with ``attribution`` (shape
252 ``[component, *batch_and_pos]``) and aligned ``labels``.
254 Raises:
255 ValueError: If ``unit`` is invalid, ``answer_tokens`` is ``None``,
256 neither ``input`` nor ``cache`` is provided, or a
257 ``TransformerBridge`` is passed without compatibility mode enabled
258 or with an active learned final-norm weight or bias.
259 NotImplementedError: If a ``TransformerBridge`` reports a hybrid block
260 layout (Mamba/SSM/Mixer/LinearAttention).
261 """
262 if unit not in _VALID_UNITS:
263 raise ValueError(f"unit must be one of {_VALID_UNITS}, got {unit!r}")
264 if answer_tokens is None:
265 raise ValueError("answer_tokens is required")
267 _validate_bridge_compatibility(model)
269 if cache is None:
270 if input is None:
271 raise ValueError("provide either `input` to run the model, or a precomputed `cache`")
272 _, cache = model.run_with_cache(input)
274 residual_stack, labels = _residual_stack_and_labels(cache, unit, pos)
276 # logit_attrs applies the final LayerNorm scaling (with the same pos slice)
277 # and dots each component against the (correct - incorrect) unembed direction.
278 attribution = cache.logit_attrs(
279 residual_stack,
280 tokens=answer_tokens,
281 incorrect_tokens=incorrect_tokens,
282 pos_slice=pos,
283 has_batch_dim=cache.has_batch_dim,
284 )
286 return DirectLogitAttribution(attribution=attribution, labels=labels, unit=unit)