Coverage for transformer_lens/tools/analysis/direct_logit_attribution.py: 93%

1"""Direct Logit Attribution (DLA). 

2 

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. 

14 

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. 

22 

23Example:: 

24 

25 from transformer_lens import TransformerBridge 

26 from transformer_lens.tools.analysis import direct_logit_attribution 

27 

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""" 

40 

41from dataclasses import dataclass 

42from typing import List, Optional, Union 

43 

44import torch 

45from jaxtyping import Float 

46 

47from transformer_lens.ActivationCache import ActivationCache 

48from transformer_lens.utilities import SliceInput 

49 

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] 

57 

58# Which structural unit the residual stream is decomposed into. 

59Unit = str # one of: "component", "layer", "head" 

60 

61_VALID_UNITS = ("component", "layer", "head") 

62 

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") 

67 

68 

69@dataclass 

70class DirectLogitAttribution: 

71 """Result of a :func:`direct_logit_attribution` call. 

72 

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 """ 

87 

88 attribution: Float[torch.Tensor, "component *batch_and_pos"] 

89 labels: List[str] 

90 unit: Unit 

91 

92 def top(self, k: int = 5) -> List[tuple]: 

93 """Return the ``k`` highest-attribution ``(label, value)`` pairs. 

94 

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())] 

104 

105 

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. 

112 

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}") 

130 

131 

132def _validate_bridge_compatibility(model) -> None: 

133 """Reject Bridge inputs that DLA can't produce correct numbers for. 

134 

135 The compatibility-mode check catches a silent- 

136 correctness footgun: without folded LN, the projection direction in 

137 ``logit_attrs`` is wrong on a Bridge. The hybrid-arch check catches Mamba/ 

138 SSM blocks early with a clear error rather than letting ``decompose_resid`` 

139 raise a confusing KeyError downstream. 

140 """ 

141 # Lazy import — keeps the module importable without dragging in the bridge. 

142 from transformer_lens.model_bridge import TransformerBridge 

143 

144 if not isinstance(model, TransformerBridge): 144 ↛ 145line 144 didn't jump to line 145 because the condition on line 144 was never true

145 return 

146 

147 if not getattr(model, "compatibility_mode", False): 

148 raise ValueError( 

149 "DLA on a TransformerBridge requires compatibility mode so that LayerNorm " 

150 "weights are folded into W_U. Call `model.enable_compatibility_mode()` " 

151 "after loading the bridge, then re-run DLA." 

152 ) 

153 

154 layer_types = model.layer_types() 

155 hybrid = [lt for lt in layer_types if any(p in _HYBRID_VARIANT_NAMES for p in lt.split("+"))] 

156 if hybrid: 

157 raise NotImplementedError( 

158 f"DLA does not yet support hybrid architectures (found block types {hybrid}). " 

159 f"Only standard attention + MLP transformers (e.g. GPT-2, LLaMA, Pythia) are " 

160 f"supported; hybrid support requires extending ActivationCache.decompose_resid." 

161 ) 

162 

163 

164def direct_logit_attribution( 

165 model, 

166 input: Union[str, List[str], torch.Tensor, None] = None, 

167 answer_tokens: Optional[TokenInput] = None, 

168 incorrect_tokens: Optional[TokenInput] = None, 

169 *, 

170 unit: Unit = "component", 

171 pos: SliceInput = -1, 

172 cache: Optional[ActivationCache] = None, 

173) -> DirectLogitAttribution: 

174 """Compute Direct Logit Attribution for a prompt. 

175 

176 Decomposes the contribution of model components to the logit of 

177 ``answer_tokens`` (or, if ``incorrect_tokens`` is given, to the logit 

178 *difference* ``answer - incorrect`` along the ``W_U`` direction, which is 

179 usually what you want for circuit analysis). 

180 

181 The model is run once with caching unless a precomputed ``cache`` is passed. 

182 

