Coverage for transformer_lens/FactoredMatrix.py: 97%
172 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-10 05:10 +0000
1"""Factored Matrix.
3Utilities for representing a matrix as a product of two matrices, and for efficient calculation of
4eigenvalues, norm and SVD.
5"""
7from __future__ import annotations
9from functools import cached_property
10from typing import Any, List, Protocol, Tuple, Union, cast, overload, runtime_checkable
12import numpy as np
13import torch
14from jaxtyping import Complex, Float
16import transformer_lens.utilities.tensors as tensor_utils
17from transformer_lens.utilities.slice import to_python_int
20@runtime_checkable
21class TensorLike(Protocol):
22 """Minimal tensor protocol that FactoredMatrix accepts in place of torch.Tensor.
24 Allows duck-typed inputs (e.g. jaxtyping wrappers, custom array types) that
25 aren't torch.Tensor subclasses but support the operations FactoredMatrix uses
26 when constructing, multiplying, and broadcasting its A and B factors.
27 """
29 @property
30 def ndim(self) -> int:
31 ...
33 @property
34 def shape(self) -> Any:
35 ...
37 def size(self, dim: int) -> int:
38 ...
40 def unsqueeze(self, dim: int) -> Any:
41 ...
43 def squeeze(self, dim: int) -> Any:
44 ...
46 def broadcast_to(self, shape: Any) -> Any:
47 ...
49 def __matmul__(self, other: Any) -> Any:
50 ...
53class FactoredMatrix:
54 """
55 Class to represent low rank factored matrices, where the matrix is represented as a product of two matrices. Has utilities for efficient calculation of eigenvalues, norm and SVD.
56 """
58 def __init__(
59 self,
60 A: Union[Float[torch.Tensor, "... ldim mdim"], TensorLike],
61 B: Union[Float[torch.Tensor, "... mdim rdim"], TensorLike],
62 ):
63 """Construct a FactoredMatrix from factors A and B.
65 A and B may be torch.Tensor or TensorLike duck types. TensorLike inputs
66 are only fully supported by matmul-family operations (``@``, ``AB``,
67 ``BA``); operations like ``svd()``, ``norm()``, ``transpose()``,
68 ``__getitem__``, and eigenvalue methods require both factors to be
69 actual torch.Tensor and will raise AttributeError on TensorLike inputs.
70 """
71 # Cast to Tensor for type-checker purposes. At runtime A and B may be
72 # TensorLike duck types; the class methods trust the protocol.
73 self.A: torch.Tensor = cast(torch.Tensor, A)
74 self.B: torch.Tensor = cast(torch.Tensor, B)
75 assert self.A.size(-1) == self.B.size(
76 -2
77 ), f"Factored matrix must match on inner dimension, shapes were a: {self.A.shape}, b:{self.B.shape}"
78 self.ldim = self.A.size(-2)
79 self.rdim = self.B.size(-1)
80 self.mdim = self.B.size(-2)
81 self.has_leading_dims = (self.A.ndim > 2) or (self.B.ndim > 2)
82 try:
83 self.shape = torch.broadcast_shapes(self.A.shape[:-2], self.B.shape[:-2]) + (
84 self.ldim,
85 self.rdim,
86 )
87 except RuntimeError as e:
88 raise RuntimeError(
89 f"Shape mismatch: Cannot broadcast leading dimensions. A has shape {self.A.shape}, B has shape {self.B.shape}. {str(e)}"
90 ) from e
91 try:
92 self.A = self.A.broadcast_to(self.shape[:-2] + (self.ldim, self.mdim))
93 self.B = self.B.broadcast_to(self.shape[:-2] + (self.mdim, self.rdim))
94 except RuntimeError as e:
95 raise RuntimeError(
96 f"Shape mismatch: Cannot broadcast tensors. A has shape {self.A.shape}, B has shape {self.B.shape}, expected broadcast shape {self.shape}. {str(e)}"
97 ) from e
99 @overload
100 def __matmul__(
101 self,
102 other: Union[
103 Float[torch.Tensor, "... rdim new_rdim"],
104 "FactoredMatrix",
105 ],
106 ) -> "FactoredMatrix":
107 ...
109 @overload
110 def __matmul__( # type: ignore
111 self,
112 other: Float[torch.Tensor, "rdim"],
113 ) -> Float[torch.Tensor, "... ldim"]:
114 ...
116 def __matmul__(
117 self,
118 other: Union[
119 Float[torch.Tensor, "... rdim new_rdim"],
120 Float[torch.Tensor, "rdim"],
121 "FactoredMatrix",
122 TensorLike,
123 ],
124 ) -> Union["FactoredMatrix", Float[torch.Tensor, "... ldim"], TensorLike]:
125 if isinstance(other, FactoredMatrix):
126 return (self @ other.A) @ other.B
127 else:
128 if other.ndim < 2:
129 # It's a vector, so we collapse the factorisation and just return a vector
130 # Squeezing/Unsqueezing is to preserve broadcasting working nicely
131 return (self.A @ (self.B @ other.unsqueeze(-1))).squeeze(-1)
132 else:
133 assert (
134 other.size(-2) == self.rdim
135 ), f"Right matrix must match on inner dimension, shapes were self: {self.shape}, other:{other.shape}"
136 if self.rdim > self.mdim:
137 # other is Tensor or TensorLike; runtime delegates to
138 # the appropriate __matmul__/__rmatmul__ overload.
139 return FactoredMatrix(self.A, self.B @ cast(torch.Tensor, other))
140 else:
141 return FactoredMatrix(self.AB, other)
143 @overload
144 def __rmatmul__( # type: ignore
145 self,
146 other: Union[
147 Float[torch.Tensor, "... new_rdim ldim"],
148 "FactoredMatrix",
149 ],
150 ) -> "FactoredMatrix":
151 ...
153 @overload
154 def __rmatmul__( # type: ignore
155 self,
156 other: Float[torch.Tensor, "ldim"],
157 ) -> Float[torch.Tensor, "... rdim"]:
158 ...
160 def __rmatmul__( # type: ignore
161 self,
162 other: Union[
163 Float[torch.Tensor, "... new_rdim ldim"],
164 Float[torch.Tensor, "ldim"],
165 "FactoredMatrix",
166 TensorLike,
167 ],
168 ) -> Union["FactoredMatrix", Float[torch.Tensor, "... rdim"], TensorLike]:
169 if isinstance(other, FactoredMatrix): 169 ↛ 170line 169 didn't jump to line 170 because the condition on line 169 was never true
170 return other.A @ (other.B @ self)
171 else:
172 assert (
173 other.size(-1) == self.ldim
174 ), f"Left matrix must match on inner dimension, shapes were self: {self.shape}, other:{other.shape}"
175 if other.ndim < 2:
176 # It's a vector, so we collapse the factorisation and just return a vector
177 return ((other.unsqueeze(-2) @ self.A) @ self.B).squeeze(-2)
178 elif self.ldim > self.mdim:
179 return FactoredMatrix(other @ self.A, self.B)
180 else:
181 return FactoredMatrix(other, self.AB)
183 def __mul__(self, scalar: Union[int, float, torch.Tensor]) -> FactoredMatrix:
184 """
185 Left scalar multiplication. Scalar multiplication distributes over matrix multiplication, so we can just multiply one of the factor matrices by the scalar.
186 """
187 if isinstance(scalar, torch.Tensor):
188 assert (
189 scalar.numel() == 1
190 ), f"Tensor must be a scalar for use with * but was of shape {scalar.shape}. For matrix multiplication, use @ instead."
191 return FactoredMatrix(self.A * scalar, self.B)
193 def __rmul__(self, scalar: Union[int, float, torch.Tensor]) -> FactoredMatrix: # type: ignore
194 """
195 Right scalar multiplication. For scalar multiplication from the right, we can reuse the __mul__ method.
196 """
197 return self * scalar
199 @property
200 def AB(self) -> Union[Float[torch.Tensor, "*leading_dims ldim rdim"], TensorLike]:
201 """The product matrix - expensive to compute, and can consume a lot of GPU memory.
203 Returns a TensorLike when A or B is a non-Tensor TensorLike duck type.
204 """
205 return self.A @ self.B
207 @property
208 def BA(self) -> Float[torch.Tensor, "*leading_dims rdim ldim"]:
209 """The reverse product. Only makes sense when ldim==rdim"""
210 assert (
211 self.rdim == self.ldim
212 ), f"Can only take ba if ldim==rdim, shapes were self: {self.shape}"
213 return self.B @ self.A
215 @property
216 def T(self) -> FactoredMatrix:
217 return FactoredMatrix(self.B.transpose(-2, -1), self.A.transpose(-2, -1))
219 @cached_property
220 def _svd_cached(
221 self,
222 ) -> Tuple[
223 Float[torch.Tensor, "*leading_dims ldim mdim"],
224 Float[torch.Tensor, "*leading_dims mdim"],
225 Float[torch.Tensor, "*leading_dims rdim mdim"],
226 ]:
227 # Cache on the instance (frees with it) rather than class-level — fixes the lru_cache leak.
228 Ua, Sa, Vha = torch.linalg.svd(self.A, full_matrices=False)
229 Ub, Sb, Vhb = torch.linalg.svd(self.B, full_matrices=False)
230 Va = tensor_utils.transpose(Vha)
231 Vb = tensor_utils.transpose(Vhb)
232 middle = Sa[..., :, None] * tensor_utils.transpose(Va) @ Ub * Sb[..., None, :]
233 Um, Sm, Vhm = torch.linalg.svd(middle, full_matrices=False)
234 Vm = tensor_utils.transpose(Vhm)
235 return Ua @ Um, Sm, Vb @ Vm
237 def svd(
238 self,
239 ) -> Tuple[
240 Float[torch.Tensor, "*leading_dims ldim mdim"],
241 Float[torch.Tensor, "*leading_dims mdim"],
242 Float[torch.Tensor, "*leading_dims rdim mdim"],
243 ]:
244 """Singular Value Decomposition: returns ``(U, S, V)`` such that ``U @ S.diag() @ V.transpose(-2, -1) == M``."""
245 return self._svd_cached
247 @property
248 def U(self) -> Float[torch.Tensor, "*leading_dims ldim mdim"]:
249 return self.svd()[0]
251 @property
252 def S(self) -> Float[torch.Tensor, "*leading_dims mdim"]:
253 return self.svd()[1]
255 @property
256 def V(self) -> Float[torch.Tensor, "*leading_dims rdim mdim"]:
257 """Right singular vectors. ``M == U @ S.diag() @ V.transpose(-2, -1)``."""
258 return self.svd()[2]
260 @property
261 def Vh(self) -> Float[torch.Tensor, "*leading_dims rdim mdim"]:
262 """Deprecated alias for :attr:`V` — historically misnamed; returns V, not its conjugate transpose."""
263 import warnings
265 warnings.warn(
266 "FactoredMatrix.Vh has always returned V (right singular vectors), not Vh. "
267 "Use .V for the canonical name; for the actual Hermitian transpose use "
268 ".V.transpose(-2, -1). The .Vh alias will be removed in a future release.",
269 DeprecationWarning,
270 stacklevel=2,
271 )
272 return self.svd()[2]
274 @property
275 def eigenvalues(self) -> Complex[torch.Tensor, "*leading_dims mdim"]:
276 """
277 Eigenvalues of AB are the same as for BA (apart from trailing zeros), because if BAv=kv ABAv = A(BAv)=kAv,
278 so Av is an eigenvector of AB with eigenvalue k.
279 """
280 input_matrix = self.BA
281 if input_matrix.dtype in [torch.bfloat16, torch.float16]:
282 # Cast to float32 because eig is not implemented for 16-bit on CPU/CUDA
283 input_matrix = input_matrix.to(torch.float32)
284 return torch.linalg.eig(input_matrix).eigenvalues
286 def _convert_to_slice(self, sequence: Union[Tuple, List], idx: int) -> Tuple:
287 """
288 e.g. if sequence = (1, 2, 3) and idx = 1, return (1, slice(2, 3), 3). This only edits elements if they are ints.
289 """
290 if isinstance(idx, int): 290 ↛ 304line 290 didn't jump to line 304 because the condition on line 290 was always true
291 sequence = list(sequence)
292 entry = sequence[idx]
293 # bool keeps its legacy int-index treatment here, so check plain int
294 # before the int-like coercion.
295 value = int(entry) if isinstance(entry, int) else to_python_int(entry)
296 if value is not None:
297 # `value + 1` selects the single requested element, except when
298 # value == -1: there `value + 1 == 0` yields the empty slice(-1, 0).
299 # Use `None` as the stop so the final element is kept.
300 stop = value + 1 if value != -1 else None
301 sequence[idx] = slice(value, stop)
302 sequence = tuple(sequence)
304 return sequence
306 def __getitem__(self, idx: Union[int, Tuple]) -> FactoredMatrix:
307 """Index leading dimensions and matrix rows/columns without forming the product.
309 Ellipsis expands across the product's dimensions, excluding new axes (``None``).
310 Integer row and column indices retain singleton dimensions in the returned factors.
311 New axes next to matrix rows/columns reshape the factors without forming the product.
312 """
313 if not isinstance(idx, tuple): 313 ↛ 316line 313 didn't jump to line 316 because the condition on line 313 was always true
314 idx = (idx,)
316 def index_ndim(item: Any) -> int:
317 if item is None or item is Ellipsis:
318 return 0
319 # Infer nested sequence masks as PyTorch does, leaving the actual index unchanged.
320 if isinstance(item, (list, tuple)):
321 sequence = torch.as_tensor(item, device="cpu")
322 return sequence.ndim if sequence.dtype == torch.bool else 1
323 # Boolean masks consume one input axis per mask dimension, unlike integer arrays.
324 if isinstance(item, torch.Tensor) and item.dtype in (torch.bool, torch.uint8):
325 return item.ndim
326 if isinstance(item, np.ndarray) and item.dtype in (np.bool_, np.uint8):
327 return item.ndim
328 return 1
330 ellipsis_positions = [i for i, item in enumerate(idx) if item is Ellipsis]
331 if len(ellipsis_positions) > 1:
332 raise IndexError("An index can only have a single ellipsis")
333 if ellipsis_positions:
334 consumed = sum(index_ndim(item) for item in idx)
335 position = ellipsis_positions[0]
336 # Expand against AB's axes before mapping row/column indices onto the factors.
337 idx = idx[:position] + (slice(None),) * (self.ndim - consumed) + idx[position + 1 :]
338 factor_indices = []
339 matrix_newaxes = []
340 length = 0
341 for item in idx:
342 if item is None and length >= self.ndim - 2:
343 matrix_newaxes.append(length - (self.ndim - 2))
344 else:
345 factor_indices.append(item)
346 length += index_ndim(item)
347 idx = tuple(factor_indices)
348 if length <= len(self.shape) - 2:
349 result = FactoredMatrix(self.A[idx], self.B[idx])
350 elif length == len(self.shape) - 1:
351 idx = self._convert_to_slice(idx, -1)
352 result = FactoredMatrix(self.A[idx], self.B[idx[:-1]])
353 elif length == len(self.shape):
354 idx = self._convert_to_slice(idx, -1)
355 idx = self._convert_to_slice(idx, -2)
356 result = FactoredMatrix(self.A[idx[:-1]], self.B[idx[:-2] + (slice(None), idx[-1])])
357 else:
358 raise ValueError(
359 f"{idx} is too long an index for a FactoredMatrix with shape {self.shape}"
360 )
362 for position in matrix_newaxes:
363 if position == 0:
364 axis = result.ndim - 2
365 result = FactoredMatrix(result.A.unsqueeze(axis), result.B.unsqueeze(axis))
366 elif position == 1:
367 # Promote the row axis to a leading dimension, keeping the contraction intact.
368 result = FactoredMatrix(result.A.unsqueeze(-2), result.B.unsqueeze(-3))
369 else:
370 # A trailing axis makes the old column axis the new row axis.
371 result = FactoredMatrix(
372 result.B.transpose(-2, -1).unsqueeze(-3), result.A.unsqueeze(-1)
373 )
374 return result
376 def norm(self) -> Float[torch.Tensor, "*leading_dims"]:
377 """
378 Frobenius norm is sqrt(sum of squared singular values)
379 """
380 return self.S.pow(2).sum(-1).sqrt()
382 def __repr__(self):
383 return f"FactoredMatrix: Shape({self.shape}), Hidden Dim({self.mdim})"
385 def make_even(self) -> FactoredMatrix:
386 """
387 Returns the factored form of (U @ S.sqrt().diag(), S.sqrt().diag() @ V.T) where U, S, V are the SVD of the matrix. This is an equivalent factorisation, but more even - each half has half the singular values, and orthogonal rows/cols
388 """
389 return FactoredMatrix(
390 self.U * self.S.sqrt()[..., None, :],
391 self.S.sqrt()[..., :, None] * tensor_utils.transpose(self.V),
392 )
394 def get_corner(self, k=3):
395 return tensor_utils.get_corner(self.A[..., :k, :] @ self.B[..., :, :k], k)
397 @property
398 def ndim(self) -> int:
399 return len(self.shape)
401 def collapse_l(self) -> Float[torch.Tensor, "*leading_dims mdim rdim"]:
402 """
403 Collapses the left side of the factorization by removing the orthogonal factor (given by self.U). Returns a (..., mdim, rdim) tensor
404 """
405 return self.S[..., :, None] * tensor_utils.transpose(self.V)
407 def collapse_r(self) -> Float[torch.Tensor, "*leading_dims ldim mdim"]:
408 """
409 Analogous to collapse_l, returns a (..., ldim, mdim) tensor
410 """
411 return self.U * self.S[..., None, :]
413 def unsqueeze(self, k: int) -> FactoredMatrix:
414 return FactoredMatrix(self.A.unsqueeze(k), self.B.unsqueeze(k))
416 @property
417 def pair(
418 self,
419 ) -> Tuple[
420 Float[torch.Tensor, "*leading_dims ldim mdim"],
421 Float[torch.Tensor, "*leading_dims mdim rdim"],
422 ]:
423 return (self.A, self.B)