Coverage for transformer_lens/ActivationCache.py: 96%
477 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +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
730 return_labels:
731 Whether to return a list of labels for the residual stream components.
732 Useful for labelling graphs.
734 Returns:
735 A tensor of the accumulated residual streams. If `return_labels` is True, also returns
736 a list of labels for the components (as a tuple in the form `(components, labels)`).
737 """
738 if not isinstance(pos_slice, Slice):
739 pos_slice = Slice(pos_slice)
740 pos_slice = cast(Slice, pos_slice) # mypy can't seem to infer this
741 if layer is None or layer == -1:
742 # Default to the residual stream immediately pre unembed
743 layer = self.model.cfg.n_layers
744 assert isinstance(layer, int)
746 incl_attn = mode != "mlp"
747 incl_mlp = mode != "attn" and not self.model.cfg.attn_only
748 components_list = []
749 labels = []
750 if incl_embeds:
751 if self.has_embed: 751 ↛ 754line 751 didn't jump to line 754 because the condition on line 751 was always true
752 components_list = [self["hook_embed"]]
753 labels.append("embed")
754 if self.has_pos_embed: 754 ↛ 758line 754 didn't jump to line 758 because the condition on line 754 was always true
755 components_list.append(self["hook_pos_embed"])
756 labels.append("pos_embed")
758 for l in range(layer):
759 if incl_attn:
760 components_list.append(self[("attn_out", l)])
761 labels.append(f"{l}_attn_out")
762 if incl_mlp:
763 components_list.append(self[("mlp_out", l)])
764 labels.append(f"{l}_mlp_out")
765 if mlp_input and incl_attn:
766 components_list.append(self[("attn_out", layer)])
767 labels.append(f"{layer}_attn_out")
768 components_list = [pos_slice.apply(c, dim=-2) for c in components_list]
769 components = torch.stack(components_list, dim=0)
770 if apply_ln:
771 components = self.apply_ln_to_stack(
772 components, layer, pos_slice=pos_slice, mlp_input=mlp_input
773 )
774 if return_labels:
775 return components, labels
776 else:
777 return components
779 def compute_head_results(
780 self,
781 ):
782 """Compute Head Results.
784 Computes and caches the results for each attention head, ie the amount contributed to the
785 residual stream from that head. attn_out for a layer is the sum of head results plus b_O.
786 Intended use is to enable use_attn_results when running and caching the model, but this can
787 be useful if you forget.
789 Existing per-head results are preserved, including those without a batch dimension.
790 Missing or invalid results are recomputed from cached ``hook_z`` and the model's
791 ``W_O`` weights. A per-head contraction avoids materializing the broadcast product
792 over both ``d_head`` and ``d_model``. As in the bridge forward pass, the contraction
793 follows active autocast settings.
795 TransformerBridge exposes ``blocks[i].attn.W_O`` via its component-mapping
796 compatibility shim.
797 """
798 if self._ensure_head_results():
799 logging.warning("Tried to compute head results when they were already cached")
801 def _ensure_head_results(self) -> bool:
802 """Recompute missing or invalid per-head results; True if all were already cached."""
803 expected_ndim = 4 if self.has_batch_dim else 3
804 all_cached = True
805 for layer in range(self.model.cfg.n_layers):
806 result_key = f"blocks.{layer}.attn.hook_result"
807 cached_result = self.cache_dict.get(result_key)
808 # Batchless per-head results are valid 3D tensors.
809 if (
810 isinstance(cached_result, torch.Tensor)
811 and cached_result.ndim == expected_ndim
812 and cached_result.shape[-2] == self.model.cfg.n_heads
813 ):
814 continue
815 all_cached = False
817 z = self[("z", layer, "attn")]
818 weights = self.model.blocks[layer].attn.W_O
819 if z.dtype != weights.dtype:
820 # Preserve broadcast promotion for caches collected under autocast.
821 dtype = torch.promote_types(z.dtype, weights.dtype)
822 z, weights = z.to(dtype), weights.to(dtype)
823 self.cache_dict[result_key] = torch.einsum("...hd,hdm->...hm", z, weights)
824 return all_cached
826 def ssm_layers(self, mixer_type: Optional[Union[type, Tuple[type, ...]]] = None) -> List[int]:
827 """Return the block indices whose mixer is an SSM / recurrent mixer.
829 Family-agnostic and purely structural: finds each block's *realized* SSM
830 mixer in its variant slot (``.mixer`` / ``.linear_attn`` / …) via
831 ``find_ssm_mixer``, which excludes a hybrid's passthrough ``.mixer`` on
832 non-SSM layers (e.g. NemotronH attention/MLP/MoE) — no dependence on
833 ``cfg.layers_block_type``.
835 Args:
836 mixer_type: Optional concrete bridge class to filter to (e.g.
837 ``SSM2MixerBridge`` for only Mamba-2 layers).
839 Returns:
840 Ascending list of SSM block indices.
841 """
842 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
843 find_ssm_mixer,
844 )
846 bridge = self.model
847 layers: List[int] = []
848 for i, block in enumerate(bridge.blocks):
849 mixer = find_ssm_mixer(block)
850 if mixer is None:
851 continue
852 if mixer_type is not None and not isinstance(mixer, mixer_type):
853 continue
854 layers.append(i)
855 return layers
857 def _over_ssm_layers(
858 self,
859 fn: Callable[[Any, int], torch.Tensor],
860 mixer_type: Optional[Union[type, Tuple[type, ...]]] = None,
861 ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]:
862 """Apply ``fn(mixer, layer_idx)`` over every SSM layer; stack or dict.
864 The single SSM layer-enumeration used by both ``compute_ssm_state`` and
865 ``compute_ssm_effective_attention``. Returns a stacked tensor (dim 0 =
866 layer) when *every* block is an SSM layer, else a ``{layer_idx: result}``
867 dict over the SSM layers (heterogeneous hybrids).
868 """
869 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
870 find_ssm_mixer,
871 )
873 bridge = self.model
874 indices = self.ssm_layers(mixer_type=mixer_type)
875 if not indices: 875 ↛ 876line 875 didn't jump to line 876 because the condition on line 875 was never true
876 raise RuntimeError(
877 "No SSM mixer layers found. Use an SSM / hybrid bridge and "
878 "run_with_cache first (gated-delta-net needs use_cache=False)."
879 )
880 results: Dict[int, torch.Tensor] = {
881 i: fn(cast(Any, find_ssm_mixer(bridge.blocks[i])), i) for i in indices
882 }
883 if indices == list(range(len(bridge.blocks))):
884 return torch.stack([results[i] for i in indices], dim=0)
885 return results
887 def compute_ssm_effective_attention(
888 self,
889 layer: Optional[int] = None,
890 include_dt_scaling: bool = False,
891 per_state_coord: bool = False,
892 ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]:
893 """Materialize SSM effective attention for one or all SSM layers, any family.
895 The single discovery surface for effective attention across Mamba-1,
896 Mamba-2, and gated-delta-net layers; dispatches to each layer's mixer
897 ``compute_effective_attention`` regardless of family. Family-specific
898 options are forwarded only to mixers whose signature accepts them.
900 Args:
901 layer: Specific block index, or None for every SSM layer.
902 include_dt_scaling: Forwarded to Mamba-1/Mamba-2 mixers (the
903 reconstruction form); gated-delta-net does not accept it.
904 per_state_coord: Forwarded to Mamba-1 only (per-state-coordinate
905 matrices). Setting it True for another family raises ValueError.
907 Returns:
908 A per-layer matrix for a single ``layer``; for ``layer=None`` a
909 stacked tensor (dim 0 = layer) when every block is an SSM layer, else
910 a ``{layer_idx: matrix}`` dict over the SSM layers.
912 Raises:
913 TypeError: If the requested ``layer`` has no SSM mixer.
914 ValueError: If an option is set that the target mixer does not support.
915 RuntimeError: If ``layer=None`` finds no SSM layers.
916 """
917 import inspect
919 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
920 find_ssm_mixer,
921 )
923 def _call(mixer: Any, layer_idx: int) -> torch.Tensor:
924 fn = cast(Any, mixer).compute_effective_attention
925 params = inspect.signature(fn).parameters
926 kwargs: Dict[str, Any] = {}
927 for name, value in (
928 ("include_dt_scaling", include_dt_scaling),
929 ("per_state_coord", per_state_coord),
930 ):
931 if name in params:
932 kwargs[name] = value
933 elif value:
934 raise ValueError(
935 f"{type(mixer).__name__}.compute_effective_attention does not "
936 f"support {name}=True."
937 )
938 result: torch.Tensor = fn(cache=self, layer_idx=layer_idx, **kwargs)
939 return result
941 if layer is not None:
942 single = find_ssm_mixer(self.model.blocks[layer])
943 if single is None:
944 raise TypeError(f"Block {layer} has no SSM mixer (no compute_effective_attention).")
945 return _call(single, layer)
946 return self._over_ssm_layers(_call)
948 def compute_ssm_state(
949 self,
950 layer: Optional[int] = None,
951 time_step: Optional[int] = None,
952 ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]:
953 """Reconstruct the recurrent SSM state ``S`` from this cache.
955 The single discovery surface for recurrent state across families —
956 ``SSMMixerBridge`` (Mamba-1), ``SSM2MixerBridge`` (Mamba-2) and
957 ``GatedDeltaNetBridge`` (gated delta rule) each reconstruct it — mirroring
958 ``compute_head_results``. Read-only post-hoc reconstruction from cached
959 hooks (no forward re-run); requires an SSM / SSM-hybrid bridge cached via
960 ``run_with_cache`` (gated-delta-net additionally needs ``use_cache=False``
961 so its interior hooks fire). See the mixer's ``compute_ssm_state`` for the
962 recurrence, shapes, and the ``time_step`` memory bound; state shape is
963 family-specific, so ``layer=None`` returns a dict except when every block
964 shares one mixer type.
966 Args:
967 layer: Specific block index, or None for every SSM-state layer.
968 time_step: Optional single position (memory-bounded); None for all.
970 Returns:
971 A per-layer state tensor for a single ``layer``; for ``layer=None`` a
972 stacked tensor (dim 0 = layer) when every block is an SSM-state layer,
973 else a ``{layer_idx: state}`` dict over those layers.
975 Raises:
976 TypeError: If the requested ``layer`` has no state-reconstructing mixer.
977 RuntimeError: If ``layer=None`` finds no such layers.
978 """
979 from transformer_lens.model_bridge.generalized_components import (
980 GatedDeltaNetBridge,
981 SSM2MixerBridge,
982 SSMMixerBridge,
983 )
984 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
985 find_ssm_mixer,
986 )
988 # Every recurrent family reconstructs S_t from its cached interior hooks.
989 state_mixers = (SSMMixerBridge, SSM2MixerBridge, GatedDeltaNetBridge)
991 def _call(mixer: Any, layer_idx: int) -> torch.Tensor:
992 state: torch.Tensor = cast(Any, mixer).compute_ssm_state(
993 self, layer_idx=layer_idx, time_step=time_step
994 )
995 return state
997 if layer is not None:
998 single = find_ssm_mixer(self.model.blocks[layer])
999 if not isinstance(single, state_mixers):
1000 raise TypeError(
1001 f"Block {layer} has no state-reconstructing SSM mixer; "
1002 "compute_ssm_state supports Mamba-1 / Mamba-2 / gated-delta-net."
1003 )
1004 return _call(single, layer)
1005 return self._over_ssm_layers(_call, mixer_type=state_mixers)
1007 def stack_head_results(
1008 self,
1009 layer: int = -1,
1010 return_labels: bool = False,
1011 incl_remainder: bool = False,
1012 pos_slice: Union[Slice, SliceInput] = None,
1013 apply_ln: bool = False,
1014 ) -> Union[
1015 Float[torch.Tensor, "num_components *batch_and_pos_dims d_model"],
1016 Tuple[Float[torch.Tensor, "num_components *batch_and_pos_dims d_model"], List[str]],
1017 ]:
1018 """Stack Head Results.
1020 Returns a stack of all head results (ie residual stream contribution) up to layer L. A good
1021 way to decompose the outputs of attention layers into attribution by specific heads. Note
1022 that the num_components axis has length layer x n_heads ((layer head_index) in einops
1023 notation).
1025 Args:
1026 layer:
1027 Layer index - heads at all layers strictly before this are included. layer must be
1028 in [1, n_layers-1], or any of (n_layers, -1, None), which all mean the final layer.
1029 return_labels:
1030 Whether to also return a list of labels of the form "L0H0" for the heads.
1031 incl_remainder:
1032 Whether to return a final term which is "the rest of the residual stream".
1033 pos_slice:
1034 A slice object to apply to the pos dimension. Defaults to None, do nothing.
1035 apply_ln:
1036 Whether to apply LayerNorm to the stack.
1037 """
1038 if not isinstance(pos_slice, Slice):
1039 pos_slice = Slice(pos_slice)
1040 pos_slice = cast(Slice, pos_slice) # mypy can't seem to infer this
1041 if layer is None or layer == -1:
1042 # Default to the residual stream immediately pre unembed
1043 layer = self.model.cfg.n_layers
1045 # Idempotent; cleans up stale Bridge entries without the 'already cached' warning,
1046 # which is only meaningful for direct compute_head_results callers.
1047 self._ensure_head_results()
1049 components: Any = []
1050 labels = []
1051 for l in range(layer):
1052 # Note that this has shape batch x pos x head_index x d_model
1053 components.append(pos_slice.apply(self[("result", l, "attn")], dim=-3))
1054 labels.extend([f"L{l}H{h}" for h in range(self.model.cfg.n_heads)])
1055 if components:
1056 components = torch.cat(components, dim=-2)
1057 components = einops.rearrange(
1058 components,
1059 "... concat_head_index d_model -> concat_head_index ... d_model",
1060 )
1061 if incl_remainder:
1062 remainder = pos_slice.apply(
1063 self[("resid_post", layer - 1)], dim=-2
1064 ) - components.sum(dim=0)
1065 components = torch.cat([components, remainder[None]], dim=0)
1066 labels.append("remainder")
1067 elif incl_remainder:
1068 # There are no components, so the remainder is the entire thing.
1069 components = torch.cat(
1070 [pos_slice.apply(self[("resid_post", layer - 1)], dim=-2)[None]], dim=0
1071 )
1072 labels.append("remainder")
1073 else:
1074 # If this is called with layer 0, we return an empty tensor of the right shape to be
1075 # stacked correctly. This uses the shape of hook_embed, which is pretty janky since it
1076 # assumes embed is in the cache. But it's hard to explicitly code the shape, since it
1077 # depends on the pos slice, whether we have a batch dim, etc. And it's pretty messy!
1078 components = torch.zeros(
1079 0,
1080 *pos_slice.apply(self["hook_embed"], dim=-2).shape,
1081 device=self.model.cfg.device,
1082 )
1084 if apply_ln:
1085 components = self.apply_ln_to_stack(components, layer, pos_slice=pos_slice)
1087 if return_labels:
1088 return components, labels
1089 else:
1090 return components
1092 def stack_activation(
1093 self,
1094 activation_name: str,
1095 layer: int = -1,
1096 sublayer_type: Optional[str] = None,
1097 ) -> Float[torch.Tensor, "layers_covered ..."]:
1098 """Stack Activations.
1100 Flexible way to stack activations with a given name.
1102 Args:
1103 activation_name:
1104 The name of the activation to be stacked
1105 layer:
1106 'Layer index - heads' at all layers strictly before this are included. layer must be
1107 in [1, n_layers-1], or any of (n_layers, -1, None), which all mean the final layer.
1108 sublayer_type:
1109 The sub layer type of the activation, passed to utils.get_act_name. Can normally be
1110 inferred.
1111 incl_remainder:
1112 Whether to return a final term which is "the rest of the residual stream".
1113 """
1114 if layer is None or layer == -1:
1115 # Default to the residual stream immediately pre unembed
1116 layer = self.model.cfg.n_layers
1118 components = []
1119 for l in range(layer):
1120 components.append(self[(activation_name, l, sublayer_type)])
1122 return torch.stack(components, dim=0)
1124 def get_neuron_results(
1125 self,
1126 layer: int,
1127 neuron_slice: Union[Slice, SliceInput] = None,
1128 pos_slice: Union[Slice, SliceInput] = None,
1129 project_output_onto: Optional[torch.Tensor] = None,
1130 ) -> torch.Tensor:
1131 """Get Neuron Results.
1133 Get the results of for neurons in a specific layer (i.e, how much each neuron contributes to
1134 the residual stream). Does it for the subset of neurons specified by neuron_slice, defaults
1135 to all of them. Does *not* cache these because it's expensive in space and cheap to compute.
1137 Args:
1138 layer:
1139 Layer index.
1140 neuron_slice:
1141 Slice of the neuron.
1142 pos_slice:
1143 Slice of the positions.
1144 project_output_onto:
1145 Optional ``[d_model]`` or ``[d_model, num_outputs]`` projection. Contracted with
1146 ``W_out`` *before* the per-neuron expansion so the ``[..., d_mlp, d_model]``
1147 intermediate is never materialized.
1149 Returns:
1150 Last-dim is ``d_model`` (default), ``num_outputs`` (2D projection), or squeezed
1151 (1D projection).
1152 """
1153 if not isinstance(neuron_slice, Slice):
1154 neuron_slice = Slice(neuron_slice)
1155 if not isinstance(pos_slice, Slice):
1156 pos_slice = Slice(pos_slice)
1158 neuron_acts = self[("post", layer, "mlp")]
1159 block = self.model.blocks[layer]
1160 W_out = block.mlp.W_out
1161 if pos_slice is not None: 1161 ↛ 1165line 1161 didn't jump to line 1165 because the condition on line 1161 was always true
1162 # Note - order is important, as Slice.apply *may* collapse a dimension, so this ensures
1163 # that position dimension is -2 when we apply position slice
1164 neuron_acts = pos_slice.apply(neuron_acts, dim=-2)
1165 if neuron_slice is not None: 1165 ↛ 1168line 1165 didn't jump to line 1168 because the condition on line 1165 was always true
1166 neuron_acts = neuron_slice.apply(neuron_acts, dim=-1)
1167 W_out = neuron_slice.apply(W_out, dim=0)
1168 if project_output_onto is None:
1169 return neuron_acts[..., None] * W_out
1170 # W_out: [d_mlp, d_model]; project: [d_model] or [d_model, n_outs]
1171 projected = W_out @ project_output_onto
1172 if project_output_onto.ndim == 1:
1173 return neuron_acts * projected
1174 return neuron_acts[..., None] * projected
1176 def _get_cached_ln_scale(
1177 self,
1178 layer: Optional[int],
1179 mlp_input: bool,
1180 pos_slice: Slice,
1181 batch_slice: Optional[Slice] = None,
1182 ) -> torch.Tensor:
1183 """Look up the cached LN scale and apply pos/batch slicing. Surfaces a clearer error
1184 when the expected hook isn't in the cache (some non-decoder-only architectures expose
1185 LN scale at a different path or not at all).
1186 """
1187 if layer == self.model.cfg.n_layers or layer is None:
1188 key = "ln_final.hook_scale"
1189 else:
1190 key = f"blocks.{layer}.ln{2 if mlp_input else 1}.hook_scale"
1191 try:
1192 scale = self[key]
1193 except KeyError as e:
1194 raise KeyError(
1195 f"Cached LN scale not found at '{key}'. apply_ln operations require the model "
1196 f"to have cached this hook (some non-decoder-only architectures expose LN scale "
1197 f"under different module paths)."
1198 ) from e
1199 scale = pos_slice.apply(scale, dim=-2)
1200 if batch_slice is not None and self.has_batch_dim:
1201 scale = batch_slice.apply(scale)
1202 return scale
1204 def _stack_neuron_results_apply_ln_projected(
1205 self,
1206 layer: int,
1207 pos_slice: Slice,
1208 neuron_slice: Slice,
1209 project_2d: torch.Tensor,
1210 mlp_input: bool = False,
1211 ) -> torch.Tensor:
1212 """LN-applied neuron stack with projection folded in — no d_mlp×d_model intermediate.
1214 Analytical formula (LN models, cached scale ``s``):
1215 ``LN_s(a_n * W_out_n) @ p = (a_n / s) * (W_out_n @ p - mean(W_out_n) * sum_p)``
1216 RMS models drop the ``mean(W_out_n) * sum_p`` term (no centering). ``s`` is the ln2
1217 scale of ``layer`` when ``mlp_input`` is set, else the ln1 scale.
1218 """
1219 scale = self._get_cached_ln_scale(layer, mlp_input=mlp_input, pos_slice=pos_slice)
1221 apply_centering = self.model.cfg.normalization_type in ["LN", "LNPre"]
1222 sum_p = project_2d.sum(dim=0) if apply_centering else None # [n_outs]
1224 components: list = []
1225 for l in range(layer):
1226 block = self.model.blocks[l]
1227 W_out_l = block.mlp.W_out # [d_mlp, d_model]
1228 W_out_l_sliced = neuron_slice.apply(W_out_l, dim=0)
1229 W_proj_l = W_out_l_sliced @ project_2d # [d_mlp, n_outs]
1230 if apply_centering:
1231 assert sum_p is not None # set when apply_centering, narrow for mypy
1232 W_means_l = W_out_l_sliced.mean(dim=-1) # [d_mlp]
1233 lin_form_l = W_proj_l - W_means_l[:, None] * sum_p[None, :]
1234 else:
1235 lin_form_l = W_proj_l
1236 a_l = self[("post", l, "mlp")]
1237 a_l = pos_slice.apply(a_l, dim=-2)
1238 a_l = neuron_slice.apply(a_l, dim=-1)
1239 # (a_l / s)[..., None] is [..., d_mlp, 1]; broadcast with lin_form_l [d_mlp, n_outs]
1240 components.append((a_l / scale)[..., None] * lin_form_l)
1241 if not components: 1241 ↛ 1242line 1241 didn't jump to line 1242 because the condition on line 1241 was never true
1242 empty_src = pos_slice.apply(self["hook_embed"], dim=-2)
1243 return torch.zeros(
1244 0, *empty_src.shape[:-1], project_2d.shape[-1], device=self.model.cfg.device
1245 )
1246 stacked = torch.cat(components, dim=-2)
1247 return einops.rearrange(
1248 stacked, "... concat_neuron_index n_outs -> concat_neuron_index ... n_outs"
1249 )
1251 def stack_neuron_results(
1252 self,
1253 layer: int,
1254 pos_slice: Union[Slice, SliceInput] = None,
1255 neuron_slice: Union[Slice, SliceInput] = None,
1256 return_labels: bool = False,
1257 incl_remainder: bool = False,
1258 apply_ln: bool = False,
1259 project_output_onto: Optional[torch.Tensor] = None,
1260 mlp_input: bool = False,
1261 ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[str]]]:
1262 """Stack Neuron Results
1264 Returns a stack of all neuron results (ie residual stream contribution) up to layer L - ie
1265 the amount each individual neuron contributes to the residual stream. Also returns a list of
1266 labels of the form "L0N0" for the neurons. A good way to decompose the outputs of MLP layers
1267 into attribution by specific neurons.
1269 Note that doing this for all neurons is SUPER expensive on GPU memory and only works for
1270 small models or short inputs. Pass ``project_output_onto`` to fold the projection into the
1271 per-neuron expansion and avoid the ``[..., d_mlp, d_model]`` intermediate.
1273 Args:
1274 layer:
1275 Layer index - heads at all layers strictly before this are included. layer must be
1276 in [1, n_layers]. ``layer=0`` follows the negative-index convention: the neuron
1277 stack is empty and with ``incl_remainder=True`` the remainder is the final layer's
1278 residual stream (``resid_post`` of layer ``n_layers - 1``), i.e. the whole model.
1279 pos_slice:
1280 Slice of the positions.
1281 neuron_slice:
1282 Slice of the neurons. An int or an integer-mode ``Slice(n)`` selects a single
1283 neuron and is treated like ``[n]``.
1284 return_labels:
1285 Whether to also return a list of labels of the form "L0N0" for the neurons.
1286 incl_remainder:
1287 Whether to return a final term which is "the rest of the residual stream".
1288 apply_ln:
1289 Whether to apply LayerNorm to the stack.
1290 project_output_onto:
1291 Optional ``[d_model]`` or ``[d_model, num_outputs]`` tensor. When set, each
1292 component's last d_model dim is replaced by the projection (memory-efficient for
1293 direction analyses; see ``get_neuron_results``). Combined with ``apply_ln=True``,
1294 the projection is folded into the analytical cached-scale LN so the
1295 ``[..., d_mlp, d_model]`` intermediate is still never materialized.
1296 mlp_input:
1297 Treat the stack as the input to ``layer``'s MLP rather than its attention, as in
1298 ``decompose_resid``: with ``incl_remainder=True`` the remainder fills the stack up to
1299 ``resid_mid`` instead of ``resid_pre``, and with ``apply_ln=True`` the LN scale of
1300 ln2 is used instead of ln1. The neurons stacked are unchanged.
1301 """
1302 if layer is None or layer == -1:
1303 # Default to the residual stream immediately pre unembed
1304 layer = self.model.cfg.n_layers
1306 # layer=0 is the documented whole-model case (remainder reads resid_post of layer -1, the
1307 # final layer). Anything below it would wrap the remainder target to an arbitrary late
1308 # layer via __getitem__, silently returning the wrong residual stream.
1309 if incl_remainder and layer < 0:
1310 raise ValueError(
1311 f"layer must be >= 0 with incl_remainder=True (-1/None mean n_layers); got {layer}."
1312 )
1314 # Layers are concatenated along the neuron axis, so it must not collapse. unwrap turns an
1315 # int into [n] but leaves a Slice as is, so an integer-mode Slice(n) is rebuilt as [n]
1316 # here. A new Slice is made, so the caller's Slice still collapses elsewhere.
1317 neuron_slice = Slice.unwrap(neuron_slice)
1318 if neuron_slice.mode == "int":
1319 assert isinstance(neuron_slice.slice, int) # narrow for mypy
1320 neuron_slice = Slice([neuron_slice.slice])
1321 if not isinstance(pos_slice, Slice):
1322 pos_slice = Slice(pos_slice)
1324 project_2d, squeeze_projected = _normalize_projection_to_2d(project_output_onto)
1326 d_mlp = self.model.cfg.d_mlp
1327 assert d_mlp is not None, "model.cfg.d_mlp must be set"
1328 neuron_labels = neuron_slice.apply(torch.arange(d_mlp), dim=0)
1330 labels = [f"L{l}N{h}" for l in range(layer) for h in neuron_labels]
1331 # The remainder fills the stack up to the residual stream that layer's LN normalizes:
1332 # resid_mid (the MLP input) with mlp_input, else resid_pre (resid_post of layer - 1).
1333 remainder_target = ("resid_mid", layer) if mlp_input else ("resid_post", layer - 1)
1334 components: Any
1335 ln_folded = apply_ln and project_2d is not None
1336 if ln_folded:
1337 assert project_2d is not None # narrow for mypy
1338 # Analytical LN+projection — no d_mlp×d_model intermediate.
1339 components = self._stack_neuron_results_apply_ln_projected(
1340 layer, pos_slice, neuron_slice, project_2d, mlp_input=mlp_input
1341 )
1342 if incl_remainder:
1343 # Linearity of cached-scale LN: remainder is LN_s(target) @ p - sum(neurons).
1344 target = pos_slice.apply(self[remainder_target], dim=-2)
1345 target_ln = self.apply_ln_to_stack(
1346 target[None], layer, pos_slice=pos_slice, mlp_input=mlp_input
1347 )[0]
1348 remainder = target_ln @ project_2d
1349 if components.shape[0] > 0: 1349 ↛ 1351line 1349 didn't jump to line 1351 because the condition on line 1349 was always true
1350 remainder = remainder - components.sum(dim=0)
1351 components = torch.cat([components, remainder[None]], dim=0)
1352 labels.append("remainder")
1353 else:
1354 per_layer: list = []
1355 for l in range(layer):
1356 per_layer.append(
1357 self.get_neuron_results(
1358 l,
1359 pos_slice=pos_slice,
1360 neuron_slice=neuron_slice,
1361 project_output_onto=project_2d,
1362 )
1363 )
1364 if per_layer:
1365 components = torch.cat(per_layer, dim=-2)
1366 components = einops.rearrange(
1367 components,
1368 "... concat_neuron_index d_model -> concat_neuron_index ... d_model",
1369 )
1370 if incl_remainder:
1371 remainder_full = pos_slice.apply(self[remainder_target], dim=-2)
1372 if project_2d is not None:
1373 remainder_full = remainder_full @ project_2d
1374 remainder = remainder_full - components.sum(dim=0)
1375 components = torch.cat([components, remainder[None]], dim=0)
1376 labels.append("remainder")
1377 elif incl_remainder:
1378 remainder_full = pos_slice.apply(self[remainder_target], dim=-2)
1379 if project_2d is not None: 1379 ↛ 1380line 1379 didn't jump to line 1380 because the condition on line 1379 was never true
1380 remainder_full = remainder_full @ project_2d
1381 components = torch.cat([remainder_full[None]], dim=0)
1382 labels.append("remainder")
1383 else:
1384 empty_shape_src = pos_slice.apply(self["hook_embed"], dim=-2)
1385 if project_2d is not None: 1385 ↛ 1386line 1385 didn't jump to line 1386 because the condition on line 1385 was never true
1386 empty_shape_src = empty_shape_src @ project_2d
1387 components = torch.zeros(0, *empty_shape_src.shape, device=self.model.cfg.device)
1389 if apply_ln:
1390 components = self.apply_ln_to_stack(
1391 components, layer, pos_slice=pos_slice, mlp_input=mlp_input
1392 )
1394 if squeeze_projected:
1395 components = components.squeeze(-1)
1397 if return_labels:
1398 return components, labels
1399 else:
1400 return components
1402 def apply_ln_to_stack(
1403 self,
1404 residual_stack: Float[torch.Tensor, "num_components *batch_and_pos_dims d_model"],
1405 layer: Optional[int] = None,
1406 mlp_input: bool = False,
1407 pos_slice: Union[Slice, SliceInput] = None,
1408 batch_slice: Union[Slice, SliceInput] = None,
1409 has_batch_dim: bool = True,
1410 recompute_ln: bool = False,
1411 ) -> Float[torch.Tensor, "num_components *batch_and_pos_dims_out d_model"]:
1412 """Apply Layer Norm to a Stack.
1414 Takes a stack of components of the residual stream (eg outputs of decompose_resid or
1415 accumulated_resid), treats them as the input to a specific layer, and applies the layer norm
1416 scaling of that layer to them, using the cached scale factors - simulating what that
1417 component of the residual stream contributes to that layer's input.
1419 The layernorm scale is global across the entire residual stream for each layer, batch
1420 element and position, which is why we need to use the cached scale factors rather than just
1421 applying a new LayerNorm.
1423 When recompute_ln=True and the target layer is the final layer (unembed), each
1424 component is normalized using stats recomputed from that component; use this for logit lens
1425 analysis. When recompute_ln=False, a single cached scale is used for all components.
1427 If the model does not use LayerNorm or RMSNorm, it returns the residual stack unchanged.
1429 Args:
1430 residual_stack:
1431 A tensor, whose final dimension is d_model. The other trailing dimensions are
1432 assumed to be the same as the stored hook_scale - which may or may not include batch
1433 or position dimensions.
1434 layer:
1435 The layer we're taking the input to. In [0, n_layers], n_layers means the unembed.
1436 None maps to the n_layers case, ie the unembed.
1437 mlp_input:
1438 Whether the input is to the MLP or attn (ie ln2 vs ln1). Defaults to False, ie ln1.
1439 If layer==n_layers, must be False, and we use ln_final
1440 pos_slice:
1441 The slice to take of positions, if residual_stack is not over the full context, None
1442 means do nothing. It is assumed that pos_slice has already been applied to
1443 residual_stack, and this is only applied to the scale. See utils.Slice for details.
1444 Defaults to None, do nothing.
1445 batch_slice:
1446 The slice to take on the batch dimension. Defaults to None, do nothing.
1447 has_batch_dim:
1448 Whether residual_stack has a batch dimension.
1449 recompute_ln:
1450 If True and target layer is the unembed (final layer), apply the final layer norm
1451 to each component with statistics recomputed from that component. Defaults to False.
1453 """
1454 if not isinstance(pos_slice, Slice):
1455 pos_slice = Slice(pos_slice)
1456 if not isinstance(batch_slice, Slice):
1457 batch_slice = Slice(batch_slice)
1459 if layer is None or layer == -1:
1460 # Default to the residual stream immediately pre unembed
1461 layer = self.model.cfg.n_layers
1463 if has_batch_dim:
1464 # Apply batch slice to the stack
1465 residual_stack = batch_slice.apply(residual_stack, dim=1)
1467 # Checked only after slicing so callers like logit_attrs still get their
1468 # batch_slice honoured on models without LayerNorm/RMSNorm.
1469 if self.model.cfg.normalization_type not in ["LN", "LNPre", "RMS", "RMSPre"]:
1470 return residual_stack
1472 # Logit lens: apply final layer norm to each component with recomputed statistics
1473 if recompute_ln and layer == self.model.cfg.n_layers and hasattr(self.model, "ln_final"):
1474 ln_final = self.model.ln_final
1475 results = []
1476 for i in range(residual_stack.shape[0]):
1477 x = residual_stack[i]
1478 original_shape = x.shape
1479 # ln_final expects (batch, pos, d_model); restore missing structural dimensions
1480 if not has_batch_dim:
1481 x = x.unsqueeze(0)
1482 if x.ndim == 2:
1483 x = x.unsqueeze(1)
1484 out = ln_final(x)
1485 results.append(out.reshape(original_shape))
1486 return torch.stack(results, dim=0)
1488 # Center the stack onlny if the model uses LayerNorm
1489 if self.model.cfg.normalization_type in ["LN", "LNPre"]:
1490 residual_stack = residual_stack - residual_stack.mean(dim=-1, keepdim=True)
1492 # Shape is [batch, position, 1] or [position, 1]; final dim is a dummy for broadcasting.
1493 scale = self._get_cached_ln_scale(layer, mlp_input, pos_slice, batch_slice)
1495 return residual_stack / scale
1497 def get_full_resid_decomposition(
1498 self,
1499 layer: Optional[int] = None,
1500 mlp_input: bool = False,
1501 expand_neurons: bool = True,
1502 apply_ln: bool = False,
1503 pos_slice: Union[Slice, SliceInput] = None,
1504 return_labels: bool = False,
1505 project_output_onto: Optional[torch.Tensor] = None,
1506 ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[str]]]:
1507 """Get the full Residual Decomposition.
1509 Decomposes the residual stream that is input into some layer into its
1510 constituent components: every attention head result, every neuron (or
1511 MLP layer) result, the embeddings, and the accumulated biases.
1513 The returned tensor stacks components along ``dim=0`` in this order:
1515 1. Attention head results, layer-by-layer (``L * n_heads`` rows)
1516 2. Neuron / MLP results (only if ``cfg.attn_only=False`` and
1517 ``layer > 0``; ``L * d_mlp`` rows when ``expand_neurons=True``,
1518 else ``L`` rows)
1519 3. ``embed`` (1 row, if the model has token embeddings)
1520 4. ``pos_embed`` (1 row, if the model has positional embeddings)
1521 5. ``bias`` (1 row, the accumulated layer biases)
1523 ``return_labels=True`` returns a list of strings in the same order, so
1524 ``labels[i]`` always names ``stack[i]``. If you need to extract a
1525 specific component, slice by label rather than by hard-coded index —
1526 the row counts depend on ``layer``, ``expand_neurons``,
1527 ``cfg.attn_only``, and whether the model has positional embeddings.
1529 Args:
1530 layer:
1531 The layer we're inputting into. layer is in [0, n_layers], if layer==n_layers (or
1532 None) we're inputting into the unembed (the entire stream), if layer==0 then it's
1533 just embed and pos_embed
1534 mlp_input:
1535 Are we inputting to the MLP in that layer or the attn? Must be False for final
1536 layer, since that's the unembed.
1537 expand_neurons:
1538 Whether to expand the MLP outputs to give every neuron's result or just return the
1539 MLP layer outputs.
1540 apply_ln:
1541 Whether to apply LayerNorm to the stack.
1542 pos_slice:
1543 Slice of the positions to take.
1544 return_labels:
1545 Whether to return the labels.
1546 project_output_onto:
1547 Optional ``[d_model]`` or ``[d_model, num_outputs]`` projection. Folded in
1548 *before* the per-neuron expansion, so the ``[..., d_mlp, d_model]`` intermediate
1549 is never materialized (memory saving applies only with ``expand_neurons=True``).
1550 Combined with ``apply_ln=True``, the projection is fused into the analytical
1551 cached-scale LN so the same memory benefit holds. Output last-dim is squeezed
1552 for a 1D projection; ``num_outputs`` for 2D.
1553 """
1554 if layer is None or layer == -1:
1555 # Default to the residual stream immediately pre unembed
1556 layer = self.model.cfg.n_layers
1557 assert layer is not None # keep mypy happy
1559 if not isinstance(pos_slice, Slice):
1560 pos_slice = Slice(pos_slice)
1562 project_2d, squeeze_projected = _normalize_projection_to_2d(project_output_onto)
1563 # When both apply_ln and projection are requested, LN is applied per-component (in
1564 # d_model space for the small ones, analytically for neurons) before projection, so the
1565 # final apply_ln_to_stack call is skipped — last-dim is already n_outs.
1566 ln_folded = apply_ln and project_2d is not None
1568 def _ln_then_project(stack: torch.Tensor) -> torch.Tensor:
1569 stack = self.apply_ln_to_stack(stack, layer, pos_slice=pos_slice, mlp_input=mlp_input)
1570 return stack @ project_2d if project_2d is not None else stack
1572 head_stack, head_labels = self.stack_head_results(
1573 layer + (1 if mlp_input else 0), pos_slice=pos_slice, return_labels=True
1574 )
1575 if ln_folded:
1576 head_stack = _ln_then_project(head_stack)
1577 elif project_2d is not None:
1578 head_stack = head_stack @ project_2d
1579 labels = head_labels
1580 components = [head_stack]
1581 if not self.model.cfg.attn_only and layer > 0:
1582 if expand_neurons:
1583 # Only ask stack_neuron_results to apply LN when we want the fused analytical
1584 # path (ln_folded). For the unfolded case the outer apply_ln_to_stack handles it.
1585 neuron_stack, neuron_labels = self.stack_neuron_results(
1586 layer,
1587 pos_slice=pos_slice,
1588 return_labels=True,
1589 apply_ln=ln_folded,
1590 project_output_onto=project_2d,
1591 mlp_input=mlp_input,
1592 )
1593 labels.extend(neuron_labels)
1594 components.append(neuron_stack)
1595 else:
1596 # Get the stack of just the MLP outputs
1597 # mlp_input included for completeness, but it doesn't actually matter, since it's
1598 # just for MLP outputs
1599 mlp_stack, mlp_labels = self.decompose_resid(
1600 layer,
1601 mlp_input=mlp_input,
1602 pos_slice=pos_slice,
1603 incl_embeds=False,
1604 mode="mlp",
1605 return_labels=True,
1606 )
1607 if ln_folded: 1607 ↛ 1608line 1607 didn't jump to line 1608 because the condition on line 1607 was never true
1608 mlp_stack = _ln_then_project(mlp_stack)
1609 elif project_2d is not None: 1609 ↛ 1610line 1609 didn't jump to line 1610 because the condition on line 1609 was never true
1610 mlp_stack = mlp_stack @ project_2d
1611 labels.extend(mlp_labels)
1612 components.append(mlp_stack)
1614 if self.has_embed: 1614 ↛ 1622line 1614 didn't jump to line 1622 because the condition on line 1614 was always true
1615 embed = pos_slice.apply(self["embed"], -2)[None]
1616 if ln_folded:
1617 embed = _ln_then_project(embed)
1618 elif project_2d is not None:
1619 embed = embed @ project_2d
1620 labels.append("embed")
1621 components.append(embed)
1622 if self.has_pos_embed: 1622 ↛ 1631line 1622 didn't jump to line 1631 because the condition on line 1622 was always true
1623 pos_embed = pos_slice.apply(self["pos_embed"], -2)[None]
1624 if ln_folded:
1625 pos_embed = _ln_then_project(pos_embed)
1626 elif project_2d is not None:
1627 pos_embed = pos_embed @ project_2d
1628 labels.append("pos_embed")
1629 components.append(pos_embed)
1630 # If we didn't expand the neurons, the MLP biases are already included in the MLP outputs.
1631 bias_full = self.model.accumulated_bias(layer, mlp_input, include_mlp_biases=expand_neurons)
1632 if ln_folded:
1633 # Expand bias to per-position d_model shape so LN can center, then project.
1634 expand_shape: tuple = (1,) + tuple(head_stack.shape[1:-1]) + (self.model.cfg.d_model,)
1635 bias = _ln_then_project(bias_full.expand(expand_shape))
1636 else:
1637 if project_2d is not None:
1638 # Bias is [d_model], so project post-hoc for shape compatibility — no memory win here.
1639 bias_full = bias_full @ project_2d
1640 bias = bias_full.expand((1,) + head_stack.shape[1:])
1641 labels.append("bias")
1642 components.append(bias)
1643 residual_stack = torch.cat(components, dim=0)
1644 if apply_ln and not ln_folded:
1645 residual_stack = self.apply_ln_to_stack(
1646 residual_stack, layer, pos_slice=pos_slice, mlp_input=mlp_input
1647 )
1649 if squeeze_projected:
1650 residual_stack = residual_stack.squeeze(-1)
1652 if return_labels:
1653 return residual_stack, labels
1654 else:
1655 return residual_stack