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