Coverage for transformer_lens/ActivationCache.py: 98%
498 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +0000
1"""Activation Cache.
3The :class:`ActivationCache` is at the core of Transformer Lens. It is a wrapper that stores all
4important activations from a forward pass of the model, and provides a variety of helper functions
5to investigate them.
7Getting Started:
9When reading these docs for the first time, we recommend reading the main :class:`ActivationCache`
10class first, including the examples, and then skimming the available methods. You can then refer
11back to these docs depending on what you need to do.
12"""
14from __future__ import annotations
16import logging
17from collections import Counter
18from typing import (
19 TYPE_CHECKING,
20 Any,
21 Callable,
22 Dict,
23 Iterator,
24 List,
25 Optional,
26 Tuple,
27 Union,
28 cast,
29)
31import einops
32import torch
33from jaxtyping import Float, Int
34from typing_extensions import Literal
36import transformer_lens.utilities as utils
37from transformer_lens.utilities import Slice, SliceInput, warn_if_mps
39if TYPE_CHECKING:
40 from transformer_lens.model_protocol import TransformerLensModelWithWeights
43def _normalize_projection_to_2d(
44 project: Optional[torch.Tensor],
45) -> Tuple[Optional[torch.Tensor], bool]:
46 """Return ``(project_2d, squeeze_at_end)`` — 1D projections are reshaped to 2D for uniform internal handling and squeezed back at the user-facing return."""
47 if project is None:
48 return None, False
49 if project.ndim == 1:
50 return project.unsqueeze(-1), True
51 return project, False
54class ActivationCache:
55 """Activation Cache.
57 A wrapper that stores all important activations from a forward pass of the model, and provides a
58 variety of helper functions to investigate them.
60 The :class:`ActivationCache` is at the core of Transformer Lens. It is a wrapper that stores all
61 important activations from a forward pass of the model, and provides a variety of helper
62 functions to investigate them. The common way to access it is to run the model with
63 :meth:`transformer_lens.model_bridge.TransformerBridge.run_with_cache`.
65 Examples:
67 When investigating a particular behaviour of a model, a very common first step is to try and
68 understand which components of the model are most responsible for that behaviour. For example,
69 if you're investigating the prompt "Why did the chicken cross the" -> " road", you might want to
70 understand if there is a specific sublayer (mlp or multi-head attention) that is responsible for
71 the model predicting "road". This kind of analysis commonly falls under the category of "logit
72 attribution" or "direct logit attribution" (DLA).
74 >>> from transformer_lens.model_bridge import TransformerBridge
75 >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M")
76 >>> model.enable_compatibility_mode()
78 >>> _logits, cache = model.run_with_cache("Why did the chicken cross the")
79 >>> residual_stream, labels = cache.decompose_resid(return_labels=True, mode="attn")
80 >>> print(labels[0:3])
81 ['embed', 'pos_embed', '0_attn_out']
83 >>> answer = " road" # Note the proceeding space to match the model's tokenization
84 >>> logit_attrs = cache.logit_attrs(residual_stream, answer)
85 >>> print(logit_attrs.shape) # Attention layers
86 torch.Size([10, 1, 7])
88 >>> most_important_component_idx = torch.argmax(logit_attrs)
89 >>> print(labels[most_important_component_idx])
90 3_attn_out
92 You can also dig in with more granularity, using :meth:`get_full_resid_decomposition` to get the
93 residual stream by individual component (mlp neurons and individual attention heads). This
94 creates a larger residual stack, but the approach of using :meth"`logit_attrs` remains the same.
96 Equally you might want to find out if the model struggles to construct such excellent jokes
97 until the very last layers, or if it is trivial and the first few layers are enough. This kind
98 of analysis is called "logit lens", and you can find out more about how to do that with
99 :meth:`ActivationCache.accumulated_resid`.
101 Warning:
103 :class:`ActivationCache` is designed to be used with
104 :class:`transformer_lens.model_bridge.TransformerBridge`. Advanced helpers expect the model to
105 expose the TransformerLens weight-processing interface and generally expect a complete cache;
106 some internal methods may break with other models or partial caches.
108 The biggest footgun and source of bugs in this code will be keeping track of indexes,
109 dimensions, and the numbers of each. There are several kinds of activations:
111 * Internal attn head vectors: q, k, v, z. Shape [batch, pos, head_index, d_head].
112 * Internal attn pattern style results: pattern (post softmax), attn_scores (pre-softmax). Shape
113 [batch, head_index, query_pos, key_pos].
114 * Attn head results: result. Shape [batch, pos, head_index, d_model].
115 * Internal MLP vectors: pre, post, mid (only used for solu_ln - the part between activation +
116 layernorm). Shape [batch, pos, d_mlp].
117 * Residual stream vectors: resid_pre, resid_mid, resid_post, attn_out, mlp_out, embed,
118 pos_embed, normalized (output of each LN or LNPre). Shape [batch, pos, d_model].
119 * LayerNorm Scale: scale. Shape [batch, pos, 1].
121 Sometimes the batch dimension will be missing because we applied `remove_batch_dim` (used when
122 batch_size=1), and as such all library functions *should* be robust to that.
124 Type annotations are in the following form:
126 * layers_covered is the number of layers queried in functions that stack the residual stream.
127 * batch_and_pos_dims is the set of dimensions from batch and pos - by default this is ["batch",
128 "pos"], but is only ["pos"] if we've removed the batch dimension and is [()] if we've removed
129 batch dimension and are applying a pos slice which indexes a specific position.
131 Args:
132 cache_dict:
133 A dictionary of cached activations from a model run.
134 model:
135 The model that the activations are from.
136 has_batch_dim:
137 Whether the activations have a batch dimension.
138 """
140 def __init__(
141 self,
142 cache_dict: Dict[str, torch.Tensor],
143 model: Any,
144 has_batch_dim: bool = True,
145 ):
146 self.cache_dict = cache_dict
147 # Advanced helpers (LN folding, residual-direction projection) need the
148 # weight-processing surface, which TransformerBridge exposes.
149 self.model = cast("TransformerLensModelWithWeights", model)
150 self.has_batch_dim = has_batch_dim
151 self.has_embed = "hook_embed" in self.cache_dict
152 self.has_pos_embed = "hook_pos_embed" in self.cache_dict
154 # Note: model reference prevents garbage collection. Set cache.model = None if unneeded.
156 def _batch_size(self) -> int:
157 """The cache's batch size: the most common leading dim across entries.
159 Caches may hold non-batch entries alongside genuinely batched
160 activations — broadcast entries with a leading dim of 1 (e.g. the
161 bridge's position-index inputs) or position-indexed entries whose
162 leading dim is the sequence length (e.g. T5's relative position bias).
163 The batched activations vastly outnumber both, so the mode is the
164 reliable signal where max/min are not.
165 """
166 counts = Counter(v.size(0) for v in self.cache_dict.values() if v.ndim > 0)
167 return counts.most_common(1)[0][0] if counts else 1
169 def remove_batch_dim(self, batch_size: Optional[int] = None) -> ActivationCache:
170 """Remove the Batch Dimension (if a single batch item).
172 Args:
173 batch_size:
174 The true batch size, when the caller knows it. By default it is inferred from
175 the most common leading dimension of the cached entries, which is wrong for a
176 cache holding only flattened or position-indexed activations.
178 Returns:
179 The ActivationCache with the batch dimension removed.
180 """
181 if self.has_batch_dim:
182 if batch_size is None:
183 batch_size = self._batch_size()
184 assert (
185 batch_size == 1
186 ), f"Cannot remove batch dimension from cache with batch size {batch_size}"
187 for key in self.cache_dict:
188 if self.cache_dict[key].ndim > 0 and self.cache_dict[key].size(0) == 1:
189 self.cache_dict[key] = self.cache_dict[key][0]
190 self.has_batch_dim = False
191 else:
192 logging.warning("Tried removing batch dimension after already having removed it.")
193 return self
195 def __repr__(self) -> str:
196 """Representation of the ActivationCache.
198 Special method that returns a string representation of an object. It's normally used to give
199 a string that can be used to recreate the object, but here we just return a string that
200 describes the object.
201 """
202 return f"ActivationCache with keys {list(self.cache_dict.keys())}"
204 def __getitem__(self, key) -> torch.Tensor:
205 """Retrieve Cached Activations by Key or Shorthand.
207 Enables direct access to cached activations via dictionary-style indexing using keys or
208 shorthand naming conventions.
210 It also supports tuples for advanced indexing, with the dimension order as (name, layer_index, layer_type).
211 See :func:`transformer_lens.utilities.get_act_name` for how shorthand is converted to a full name.
214 Args:
215 key:
216 The key or shorthand name for the activation to retrieve.
218 Returns:
219 The cached activation tensor corresponding to the given key.
220 """
221 if key in self.cache_dict:
222 return self.cache_dict[key]
223 elif type(key) == str:
224 return self.cache_dict[utils.get_act_name(key)]
225 else:
226 if len(key) > 1 and key[1] is not None:
227 if key[1] < 0:
228 # Supports negative indexing on the layer dimension
229 key = (key[0], self.model.cfg.n_layers + key[1], *key[2:])
230 return self.cache_dict[utils.get_act_name(*key)]
232 def __len__(self) -> int:
233 """Length of the ActivationCache.
235 Special method that returns the length of an object (in this case the number of different
236 activations in the cache).
237 """
238 return len(self.cache_dict)
240 def to(self, device: Union[str, torch.device]) -> ActivationCache:
241 """Move the Cache to a Device.
243 Mostly useful for moving the cache to the CPU after model computation finishes to save GPU
244 memory. Note however that operations will be much slower on the CPU. Note also that some
245 methods will break unless the model is also moved to the same device, eg
246 `compute_head_results`.
248 Args:
249 device:
250 The device to move the cache to (e.g. `torch.device.cpu`).
252 """
253 warn_if_mps(device)
254 self.cache_dict = {key: value.to(device) for key, value in self.cache_dict.items()}
255 return self
257 def toggle_autodiff(self, mode: bool = False):
258 """Toggle Autodiff Globally.
260 Applies `torch.set_grad_enabled(mode)` to the global state (not just TransformerLens).
262 Warning:
264 This is pretty dangerous, since autodiff is global state - this turns off torch's
265 ability to take gradients completely and it's easy to get a bunch of errors if you don't
266 realise what you're doing.
268 But autodiff consumes a LOT of GPU memory (since every intermediate activation is cached
269 until all downstream activations are deleted - this means that computing the loss and
270 storing it in a list will keep every activation sticking around!). So often when you're
271 analysing a model's activations, and don't need to do any training, autodiff is more trouble
272 than its worth.
274 If you don't want to mess with global state, using torch.inference_mode as a context manager
275 or decorator achieves similar effects:
277 >>> with torch.inference_mode():
278 ... y = torch.Tensor([1., 2, 3])
279 >>> y.requires_grad
280 False
281 """
282 logging.warning("Changed the global state, set autodiff to %s", mode)
283 torch.set_grad_enabled(mode)
285 def keys(self):
286 """Keys of the ActivationCache.
288 Examples:
290 >>> from transformer_lens.model_bridge import TransformerBridge
291 >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M")
292 >>> model.enable_compatibility_mode()
293 >>> _logits, cache = model.run_with_cache("Some prompt")
294 >>> list(cache.keys())[0:8]
295 ['embed.hook_in', 'hook_embed', 'embed.hook_out', 'pos_embed.hook_in',
296 'hook_pos_embed', 'pos_embed.hook_out', 'blocks.0.hook_in',
297 'blocks.0.hook_resid_pre']
299 Returns:
300 List of all keys.
301 """
302 return self.cache_dict.keys()
304 def values(self):
305 """Values of the ActivationCache.
307 Returns:
308 List of all values.
309 """
310 return self.cache_dict.values()
312 def items(self):
313 """Items of the ActivationCache.
315 Returns:
316 List of all items ((key, value) tuples).
317 """
318 return self.cache_dict.items()
320 def __iter__(self) -> Iterator[str]:
321 """ActivationCache Iterator.
323 Special method that returns an iterator over the keys in the ActivationCache. Allows looping over the
324 cache.
326 Examples:
328 >>> from transformer_lens.model_bridge import TransformerBridge
329 >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M")
330 >>> model.enable_compatibility_mode()
331 >>> _logits, cache = model.run_with_cache("Some prompt")
332 >>> cache_interesting_names = []
333 >>> for key in cache:
334 ... if not key.startswith("blocks.") or key.startswith("blocks.0"):
335 ... cache_interesting_names.append(key)
336 >>> print(cache_interesting_names[0:8])
337 ['embed.hook_in', 'hook_embed', 'embed.hook_out', 'pos_embed.hook_in',
338 'hook_pos_embed', 'pos_embed.hook_out', 'blocks.0.hook_in',
339 'blocks.0.hook_resid_pre']
341 Returns:
342 Iterator over the cache.
343 """
344 return self.cache_dict.__iter__()
346 def apply_slice_to_batch_dim(self, batch_slice: Union[Slice, SliceInput]) -> ActivationCache:
347 """Apply a Slice to the Batch Dimension.
349 Args:
350 batch_slice:
351 The slice to apply to the batch dimension.
353 Returns:
354 The ActivationCache with the batch dimension sliced.
355 """
356 if not isinstance(batch_slice, Slice):
357 batch_slice = Slice(batch_slice)
358 batch_slice = cast(Slice, batch_slice) # mypy can't seem to infer this
359 assert (
360 self.has_batch_dim or batch_slice.mode == "empty"
361 ), "Cannot index into a cache without a batch dim"
362 still_has_batch_dim = (batch_slice.mode != "int") and self.has_batch_dim
363 batch_size = self._batch_size()
364 # Broadcast entries (leading dim 1 when the true batch is larger) are not
365 # batched — leave them untouched so slicing can't index out of bounds.
366 new_cache_dict = {
367 name: (
368 batch_slice.apply(param, dim=0)
369 if param.ndim > 0 and param.size(0) == batch_size
370 else param
371 )
372 for name, param in self.cache_dict.items()
373 }
374 return ActivationCache(new_cache_dict, self.model, has_batch_dim=still_has_batch_dim)
376 def accumulated_resid(
377 self,
378 layer: Optional[int] = None,
379 incl_mid: bool = False,
380 apply_ln: bool = False,
381 pos_slice: Optional[Union[Slice, SliceInput]] = None,
382 mlp_input: bool = False,
383 return_labels: bool = False,
384 ) -> Union[
385 Float[torch.Tensor, "layers_covered *batch_and_pos_dims d_model"],
386 Tuple[Float[torch.Tensor, "layers_covered *batch_and_pos_dims d_model"], List[str]],
387 ]:
388 """Accumulated Residual Stream.
390 Returns the accumulated residual stream at each layer/sub-layer. This is useful for `Logit
391 Lens <https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens>`
392 style analysis, where it can be thought of as what the model "believes" at each point in the
393 residual stream.
395 :meth:`logit_lens` does the whole projection below in one call (final norm, ``W_U``,
396 ``b_U`` and the adapter's output transform, chunked over the vocabulary); the recipe
397 here shows what it computes.
399 To project this into the vocabulary space, remember that there is a final layer norm in most
400 decoder-only transformers. Therefore, you need to first apply the final layer norm (which
401 can be done with `apply_ln`), and then multiply by the unembedding matrix (:math:`W_U`)
402 and optionally add the unembedding bias (:math:`b_U`).
404 **Note on bias terms:** There are two valid approaches for the final projection:
406 1. **With bias terms:** Use `model.unembed(normalized_resid)` which applies both :math:`W_U`
407 and :math:`b_U` (equivalent to `normalized_resid @ model.W_U + model.b_U`). This works
408 correctly with both `fold_ln=True` and `fold_ln=False` settings, as the biases are
409 handled consistently.
410 2. **Without bias terms:** Use only `normalized_resid @ model.W_U`. If taking this approach,
411 you should instantiate the model with `fold_ln=True`, which folds the layer norm scaling
412 into :math:`W_U` and the layer norm bias into :math:`b_U`. Since `apply_ln=True` will
413 apply the (now parameter-free) layer norm, and you skip :math:`b_U`, no bias terms are
414 included. With `fold_ln=False`, the layer norm bias would still be applied, which is
415 typically not desired when excluding bias terms.
417 Both approaches are commonly used in the literature and are valid interpretability choices.
419 If you instead want to look at contributions to the residual stream from each component
420 (e.g. for direct logit attribution), see :meth:`decompose_resid` instead, or
421 :meth:`get_full_resid_decomposition` if you want contributions broken down further into each
422 MLP neuron.
424 Examples:
426 Logit Lens analysis can be done as follows:
428 >>> from transformer_lens.model_bridge import TransformerBridge
429 >>> import torch
430 >>> import pandas as pd
432 >>> model = TransformerBridge.boot_transformers("roneneldan/TinyStories-1M", device="cpu")
433 >>> model.enable_compatibility_mode()
435 >>> prompt = "Why did the chicken cross the"
436 >>> answer = " road"
437 >>> logits, cache = model.run_with_cache("Why did the chicken cross the")
438 >>> answer_token = model.to_single_token(answer)
439 >>> print(answer_token)
440 2975
442 >>> accum_resid, labels = cache.accumulated_resid(return_labels=True, apply_ln=True)
443 >>> last_token_accum = accum_resid[:, 0, -1, :] # layer, batch, pos, d_model
444 >>> print(last_token_accum.shape) # layer, d_model
445 torch.Size([9, 64])
448 >>> W_U = model.W_U
449 >>> print(W_U.shape)
450 torch.Size([64, 50257])
452 >>> # Project to vocabulary without unembedding bias
453 >>> layers_logits = last_token_accum @ W_U # layer, d_vocab
454 >>> print(layers_logits.shape)
455 torch.Size([9, 50257])
457 >>> # The unembedding bias can be added on top when the model carries one
458 >>> # (the rank table below stays on the bias-free logits):
459 >>> b_U = getattr(model, "b_U", None)
460 >>> with_bias = layers_logits + b_U if b_U is not None else layers_logits
461 >>> print(with_bias.shape)
462 torch.Size([9, 50257])
464 >>> # Get the rank of the correct answer by layer
465 >>> sorted_indices = torch.argsort(layers_logits, dim=1, descending=True)
466 >>> rank_answer = (sorted_indices == 2975).nonzero(as_tuple=True)[1]
467 >>> print(pd.Series(rank_answer, index=labels))
468 0_pre 4442
469 1_pre 382
470 2_pre 982
471 3_pre 1160
472 4_pre 408
473 5_pre 145
474 6_pre 78
475 7_pre 387
476 final_post 6
477 dtype: int64
479 Args:
480 layer:
481 The layer to take components up to - by default includes resid_pre for that layer
482 and excludes resid_mid and resid_post for that layer. If set as `n_layers`, `-1` or
483 `None` it will return all residual streams, including the final one (i.e.
484 immediately pre logits). The indices are taken such that this gives the accumulated
485 streams up to the input to layer l.
486 incl_mid:
487 Whether to return `resid_mid` for all previous layers.
488 apply_ln:
489 Whether to apply the final layer norm to the stack. When True, applies
490 `model.ln_final`, which recomputes normalization statistics (mean and
491 variance/RMS) for each intermediate state in the stack, transforming the
492 activations into the format expected by the unembedding layer.
493 pos_slice:
494 A slice object to apply to the pos dimension. Defaults to None, do nothing.
495 mlp_input:
496 Whether to include resid_mid for the current layer. This essentially gives the MLP
497 input rather than the attention input.
498 return_labels:
499 Whether to return a list of labels for the residual stream components. Useful for
500 labelling graphs.
502 Returns:
503 A tensor of the accumulated residual streams. If `return_labels` is True, also returns a
504 list of labels for the components (as a tuple in the form `(components, labels)`).
505 """
506 if not isinstance(pos_slice, Slice):
507 pos_slice = Slice(pos_slice)
508 if layer is None or layer == -1:
509 # Default to the residual stream immediately pre unembed
510 layer = self.model.cfg.n_layers
511 assert isinstance(layer, int)
512 labels = []
513 components_list = []
514 for l in range(layer + 1):
515 if l == self.model.cfg.n_layers:
516 components_list.append(self[("resid_post", self.model.cfg.n_layers - 1)])
517 labels.append("final_post")
518 continue
519 components_list.append(self[("resid_pre", l)])
520 labels.append(f"{l}_pre")
521 if (incl_mid and l < layer) or (mlp_input and l == layer):
522 components_list.append(self[("resid_mid", l)])
523 labels.append(f"{l}_mid")
524 components_list = [pos_slice.apply(c, dim=-2) for c in components_list]
525 components = torch.stack(components_list, dim=0)
526 if apply_ln:
527 recompute_ln = layer == self.model.cfg.n_layers
528 components = self.apply_ln_to_stack(
529 components,
530 layer,
531 pos_slice=pos_slice,
532 mlp_input=mlp_input,
533 has_batch_dim=self.has_batch_dim,
534 recompute_ln=recompute_ln,
535 )
536 if return_labels:
537 return components, labels
538 else:
539 return components
541 def logit_lens(self, **kwargs: Any) -> Any:
542 """One-call logit lens over this cache's accumulated residual stream.
544 Delegates to :func:`transformer_lens.tools.analysis.logit_lens.logit_lens`
545 with this cache; see it for the keyword arguments (``layers``,
546 ``positions``, ``targets``, ``top_k``, ``vocab``, ``return_type`` ...).
547 The final stack entry reproduces the model's logits through the real
548 final norm, ``W_U``, ``b_U`` and the adapter's output transform.
549 """
550 # Lazy import: tools.analysis imports this module.
551 from transformer_lens.tools.analysis.logit_lens import logit_lens
553 return logit_lens(self.model, cache=self, **kwargs)
555 def logit_attrs(
556 self,
557 residual_stack: Float[torch.Tensor, "num_components *batch_and_pos_dims d_model"],
558 tokens: Union[
559 str,
560 int,
561 Int[torch.Tensor, ""],
562 Int[torch.Tensor, "batch"],
563 Int[torch.Tensor, "batch position"],
564 ],
565 incorrect_tokens: Optional[
566 Union[
567 str,
568 int,
569 Int[torch.Tensor, ""],
570 Int[torch.Tensor, "batch"],
571 Int[torch.Tensor, "batch position"],
572 ]
573 ] = None,
574 pos_slice: Union[Slice, SliceInput] = None,
575 batch_slice: Union[Slice, SliceInput] = None,
576 has_batch_dim: bool = True,
577 ) -> Float[torch.Tensor, "num_components *batch_and_pos_dims_out"]:
578 """Logit Attributions.
580 Takes a residual stack (typically the residual stream decomposed by components), and
581 calculates how much each item in the stack "contributes" to specific tokens.
583 It does this by:
584 1. Getting the residual directions of the tokens (i.e. reversing the unembed)
585 2. Taking the dot product of each item in the residual stack, with the token residual
586 directions.
588 Note that if incorrect tokens are provided, it instead takes the difference between the
589 correct and incorrect tokens (to calculate the residual directions). This is useful as
590 sometimes we want to know e.g. which components are most responsible for selecting the
591 correct token rather than an incorrect one. For example in the `Interpretability in the Wild
592 paper <https://arxiv.org/abs/2211.00593>` prompts such as "John and Mary went to the shops,
593 John gave a bag to" were investigated, and it was therefore useful to calculate attribution
594 for the :math:`\\text{Mary} - \\text{John}` residual direction.
596 Warning:
598 Choosing the correct `tokens` and `incorrect_tokens` is both important and difficult. When
599 investigating specific components it's also useful to look at it's impact on all tokens
600 (i.e. :math:`\\text{final_ln}(\\text{residual_stack_item}) W_U`).
602 Args:
603 residual_stack:
604 Stack of components of residual stream to get logit attributions for.
605 tokens:
606 Tokens to compute logit attributions on. A [batch] tensor selects one token
607 per example across all retained positions. A 1-D tensor whose length differs
608 from the batch size selects one token per position, shared across examples.
609 Use [1, position] to share per-position targets when batch and position sizes
610 are equal; [batch, position] selects targets for each example and position.
611 Token tensors should already match any batch/position slices.
612 incorrect_tokens:
613 If provided, compute attributions on logit difference between tokens and
614 incorrect_tokens. Must have the same shape as tokens.
615 pos_slice:
616 The slice to apply layer norm scaling on. Defaults to None, do nothing.
617 batch_slice:
618 The slice to take on the batch dimension during layer norm scaling. Defaults to
619 None, do nothing.
620 has_batch_dim:
621 Whether residual_stack has a batch dimension. Defaults to True.
623 Returns:
624 A tensor of the logit attributions or logit difference attributions if incorrect_tokens
625 was provided.
626 """
627 if not isinstance(pos_slice, Slice):
628 pos_slice = Slice(pos_slice)
630 if not isinstance(batch_slice, Slice):
631 batch_slice = Slice(batch_slice)
633 # Convert tokens to tensor for shape checking, but pass original to tokens_to_residual_directions
634 tokens_for_shape_check = tokens
636 if isinstance(tokens_for_shape_check, str):
637 tokens_for_shape_check = torch.as_tensor(
638 self.model.to_single_token(tokens_for_shape_check)
639 )
640 elif isinstance(tokens_for_shape_check, int):
641 tokens_for_shape_check = torch.as_tensor(tokens_for_shape_check)
643 logit_directions = self.model.tokens_to_residual_directions(tokens)
645 if incorrect_tokens is not None:
646 # Convert incorrect_tokens to tensor for shape checking, but pass original to tokens_to_residual_directions
647 incorrect_tokens_for_shape_check = incorrect_tokens
649 if isinstance(incorrect_tokens_for_shape_check, str):
650 incorrect_tokens_for_shape_check = torch.as_tensor(
651 self.model.to_single_token(incorrect_tokens_for_shape_check)
652 )
653 elif isinstance(incorrect_tokens_for_shape_check, int):
654 incorrect_tokens_for_shape_check = torch.as_tensor(incorrect_tokens_for_shape_check)
656 if tokens_for_shape_check.shape != incorrect_tokens_for_shape_check.shape:
657 raise ValueError(
658 f"tokens and incorrect_tokens must have the same shape! \
659 (tokens.shape={tokens_for_shape_check.shape}, \
660 incorrect_tokens.shape={incorrect_tokens_for_shape_check.shape})"
661 )
663 # If incorrect_tokens was provided, take the logit difference
664 logit_directions = logit_directions - self.model.tokens_to_residual_directions(
665 incorrect_tokens
666 )
668 scaled_residual_stack = self.apply_ln_to_stack(
669 residual_stack,
670 layer=-1,
671 pos_slice=pos_slice,
672 batch_slice=batch_slice,
673 has_batch_dim=has_batch_dim,
674 )
676 if (
677 has_batch_dim
678 and batch_slice.mode != "int"
679 and scaled_residual_stack.ndim == 4
680 and logit_directions.ndim == 2
681 and logit_directions.shape[0] == scaled_residual_stack.shape[1]
682 ):
683 # Per-example directions must broadcast over positions, not across the batch.
684 logit_directions = logit_directions.unsqueeze(-2)
686 logit_attrs = (scaled_residual_stack * logit_directions).sum(dim=-1)
687 return logit_attrs
689 def decompose_resid(
690 self,
691 layer: Optional[int] = None,
692 mlp_input: bool = False,
693 mode: Literal["all", "mlp", "attn"] = "all",
694 apply_ln: bool = False,
695 pos_slice: Union[Slice, SliceInput] = None,
696 incl_embeds: bool = True,
697 return_labels: bool = False,
698 ) -> Union[
699 Float[torch.Tensor, "layers_covered *batch_and_pos_dims d_model"],
700 Tuple[Float[torch.Tensor, "layers_covered *batch_and_pos_dims d_model"], List[str]],
701 ]:
702 """Decompose the Residual Stream.
704 Decomposes the residual stream input to layer L into a stack of the output of previous
705 layers. The sum of these is the input to layer L (plus embedding and pos embedding). This is
706 useful for attributing model behaviour to different components of the residual stream
708 Args:
709 layer:
710 The layer to take components up to - by default includes
711 resid_pre for that layer and excludes resid_mid and resid_post for that layer.
712 layer==n_layers means to return all layer outputs incl in the final layer, layer==0
713 means just embed and pos_embed. The indices are taken such that this gives the
714 accumulated streams up to the input to layer l
715 mlp_input:
716 Whether to include attn_out for the current
717 layer - essentially decomposing the residual stream that's input to the MLP input
718 rather than the Attn input.
719 mode:
720 Values are "all", "mlp" or "attn". "all" returns all
721 components, "mlp" returns only the MLP components, and "attn" returns only the
722 attention components. Defaults to "all".
723 apply_ln:
724 Whether to apply LayerNorm to the stack.
725 pos_slice:
726 A slice object to apply to the pos dimension.
727 Defaults to None, do nothing.
728 incl_embeds:
729 Whether to include embed & pos_embed, whose hooks expose their
730 contributions to the residual stream.
731 return_labels:
732 Whether to return a list of labels for the residual stream components.
733 Useful for labelling graphs.
735 Returns:
736 A tensor of the accumulated residual streams. If `return_labels` is True, also returns
737 a list of labels for the components (as a tuple in the form `(components, labels)`).
738 """
739 if not isinstance(pos_slice, Slice):
740 pos_slice = Slice(pos_slice)
741 pos_slice = cast(Slice, pos_slice) # mypy can't seem to infer this
742 if layer is None or layer == -1:
743 # Default to the residual stream immediately pre unembed
744 layer = self.model.cfg.n_layers
745 assert isinstance(layer, int)
747 incl_attn = mode != "mlp"
748 incl_mlp = mode != "attn" and not self.model.cfg.attn_only
749 components_list = []
750 labels = []
751 if incl_embeds:
752 if self.has_embed: 752 ↛ 755line 752 didn't jump to line 755 because the condition on line 752 was always true
753 components_list = [self["hook_embed"]]
754 labels.append("embed")
755 if self.has_pos_embed:
756 components_list.append(self["hook_pos_embed"])
757 labels.append("pos_embed")
759 for l in range(layer):
760 if incl_attn:
761 components_list.append(self[("attn_out", l)])
762 labels.append(f"{l}_attn_out")
763 if incl_mlp:
764 components_list.append(self[("mlp_out", l)])
765 labels.append(f"{l}_mlp_out")
766 if mlp_input and incl_attn:
767 components_list.append(self[("attn_out", layer)])
768 labels.append(f"{layer}_attn_out")
769 components_list = [pos_slice.apply(c, dim=-2) for c in components_list]
770 components = torch.stack(components_list, dim=0)
771 if apply_ln:
772 components = self.apply_ln_to_stack(
773 components, layer, pos_slice=pos_slice, mlp_input=mlp_input
774 )
775 if return_labels:
776 return components, labels
777 else:
778 return components
780 def compute_head_results(
781 self,
782 ):
783 """Compute Head Results.
785 Computes and caches the results for each attention head, ie the amount contributed to the
786 residual stream from that head. attn_out for a layer is the sum of head results plus b_O.
787 Intended use is to enable use_attn_results when running and caching the model, but this can
788 be useful if you forget.
790 Existing per-head results are preserved, including those without a batch dimension.
791 Missing or invalid results are recomputed from cached ``hook_z`` and the model's
792 ``W_O`` weights. A per-head contraction avoids materializing the broadcast product
793 over both ``d_head`` and ``d_model``. As in the bridge forward pass, the contraction
794 follows active autocast settings.
796 TransformerBridge exposes ``blocks[i].attn.W_O`` via its component-mapping
797 compatibility shim.
798 """
799 if self._ensure_head_results():
800 logging.warning("Tried to compute head results when they were already cached")
802 def _ensure_head_results(self) -> bool:
803 """Recompute missing or invalid per-head results; True if all were already cached."""
804 expected_ndim = 4 if self.has_batch_dim else 3
805 all_cached = True
806 for layer in range(self.model.cfg.n_layers):
807 result_key = f"blocks.{layer}.attn.hook_result"
808 cached_result = self.cache_dict.get(result_key)
809 # Batchless per-head results are valid 3D tensors.
810 if (
811 isinstance(cached_result, torch.Tensor)
812 and cached_result.ndim == expected_ndim
813 and cached_result.shape[-2] == self.model.cfg.n_heads
814 ):
815 continue
816 all_cached = False
818 z = self[("z", layer, "attn")]
819 weights = self.model.blocks[layer].attn.W_O
820 if z.dtype != weights.dtype:
821 # Preserve broadcast promotion for caches collected under autocast.
822 dtype = torch.promote_types(z.dtype, weights.dtype)
823 z, weights = z.to(dtype), weights.to(dtype)
824 result = torch.einsum("...hd,hdm->...hm", z, weights)
825 scale = getattr(self.model.blocks[layer].attn, "_result_contribution_scale", 1.0)
826 self.cache_dict[result_key] = result if scale == 1.0 else result * scale
827 return all_cached
829 def ssm_layers(self, mixer_type: Optional[Union[type, Tuple[type, ...]]] = None) -> List[int]:
830 """Return the block indices whose mixer is an SSM / recurrent mixer.
832 Family-agnostic and purely structural: finds each block's *realized* SSM
833 mixer in its variant slot (``.mixer`` / ``.linear_attn`` / …) via
834 ``find_ssm_mixer``, which excludes a hybrid's passthrough ``.mixer`` on
835 non-SSM layers (e.g. NemotronH attention/MLP/MoE) — no dependence on
836 ``cfg.layers_block_type``.
838 Args:
839 mixer_type: Optional concrete bridge class to filter to (e.g.
840 ``SSM2MixerBridge`` for only Mamba-2 layers).
842 Returns:
843 Ascending list of SSM block indices.
844 """
845 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
846 find_ssm_mixer,
847 )
849 bridge = self.model
850 layers: List[int] = []
851 for i, block in enumerate(bridge.blocks):
852 mixer = find_ssm_mixer(block)
853 if mixer is None:
854 continue
855 if mixer_type is not None and not isinstance(mixer, mixer_type):
856 continue
857 layers.append(i)
858 return layers
860 def _over_ssm_layers(
861 self,
862 fn: Callable[[Any, int], torch.Tensor],
863 mixer_type: Optional[Union[type, Tuple[type, ...]]] = None,
864 ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]:
865 """Apply ``fn(mixer, layer_idx)`` over every SSM layer; stack or dict.
867 The single SSM layer-enumeration used by both ``compute_ssm_state`` and
868 ``compute_ssm_effective_attention``. Returns a stacked tensor (dim 0 =
869 layer) when *every* block is an SSM layer, else a ``{layer_idx: result}``
870 dict over the SSM layers (heterogeneous hybrids).
871 """
872 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
873 find_ssm_mixer,
874 )
876 bridge = self.model
877 indices = self.ssm_layers(mixer_type=mixer_type)
878 if not indices: 878 ↛ 879line 878 didn't jump to line 879 because the condition on line 878 was never true
879 raise RuntimeError(
880 "No SSM mixer layers found. Use an SSM / hybrid bridge and "
881 "run_with_cache first (gated-delta-net needs use_cache=False)."
882 )
883 results: Dict[int, torch.Tensor] = {
884 i: fn(cast(Any, find_ssm_mixer(bridge.blocks[i])), i) for i in indices
885 }
886 if indices == list(range(len(bridge.blocks))):
887 return torch.stack([results[i] for i in indices], dim=0)
888 return results
890 def compute_ssm_effective_attention(
891 self,
892 layer: Optional[int] = None,
893 include_dt_scaling: bool = False,
894 per_state_coord: bool = False,
895 ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]:
896 """Materialize SSM effective attention for one or all SSM layers, any family.
898 The single discovery surface for effective attention across Mamba-1,
899 Mamba-2, and gated-delta-net layers; dispatches to each layer's mixer
900 ``compute_effective_attention`` regardless of family. Family-specific
901 options are forwarded only to mixers whose signature accepts them.
903 Args:
904 layer: Specific block index, or None for every SSM layer.
905 include_dt_scaling: Forwarded to Mamba-1/Mamba-2 mixers (the
906 reconstruction form); gated-delta-net does not accept it.
907 per_state_coord: Forwarded to Mamba-1 only (per-state-coordinate
908 matrices). Setting it True for another family raises ValueError.
910 Returns:
911 A per-layer matrix for a single ``layer``; for ``layer=None`` a
912 stacked tensor (dim 0 = layer) when every block is an SSM layer, else
913 a ``{layer_idx: matrix}`` dict over the SSM layers.
915 Raises:
916 TypeError: If the requested ``layer`` has no SSM mixer.
917 ValueError: If an option is set that the target mixer does not support.
918 RuntimeError: If ``layer=None`` finds no SSM layers.
919 """
920 import inspect
922 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
923 find_ssm_mixer,
924 )
926 def _call(mixer: Any, layer_idx: int) -> torch.Tensor:
927 fn = cast(Any, mixer).compute_effective_attention
928 params = inspect.signature(fn).parameters
929 kwargs: Dict[str, Any] = {}
930 for name, value in (
931 ("include_dt_scaling", include_dt_scaling),
932 ("per_state_coord", per_state_coord),
933 ):
934 if name in params:
935 kwargs[name] = value
936 elif value:
937 raise ValueError(
938 f"{type(mixer).__name__}.compute_effective_attention does not "
939 f"support {name}=True."
940 )
941 result: torch.Tensor = fn(cache=self, layer_idx=layer_idx, **kwargs)
942 return result
944 if layer is not None:
945 single = find_ssm_mixer(self.model.blocks[layer])
946 if single is None:
947 raise TypeError(f"Block {layer} has no SSM mixer (no compute_effective_attention).")
948 return _call(single, layer)
949 return self._over_ssm_layers(_call)
951 def compute_ssm_state(
952 self,
953 layer: Optional[int] = None,
954 time_step: Optional[int] = None,
955 ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]:
956 """Reconstruct the recurrent SSM state ``S`` from this cache.
958 The single discovery surface for recurrent state across families —
959 ``SSMMixerBridge`` (Mamba-1), ``SSM2MixerBridge`` (Mamba-2) and
960 ``GatedDeltaNetBridge`` (gated delta rule) each reconstruct it — mirroring
961 ``compute_head_results``. Read-only post-hoc reconstruction from cached
962 hooks (no forward re-run); requires an SSM / SSM-hybrid bridge cached via
963 ``run_with_cache`` (gated-delta-net additionally needs ``use_cache=False``
964 so its interior hooks fire). See the mixer's ``compute_ssm_state`` for the
965 recurrence, shapes, and the ``time_step`` memory bound; state shape is
966 family-specific, so ``layer=None`` returns a dict except when every block
967 shares one mixer type.
969 Args:
970 layer: Specific block index, or None for every SSM-state layer.
971 time_step: Optional single position (memory-bounded); None for all.
973 Returns:
974 A per-layer state tensor for a single ``layer``; for ``layer=None`` a
975 stacked tensor (dim 0 = layer) when every block is an SSM-state layer,
976 else a ``{layer_idx: state}`` dict over those layers.
978 Raises:
979 TypeError: If the requested ``layer`` has no state-reconstructing mixer.
980 RuntimeError: If ``layer=None`` finds no such layers.
981 """
982 from transformer_lens.model_bridge.generalized_components import (
983 GatedDeltaNetBridge,
984 SSM2MixerBridge,
985 SSMMixerBridge,
986 )
987 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
988 find_ssm_mixer,
989 )
991 # Every recurrent family reconstructs S_t from its cached interior hooks.
992 state_mixers = (SSMMixerBridge, SSM2MixerBridge, GatedDeltaNetBridge)
994 def _call(mixer: Any, layer_idx: int) -> torch.Tensor:
995 state: torch.Tensor = cast(Any, mixer).compute_ssm_state(
996 self, layer_idx=layer_idx, time_step=time_step
997 )
998 return state
1000 if layer is not None:
1001 single = find_ssm_mixer(self.model.blocks[layer])
1002 if not isinstance(single, state_mixers):
1003 raise TypeError(
1004 f"Block {layer} has no state-reconstructing SSM mixer; "
1005 "compute_ssm_state supports Mamba-1 / Mamba-2 / gated-delta-net."
1006 )
1007 return _call(single, layer)
1008 return self._over_ssm_layers(_call, mixer_type=state_mixers)
1010 def stack_head_results(
1011 self,
1012 layer: int = -1,
1013 return_labels: bool = False,
1014 incl_remainder: bool = False,
1015 pos_slice: Union[Slice, SliceInput] = None,
1016 apply_ln: bool = False,
1017 ) -> Union[
1018 Float[torch.Tensor, "num_components *batch_and_pos_dims d_model"],
1019 Tuple[Float[torch.Tensor, "num_components *batch_and_pos_dims d_model"], List[str]],
1020 ]:
1021 """Stack Head Results.
1023 Returns a stack of all head results (ie residual stream contribution) up to layer L. A good
1024 way to decompose the outputs of attention layers into attribution by specific heads. Note
1025 that the num_components axis has length layer x n_heads ((layer head_index) in einops
1026 notation).
1028 Args:
1029 layer:
1030 Layer index - heads at all layers strictly before this are included. layer must be
1031 in [1, n_layers-1], or any of (n_layers, -1, None), which all mean the final layer.
1032 return_labels:
1033 Whether to also return a list of labels of the form "L0H0" for the heads.
1034 incl_remainder:
1035 Whether to return a final term which is "the rest of the residual stream".
1036 pos_slice:
1037 A slice object to apply to the pos dimension. Defaults to None, do nothing.
1038 apply_ln:
1039 Whether to apply LayerNorm to the stack.
1040 """
1041 if not isinstance(pos_slice, Slice):
1042 pos_slice = Slice(pos_slice)
1043 pos_slice = cast(Slice, pos_slice) # mypy can't seem to infer this
1044 if layer is None or layer == -1:
1045 # Default to the residual stream immediately pre unembed
1046 layer = self.model.cfg.n_layers
1048 # Idempotent; cleans up stale Bridge entries without the 'already cached' warning,
1049 # which is only meaningful for direct compute_head_results callers.
1050 self._ensure_head_results()
1052 components: Any = []
1053 labels = []
1054 for l in range(layer):
1055 # Note that this has shape batch x pos x head_index x d_model
1056 components.append(pos_slice.apply(self[("result", l, "attn")], dim=-3))
1057 labels.extend([f"L{l}H{h}" for h in range(self.model.cfg.n_heads)])
1058 if components:
1059 components = torch.cat(components, dim=-2)
1060 components = einops.rearrange(
1061 components,
1062 "... concat_head_index d_model -> concat_head_index ... d_model",
1063 )
1064 if incl_remainder:
1065 remainder = pos_slice.apply(
1066 self[("resid_post", layer - 1)], dim=-2
1067 ) - components.sum(dim=0)
1068 components = torch.cat([components, remainder[None]], dim=0)
1069 labels.append("remainder")
1070 elif incl_remainder:
1071 # There are no components, so the remainder is the entire thing.
1072 components = torch.cat(
1073 [pos_slice.apply(self[("resid_post", layer - 1)], dim=-2)[None]], dim=0
1074 )
1075 labels.append("remainder")
1076 else:
1077 # If this is called with layer 0, we return an empty tensor of the right shape to be
1078 # stacked correctly. This uses the shape of hook_embed, which is pretty janky since it
1079 # assumes embed is in the cache. But it's hard to explicitly code the shape, since it
1080 # depends on the pos slice, whether we have a batch dim, etc. And it's pretty messy!
1081 empty_src = pos_slice.apply(self["hook_embed"], dim=-2)
1082 components = torch.zeros(
1083 0,
1084 *empty_src.shape,
1085 dtype=empty_src.dtype,
1086 device=empty_src.device,
1087 )
1089 if apply_ln:
1090 components = self.apply_ln_to_stack(components, layer, pos_slice=pos_slice)
1092 if return_labels:
1093 return components, labels
1094 else:
1095 return components
1097 def stack_activation(
1098 self,
1099 activation_name: str,
1100 layer: int = -1,
1101 sublayer_type: Optional[str] = None,
1102 ) -> Float[torch.Tensor, "layers_covered ..."]:
1103 """Stack Activations.
1105 Flexible way to stack activations with a given name.
1107 Args:
1108 activation_name:
1109 The name of the activation to be stacked
1110 layer:
1111 'Layer index - heads' at all layers strictly before this are included. layer must be
1112 in [1, n_layers-1], or any of (n_layers, -1, None), which all mean the final layer.
1113 sublayer_type:
1114 The sub layer type of the activation, passed to utils.get_act_name. Can normally be
1115 inferred.
1116 incl_remainder:
1117 Whether to return a final term which is "the rest of the residual stream".
1118 """
1119 if layer is None or layer == -1:
1120 # Default to the residual stream immediately pre unembed
1121 layer = self.model.cfg.n_layers
1123 components = []
1124 for l in range(layer):
1125 components.append(self[(activation_name, l, sublayer_type)])
1127 return torch.stack(components, dim=0)
1129 def get_neuron_results(
1130 self,
1131 layer: int,
1132 neuron_slice: Union[Slice, SliceInput] = None,
1133 pos_slice: Union[Slice, SliceInput] = None,
1134 project_output_onto: Optional[torch.Tensor] = None,
1135 ) -> torch.Tensor:
1136 """Get Neuron Results.
1138 Get the results of for neurons in a specific layer (i.e, how much each neuron contributes to
1139 the residual stream). Does it for the subset of neurons specified by neuron_slice, defaults
1140 to all of them. Does *not* cache these because it's expensive in space and cheap to compute.
1142 Args:
1143 layer:
1144 Layer index.
1145 neuron_slice:
1146 Slice of the neuron.
1147 pos_slice:
1148 Slice of the positions.
1149 project_output_onto:
1150 Optional ``[d_model]`` or ``[d_model, num_outputs]`` projection. Contracted with
1151 ``W_out`` *before* the per-neuron expansion so the ``[..., d_mlp, d_model]``
1152 intermediate is never materialized.
1154 Returns:
1155 Last-dim is ``d_model`` (default), ``num_outputs`` (2D projection), or squeezed
1156 (1D projection).
1157 """
1158 if not isinstance(neuron_slice, Slice):
1159 neuron_slice = Slice(neuron_slice)
1160 if not isinstance(pos_slice, Slice):
1161 pos_slice = Slice(pos_slice)
1163 neuron_acts = self[("post", layer, "mlp")]
1164 block = self.model.blocks[layer]
1165 W_out = block.mlp.W_out
1166 contribution_scale = getattr(block, "residual_contribution_scale", 1.0)
1167 if pos_slice is not None: 1167 ↛ 1171line 1167 didn't jump to line 1171 because the condition on line 1167 was always true
1168 # Note - order is important, as Slice.apply *may* collapse a dimension, so this ensures
1169 # that position dimension is -2 when we apply position slice
1170 neuron_acts = pos_slice.apply(neuron_acts, dim=-2)
1171 if neuron_slice is not None: 1171 ↛ 1174line 1171 didn't jump to line 1174 because the condition on line 1171 was always true
1172 neuron_acts = neuron_slice.apply(neuron_acts, dim=-1)
1173 W_out = neuron_slice.apply(W_out, dim=0)
1174 if contribution_scale != 1.0:
1175 neuron_acts = neuron_acts * contribution_scale
1176 if project_output_onto is None:
1177 return neuron_acts[..., None] * W_out
1178 # W_out: [d_mlp, d_model]; project: [d_model] or [d_model, n_outs]
1179 projected = W_out @ project_output_onto
1180 if project_output_onto.ndim == 1:
1181 return neuron_acts * projected
1182 return neuron_acts[..., None] * projected
1184 def _uses_ln_centering(self, layer: Optional[int]) -> bool:
1185 """Use the target norm's semantics, not just the block normalization family."""
1186 is_final = layer is None or layer == self.model.cfg.n_layers
1187 return self.model.cfg.normalization_type in ["LN", "LNPre"] and not (
1188 is_final and getattr(self.model.cfg, "final_rms", False)
1189 )
1191 def _get_cached_ln_scale(
1192 self,
1193 layer: Optional[int],
1194 mlp_input: bool,
1195 pos_slice: Slice,
1196 batch_slice: Optional[Slice] = None,
1197 ) -> torch.Tensor:
1198 """Look up the cached LN scale and apply pos/batch slicing. Surfaces a clearer error
1199 when the expected hook isn't in the cache (some non-decoder-only architectures expose
1200 LN scale at a different path or not at all).
1201 """
1202 if layer == self.model.cfg.n_layers or layer is None:
1203 key = "ln_final.hook_scale"
1204 else:
1205 key = f"blocks.{layer}.ln{2 if mlp_input else 1}.hook_scale"
1206 try:
1207 scale = self[key]
1208 except KeyError as e:
1209 norm = f"ln{2 if mlp_input else 1}"
1210 blocks = getattr(self.model, "blocks", None)
1211 block = blocks[layer] if blocks is not None and layer is not None else None
1212 if ( 1212 ↛ 1223line 1212 didn't jump to line 1223 because the condition on line 1212 was always true
1213 block is not None
1214 and getattr(block, norm, None) is None
1215 and getattr(block, f"{norm}_post", None) is not None
1216 ):
1217 sublayer = "MLP" if mlp_input else "attention"
1218 raise KeyError(
1219 f"Layer {layer} has no {norm} (pre-{sublayer} norm): this architecture "
1220 f"normalizes sublayer outputs ({norm}_post) instead, so a mid-layer cached-scale "
1221 f"LN is undefined. Use apply_ln=False, or apply LN at the final layer."
1222 ) from e
1223 raise KeyError(
1224 f"Cached LN scale not found at '{key}'. apply_ln operations require the model "
1225 f"to have cached this hook (some non-decoder-only architectures expose LN scale "
1226 f"under different module paths)."
1227 ) from e
1228 scale = pos_slice.apply(scale, dim=-2)
1229 if batch_slice is not None and self.has_batch_dim:
1230 scale = batch_slice.apply(scale)
1231 return scale
1233 def _stack_neuron_results_apply_ln_projected(
1234 self,
1235 layer: int,
1236 pos_slice: Slice,
1237 neuron_slice: Slice,
1238 project_2d: torch.Tensor,
1239 mlp_input: bool = False,
1240 ) -> torch.Tensor:
1241 """LN-applied neuron stack with projection folded in — no d_mlp×d_model intermediate.
1243 Analytical formula (LN models, cached scale ``s``):
1244 ``LN_s(a_n * W_out_n) @ p = (a_n / s) * (W_out_n @ p - mean(W_out_n) * sum_p)``
1245 RMS models drop the ``mean(W_out_n) * sum_p`` term (no centering). ``s`` is the ln2
1246 scale of ``layer`` when ``mlp_input`` is set, else the ln1 scale.
1247 """
1248 scale = self._get_cached_ln_scale(layer, mlp_input=mlp_input, pos_slice=pos_slice)
1250 apply_centering = self._uses_ln_centering(layer)
1251 sum_p = project_2d.sum(dim=0) if apply_centering else None # [n_outs]
1253 components: list = []
1254 for l in range(layer):
1255 block = self.model.blocks[l]
1256 W_out_l = block.mlp.W_out # [d_mlp, d_model]
1257 contribution_scale = getattr(block, "residual_contribution_scale", 1.0)
1258 W_out_l_sliced = neuron_slice.apply(W_out_l, dim=0)
1259 W_proj_l = W_out_l_sliced @ project_2d # [d_mlp, n_outs]
1260 if apply_centering:
1261 assert sum_p is not None # set when apply_centering, narrow for mypy
1262 W_means_l = W_out_l_sliced.mean(dim=-1) # [d_mlp]
1263 lin_form_l = W_proj_l - W_means_l[:, None] * sum_p[None, :]
1264 else:
1265 lin_form_l = W_proj_l
1266 if contribution_scale != 1.0:
1267 lin_form_l = lin_form_l * contribution_scale
1268 a_l = self[("post", l, "mlp")]
1269 a_l = pos_slice.apply(a_l, dim=-2)
1270 a_l = neuron_slice.apply(a_l, dim=-1)
1271 # (a_l / s)[..., None] is [..., d_mlp, 1]; broadcast with lin_form_l [d_mlp, n_outs]
1272 components.append((a_l / scale)[..., None] * lin_form_l)
1273 if not components:
1274 empty_src = pos_slice.apply(self["hook_embed"], dim=-2)
1275 return torch.zeros(
1276 0,
1277 *empty_src.shape[:-1],
1278 project_2d.shape[-1],
1279 dtype=torch.promote_types(
1280 torch.promote_types(empty_src.dtype, scale.dtype), project_2d.dtype
1281 ),
1282 device=empty_src.device,
1283 )
1284 stacked = torch.cat(components, dim=-2)
1285 return einops.rearrange(
1286 stacked, "... concat_neuron_index n_outs -> concat_neuron_index ... n_outs"
1287 )
1289 def stack_neuron_results(
1290 self,
1291 layer: int,
1292 pos_slice: Union[Slice, SliceInput] = None,
1293 neuron_slice: Union[Slice, SliceInput] = None,
1294 return_labels: bool = False,
1295 incl_remainder: bool = False,
1296 apply_ln: bool = False,
1297 project_output_onto: Optional[torch.Tensor] = None,
1298 mlp_input: bool = False,
1299 ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[str]]]:
1300 """Stack Neuron Results
1302 Returns a stack of all neuron results (ie residual stream contribution) up to layer L - ie
1303 the amount each individual neuron contributes to the residual stream. Also returns a list of
1304 labels of the form "L0N0" for the neurons. A good way to decompose the outputs of MLP layers
1305 into attribution by specific neurons.
1307 Note that doing this for all neurons is SUPER expensive on GPU memory and only works for
1308 small models or short inputs. Pass ``project_output_onto`` to fold the projection into the
1309 per-neuron expansion and avoid the ``[..., d_mlp, d_model]`` intermediate.
1311 Args:
1312 layer:
1313 Layer index - heads at all layers strictly before this are included. layer must be
1314 in [1, n_layers]. ``layer=0`` follows the negative-index convention: the neuron
1315 stack is empty and with ``incl_remainder=True`` the remainder is the final layer's
1316 residual stream (``resid_post`` of layer ``n_layers - 1``), i.e. the whole model.
1317 pos_slice:
1318 Slice of the positions.
1319 neuron_slice:
1320 Slice of the neurons. An int or an integer-mode ``Slice(n)`` selects a single
1321 neuron and is treated like ``[n]``.
1322 return_labels:
1323 Whether to also return a list of labels of the form "L0N0" for the neurons.
1324 incl_remainder:
1325 Whether to return a final term which is "the rest of the residual stream".
1326 apply_ln:
1327 Whether to apply LayerNorm to the stack.
1328 project_output_onto:
1329 Optional ``[d_model]`` or ``[d_model, num_outputs]`` tensor. When set, each
1330 component's last d_model dim is replaced by the projection (memory-efficient for
1331 direction analyses; see ``get_neuron_results``). Combined with ``apply_ln=True``,
1332 the projection is folded into the analytical cached-scale LN so the
1333 ``[..., d_mlp, d_model]`` intermediate is still never materialized.
1334 mlp_input:
1335 Treat the stack as the input to ``layer``'s MLP rather than its attention, as in
1336 ``decompose_resid``: with ``incl_remainder=True`` the remainder fills the stack up to
1337 ``resid_mid`` instead of ``resid_pre``, and with ``apply_ln=True`` the LN scale of
1338 ln2 is used instead of ln1. The neurons stacked are unchanged.
1339 """
1340 if layer is None or layer == -1:
1341 # Default to the residual stream immediately pre unembed
1342 layer = self.model.cfg.n_layers
1344 # layer=0 is the documented whole-model case (remainder reads resid_post of layer -1, the
1345 # final layer). Anything below it would wrap the remainder target to an arbitrary late
1346 # layer via __getitem__, silently returning the wrong residual stream.
1347 if incl_remainder and layer < 0:
1348 raise ValueError(
1349 f"layer must be >= 0 with incl_remainder=True (-1/None mean n_layers); got {layer}."
1350 )
1352 # Layers are concatenated along the neuron axis, so it must not collapse. unwrap turns an
1353 # int into [n] but leaves a Slice as is, so an integer-mode Slice(n) is rebuilt as [n]
1354 # here. A new Slice is made, so the caller's Slice still collapses elsewhere.
1355 neuron_slice = Slice.unwrap(neuron_slice)
1356 if neuron_slice.mode == "int":
1357 assert isinstance(neuron_slice.slice, int) # narrow for mypy
1358 neuron_slice = Slice([neuron_slice.slice])
1359 if not isinstance(pos_slice, Slice):
1360 pos_slice = Slice(pos_slice)
1362 project_2d, squeeze_projected = _normalize_projection_to_2d(project_output_onto)
1364 d_mlp = self.model.cfg.d_mlp
1365 assert d_mlp is not None, "model.cfg.d_mlp must be set"
1366 neuron_labels = neuron_slice.apply(torch.arange(d_mlp), dim=0)
1368 labels = [f"L{l}N{h}" for l in range(layer) for h in neuron_labels]
1369 # The remainder fills the stack up to the residual stream that layer's LN normalizes:
1370 # resid_mid (the MLP input) with mlp_input, else resid_pre (resid_post of layer - 1).
1371 remainder_target = ("resid_mid", layer) if mlp_input else ("resid_post", layer - 1)
1372 components: Any
1373 ln_folded = apply_ln and project_2d is not None
1374 if ln_folded:
1375 assert project_2d is not None # narrow for mypy
1376 # Analytical LN+projection — no d_mlp×d_model intermediate.
1377 components = self._stack_neuron_results_apply_ln_projected(
1378 layer, pos_slice, neuron_slice, project_2d, mlp_input=mlp_input
1379 )
1380 if incl_remainder:
1381 # Linearity of cached-scale LN: remainder is LN_s(target) @ p - sum(neurons).
1382 target = pos_slice.apply(self[remainder_target], dim=-2)
1383 target_ln = self.apply_ln_to_stack(
1384 target[None], layer, pos_slice=pos_slice, mlp_input=mlp_input
1385 )[0]
1386 remainder = target_ln @ project_2d
1387 if components.shape[0] > 0: 1387 ↛ 1389line 1387 didn't jump to line 1389 because the condition on line 1387 was always true
1388 remainder = remainder - components.sum(dim=0)
1389 components = torch.cat([components, remainder[None]], dim=0)
1390 labels.append("remainder")
1391 else:
1392 per_layer: list = []
1393 for l in range(layer):
1394 per_layer.append(
1395 self.get_neuron_results(
1396 l,
1397 pos_slice=pos_slice,
1398 neuron_slice=neuron_slice,
1399 project_output_onto=project_2d,
1400 )
1401 )
1402 if per_layer:
1403 components = torch.cat(per_layer, dim=-2)
1404 components = einops.rearrange(
1405 components,
1406 "... concat_neuron_index d_model -> concat_neuron_index ... d_model",
1407 )
1408 if incl_remainder:
1409 remainder_full = pos_slice.apply(self[remainder_target], dim=-2)
1410 if project_2d is not None:
1411 remainder_full = remainder_full @ project_2d
1412 remainder = remainder_full - components.sum(dim=0)
1413 components = torch.cat([components, remainder[None]], dim=0)
1414 labels.append("remainder")
1415 elif incl_remainder:
1416 remainder_full = pos_slice.apply(self[remainder_target], dim=-2)
1417 if project_2d is not None: 1417 ↛ 1418line 1417 didn't jump to line 1418 because the condition on line 1417 was never true
1418 remainder_full = remainder_full @ project_2d
1419 components = torch.cat([remainder_full[None]], dim=0)
1420 labels.append("remainder")
1421 else:
1422 empty_src = pos_slice.apply(self["hook_embed"], dim=-2)
1423 empty_shape = tuple(empty_src.shape)
1424 empty_dtype = empty_src.dtype
1425 if project_2d is not None:
1426 # Shape only: a real matmul would demand the cache and direction share a dtype.
1427 empty_shape = (*empty_shape[:-1], project_2d.shape[-1])
1428 empty_dtype = torch.promote_types(empty_dtype, project_2d.dtype)
1429 components = torch.zeros(
1430 0, *empty_shape, dtype=empty_dtype, device=empty_src.device
1431 )
1433 if apply_ln:
1434 components = self.apply_ln_to_stack(
1435 components, layer, pos_slice=pos_slice, mlp_input=mlp_input
1436 )
1438 if squeeze_projected:
1439 components = components.squeeze(-1)
1441 if return_labels:
1442 return components, labels
1443 else:
1444 return components
1446 def apply_ln_to_stack(
1447 self,
1448 residual_stack: Float[torch.Tensor, "num_components *batch_and_pos_dims d_model"],
1449 layer: Optional[int] = None,
1450 mlp_input: bool = False,
1451 pos_slice: Union[Slice, SliceInput] = None,
1452 batch_slice: Union[Slice, SliceInput] = None,
1453 has_batch_dim: bool = True,
1454 recompute_ln: bool = False,
1455 ) -> Float[torch.Tensor, "num_components *batch_and_pos_dims_out d_model"]:
1456 """Apply Layer Norm to a Stack.
1458 Takes a stack of components of the residual stream (eg outputs of decompose_resid or
1459 accumulated_resid), treats them as the input to a specific layer, and applies the layer norm
1460 scaling of that layer to them, using the cached scale factors - simulating what that
1461 component of the residual stream contributes to that layer's input.
1463 The layernorm scale is global across the entire residual stream for each layer, batch
1464 element and position, which is why we need to use the cached scale factors rather than just
1465 applying a new LayerNorm.
1467 When recompute_ln=True and the target layer is the final layer (unembed), each
1468 component is normalized using stats recomputed from that component; use this for logit lens
1469 analysis. When recompute_ln=False, a single cached scale is used for all components.
1471 If the model does not use LayerNorm or RMSNorm, it returns the residual stack unchanged.
1473 Args:
1474 residual_stack:
1475 A tensor, whose final dimension is d_model. The other trailing dimensions are
1476 assumed to be the same as the stored hook_scale - which may or may not include batch
1477 or position dimensions.
1478 layer:
1479 The layer we're taking the input to. In [0, n_layers], n_layers means the unembed.
1480 None maps to the n_layers case, ie the unembed.
1481 mlp_input:
1482 Whether the input is to the MLP or attn (ie ln2 vs ln1). Defaults to False, ie ln1.
1483 If layer==n_layers, must be False, and we use ln_final
1484 pos_slice:
1485 The slice to take of positions, if residual_stack is not over the full context, None
1486 means do nothing. It is assumed that pos_slice has already been applied to
1487 residual_stack, and this is only applied to the scale. See utils.Slice for details.
1488 Defaults to None, do nothing.
1489 batch_slice:
1490 The slice to take on the batch dimension. Defaults to None, do nothing.
1491 has_batch_dim:
1492 Whether residual_stack has a batch dimension.
1493 recompute_ln:
1494 If True and target layer is the unembed (final layer), apply the final layer norm
1495 to each component with statistics recomputed from that component. Defaults to False.
1497 """
1498 if not isinstance(pos_slice, Slice):
1499 pos_slice = Slice(pos_slice)
1500 if not isinstance(batch_slice, Slice):
1501 batch_slice = Slice(batch_slice)
1503 if layer is None or layer == -1:
1504 # Default to the residual stream immediately pre unembed
1505 layer = self.model.cfg.n_layers
1507 if has_batch_dim:
1508 # Apply batch slice to the stack
1509 residual_stack = batch_slice.apply(residual_stack, dim=1)
1511 # Checked only after slicing so callers like logit_attrs still get their
1512 # batch_slice honoured on models without LayerNorm/RMSNorm.
1513 if self.model.cfg.normalization_type not in ["LN", "LNPre", "RMS", "RMSPre"]:
1514 return residual_stack
1516 # Logit lens: apply final layer norm to each component with recomputed statistics
1517 if recompute_ln and layer == self.model.cfg.n_layers and hasattr(self.model, "ln_final"):
1518 ln_final = self.model.ln_final
1519 results = []
1520 for i in range(residual_stack.shape[0]):
1521 x = residual_stack[i]
1522 original_shape = x.shape
1523 # ln_final expects (batch, pos, d_model); restore missing structural dimensions
1524 if not has_batch_dim:
1525 x = x.unsqueeze(0)
1526 if x.ndim == 2:
1527 x = x.unsqueeze(1)
1528 out = ln_final(x)
1529 results.append(out.reshape(original_shape))
1530 return torch.stack(results, dim=0)
1532 # Final RMSNorm must not inherit the block LayerNorm centering convention.
1533 if self._uses_ln_centering(layer):
1534 residual_stack = residual_stack - residual_stack.mean(dim=-1, keepdim=True)
1536 # Shape is [batch, position, 1] or [position, 1]; final dim is a dummy for broadcasting.
1537 scale = self._get_cached_ln_scale(layer, mlp_input, pos_slice, batch_slice)
1539 return residual_stack / scale
1541 def get_full_resid_decomposition(
1542 self,
1543 layer: Optional[int] = None,
1544 mlp_input: bool = False,
1545 expand_neurons: bool = True,
1546 apply_ln: bool = False,
1547 pos_slice: Union[Slice, SliceInput] = None,
1548 return_labels: bool = False,
1549 project_output_onto: Optional[torch.Tensor] = None,
1550 ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[str]]]:
1551 """Get the full Residual Decomposition.
1553 Decomposes the residual stream that is input into some layer into its
1554 constituent components: every attention head result, every neuron (or
1555 MLP layer) result, the embeddings, and the accumulated biases.
1557 The returned tensor stacks components along ``dim=0`` in this order:
1559 1. Attention head results, layer-by-layer (``L * n_heads`` rows)
1560 2. Neuron / MLP results (only if ``cfg.attn_only=False`` and
1561 ``layer > 0``; ``L * d_mlp`` rows when ``expand_neurons=True``,
1562 else ``L`` rows)
1563 3. ``embed`` (1 row, including any post-embedding runtime multiplier)
1564 4. ``pos_embed`` (1 row, if the model has positional embeddings)
1565 5. ``bias`` (1 row, the accumulated layer biases)
1567 ``return_labels=True`` returns a list of strings in the same order, so
1568 ``labels[i]`` always names ``stack[i]``. If you need to extract a
1569 specific component, slice by label rather than by hard-coded index —
1570 the row counts depend on ``layer``, ``expand_neurons``,
1571 ``cfg.attn_only``, and whether the model has positional embeddings.
1573 Args:
1574 layer:
1575 The layer we're inputting into. layer is in [0, n_layers], if layer==n_layers (or
1576 None) we're inputting into the unembed (the entire stream), if layer==0 then it's
1577 just embed and pos_embed
1578 mlp_input:
1579 Are we inputting to the MLP in that layer or the attn? Must be False for final
1580 layer, since that's the unembed.
1581 expand_neurons:
1582 Whether to expand the MLP outputs to give every neuron's result or just return the
1583 MLP layer outputs.
1584 apply_ln:
1585 Whether to apply LayerNorm to the stack.
1586 pos_slice:
1587 Slice of the positions to take.
1588 return_labels:
1589 Whether to return the labels.
1590 project_output_onto:
1591 Optional ``[d_model]`` or ``[d_model, num_outputs]`` projection. Folded in
1592 *before* the per-neuron expansion, so the ``[..., d_mlp, d_model]`` intermediate
1593 is never materialized (memory saving applies only with ``expand_neurons=True``).
1594 Combined with ``apply_ln=True``, the projection is fused into the analytical
1595 cached-scale LN so the same memory benefit holds. Output last-dim is squeezed
1596 for a 1D projection; ``num_outputs`` for 2D.
1597 """
1598 if layer is None or layer == -1:
1599 # Default to the residual stream immediately pre unembed
1600 layer = self.model.cfg.n_layers
1601 assert layer is not None # keep mypy happy
1603 if not isinstance(pos_slice, Slice):
1604 pos_slice = Slice(pos_slice)
1606 project_2d, squeeze_projected = _normalize_projection_to_2d(project_output_onto)
1607 # When both apply_ln and projection are requested, LN is applied per-component (in
1608 # d_model space for the small ones, analytically for neurons) before projection, so the
1609 # final apply_ln_to_stack call is skipped — last-dim is already n_outs.
1610 ln_folded = apply_ln and project_2d is not None
1612 def _ln_then_project(stack: torch.Tensor) -> torch.Tensor:
1613 stack = self.apply_ln_to_stack(stack, layer, pos_slice=pos_slice, mlp_input=mlp_input)
1614 return stack @ project_2d if project_2d is not None else stack
1616 head_stack, head_labels = self.stack_head_results(
1617 layer + (1 if mlp_input else 0), pos_slice=pos_slice, return_labels=True
1618 )
1619 if ln_folded:
1620 head_stack = _ln_then_project(head_stack)
1621 elif project_2d is not None:
1622 head_stack = head_stack @ project_2d
1623 labels = head_labels
1624 components = [head_stack]
1625 if not self.model.cfg.attn_only and layer > 0:
1626 if expand_neurons:
1627 # Only ask stack_neuron_results to apply LN when we want the fused analytical
1628 # path (ln_folded). For the unfolded case the outer apply_ln_to_stack handles it.
1629 neuron_stack, neuron_labels = self.stack_neuron_results(
1630 layer,
1631 pos_slice=pos_slice,
1632 return_labels=True,
1633 apply_ln=ln_folded,
1634 project_output_onto=project_2d,
1635 mlp_input=mlp_input,
1636 )
1637 labels.extend(neuron_labels)
1638 components.append(neuron_stack)
1639 else:
1640 # Get the stack of just the MLP outputs
1641 # mlp_input included for completeness, but it doesn't actually matter, since it's
1642 # just for MLP outputs
1643 mlp_stack, mlp_labels = self.decompose_resid(
1644 layer,
1645 mlp_input=mlp_input,
1646 pos_slice=pos_slice,
1647 incl_embeds=False,
1648 mode="mlp",
1649 return_labels=True,
1650 )
1651 if ln_folded:
1652 mlp_stack = _ln_then_project(mlp_stack)
1653 elif project_2d is not None: 1653 ↛ 1654line 1653 didn't jump to line 1654 because the condition on line 1653 was never true
1654 mlp_stack = mlp_stack @ project_2d
1655 labels.extend(mlp_labels)
1656 components.append(mlp_stack)
1658 if self.has_embed: 1658 ↛ 1666line 1658 didn't jump to line 1666 because the condition on line 1658 was always true
1659 embed = pos_slice.apply(self["hook_embed"], -2)[None]
1660 if ln_folded:
1661 embed = _ln_then_project(embed)
1662 elif project_2d is not None:
1663 embed = embed @ project_2d
1664 labels.append("embed")
1665 components.append(embed)
1666 if self.has_pos_embed:
1667 pos_embed = pos_slice.apply(self["pos_embed"], -2)[None]
1668 if ln_folded:
1669 pos_embed = _ln_then_project(pos_embed)
1670 elif project_2d is not None:
1671 pos_embed = pos_embed @ project_2d
1672 labels.append("pos_embed")
1673 components.append(pos_embed)
1674 # If we didn't expand the neurons, the MLP biases are already included in the MLP outputs.
1675 bias_full = self.model.accumulated_bias(layer, mlp_input, include_mlp_biases=expand_neurons)
1676 if ln_folded:
1677 # Expand bias to per-position d_model shape so LN can center, then project.
1678 expand_shape: tuple = (1,) + tuple(head_stack.shape[1:-1]) + (self.model.cfg.d_model,)
1679 bias = _ln_then_project(bias_full.expand(expand_shape))
1680 else:
1681 if project_2d is not None:
1682 # Bias is [d_model], so project post-hoc for shape compatibility — no memory win here.
1683 bias_full = bias_full @ project_2d
1684 bias = bias_full.expand((1,) + head_stack.shape[1:])
1685 labels.append("bias")
1686 components.append(bias)
1687 residual_stack = torch.cat(components, dim=0)
1688 if apply_ln and not ln_folded:
1689 residual_stack = self.apply_ln_to_stack(
1690 residual_stack, layer, pos_slice=pos_slice, mlp_input=mlp_input
1691 )
1693 if squeeze_projected:
1694 residual_stack = residual_stack.squeeze(-1)
1696 if return_labels:
1697 return residual_stack, labels
1698 else:
1699 return residual_stack