Coverage for transformer_lens/SVDInterpreter.py: 91%

1"""SVD Interpreter. 

2 

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""" 

7 

8from typing import NoReturn, Optional, Union 

9 

10import torch 

11from typing_extensions import Literal 

12 

13from transformer_lens.FactoredMatrix import FactoredMatrix, upcast_for_linalg 

14from transformer_lens.model_protocol import TransformerLensModel 

15 

16OUTPUT_EMBEDDING = "unembed.W_U" 

17VECTOR_TYPES = ["OV", "w_in", "w_out"] 

18 

19 

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()} 

35 

36 @torch.no_grad() 

37 def get_singular_vectors( 

38 self, 

39 vector_type: Union[Literal["OV"], Literal["w_in"], Literal["w_out"]], 

40 layer_index: int, 

41 num_vectors: int = 10, 

42 head_index: Optional[int] = None, 

43 ) -> torch.Tensor: 

44 """Gets the singular vectors for a given vector type, layer, and optionally head. 

45 

46 This tensor can then be plotted using Neel's PySvelte, as demonstrated in the demo for this 

47 feature. The demo also points out some "gotchas" in this feature - numerical instability 

48 means inconsistency across devices, and default weight processing doesn't 

49 replicate the original SVD post very well. So I'd recommend checking out the demo if you 

50 want to use this! 

51 

52 Runs under ``torch.no_grad()``, so the returned tensor is always detached from autograd. 

53 

54 Example: 

55 

56 .. code-block:: python 

57 

58 from transformer_lens import SVDInterpreter 

59 from transformer_lens.model_bridge import TransformerBridge 

60 

61 model = TransformerBridge.boot_transformers('gpt2-medium') 

62 svd_interpreter = SVDInterpreter(model) 

63 

64 ov = svd_interpreter.get_singular_vectors('OV', layer_index=22, head_index=10) 

65 

66 all_tokens = [model.to_str_tokens(np.array([i])) for i in range(model.cfg.d_vocab)] 

67 all_tokens = [all_tokens[i][0] for i in range(model.cfg.d_vocab)] 

68 

69 def plot_matrix(matrix, tokens, k=10, filter="topk"): 

70 pysvelte.TopKTable( 

71 tokens=all_tokens, 

72 activations=matrix, 

73 obj_type="SVD direction", 

74 k=k, 

75 filter=filter 

76 ).show() 

77 

78 plot_matrix(ov, all_tokens) 

79 

80 Args: 

81 vector_type: Type of the vector: 

82 - "OV": Singular vectors of the OV matrix for a particular layer and head. 

83 - "w_in": Singular vectors of the w_in matrix for a particular layer. 

84 - "w_out": Singular vectors of the w_out matrix for a particular layer. 

85 layer_index: The index of the layer. 

86 num_vectors: Number of vectors. 

87 head_index: Index of the head. 

88 

89 Raises: 

90 NotImplementedError: If the requested layer does not expose a single dense MLP weight, 

91 such as a sparse-MoE layer that requires an expert-aware interpretation. 

92 """ 

93 

94 if head_index is None: 

95 assert vector_type in [ 

96 "w_in", 

97 "w_out", 

98 ], f"Head index optional only for w_in and w_out, got {vector_type}" 

99 

100 matrix: Union[FactoredMatrix, torch.Tensor] 

101 if vector_type == "OV": 

102 assert head_index is not None # keep mypy happy 

103 matrix = self._get_OV_matrix(layer_index, head_index) 

104 V = matrix.V.T 

105 

106 elif vector_type == "w_in": 

107 matrix = self._get_w_in_matrix(layer_index) 

108 _, _, V = torch.linalg.svd(upcast_for_linalg(matrix)) 

109 

110 elif vector_type == "w_out": 110 ↛ 115line 110 didn't jump to line 115 because the condition on line 110 was always true

111 matrix = self._get_w_out_matrix(layer_index) 

112 _, _, V = torch.linalg.svd(upcast_for_linalg(matrix)) 

113 