183 Note that DLA attributes only the part of a logit that comes from the 

184 residual stream through the unembedding direction; the unembedding bias 

185 ``b_U`` is a per-token constant that no component produces. So a complete 

186 decomposition reconstructs ``logit[token] - b_U[token]`` rather than the raw 

187 logit. 

188 

189 On a ``TransformerBridge``, compatibility mode must be enabled (so the final 

190 LayerNorm is folded into ``W_U``) — otherwise the projection direction is 

191 wrong and DLA returns silently incorrect numbers. Hybrid architectures 

192 (Mamba/SSM/Mixer/LinearAttention) are not yet supported because 

193 ``decompose_resid`` only understands the ``attn_out + mlp_out`` block layout; 

194 both conditions raise an explicit error at call time. 

195 

196 Args: 

197 model: 

198 A ``TransformerBridge`` with ``enable_compatibility_mode()`` 

199 already called. 

200 input: 

201 Prompt to run — a string, list of strings, or token tensor. Optional 

202 only when a precomputed ``cache`` is supplied. 

203 answer_tokens: 

204 The correct token(s) to attribute, as a string, id, or tensor. A 

205 string is converted with ``model.to_single_token``. 

206 incorrect_tokens: 

207 Optional baseline token(s). When given, attribution is computed for 

208 the ``answer - incorrect`` residual direction. Must broadcast to the 

209 same shape as ``answer_tokens``. 

210 unit: 

211 Decomposition granularity: 

212 

213 - ``"component"`` (default): embedding + each layer's attention and 

214 MLP output (via ``decompose_resid``). 

215 - ``"layer"``: cumulative residual stream after each sublayer, i.e. 

216 logit-lens trajectory (via ``accumulated_resid``). 

217 - ``"head"``: each attention head individually, plus a remainder 

218 term for everything else (via ``stack_head_results``). 

219 pos: 

220 Sequence position(s) to attribute. Defaults to ``-1`` (the final 

221 token, the usual choice for next-token DLA). Pass ``None`` to keep 

222 every position (the result then has a trailing position axis). 

223 cache: 

224 Optional precomputed ``ActivationCache`` to reuse instead of running 

225 the model again. 

226 

227 Returns: 

228 A :class:`DirectLogitAttribution` with ``attribution`` (shape 

229 ``[component, *batch_and_pos]``) and aligned ``labels``. 

230 

231 Raises: 

232 ValueError: If ``unit`` is invalid, ``answer_tokens`` is ``None``, 

233 neither ``input`` nor ``cache`` is provided, or a 

234 ``TransformerBridge`` is passed without compatibility mode enabled. 

235 NotImplementedError: If a ``TransformerBridge`` reports a hybrid block 

236 layout (Mamba/SSM/Mixer/LinearAttention). 

237 """ 

238 if unit not in _VALID_UNITS: 

239 raise ValueError(f"unit must be one of {_VALID_UNITS}, got {unit!r}") 

240 if answer_tokens is None: 

241 raise ValueError("answer_tokens is required") 

242 

243 _validate_bridge_compatibility(model) 

244 

245 if cache is None: 

246 if input is None: 

247 raise ValueError("provide either `input` to run the model, or a precomputed `cache`") 

248 _, cache = model.run_with_cache(input) 

249 

250 residual_stack, labels = _residual_stack_and_labels(cache, unit, pos) 

251 

252 # logit_attrs applies the final LayerNorm scaling (with the same pos slice) 

253 # and dots each component against the (correct - incorrect) unembed direction. 

254 attribution = cache.logit_attrs( 

255 residual_stack, 

256 tokens=answer_tokens, 

257 incorrect_tokens=incorrect_tokens, 

258 pos_slice=pos, 

259 has_batch_dim=cache.has_batch_dim, 

260 ) 

261 

262 return DirectLogitAttribution(attribution=attribution, labels=labels, unit=unit)