Coverage for transformer_lens/SVDInterpreter.py: 91%
63 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""SVD Interpreter.
3Module for getting the singular vectors of the OV, w_in, and w_out matrices of a
4:class:`transformer_lens.model_bridge.TransformerBridge` (or any model exposing
5the TransformerLens weight surface).
6"""
8from typing import NoReturn, Optional, Union
10import torch
11from typing_extensions import Literal
13from transformer_lens.FactoredMatrix import FactoredMatrix
14from transformer_lens.model_protocol import TransformerLensModel
16OUTPUT_EMBEDDING = "unembed.W_U"
17VECTOR_TYPES = ["OV", "w_in", "w_out"]
20class SVDInterpreter:
21 # Base protocol at runtime: beartype validates via getattr_static, which
22 # cannot see nn.Module instance submodules, so the WithWeights surface
23 # would spuriously reject legacy models. Everything touched here (cfg,
24 # tl_parameters/named_parameters fallback) is on the base surface.
25 def __init__(self, model: TransformerLensModel):
26 self.model = model
27 self.cfg = model.cfg
28 # Use tl_parameters() for TransformerBridge (returns TL-style dict); other
29 # nn.Module models with TL-style parameter names use named_parameters().
30 if hasattr(model, "tl_parameters"): 30 ↛ 33line 30 didn't jump to line 33 because the condition on line 30 was always true
31 self.params = model.tl_parameters()
32 else:
33 assert isinstance(model, torch.nn.Module) # named_parameters() fallback
34 self.params = {name: param for name, param in model.named_parameters()}
36 def get_singular_vectors(
37 self,
38 vector_type: Union[Literal["OV"], Literal["w_in"], Literal["w_out"]],
39 layer_index: int,
40 num_vectors: int = 10,
41 head_index: Optional[int] = None,
42 ) -> torch.Tensor:
43 """Gets the singular vectors for a given vector type, layer, and optionally head.
45 This tensor can then be plotted using Neel's PySvelte, as demonstrated in the demo for this
46 feature. The demo also points out some "gotchas" in this feature - numerical instability
47 means inconsistency across devices, and default weight processing doesn't
48 replicate the original SVD post very well. So I'd recommend checking out the demo if you
49 want to use this!
51 Example:
53 .. code-block:: python
55 from transformer_lens import SVDInterpreter
56 from transformer_lens.model_bridge import TransformerBridge
58 model = TransformerBridge.boot_transformers('gpt2-medium')
59 svd_interpreter = SVDInterpreter(model)
61 ov = svd_interpreter.get_singular_vectors('OV', layer_index=22, head_index=10)
63 all_tokens = [model.to_str_tokens(np.array([i])) for i in range(model.cfg.d_vocab)]
64 all_tokens = [all_tokens[i][0] for i in range(model.cfg.d_vocab)]
66 def plot_matrix(matrix, tokens, k=10, filter="topk"):
67 pysvelte.TopKTable(
68 tokens=all_tokens,
69 activations=matrix,
70 obj_type="SVD direction",
71 k=k,
72 filter=filter
73 ).show()
75 plot_matrix(ov, all_tokens)
77 Args:
78 vector_type: Type of the vector:
79 - "OV": Singular vectors of the OV matrix for a particular layer and head.
80 - "w_in": Singular vectors of the w_in matrix for a particular layer.
81 - "w_out": Singular vectors of the w_out matrix for a particular layer.
82 layer_index: The index of the layer.
83 num_vectors: Number of vectors.
84 head_index: Index of the head.
86 Raises:
87 NotImplementedError: If the requested layer does not expose a single dense MLP weight,
88 such as a sparse-MoE layer that requires an expert-aware interpretation.
89 """
91 if head_index is None:
92 assert vector_type in [
93 "w_in",
94 "w_out",
95 ], f"Head index optional only for w_in and w_out, got {vector_type}"
97 matrix: Union[FactoredMatrix, torch.Tensor]
98 if vector_type == "OV":
99 assert head_index is not None # keep mypy happy
100 matrix = self._get_OV_matrix(layer_index, head_index)
101 V = matrix.V.T
103 elif vector_type == "w_in":
104 matrix = self._get_w_in_matrix(layer_index)
105 _, _, V = torch.linalg.svd(matrix)
107 elif vector_type == "w_out": 107 ↛ 112line 107 didn't jump to line 112 because the condition on line 107 was always true
108 matrix = self._get_w_out_matrix(layer_index)
109 _, _, V = torch.linalg.svd(matrix)
111 else:
112 raise ValueError(f"Vector type must be in {VECTOR_TYPES}, instead got {vector_type}")
114 return self._get_singular_vectors_from_matrix(V, self.params[OUTPUT_EMBEDDING], num_vectors)
116 def _get_singular_vectors_from_matrix(
117 self,
118 V: Union[torch.Tensor, FactoredMatrix],
119 embedding: torch.Tensor,
120 num_vectors: int = 10,
121 ) -> torch.Tensor:
122 """Returns the top num_vectors singular vectors from a matrix."""
124 vectors_list = []
125 for i in range(num_vectors):
126 activations = V[i, :].float() @ embedding # type: ignore
127 vectors_list.append(activations)
129 vectors = torch.stack(vectors_list, dim=1).unsqueeze(1)
130 assert vectors.shape == (
131 self.cfg.d_vocab,
132 1,
133 num_vectors,
134 ), f"Vectors shape should be {self.cfg.d_vocab, 1, num_vectors} but got {vectors.shape}"
135 return vectors
137 def _get_OV_matrix(self, layer_index: int, head_index: int) -> FactoredMatrix:
138 """Gets the OV matrix for a particular layer and head."""
140 assert (
141 0 <= layer_index < self.cfg.n_layers
142 ), f"Layer index must be between 0 and {self.cfg.n_layers-1} but got {layer_index}"
143 assert (
144 0 <= head_index < self.cfg.n_heads
145 ), f"Head index must be between 0 and {self.cfg.n_heads-1} but got {head_index}"
147 W_V: torch.Tensor = self.params[f"blocks.{layer_index}.attn.W_V"]
148 W_O: torch.Tensor = self.params[f"blocks.{layer_index}.attn.W_O"]
149 W_V, W_O = W_V[head_index, :, :], W_O[head_index, :, :]
151 return FactoredMatrix(W_V, W_O)
153 def _get_w_in_matrix(self, layer_index: int) -> torch.Tensor:
154 """Gets the w_in matrix for a particular layer."""
156 assert (
157 0 <= layer_index < self.cfg.n_layers
158 ), f"Layer index must be between 0 and {self.cfg.n_layers-1} but got {layer_index}"
160 key = f"blocks.{layer_index}.mlp.W_in"
161 if key not in self.params:
162 self._raise_unsupported_mlp_weight("w_in", layer_index)
163 w_in = self.params[key].T
165 if f"blocks.{layer_index}.ln2.w" in self.params: # If fold_ln == False 165 ↛ 169line 165 didn't jump to line 169 because the condition on line 165 was always true
166 ln_2 = self.params[f"blocks.{layer_index}.ln2.w"]
167 return w_in * ln_2
169 return w_in
171 def _get_w_out_matrix(self, layer_index: int) -> torch.Tensor:
172 """Gets the w_out matrix for a particular layer."""
174 assert (
175 0 <= layer_index < self.cfg.n_layers
176 ), f"Layer index must be between 0 and {self.cfg.n_layers-1} but got {layer_index}"
178 key = f"blocks.{layer_index}.mlp.W_out"
179 if key not in self.params:
180 self._raise_unsupported_mlp_weight("w_out", layer_index)
181 return self.params[key]
183 def _raise_unsupported_mlp_weight(self, weight_name: str, layer_index: int) -> NoReturn:
184 raise NotImplementedError(
185 f"SVDInterpreter cannot analyze {weight_name} for layer {layer_index}: "
186 "the layer does not expose a single dense MLP weight. Sparse MoE layers "
187 "require an explicit expert-aware interpretation."
188 )