114 else: 

115 raise ValueError(f"Vector type must be in {VECTOR_TYPES}, instead got {vector_type}") 

116 

117 return self._get_singular_vectors_from_matrix(V, self.params[OUTPUT_EMBEDDING], num_vectors) 

118 

119 def _get_singular_vectors_from_matrix( 

120 self, 

121 V: torch.Tensor, 

122 embedding: torch.Tensor, 

123 num_vectors: int = 10, 

124 ) -> torch.Tensor: 

125 """Returns the top num_vectors singular vectors from a matrix.""" 

126 

127 embedding = embedding.to(dtype=V.dtype) 

128 vectors = (embedding.T @ V[:num_vectors].T).unsqueeze(1) 

129 assert vectors.shape == ( 

130 self.cfg.d_vocab, 

131 1, 

132 num_vectors, 

133 ), f"Vectors shape should be {self.cfg.d_vocab, 1, num_vectors} but got {vectors.shape}" 

134 return vectors 

135 

136 def _get_OV_matrix(self, layer_index: int, head_index: int) -> FactoredMatrix: 

137 """Gets the OV matrix for a particular layer and head.""" 

138 

139 assert ( 

140 0 <= layer_index < self.cfg.n_layers 

141 ), f"Layer index must be between 0 and {self.cfg.n_layers-1} but got {layer_index}" 

142 assert ( 

143 0 <= head_index < self.cfg.n_heads 

144 ), f"Head index must be between 0 and {self.cfg.n_heads-1} but got {head_index}" 

145 

146 W_V: torch.Tensor = self.params[f"blocks.{layer_index}.attn.W_V"] 

147 W_O: torch.Tensor = self.params[f"blocks.{layer_index}.attn.W_O"] 

148 W_V, W_O = W_V[head_index, :, :], W_O[head_index, :, :] 

149 

150 return FactoredMatrix(W_V, W_O) 

151 

152 def _get_w_in_matrix(self, layer_index: int) -> torch.Tensor: 

153 """Gets the w_in matrix for a particular layer.""" 

154 

155 assert ( 

156 0 <= layer_index < self.cfg.n_layers 

157 ), f"Layer index must be between 0 and {self.cfg.n_layers-1} but got {layer_index}" 

158 

159 key = f"blocks.{layer_index}.mlp.W_in" 

160 if key not in self.params: 

161 self._raise_unsupported_mlp_weight("w_in", layer_index) 

162 w_in = self.params[key].T 

163 

164 if f"blocks.{layer_index}.ln2.w" in self.params: # If fold_ln == False 164 ↛ 172line 164 didn't jump to line 172 because the condition on line 164 was always true

165 ln_2 = self.params[f"blocks.{layer_index}.ln2.w"] 

166 if getattr(self.cfg, "rmsnorm_uses_offset", False): 

167 # Offset-RMS models (Gemma family) store w as an offset from 1: 

168 # the effective scale is (1 + w), and post-fold w = 0 is identity. 

169 ln_2 = 1.0 + ln_2 

170 return w_in * ln_2 

171 

172 return w_in 

173 

174 def _get_w_out_matrix(self, layer_index: int) -> torch.Tensor: 

175 """Gets the w_out matrix for a particular layer.""" 

176 

177 assert ( 

178 0 <= layer_index < self.cfg.n_layers 

179 ), f"Layer index must be between 0 and {self.cfg.n_layers-1} but got {layer_index}" 

180 

181 key = f"blocks.{layer_index}.mlp.W_out" 

182 if key not in self.params: 

183 self._raise_unsupported_mlp_weight("w_out", layer_index) 

184 return self.params[key] 

185 

186 def _raise_unsupported_mlp_weight(self, weight_name: str, layer_index: int) -> NoReturn: 

187 raise NotImplementedError( 

188 f"SVDInterpreter cannot analyze {weight_name} for layer {layer_index}: " 

189 "the layer does not expose a single dense MLP weight. Sparse MoE layers " 

190 "require an explicit expert-aware interpretation." 

191 )