Coverage for transformer_lens/ActivationCache.py: 94%
466 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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 numpy as np
33import torch
34from jaxtyping import Float, Int
35from typing_extensions import Literal
37import transformer_lens.utilities as utils
38from transformer_lens.utilities import Slice, SliceInput, warn_if_mps
40if TYPE_CHECKING:
41 from transformer_lens.HookedTransformer import HookedTransformer
44def _normalize_projection_to_2d(
45 project: Optional[torch.Tensor],
46) -> Tuple[Optional[torch.Tensor], bool]:
47 """Return ``(project_2d, squeeze_at_end)`` — 1D projections are reshaped to 2D for uniform internal handling and squeezed back at the user-facing return."""
48 if project is None:
49 return None, False
50 if project.ndim == 1:
51 return project.unsqueeze(-1), True
52 return project, False
55class ActivationCache:
56 """Activation Cache.
58 A wrapper that stores all important activations from a forward pass of the model, and provides a
59 variety of helper functions to investigate them.
61 The :class:`ActivationCache` is at the core of Transformer Lens. It is a wrapper that stores all
62 important activations from a forward pass of the model, and provides a variety of helper
63 functions to investigate them. The common way to access it is to run the model with
64 :meth:`transformer_lens.HookedTransformer.HookedTransformer.run_with_cache`.
66 Examples:
68 When investigating a particular behaviour of a model, a very common first step is to try and
69 understand which components of the model are most responsible for that behaviour. For example,
70 if you're investigating the prompt "Why did the chicken cross the" -> " road", you might want to
71 understand if there is a specific sublayer (mlp or multi-head attention) that is responsible for
72 the model predicting "road". This kind of analysis commonly falls under the category of "logit
73 attribution" or "direct logit attribution" (DLA).
75 >>> from transformer_lens import HookedTransformer
76 >>> model = HookedTransformer.from_pretrained("tiny-stories-1M")
77 Loaded pretrained model tiny-stories-1M into HookedTransformer
79 >>> _logits, cache = model.run_with_cache("Why did the chicken cross the")
80 >>> residual_stream, labels = cache.decompose_resid(return_labels=True, mode="attn")
81 >>> print(labels[0:3])
82 ['embed', 'pos_embed', '0_attn_out']
84 >>> answer = " road" # Note the proceeding space to match the model's tokenization
85 >>> logit_attrs = cache.logit_attrs(residual_stream, answer)
86 >>> print(logit_attrs.shape) # Attention layers
87 torch.Size([10, 1, 7])
89 >>> most_important_component_idx = torch.argmax(logit_attrs)
90 >>> print(labels[most_important_component_idx])
91 3_attn_out
93 You can also dig in with more granularity, using :meth:`get_full_resid_decomposition` to get the
94 residual stream by individual component (mlp neurons and individual attention heads). This
95 creates a larger residual stack, but the approach of using :meth"`logit_attrs` remains the same.
97 Equally you might want to find out if the model struggles to construct such excellent jokes
98 until the very last layers, or if it is trivial and the first few layers are enough. This kind
99 of analysis is called "logit lens", and you can find out more about how to do that with
100 :meth:`ActivationCache.accumulated_resid`.
102 Warning:
104 :class:`ActivationCache` is designed to be used with
105 :class:`transformer_lens.HookedTransformer`, and will not work with other models. It's also
106 designed to be used with all activations of :class:`transformer_lens.HookedTransformer` being
107 cached, and some internal methods will break without that.
109 The biggest footgun and source of bugs in this code will be keeping track of indexes,
110 dimensions, and the numbers of each. There are several kinds of activations:
112 * Internal attn head vectors: q, k, v, z. Shape [batch, pos, head_index, d_head].
113 * Internal attn pattern style results: pattern (post softmax), attn_scores (pre-softmax). Shape
114 [batch, head_index, query_pos, key_pos].
115 * Attn head results: result. Shape [batch, pos, head_index, d_model].
116 * Internal MLP vectors: pre, post, mid (only used for solu_ln - the part between activation +
117 layernorm). Shape [batch, pos, d_mlp].
118 * Residual stream vectors: resid_pre, resid_mid, resid_post, attn_out, mlp_out, embed,
119 pos_embed, normalized (output of each LN or LNPre). Shape [batch, pos, d_model].
120 * LayerNorm Scale: scale. Shape [batch, pos, 1].
122 Sometimes the batch dimension will be missing because we applied `remove_batch_dim` (used when
123 batch_size=1), and as such all library functions *should* be robust to that.
125 Type annotations are in the following form:
127 * layers_covered is the number of layers queried in functions that stack the residual stream.
128 * batch_and_pos_dims is the set of dimensions from batch and pos - by default this is ["batch",
129 "pos"], but is only ["pos"] if we've removed the batch dimension and is [()] if we've removed
130 batch dimension and are applying a pos slice which indexes a specific position.
132 Args:
133 cache_dict:
134 A dictionary of cached activations from a model run.
135 model:
136 The model that the activations are from.
137 has_batch_dim:
138 Whether the activations have a batch dimension.
139 """
141 def __init__(
142 self,
143 cache_dict: Dict[str, torch.Tensor],
144 model: Any,
145 has_batch_dim: bool = True,
146 ):
147 self.cache_dict = cache_dict
148 # Helper methods require HT-internal structure; bridge users only use cache_dict.
149 self.model = cast("HookedTransformer", 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) -> ActivationCache:
170 """Remove the Batch Dimension (if a single batch item).
172 Returns:
173 The ActivationCache with the batch dimension removed.
174 """
175 if self.has_batch_dim:
176 batch_size = self._batch_size()
177 assert (
178 batch_size == 1
179 ), f"Cannot remove batch dimension from cache with batch size {batch_size}"
180 for key in self.cache_dict:
181 if self.cache_dict[key].ndim > 0 and self.cache_dict[key].size(0) == 1:
182 self.cache_dict[key] = self.cache_dict[key][0]
183 self.has_batch_dim = False
184 else:
185 logging.warning("Tried removing batch dimension after already having removed it.")
186 return self
188 def __repr__(self) -> str:
189 """Representation of the ActivationCache.
191 Special method that returns a string representation of an object. It's normally used to give
192 a string that can be used to recreate the object, but here we just return a string that
193 describes the object.
194 """
195 return f"ActivationCache with keys {list(self.cache_dict.keys())}"
197 def __getitem__(self, key) -> torch.Tensor:
198 """Retrieve Cached Activations by Key or Shorthand.
200 Enables direct access to cached activations via dictionary-style indexing using keys or
201 shorthand naming conventions.
203 It also supports tuples for advanced indexing, with the dimension order as (name, layer_index, layer_type).
204 See :func:`transformer_lens.utils.get_act_name` for how shorthand is converted to a full name.
207 Args:
208 key:
209 The key or shorthand name for the activation to retrieve.
211 Returns:
212 The cached activation tensor corresponding to the given key.
213 """
214 if key in self.cache_dict:
215 return self.cache_dict[key]
216 elif type(key) == str:
217 return self.cache_dict[utils.get_act_name(key)]
218 else:
219 if len(key) > 1 and key[1] is not None:
220 if key[1] < 0:
221 # Supports negative indexing on the layer dimension
222 key = (key[0], self.model.cfg.n_layers + key[1], *key[2:])
223 return self.cache_dict[utils.get_act_name(*key)]
225 def __len__(self) -> int:
226 """Length of the ActivationCache.
228 Special method that returns the length of an object (in this case the number of different
229 activations in the cache).
230 """
231 return len(self.cache_dict)
233 def to(self, device: Union[str, torch.device]) -> ActivationCache:
234 """Move the Cache to a Device.
236 Mostly useful for moving the cache to the CPU after model computation finishes to save GPU
237 memory. Note however that operations will be much slower on the CPU. Note also that some
238 methods will break unless the model is also moved to the same device, eg
239 `compute_head_results`.
241 Args:
242 device:
243 The device to move the cache to (e.g. `torch.device.cpu`).
245 """
246 warn_if_mps(device)
247 self.cache_dict = {key: value.to(device) for key, value in self.cache_dict.items()}
248 return self
250 def toggle_autodiff(self, mode: bool = False):
251 """Toggle Autodiff Globally.
253 Applies `torch.set_grad_enabled(mode)` to the global state (not just TransformerLens).
255 Warning:
257 This is pretty dangerous, since autodiff is global state - this turns off torch's
258 ability to take gradients completely and it's easy to get a bunch of errors if you don't
259 realise what you're doing.
261 But autodiff consumes a LOT of GPU memory (since every intermediate activation is cached
262 until all downstream activations are deleted - this means that computing the loss and
263 storing it in a list will keep every activation sticking around!). So often when you're
264 analysing a model's activations, and don't need to do any training, autodiff is more trouble
265 than its worth.
267 If you don't want to mess with global state, using torch.inference_mode as a context manager
268 or decorator achieves similar effects:
270 >>> with torch.inference_mode():
271 ... y = torch.Tensor([1., 2, 3])
272 >>> y.requires_grad
273 False
274 """
275 logging.warning("Changed the global state, set autodiff to %s", mode)
276 torch.set_grad_enabled(mode)
278 def keys(self):
279 """Keys of the ActivationCache.
281 Examples:
283 >>> from transformer_lens import HookedTransformer
284 >>> model = HookedTransformer.from_pretrained("tiny-stories-1M")
285 Loaded pretrained model tiny-stories-1M into HookedTransformer
286 >>> _logits, cache = model.run_with_cache("Some prompt")
287 >>> list(cache.keys())[0:3]
288 ['hook_embed', 'hook_pos_embed', 'blocks.0.hook_resid_pre']
290 Returns:
291 List of all keys.
292 """
293 return self.cache_dict.keys()
295 def values(self):
296 """Values of the ActivationCache.
298 Returns:
299 List of all values.
300 """
301 return self.cache_dict.values()
303 def items(self):
304 """Items of the ActivationCache.
306 Returns:
307 List of all items ((key, value) tuples).
308 """
309 return self.cache_dict.items()
311 def __iter__(self) -> Iterator[str]:
312 """ActivationCache Iterator.
314 Special method that returns an iterator over the keys in the ActivationCache. Allows looping over the
315 cache.
317 Examples:
319 >>> from transformer_lens import HookedTransformer
320 >>> model = HookedTransformer.from_pretrained("tiny-stories-1M")
321 Loaded pretrained model tiny-stories-1M into HookedTransformer
322 >>> _logits, cache = model.run_with_cache("Some prompt")
323 >>> cache_interesting_names = []
324 >>> for key in cache:
325 ... if not key.startswith("blocks.") or key.startswith("blocks.0"):
326 ... cache_interesting_names.append(key)
327 >>> print(cache_interesting_names[0:3])
328 ['hook_embed', 'hook_pos_embed', 'blocks.0.hook_resid_pre']
330 Returns:
331 Iterator over the cache.
332 """
333 return self.cache_dict.__iter__()
335 def apply_slice_to_batch_dim(self, batch_slice: Union[Slice, SliceInput]) -> ActivationCache:
336 """Apply a Slice to the Batch Dimension.
338 Args:
339 batch_slice:
340 The slice to apply to the batch dimension.
342 Returns:
343 The ActivationCache with the batch dimension sliced.
344 """
345 if not isinstance(batch_slice, Slice):
346 batch_slice = Slice(batch_slice)
347 batch_slice = cast(Slice, batch_slice) # mypy can't seem to infer this
348 assert (
349 self.has_batch_dim or batch_slice.mode == "empty"
350 ), "Cannot index into a cache without a batch dim"
351 still_has_batch_dim = (batch_slice.mode != "int") and self.has_batch_dim
352 batch_size = self._batch_size()
353 # Broadcast entries (leading dim 1 when the true batch is larger) are not
354 # batched — leave them untouched so slicing can't index out of bounds.
355 new_cache_dict = {
356 name: (
357 batch_slice.apply(param, dim=0)
358 if param.ndim > 0 and param.size(0) == batch_size
359 else param
360 )
361 for name, param in self.cache_dict.items()
362 }
363 return ActivationCache(new_cache_dict, self.model, has_batch_dim=still_has_batch_dim)
365 def accumulated_resid(
366 self,
367 layer: Optional[int] = None,
368 incl_mid: bool = False,
369 apply_ln: bool = False,
370 pos_slice: Optional[Union[Slice, SliceInput]] = None,
371 mlp_input: bool = False,
372 return_labels: bool = False,
373 ) -> Union[
374 Float[torch.Tensor, "layers_covered *batch_and_pos_dims d_model"],
375 Tuple[Float[torch.Tensor, "layers_covered *batch_and_pos_dims d_model"], List[str]],
376 ]:
377 """Accumulated Residual Stream.
379 Returns the accumulated residual stream at each layer/sub-layer. This is useful for `Logit
380 Lens <https://www.lesswrong.com/posts/AcKRB8wDpdaN6v6ru/interpreting-gpt-the-logit-lens>`
381 style analysis, where it can be thought of as what the model "believes" at each point in the
382 residual stream.
384 To project this into the vocabulary space, remember that there is a final layer norm in most
385 decoder-only transformers. Therefore, you need to first apply the final layer norm (which
386 can be done with `apply_ln`), and then multiply by the unembedding matrix (:math:`W_U`)
387 and optionally add the unembedding bias (:math:`b_U`).
389 **Note on bias terms:** There are two valid approaches for the final projection:
391 1. **With bias terms:** Use `model.unembed(normalized_resid)` which applies both :math:`W_U`
392 and :math:`b_U` (equivalent to `normalized_resid @ model.W_U + model.b_U`). This works
393 correctly with both `fold_ln=True` and `fold_ln=False` settings, as the biases are
394 handled consistently.
395 2. **Without bias terms:** Use only `normalized_resid @ model.W_U`. If taking this approach,
396 you should instantiate the model with `fold_ln=True`, which folds the layer norm scaling
397 into :math:`W_U` and the layer norm bias into :math:`b_U`. Since `apply_ln=True` will
398 apply the (now parameter-free) layer norm, and you skip :math:`b_U`, no bias terms are
399 included. With `fold_ln=False`, the layer norm bias would still be applied, which is
400 typically not desired when excluding bias terms.
402 Both approaches are commonly used in the literature and are valid interpretability choices.
404 If you instead want to look at contributions to the residual stream from each component
405 (e.g. for direct logit attribution), see :meth:`decompose_resid` instead, or
406 :meth:`get_full_resid_decomposition` if you want contributions broken down further into each
407 MLP neuron.
409 Examples:
411 Logit Lens analysis can be done as follows:
413 >>> from transformer_lens import HookedTransformer
414 >>> import torch
415 >>> import pandas as pd
417 >>> model = HookedTransformer.from_pretrained("tiny-stories-1M", device="cpu", fold_ln=True)
418 Loaded pretrained model tiny-stories-1M into HookedTransformer
420 >>> prompt = "Why did the chicken cross the"
421 >>> answer = " road"
422 >>> logits, cache = model.run_with_cache("Why did the chicken cross the")
423 >>> answer_token = model.to_single_token(answer)
424 >>> print(answer_token)
425 2975
427 >>> accum_resid, labels = cache.accumulated_resid(return_labels=True, apply_ln=True)
428 >>> last_token_accum = accum_resid[:, 0, -1, :] # layer, batch, pos, d_model
429 >>> print(last_token_accum.shape) # layer, d_model
430 torch.Size([9, 64])
433 >>> W_U = model.W_U
434 >>> print(W_U.shape)
435 torch.Size([64, 50257])
437 >>> # Project to vocabulary without unembedding bias
438 >>> layers_logits = last_token_accum @ W_U # layer, d_vocab
439 >>> print(layers_logits.shape)
440 torch.Size([9, 50257])
442 >>> # If you want to apply the unembedding bias, add b_U when present:
443 >>> # b_U = getattr(model, "b_U", None)
444 >>> # layers_logits = layers_logits + b_U if b_U is not None else layers_logits
445 >>> # print(layers_logits.shape)
446 torch.Size([9, 50257])
448 >>> # Get the rank of the correct answer by layer
449 >>> sorted_indices = torch.argsort(layers_logits, dim=1, descending=True)
450 >>> rank_answer = (sorted_indices == 2975).nonzero(as_tuple=True)[1]
451 >>> print(pd.Series(rank_answer, index=labels))
452 0_pre 4442
453 1_pre 382
454 2_pre 982
455 3_pre 1160
456 4_pre 408
457 5_pre 145
458 6_pre 78
459 7_pre 387
460 final_post 6
461 dtype: int64
463 Args:
464 layer:
465 The layer to take components up to - by default includes resid_pre for that layer
466 and excludes resid_mid and resid_post for that layer. If set as `n_layers`, `-1` or
467 `None` it will return all residual streams, including the final one (i.e.
468 immediately pre logits). The indices are taken such that this gives the accumulated
469 streams up to the input to layer l.
470 incl_mid:
471 Whether to return `resid_mid` for all previous layers.
472 apply_ln:
473 Whether to apply the final layer norm to the stack. When True, applies
474 `model.ln_final`, which recomputes normalization statistics (mean and
475 variance/RMS) for each intermediate state in the stack, transforming the
476 activations into the format expected by the unembedding layer.
477 pos_slice:
478 A slice object to apply to the pos dimension. Defaults to None, do nothing.
479 mlp_input:
480 Whether to include resid_mid for the current layer. This essentially gives the MLP
481 input rather than the attention input.
482 return_labels:
483 Whether to return a list of labels for the residual stream components. Useful for
484 labelling graphs.
486 Returns:
487 A tensor of the accumulated residual streams. If `return_labels` is True, also returns a
488 list of labels for the components (as a tuple in the form `(components, labels)`).
489 """
490 if not isinstance(pos_slice, Slice):
491 pos_slice = Slice(pos_slice)
492 if layer is None or layer == -1:
493 # Default to the residual stream immediately pre unembed
494 layer = self.model.cfg.n_layers
495 assert isinstance(layer, int)
496 labels = []
497 components_list = []
498 for l in range(layer + 1):
499 if l == self.model.cfg.n_layers:
500 components_list.append(self[("resid_post", self.model.cfg.n_layers - 1)])
501 labels.append("final_post")
502 continue
503 components_list.append(self[("resid_pre", l)])
504 labels.append(f"{l}_pre")
505 if (incl_mid and l < layer) or (mlp_input and l == layer):
506 components_list.append(self[("resid_mid", l)])
507 labels.append(f"{l}_mid")
508 components_list = [pos_slice.apply(c, dim=-2) for c in components_list]
509 components = torch.stack(components_list, dim=0)
510 if apply_ln:
511 recompute_ln = layer == self.model.cfg.n_layers
512 components = self.apply_ln_to_stack(
513 components,
514 layer,
515 pos_slice=pos_slice,
516 mlp_input=mlp_input,
517 has_batch_dim=self.has_batch_dim,
518 recompute_ln=recompute_ln,
519 )
520 if return_labels:
521 return components, labels
522 else:
523 return components
525 def logit_attrs(
526 self,
527 residual_stack: Float[torch.Tensor, "num_components *batch_and_pos_dims d_model"],
528 tokens: Union[
529 str,
530 int,
531 Int[torch.Tensor, ""],
532 Int[torch.Tensor, "batch"],
533 Int[torch.Tensor, "batch position"],
534 ],
535 incorrect_tokens: Optional[
536 Union[
537 str,
538 int,
539 Int[torch.Tensor, ""],
540 Int[torch.Tensor, "batch"],
541 Int[torch.Tensor, "batch position"],
542 ]
543 ] = None,
544 pos_slice: Union[Slice, SliceInput] = None,
545 batch_slice: Union[Slice, SliceInput] = None,
546 has_batch_dim: bool = True,
547 ) -> Float[torch.Tensor, "num_components *batch_and_pos_dims_out"]:
548 """Logit Attributions.
550 Takes a residual stack (typically the residual stream decomposed by components), and
551 calculates how much each item in the stack "contributes" to specific tokens.
553 It does this by:
554 1. Getting the residual directions of the tokens (i.e. reversing the unembed)
555 2. Taking the dot product of each item in the residual stack, with the token residual
556 directions.
558 Note that if incorrect tokens are provided, it instead takes the difference between the
559 correct and incorrect tokens (to calculate the residual directions). This is useful as
560 sometimes we want to know e.g. which components are most responsible for selecting the
561 correct token rather than an incorrect one. For example in the `Interpretability in the Wild
562 paper <https://arxiv.org/abs/2211.00593>` prompts such as "John and Mary went to the shops,
563 John gave a bag to" were investigated, and it was therefore useful to calculate attribution
564 for the :math:`\\text{Mary} - \\text{John}` residual direction.
566 Warning:
568 Choosing the correct `tokens` and `incorrect_tokens` is both important and difficult. When
569 investigating specific components it's also useful to look at it's impact on all tokens
570 (i.e. :math:`\\text{final_ln}(\\text{residual_stack_item}) W_U`).
572 Args:
573 residual_stack:
574 Stack of components of residual stream to get logit attributions for.
575 tokens:
576 Tokens to compute logit attributions on.
577 incorrect_tokens:
578 If provided, compute attributions on logit difference between tokens and
579 incorrect_tokens. Must have the same shape as tokens.
580 pos_slice:
581 The slice to apply layer norm scaling on. Defaults to None, do nothing.
582 batch_slice:
583 The slice to take on the batch dimension during layer norm scaling. Defaults to
584 None, do nothing.
585 has_batch_dim:
586 Whether residual_stack has a batch dimension. Defaults to True.
588 Returns:
589 A tensor of the logit attributions or logit difference attributions if incorrect_tokens
590 was provided.
591 """
592 if not isinstance(pos_slice, Slice):
593 pos_slice = Slice(pos_slice)
595 if not isinstance(batch_slice, Slice):
596 batch_slice = Slice(batch_slice)
598 # Convert tokens to tensor for shape checking, but pass original to tokens_to_residual_directions
599 tokens_for_shape_check = tokens
601 if isinstance(tokens_for_shape_check, str):
602 tokens_for_shape_check = torch.as_tensor(
603 self.model.to_single_token(tokens_for_shape_check)
604 )
605 elif isinstance(tokens_for_shape_check, int):
606 tokens_for_shape_check = torch.as_tensor(tokens_for_shape_check)
608 logit_directions = self.model.tokens_to_residual_directions(tokens)
610 if incorrect_tokens is not None:
611 # Convert incorrect_tokens to tensor for shape checking, but pass original to tokens_to_residual_directions
612 incorrect_tokens_for_shape_check = incorrect_tokens
614 if isinstance(incorrect_tokens_for_shape_check, str):
615 incorrect_tokens_for_shape_check = torch.as_tensor(
616 self.model.to_single_token(incorrect_tokens_for_shape_check)
617 )
618 elif isinstance(incorrect_tokens_for_shape_check, int):
619 incorrect_tokens_for_shape_check = torch.as_tensor(incorrect_tokens_for_shape_check)
621 if tokens_for_shape_check.shape != incorrect_tokens_for_shape_check.shape:
622 raise ValueError(
623 f"tokens and incorrect_tokens must have the same shape! \
624 (tokens.shape={tokens_for_shape_check.shape}, \
625 incorrect_tokens.shape={incorrect_tokens_for_shape_check.shape})"
626 )
628 # If incorrect_tokens was provided, take the logit difference
629 logit_directions = logit_directions - self.model.tokens_to_residual_directions(
630 incorrect_tokens
631 )
633 scaled_residual_stack = self.apply_ln_to_stack(
634 residual_stack,
635 layer=-1,
636 pos_slice=pos_slice,
637 batch_slice=batch_slice,
638 has_batch_dim=has_batch_dim,
639 )
641 logit_attrs = (scaled_residual_stack * logit_directions).sum(dim=-1)
642 return logit_attrs
644 def decompose_resid(
645 self,
646 layer: Optional[int] = None,
647 mlp_input: bool = False,
648 mode: Literal["all", "mlp", "attn"] = "all",
649 apply_ln: bool = False,
650 pos_slice: Union[Slice, SliceInput] = None,
651 incl_embeds: bool = True,
652 return_labels: bool = False,
653 ) -> Union[
654 Float[torch.Tensor, "layers_covered *batch_and_pos_dims d_model"],
655 Tuple[Float[torch.Tensor, "layers_covered *batch_and_pos_dims d_model"], List[str]],
656 ]:
657 """Decompose the Residual Stream.
659 Decomposes the residual stream input to layer L into a stack of the output of previous
660 layers. The sum of these is the input to layer L (plus embedding and pos embedding). This is
661 useful for attributing model behaviour to different components of the residual stream
663 Args:
664 layer:
665 The layer to take components up to - by default includes
666 resid_pre for that layer and excludes resid_mid and resid_post for that layer.
667 layer==n_layers means to return all layer outputs incl in the final layer, layer==0
668 means just embed and pos_embed. The indices are taken such that this gives the
669 accumulated streams up to the input to layer l
670 mlp_input:
671 Whether to include attn_out for the current
672 layer - essentially decomposing the residual stream that's input to the MLP input
673 rather than the Attn input.
674 mode:
675 Values are "all", "mlp" or "attn". "all" returns all
676 components, "mlp" returns only the MLP components, and "attn" returns only the
677 attention components. Defaults to "all".
678 apply_ln:
679 Whether to apply LayerNorm to the stack.
680 pos_slice:
681 A slice object to apply to the pos dimension.
682 Defaults to None, do nothing.
683 incl_embeds:
684 Whether to include embed & pos_embed
685 return_labels:
686 Whether to return a list of labels for the residual stream components.
687 Useful for labelling graphs.
689 Returns:
690 A tensor of the accumulated residual streams. If `return_labels` is True, also returns
691 a list of labels for the components (as a tuple in the form `(components, labels)`).
692 """
693 if not isinstance(pos_slice, Slice):
694 pos_slice = Slice(pos_slice)
695 pos_slice = cast(Slice, pos_slice) # mypy can't seem to infer this
696 if layer is None or layer == -1:
697 # Default to the residual stream immediately pre unembed
698 layer = self.model.cfg.n_layers
699 assert isinstance(layer, int)
701 incl_attn = mode != "mlp"
702 incl_mlp = mode != "attn" and not self.model.cfg.attn_only
703 components_list = []
704 labels = []
705 if incl_embeds:
706 if self.has_embed: 706 ↛ 709line 706 didn't jump to line 709 because the condition on line 706 was always true
707 components_list = [self["hook_embed"]]
708 labels.append("embed")
709 if self.has_pos_embed: 709 ↛ 713line 709 didn't jump to line 713 because the condition on line 709 was always true
710 components_list.append(self["hook_pos_embed"])
711 labels.append("pos_embed")
713 for l in range(layer):
714 if incl_attn:
715 components_list.append(self[("attn_out", l)])
716 labels.append(f"{l}_attn_out")
717 if incl_mlp:
718 components_list.append(self[("mlp_out", l)])
719 labels.append(f"{l}_mlp_out")
720 if mlp_input and incl_attn:
721 components_list.append(self[("attn_out", layer)])
722 labels.append(f"{layer}_attn_out")
723 components_list = [pos_slice.apply(c, dim=-2) for c in components_list]
724 components = torch.stack(components_list, dim=0)
725 if apply_ln:
726 components = self.apply_ln_to_stack(
727 components, layer, pos_slice=pos_slice, mlp_input=mlp_input
728 )
729 if return_labels:
730 return components, labels
731 else:
732 return components
734 def compute_head_results(
735 self,
736 ):
737 """Compute Head Results.
739 Computes and caches the results for each attention head, ie the amount contributed to the
740 residual stream from that head. attn_out for a layer is the sum of head results plus b_O.
741 Intended use is to enable use_attn_results when running and caching the model, but this can
742 be useful if you forget.
744 Works for both HookedTransformer and TransformerBridge — bridge exposes
745 ``blocks[i].attn.W_O`` via its component-mapping compatibility shim.
746 """
747 # Return if valid 4D results exist; replace stale 3D Bridge entries if needed
748 first_key = "blocks.0.attn.hook_result"
749 if first_key in self.cache_dict:
750 val = self.cache_dict[first_key]
751 if isinstance(val, torch.Tensor) and val.ndim >= 4: 751 ↛ 755line 751 didn't jump to line 755 because the condition on line 751 was always true
752 logging.warning("Tried to compute head results when they were already cached")
753 return
754 # Remove stale 3D entries before recomputing
755 for layer in range(self.model.cfg.n_layers):
756 key = f"blocks.{layer}.attn.hook_result"
757 if key in self.cache_dict:
758 del self.cache_dict[key]
759 for layer in range(self.model.cfg.n_layers):
760 # Note that we haven't enabled set item on this object so we need to edit the underlying
761 # cache_dict directly.
763 # Add singleton dimension to match W_O's shape for broadcasting
764 z = einops.rearrange(
765 self[("z", layer, "attn")],
766 "... head_index d_head -> ... head_index d_head 1",
767 )
769 # Element-wise multiplication of z and W_O (with shape [head_index, d_head, d_model])
770 block = self.model.blocks[layer]
771 result = z * block.attn.W_O
773 # Sum over d_head to get the contribution of each head to the residual stream
774 self.cache_dict[f"blocks.{layer}.attn.hook_result"] = result.sum(dim=-2)
776 def ssm_layers(self, mixer_type: Optional[Union[type, Tuple[type, ...]]] = None) -> List[int]:
777 """Return the block indices whose mixer is an SSM / recurrent mixer.
779 Family-agnostic and purely structural: finds each block's *realized* SSM
780 mixer in its variant slot (``.mixer`` / ``.linear_attn`` / …) via
781 ``find_ssm_mixer``, which excludes a hybrid's passthrough ``.mixer`` on
782 non-SSM layers (e.g. NemotronH attention/MLP/MoE) — no dependence on
783 ``cfg.layers_block_type``.
785 Args:
786 mixer_type: Optional concrete bridge class to filter to (e.g.
787 ``SSM2MixerBridge`` for only Mamba-2 layers).
789 Returns:
790 Ascending list of SSM block indices.
791 """
792 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
793 find_ssm_mixer,
794 )
796 bridge = self.model
797 layers: List[int] = []
798 for i, block in enumerate(bridge.blocks):
799 mixer = find_ssm_mixer(block)
800 if mixer is None:
801 continue
802 if mixer_type is not None and not isinstance(mixer, mixer_type):
803 continue
804 layers.append(i)
805 return layers
807 def _over_ssm_layers(
808 self,
809 fn: Callable[[Any, int], torch.Tensor],
810 mixer_type: Optional[Union[type, Tuple[type, ...]]] = None,
811 ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]:
812 """Apply ``fn(mixer, layer_idx)`` over every SSM layer; stack or dict.
814 The single SSM layer-enumeration used by both ``compute_ssm_state`` and
815 ``compute_ssm_effective_attention``. Returns a stacked tensor (dim 0 =
816 layer) when *every* block is an SSM layer, else a ``{layer_idx: result}``
817 dict over the SSM layers (heterogeneous hybrids).
818 """
819 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
820 find_ssm_mixer,
821 )
823 bridge = self.model
824 indices = self.ssm_layers(mixer_type=mixer_type)
825 if not indices: 825 ↛ 826line 825 didn't jump to line 826 because the condition on line 825 was never true
826 raise RuntimeError(
827 "No SSM mixer layers found. Use an SSM / hybrid bridge and "
828 "run_with_cache first (gated-delta-net needs use_cache=False)."
829 )
830 results: Dict[int, torch.Tensor] = {
831 i: fn(cast(Any, find_ssm_mixer(bridge.blocks[i])), i) for i in indices
832 }
833 if indices == list(range(len(bridge.blocks))):
834 return torch.stack([results[i] for i in indices], dim=0)
835 return results
837 def compute_ssm_effective_attention(
838 self,
839 layer: Optional[int] = None,
840 include_dt_scaling: bool = False,
841 per_state_coord: bool = False,
842 ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]:
843 """Materialize SSM effective attention for one or all SSM layers, any family.
845 The single discovery surface for effective attention across Mamba-1,
846 Mamba-2, and gated-delta-net layers; dispatches to each layer's mixer
847 ``compute_effective_attention`` regardless of family. Family-specific
848 options are forwarded only to mixers whose signature accepts them.
850 Args:
851 layer: Specific block index, or None for every SSM layer.
852 include_dt_scaling: Forwarded to Mamba-1/Mamba-2 mixers (the
853 reconstruction form); gated-delta-net does not accept it.
854 per_state_coord: Forwarded to Mamba-1 only (per-state-coordinate
855 matrices). Setting it True for another family raises ValueError.
857 Returns:
858 A per-layer matrix for a single ``layer``; for ``layer=None`` a
859 stacked tensor (dim 0 = layer) when every block is an SSM layer, else
860 a ``{layer_idx: matrix}`` dict over the SSM layers.
862 Raises:
863 TypeError: If the requested ``layer`` has no SSM mixer.
864 ValueError: If an option is set that the target mixer does not support.
865 RuntimeError: If ``layer=None`` finds no SSM layers.
866 """
867 import inspect
869 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
870 find_ssm_mixer,
871 )
873 def _call(mixer: Any, layer_idx: int) -> torch.Tensor:
874 fn = cast(Any, mixer).compute_effective_attention
875 params = inspect.signature(fn).parameters
876 kwargs: Dict[str, Any] = {}
877 for name, value in (
878 ("include_dt_scaling", include_dt_scaling),
879 ("per_state_coord", per_state_coord),
880 ):
881 if name in params:
882 kwargs[name] = value
883 elif value:
884 raise ValueError(
885 f"{type(mixer).__name__}.compute_effective_attention does not "
886 f"support {name}=True."
887 )
888 result: torch.Tensor = fn(cache=self, layer_idx=layer_idx, **kwargs)
889 return result
891 if layer is not None:
892 single = find_ssm_mixer(self.model.blocks[layer])
893 if single is None:
894 raise TypeError(f"Block {layer} has no SSM mixer (no compute_effective_attention).")
895 return _call(single, layer)
896 return self._over_ssm_layers(_call)
898 def compute_ssm_state(
899 self,
900 layer: Optional[int] = None,
901 time_step: Optional[int] = None,
902 ) -> Union[torch.Tensor, Dict[int, torch.Tensor]]:
903 """Reconstruct the recurrent SSM state ``S`` from this cache.
905 The single discovery surface for recurrent state across families —
906 ``SSMMixerBridge`` (Mamba-1), ``SSM2MixerBridge`` (Mamba-2) and
907 ``GatedDeltaNetBridge`` (gated delta rule) each reconstruct it — mirroring
908 ``compute_head_results``. Read-only post-hoc reconstruction from cached
909 hooks (no forward re-run); requires an SSM / SSM-hybrid bridge cached via
910 ``run_with_cache`` (gated-delta-net additionally needs ``use_cache=False``
911 so its interior hooks fire). See the mixer's ``compute_ssm_state`` for the
912 recurrence, shapes, and the ``time_step`` memory bound; state shape is
913 family-specific, so ``layer=None`` returns a dict except when every block
914 shares one mixer type.
916 Args:
917 layer: Specific block index, or None for every SSM-state layer.
918 time_step: Optional single position (memory-bounded); None for all.
920 Returns:
921 A per-layer state tensor for a single ``layer``; for ``layer=None`` a
922 stacked tensor (dim 0 = layer) when every block is an SSM-state layer,
923 else a ``{layer_idx: state}`` dict over those layers.
925 Raises:
926 TypeError: If the requested ``layer`` has no state-reconstructing mixer.
927 RuntimeError: If ``layer=None`` finds no such layers.
928 """
929 from transformer_lens.model_bridge.generalized_components import (
930 GatedDeltaNetBridge,
931 SSM2MixerBridge,
932 SSMMixerBridge,
933 )
934 from transformer_lens.model_bridge.generalized_components.ssm_protocol import (
935 find_ssm_mixer,
936 )
938 # Every recurrent family reconstructs S_t from its cached interior hooks.
939 state_mixers = (SSMMixerBridge, SSM2MixerBridge, GatedDeltaNetBridge)
941 def _call(mixer: Any, layer_idx: int) -> torch.Tensor:
942 state: torch.Tensor = cast(Any, mixer).compute_ssm_state(
943 self, layer_idx=layer_idx, time_step=time_step
944 )
945 return state
947 if layer is not None:
948 single = find_ssm_mixer(self.model.blocks[layer])
949 if not isinstance(single, state_mixers):
950 raise TypeError(
951 f"Block {layer} has no state-reconstructing SSM mixer; "
952 "compute_ssm_state supports Mamba-1 / Mamba-2 / gated-delta-net."
953 )
954 return _call(single, layer)
955 return self._over_ssm_layers(_call, mixer_type=state_mixers)
957 def stack_head_results(
958 self,
959 layer: int = -1,
960 return_labels: bool = False,
961 incl_remainder: bool = False,
962 pos_slice: Union[Slice, SliceInput] = None,
963 apply_ln: bool = False,
964 ) -> Union[
965 Float[torch.Tensor, "num_components *batch_and_pos_dims d_model"],
966 Tuple[Float[torch.Tensor, "num_components *batch_and_pos_dims d_model"], List[str]],
967 ]:
968 """Stack Head Results.
970 Returns a stack of all head results (ie residual stream contribution) up to layer L. A good
971 way to decompose the outputs of attention layers into attribution by specific heads. Note
972 that the num_components axis has length layer x n_heads ((layer head_index) in einops
973 notation).
975 Args:
976 layer:
977 Layer index - heads at all layers strictly before this are included. layer must be
978 in [1, n_layers-1], or any of (n_layers, -1, None), which all mean the final layer.
979 return_labels:
980 Whether to also return a list of labels of the form "L0H0" for the heads.
981 incl_remainder:
982 Whether to return a final term which is "the rest of the residual stream".
983 pos_slice:
984 A slice object to apply to the pos dimension. Defaults to None, do nothing.
985 apply_ln:
986 Whether to apply LayerNorm to the stack.
987 """
988 if not isinstance(pos_slice, Slice):
989 pos_slice = Slice(pos_slice)
990 pos_slice = cast(Slice, pos_slice) # mypy can't seem to infer this
991 if layer is None or layer == -1:
992 # Default to the residual stream immediately pre unembed
993 layer = self.model.cfg.n_layers
995 # Idempotent; cleans up stale Bridge entries
996 self.compute_head_results()
998 components: Any = []
999 labels = []
1000 for l in range(layer):
1001 # Note that this has shape batch x pos x head_index x d_model
1002 components.append(pos_slice.apply(self[("result", l, "attn")], dim=-3))
1003 labels.extend([f"L{l}H{h}" for h in range(self.model.cfg.n_heads)])
1004 if components:
1005 components = torch.cat(components, dim=-2)
1006 components = einops.rearrange(
1007 components,
1008 "... concat_head_index d_model -> concat_head_index ... d_model",
1009 )
1010 if incl_remainder:
1011 remainder = pos_slice.apply(
1012 self[("resid_post", layer - 1)], dim=-2
1013 ) - components.sum(dim=0)
1014 components = torch.cat([components, remainder[None]], dim=0)
1015 labels.append("remainder")
1016 elif incl_remainder:
1017 # There are no components, so the remainder is the entire thing.
1018 components = torch.cat(
1019 [pos_slice.apply(self[("resid_post", layer - 1)], dim=-2)[None]], dim=0
1020 )
1021 labels.append("remainder")
1022 else:
1023 # If this is called with layer 0, we return an empty tensor of the right shape to be
1024 # stacked correctly. This uses the shape of hook_embed, which is pretty janky since it
1025 # assumes embed is in the cache. But it's hard to explicitly code the shape, since it
1026 # depends on the pos slice, whether we have a batch dim, etc. And it's pretty messy!
1027 components = torch.zeros(
1028 0,
1029 *pos_slice.apply(self["hook_embed"], dim=-2).shape,
1030 device=self.model.cfg.device,
1031 )
1033 if apply_ln:
1034 components = self.apply_ln_to_stack(components, layer, pos_slice=pos_slice)
1036 if return_labels:
1037 return components, labels
1038 else:
1039 return components
1041 def stack_activation(
1042 self,
1043 activation_name: str,
1044 layer: int = -1,
1045 sublayer_type: Optional[str] = None,
1046 ) -> Float[torch.Tensor, "layers_covered ..."]:
1047 """Stack Activations.
1049 Flexible way to stack activations with a given name.
1051 Args:
1052 activation_name:
1053 The name of the activation to be stacked
1054 layer:
1055 'Layer index - heads' at all layers strictly before this are included. layer must be
1056 in [1, n_layers-1], or any of (n_layers, -1, None), which all mean the final layer.
1057 sublayer_type:
1058 The sub layer type of the activation, passed to utils.get_act_name. Can normally be
1059 inferred.
1060 incl_remainder:
1061 Whether to return a final term which is "the rest of the residual stream".
1062 """
1063 if layer is None or layer == -1:
1064 # Default to the residual stream immediately pre unembed
1065 layer = self.model.cfg.n_layers
1067 components = []
1068 for l in range(layer):
1069 components.append(self[(activation_name, l, sublayer_type)])
1071 return torch.stack(components, dim=0)
1073 def get_neuron_results(
1074 self,
1075 layer: int,
1076 neuron_slice: Union[Slice, SliceInput] = None,
1077 pos_slice: Union[Slice, SliceInput] = None,
1078 project_output_onto: Optional[torch.Tensor] = None,
1079 ) -> torch.Tensor:
1080 """Get Neuron Results.
1082 Get the results of for neurons in a specific layer (i.e, how much each neuron contributes to
1083 the residual stream). Does it for the subset of neurons specified by neuron_slice, defaults
1084 to all of them. Does *not* cache these because it's expensive in space and cheap to compute.
1086 Args:
1087 layer:
1088 Layer index.
1089 neuron_slice:
1090 Slice of the neuron.
1091 pos_slice:
1092 Slice of the positions.
1093 project_output_onto:
1094 Optional ``[d_model]`` or ``[d_model, num_outputs]`` projection. Contracted with
1095 ``W_out`` *before* the per-neuron expansion so the ``[..., d_mlp, d_model]``
1096 intermediate is never materialized.
1098 Returns:
1099 Last-dim is ``d_model`` (default), ``num_outputs`` (2D projection), or squeezed
1100 (1D projection).
1101 """
1102 if not isinstance(neuron_slice, Slice):
1103 neuron_slice = Slice(neuron_slice)
1104 if not isinstance(pos_slice, Slice):
1105 pos_slice = Slice(pos_slice)
1107 neuron_acts = self[("post", layer, "mlp")]
1108 block = self.model.blocks[layer]
1109 W_out = block.mlp.W_out
1110 if pos_slice is not None: 1110 ↛ 1114line 1110 didn't jump to line 1114 because the condition on line 1110 was always true
1111 # Note - order is important, as Slice.apply *may* collapse a dimension, so this ensures
1112 # that position dimension is -2 when we apply position slice
1113 neuron_acts = pos_slice.apply(neuron_acts, dim=-2)
1114 if neuron_slice is not None: 1114 ↛ 1117line 1114 didn't jump to line 1117 because the condition on line 1114 was always true
1115 neuron_acts = neuron_slice.apply(neuron_acts, dim=-1)
1116 W_out = neuron_slice.apply(W_out, dim=0)
1117 if project_output_onto is None:
1118 return neuron_acts[..., None] * W_out
1119 # W_out: [d_mlp, d_model]; project: [d_model] or [d_model, n_outs]
1120 projected = W_out @ project_output_onto
1121 if projected.ndim == 1:
1122 return neuron_acts * projected
1123 return neuron_acts[..., None] * projected
1125 def _get_cached_ln_scale(
1126 self,
1127 layer: Optional[int],
1128 mlp_input: bool,
1129 pos_slice: Slice,
1130 batch_slice: Optional[Slice] = None,
1131 ) -> torch.Tensor:
1132 """Look up the cached LN scale and apply pos/batch slicing. Surfaces a clearer error
1133 when the expected hook isn't in the cache (some non-decoder-only architectures expose
1134 LN scale at a different path or not at all).
1135 """
1136 if layer == self.model.cfg.n_layers or layer is None:
1137 key = "ln_final.hook_scale"
1138 else:
1139 key = f"blocks.{layer}.ln{2 if mlp_input else 1}.hook_scale"
1140 try:
1141 scale = self[key]
1142 except KeyError as e:
1143 raise KeyError(
1144 f"Cached LN scale not found at '{key}'. apply_ln operations require the model "
1145 f"to have cached this hook (some non-decoder-only architectures expose LN scale "
1146 f"under different module paths)."
1147 ) from e
1148 scale = pos_slice.apply(scale, dim=-2)
1149 if batch_slice is not None and self.has_batch_dim:
1150 scale = batch_slice.apply(scale)
1151 return scale
1153 def _stack_neuron_results_apply_ln_projected(
1154 self,
1155 layer: int,
1156 pos_slice: Slice,
1157 neuron_slice: Slice,
1158 project_2d: torch.Tensor,
1159 ) -> torch.Tensor:
1160 """LN-applied neuron stack with projection folded in — no d_mlp×d_model intermediate.
1162 Analytical formula (LN models, cached scale ``s``):
1163 ``LN_s(a_n * W_out_n) @ p = (a_n / s) * (W_out_n @ p - mean(W_out_n) * sum_p)``
1164 RMS models drop the ``mean(W_out_n) * sum_p`` term (no centering). Always uses the
1165 ln1 scale (mlp_input=False) since ``stack_neuron_results`` doesn't expose mlp_input.
1167 """
1168 scale = self._get_cached_ln_scale(layer, mlp_input=False, pos_slice=pos_slice)
1170 apply_centering = self.model.cfg.normalization_type in ["LN", "LNPre"]
1171 sum_p = project_2d.sum(dim=0) if apply_centering else None # [n_outs]
1173 components: list = []
1174 for l in range(layer):
1175 block = self.model.blocks[l]
1176 W_out_l = block.mlp.W_out # [d_mlp, d_model]
1177 W_out_l_sliced = neuron_slice.apply(W_out_l, dim=0)
1178 W_proj_l = W_out_l_sliced @ project_2d # [d_mlp, n_outs]
1179 if apply_centering: 1179 ↛ 1184line 1179 didn't jump to line 1184 because the condition on line 1179 was always true
1180 assert sum_p is not None # set when apply_centering, narrow for mypy
1181 W_means_l = W_out_l_sliced.mean(dim=-1) # [d_mlp]
1182 lin_form_l = W_proj_l - W_means_l[:, None] * sum_p[None, :]
1183 else:
1184 lin_form_l = W_proj_l
1185 a_l = self[("post", l, "mlp")]
1186 a_l = pos_slice.apply(a_l, dim=-2)
1187 a_l = neuron_slice.apply(a_l, dim=-1)
1188 # (a_l / s)[..., None] is [..., d_mlp, 1]; broadcast with lin_form_l [d_mlp, n_outs]
1189 components.append((a_l / scale)[..., None] * lin_form_l)
1190 if not components: 1190 ↛ 1191line 1190 didn't jump to line 1191 because the condition on line 1190 was never true
1191 empty_src = pos_slice.apply(self["hook_embed"], dim=-2)
1192 return torch.zeros(
1193 0, *empty_src.shape[:-1], project_2d.shape[-1], device=self.model.cfg.device
1194 )
1195 stacked = torch.cat(components, dim=-2)
1196 return einops.rearrange(
1197 stacked, "... concat_neuron_index n_outs -> concat_neuron_index ... n_outs"
1198 )
1200 def stack_neuron_results(
1201 self,
1202 layer: int,
1203 pos_slice: Union[Slice, SliceInput] = None,
1204 neuron_slice: Union[Slice, SliceInput] = None,
1205 return_labels: bool = False,
1206 incl_remainder: bool = False,
1207 apply_ln: bool = False,
1208 project_output_onto: Optional[torch.Tensor] = None,
1209 ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[str]]]:
1210 """Stack Neuron Results
1212 Returns a stack of all neuron results (ie residual stream contribution) up to layer L - ie
1213 the amount each individual neuron contributes to the residual stream. Also returns a list of
1214 labels of the form "L0N0" for the neurons. A good way to decompose the outputs of MLP layers
1215 into attribution by specific neurons.
1217 Note that doing this for all neurons is SUPER expensive on GPU memory and only works for
1218 small models or short inputs. Pass ``project_output_onto`` to fold the projection into the
1219 per-neuron expansion and avoid the ``[..., d_mlp, d_model]`` intermediate.
1221 Args:
1222 layer:
1223 Layer index - heads at all layers strictly before this are included. layer must be
1224 in [1, n_layers]
1225 pos_slice:
1226 Slice of the positions.
1227 neuron_slice:
1228 Slice of the neurons.
1229 return_labels:
1230 Whether to also return a list of labels of the form "L0H0" for the heads.
1231 incl_remainder:
1232 Whether to return a final term which is "the rest of the residual stream".
1233 apply_ln:
1234 Whether to apply LayerNorm to the stack.
1235 project_output_onto:
1236 Optional ``[d_model]`` or ``[d_model, num_outputs]`` tensor. When set, each
1237 component's last d_model dim is replaced by the projection (memory-efficient for
1238 direction analyses; see ``get_neuron_results``). Combined with ``apply_ln=True``,
1239 the projection is folded into the analytical cached-scale LN so the
1240 ``[..., d_mlp, d_model]`` intermediate is still never materialized.
1241 """
1242 if layer is None or layer == -1:
1243 # Default to the residual stream immediately pre unembed
1244 layer = self.model.cfg.n_layers
1246 if not isinstance(neuron_slice, Slice):
1247 neuron_slice = Slice(neuron_slice)
1248 if not isinstance(pos_slice, Slice):
1249 pos_slice = Slice(pos_slice)
1251 project_2d, squeeze_projected = _normalize_projection_to_2d(project_output_onto)
1253 d_mlp = self.model.cfg.d_mlp
1254 assert d_mlp is not None, "model.cfg.d_mlp must be set"
1255 neuron_labels: Union[torch.Tensor, np.ndarray] = neuron_slice.apply(
1256 torch.arange(d_mlp), dim=0
1257 )
1258 if isinstance(neuron_labels, int): 1258 ↛ 1259line 1258 didn't jump to line 1259 because the condition on line 1258 was never true
1259 neuron_labels = np.array([neuron_labels])
1261 labels = [f"L{l}N{h}" for l in range(layer) for h in neuron_labels]
1262 components: Any
1263 ln_folded = apply_ln and project_2d is not None
1264 if ln_folded:
1265 assert project_2d is not None # narrow for mypy
1266 # Analytical LN+projection — no d_mlp×d_model intermediate.
1267 components = self._stack_neuron_results_apply_ln_projected(
1268 layer, pos_slice, neuron_slice, project_2d
1269 )
1270 if incl_remainder:
1271 # Linearity of cached-scale LN: remainder is LN_s(resid_post) @ p - sum(neurons).
1272 resid_post = pos_slice.apply(self[("resid_post", layer - 1)], dim=-2)
1273 resid_post_ln = self.apply_ln_to_stack(
1274 resid_post[None], layer, pos_slice=pos_slice
1275 )[0]
1276 remainder = resid_post_ln @ project_2d
1277 if components.shape[0] > 0: 1277 ↛ 1279line 1277 didn't jump to line 1279 because the condition on line 1277 was always true
1278 remainder = remainder - components.sum(dim=0)
1279 components = torch.cat([components, remainder[None]], dim=0)
1280 labels.append("remainder")
1281 else:
1282 per_layer: list = []
1283 for l in range(layer):
1284 per_layer.append(
1285 self.get_neuron_results(
1286 l,
1287 pos_slice=pos_slice,
1288 neuron_slice=neuron_slice,
1289 project_output_onto=project_2d,
1290 )
1291 )
1292 if per_layer:
1293 components = torch.cat(per_layer, dim=-2)
1294 components = einops.rearrange(
1295 components,
1296 "... concat_neuron_index d_model -> concat_neuron_index ... d_model",
1297 )
1298 if incl_remainder:
1299 remainder_full = pos_slice.apply(self[("resid_post", layer - 1)], dim=-2)
1300 if project_2d is not None:
1301 remainder_full = remainder_full @ project_2d
1302 remainder = remainder_full - components.sum(dim=0)
1303 components = torch.cat([components, remainder[None]], dim=0)
1304 labels.append("remainder")
1305 elif incl_remainder:
1306 remainder_full = pos_slice.apply(self[("resid_post", layer - 1)], dim=-2)
1307 if project_2d is not None: 1307 ↛ 1308line 1307 didn't jump to line 1308 because the condition on line 1307 was never true
1308 remainder_full = remainder_full @ project_2d
1309 components = torch.cat([remainder_full[None]], dim=0)
1310 labels.append("remainder")
1311 else:
1312 empty_shape_src = pos_slice.apply(self["hook_embed"], dim=-2)
1313 if project_2d is not None: 1313 ↛ 1314line 1313 didn't jump to line 1314 because the condition on line 1313 was never true
1314 empty_shape_src = empty_shape_src @ project_2d
1315 components = torch.zeros(0, *empty_shape_src.shape, device=self.model.cfg.device)
1317 if apply_ln:
1318 components = self.apply_ln_to_stack(components, layer, pos_slice=pos_slice)
1320 if squeeze_projected:
1321 components = components.squeeze(-1)
1323 if return_labels:
1324 return components, labels
1325 else:
1326 return components
1328 def apply_ln_to_stack(
1329 self,
1330 residual_stack: Float[torch.Tensor, "num_components *batch_and_pos_dims d_model"],
1331 layer: Optional[int] = None,
1332 mlp_input: bool = False,
1333 pos_slice: Union[Slice, SliceInput] = None,
1334 batch_slice: Union[Slice, SliceInput] = None,
1335 has_batch_dim: bool = True,
1336 recompute_ln: bool = False,
1337 ) -> Float[torch.Tensor, "num_components *batch_and_pos_dims_out d_model"]:
1338 """Apply Layer Norm to a Stack.
1340 Takes a stack of components of the residual stream (eg outputs of decompose_resid or
1341 accumulated_resid), treats them as the input to a specific layer, and applies the layer norm
1342 scaling of that layer to them, using the cached scale factors - simulating what that
1343 component of the residual stream contributes to that layer's input.
1345 The layernorm scale is global across the entire residual stream for each layer, batch
1346 element and position, which is why we need to use the cached scale factors rather than just
1347 applying a new LayerNorm.
1349 When recompute_ln=True and the target layer is the final layer (unembed), each
1350 component is normalized using stats recomputed from that component; use this for logit lens
1351 analysis. When recompute_ln=False, a single cached scale is used for all components.
1353 If the model does not use LayerNorm or RMSNorm, it returns the residual stack unchanged.
1355 Args:
1356 residual_stack:
1357 A tensor, whose final dimension is d_model. The other trailing dimensions are
1358 assumed to be the same as the stored hook_scale - which may or may not include batch
1359 or position dimensions.
1360 layer:
1361 The layer we're taking the input to. In [0, n_layers], n_layers means the unembed.
1362 None maps to the n_layers case, ie the unembed.
1363 mlp_input:
1364 Whether the input is to the MLP or attn (ie ln2 vs ln1). Defaults to False, ie ln1.
1365 If layer==n_layers, must be False, and we use ln_final
1366 pos_slice:
1367 The slice to take of positions, if residual_stack is not over the full context, None
1368 means do nothing. It is assumed that pos_slice has already been applied to
1369 residual_stack, and this is only applied to the scale. See utils.Slice for details.
1370 Defaults to None, do nothing.
1371 batch_slice:
1372 The slice to take on the batch dimension. Defaults to None, do nothing.
1373 has_batch_dim:
1374 Whether residual_stack has a batch dimension.
1375 recompute_ln:
1376 If True and target layer is the unembed (final layer), apply the final layer norm
1377 to each component with statistics recomputed from that component. Defaults to False.
1379 """
1380 if self.model.cfg.normalization_type not in ["LN", "LNPre", "RMS", "RMSPre"]: 1380 ↛ 1382line 1380 didn't jump to line 1382 because the condition on line 1380 was never true
1381 # The model does not use LayerNorm, so we don't need to do anything.
1382 return residual_stack
1383 if not isinstance(pos_slice, Slice):
1384 pos_slice = Slice(pos_slice)
1385 if not isinstance(batch_slice, Slice):
1386 batch_slice = Slice(batch_slice)
1388 if layer is None or layer == -1:
1389 # Default to the residual stream immediately pre unembed
1390 layer = self.model.cfg.n_layers
1392 if has_batch_dim:
1393 # Apply batch slice to the stack
1394 residual_stack = batch_slice.apply(residual_stack, dim=1)
1396 # Logit lens: apply final layer norm to each component with recomputed statistics
1397 if recompute_ln and layer == self.model.cfg.n_layers and hasattr(self.model, "ln_final"):
1398 ln_final = self.model.ln_final
1399 results = []
1400 for i in range(residual_stack.shape[0]):
1401 x = residual_stack[i]
1402 original_shape = x.shape
1403 # ln_final expects (batch, pos, d_model); restore missing structural dimensions
1404 if not has_batch_dim:
1405 x = x.unsqueeze(0)
1406 if x.ndim == 2:
1407 x = x.unsqueeze(1)
1408 out = ln_final(x)
1409 results.append(out.reshape(original_shape))
1410 return torch.stack(results, dim=0)
1412 # Center the stack onlny if the model uses LayerNorm
1413 if self.model.cfg.normalization_type in ["LN", "LNPre"]:
1414 residual_stack = residual_stack - residual_stack.mean(dim=-1, keepdim=True)
1416 # Shape is [batch, position, 1] or [position, 1]; final dim is a dummy for broadcasting.
1417 scale = self._get_cached_ln_scale(layer, mlp_input, pos_slice, batch_slice)
1419 return residual_stack / scale
1421 def get_full_resid_decomposition(
1422 self,
1423 layer: Optional[int] = None,
1424 mlp_input: bool = False,
1425 expand_neurons: bool = True,
1426 apply_ln: bool = False,
1427 pos_slice: Union[Slice, SliceInput] = None,
1428 return_labels: bool = False,
1429 project_output_onto: Optional[torch.Tensor] = None,
1430 ) -> Union[torch.Tensor, Tuple[torch.Tensor, List[str]]]:
1431 """Get the full Residual Decomposition.
1433 Decomposes the residual stream that is input into some layer into its
1434 constituent components: every attention head result, every neuron (or
1435 MLP layer) result, the embeddings, and the accumulated biases.
1437 The returned tensor stacks components along ``dim=0`` in this order:
1439 1. Attention head results, layer-by-layer (``L * n_heads`` rows)
1440 2. Neuron / MLP results (only if ``cfg.attn_only=False`` and
1441 ``layer > 0``; ``L * d_mlp`` rows when ``expand_neurons=True``,
1442 else ``L`` rows)
1443 3. ``embed`` (1 row, if the model has token embeddings)
1444 4. ``pos_embed`` (1 row, if the model has positional embeddings)
1445 5. ``bias`` (1 row, the accumulated layer biases)
1447 ``return_labels=True`` returns a list of strings in the same order, so
1448 ``labels[i]`` always names ``stack[i]``. If you need to extract a
1449 specific component, slice by label rather than by hard-coded index —
1450 the row counts depend on ``layer``, ``expand_neurons``,
1451 ``cfg.attn_only``, and whether the model has positional embeddings.
1453 Args:
1454 layer:
1455 The layer we're inputting into. layer is in [0, n_layers], if layer==n_layers (or
1456 None) we're inputting into the unembed (the entire stream), if layer==0 then it's
1457 just embed and pos_embed
1458 mlp_input:
1459 Are we inputting to the MLP in that layer or the attn? Must be False for final
1460 layer, since that's the unembed.
1461 expand_neurons:
1462 Whether to expand the MLP outputs to give every neuron's result or just return the
1463 MLP layer outputs.
1464 apply_ln:
1465 Whether to apply LayerNorm to the stack.
1466 pos_slice:
1467 Slice of the positions to take.
1468 return_labels:
1469 Whether to return the labels.
1470 project_output_onto:
1471 Optional ``[d_model]`` or ``[d_model, num_outputs]`` projection. Folded in
1472 *before* the per-neuron expansion, so the ``[..., d_mlp, d_model]`` intermediate
1473 is never materialized (memory saving applies only with ``expand_neurons=True``).
1474 Combined with ``apply_ln=True``, the projection is fused into the analytical
1475 cached-scale LN so the same memory benefit holds. Output last-dim is squeezed
1476 for a 1D projection; ``num_outputs`` for 2D.
1477 """
1478 if layer is None or layer == -1:
1479 # Default to the residual stream immediately pre unembed
1480 layer = self.model.cfg.n_layers
1481 assert layer is not None # keep mypy happy
1483 if not isinstance(pos_slice, Slice):
1484 pos_slice = Slice(pos_slice)
1486 project_2d, squeeze_projected = _normalize_projection_to_2d(project_output_onto)
1487 # When both apply_ln and projection are requested, LN is applied per-component (in
1488 # d_model space for the small ones, analytically for neurons) before projection, so the
1489 # final apply_ln_to_stack call is skipped — last-dim is already n_outs.
1490 ln_folded = apply_ln and project_2d is not None
1492 def _ln_then_project(stack: torch.Tensor) -> torch.Tensor:
1493 stack = self.apply_ln_to_stack(stack, layer, pos_slice=pos_slice, mlp_input=mlp_input)
1494 return stack @ project_2d if project_2d is not None else stack
1496 head_stack, head_labels = self.stack_head_results(
1497 layer + (1 if mlp_input else 0), pos_slice=pos_slice, return_labels=True
1498 )
1499 if ln_folded:
1500 head_stack = _ln_then_project(head_stack)
1501 elif project_2d is not None:
1502 head_stack = head_stack @ project_2d
1503 labels = head_labels
1504 components = [head_stack]
1505 if not self.model.cfg.attn_only and layer > 0:
1506 if expand_neurons:
1507 # Only ask stack_neuron_results to apply LN when we want the fused analytical
1508 # path (ln_folded). For the unfolded case the outer apply_ln_to_stack handles it.
1509 neuron_stack, neuron_labels = self.stack_neuron_results(
1510 layer,
1511 pos_slice=pos_slice,
1512 return_labels=True,
1513 apply_ln=ln_folded,
1514 project_output_onto=project_2d,
1515 )
1516 labels.extend(neuron_labels)
1517 components.append(neuron_stack)
1518 else:
1519 # Get the stack of just the MLP outputs
1520 # mlp_input included for completeness, but it doesn't actually matter, since it's
1521 # just for MLP outputs
1522 mlp_stack, mlp_labels = self.decompose_resid(
1523 layer,
1524 mlp_input=mlp_input,
1525 pos_slice=pos_slice,
1526 incl_embeds=False,
1527 mode="mlp",
1528 return_labels=True,
1529 )
1530 if ln_folded: 1530 ↛ 1531line 1530 didn't jump to line 1531 because the condition on line 1530 was never true
1531 mlp_stack = _ln_then_project(mlp_stack)
1532 elif project_2d is not None: 1532 ↛ 1533line 1532 didn't jump to line 1533 because the condition on line 1532 was never true
1533 mlp_stack = mlp_stack @ project_2d
1534 labels.extend(mlp_labels)
1535 components.append(mlp_stack)
1537 if self.has_embed: 1537 ↛ 1545line 1537 didn't jump to line 1545 because the condition on line 1537 was always true
1538 embed = pos_slice.apply(self["embed"], -2)[None]
1539 if ln_folded:
1540 embed = _ln_then_project(embed)
1541 elif project_2d is not None:
1542 embed = embed @ project_2d
1543 labels.append("embed")
1544 components.append(embed)
1545 if self.has_pos_embed: 1545 ↛ 1554line 1545 didn't jump to line 1554 because the condition on line 1545 was always true
1546 pos_embed = pos_slice.apply(self["pos_embed"], -2)[None]
1547 if ln_folded:
1548 pos_embed = _ln_then_project(pos_embed)
1549 elif project_2d is not None:
1550 pos_embed = pos_embed @ project_2d
1551 labels.append("pos_embed")
1552 components.append(pos_embed)
1553 # If we didn't expand the neurons, the MLP biases are already included in the MLP outputs.
1554 bias_full = self.model.accumulated_bias(layer, mlp_input, include_mlp_biases=expand_neurons)
1555 if ln_folded:
1556 # Expand bias to per-position d_model shape so LN can center, then project.
1557 expand_shape: tuple = (1,) + tuple(head_stack.shape[1:-1]) + (self.model.cfg.d_model,)
1558 bias = _ln_then_project(bias_full.expand(expand_shape))
1559 else:
1560 if project_2d is not None:
1561 # Bias is [d_model], so project post-hoc for shape compatibility — no memory win here.
1562 bias_full = bias_full @ project_2d
1563 bias = bias_full.expand((1,) + head_stack.shape[1:])
1564 labels.append("bias")
1565 components.append(bias)
1566 residual_stack = torch.cat(components, dim=0)
1567 if apply_ln and not ln_folded:
1568 residual_stack = self.apply_ln_to_stack(
1569 residual_stack, layer, pos_slice=pos_slice, mlp_input=mlp_input
1570 )
1572 if squeeze_projected:
1573 residual_stack = residual_stack.squeeze(-1)
1575 if return_labels:
1576 return residual_stack, labels
1577 else:
1578 return residual_stack