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