Coverage for transformer_lens/tools/analysis/jacobian_lens_decomposition.py: 94%
265 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-21 19:27 +0000
1"""J-space sparse decomposition for the Jacobian lens.
3Decomposes an activation (or a steering / sparse-autoencoder direction) into a sparse
4nonnegative combination of J-lens vectors, following Gurnee et al. (2026), "Verbalizable
5Representations Form a Global Workspace in Language Models" (Transformer Circuits Thread).
7The decomposition is greedy: at each of at most ``k`` steps the atom most correlated with the
8current residual (using unit-normalised atoms, so high-norm atoms are not preferred for their
9scale alone) is added to the *selected* set, and the selected-set coefficients are updated
10under a nonnegativity constraint. ``k`` is an upper bound, not a target: selection stops early
11once no unselected atom has a materially positive residual correlation (under nonnegativity a
12negatively-correlated atom cannot reduce the residual), so fewer than ``k`` atoms may be
13selected. Two coefficient-update rules are provided via ``algorithm``:
15- ``"nonnegative_orthogonal_matching_pursuit"`` (default) -- the active-set coefficients are
16 re-solved with a float64 active-set nonnegative least-squares fit at each step, then checked
17 against the KKT conditions.
18- ``"gradient_pursuit"`` -- the directional update of Blumensath & Davies (2008): a single
19 optimal-step-size gradient step on the active coefficients, projected onto the
20 nonnegative orthant.
22Both use the same greedy selection rule and both are nonnegative; they differ in the
23coefficient update and therefore can choose different atoms at later steps. At vocabulary
24scale, correlation over all atoms costs ``O(num_atoms * d_model)`` per step, while the exact
25re-solve adds float64 linear algebra on the small selected set. The NNLS update is the default
26because it optimizes all selected coefficients jointly. ``"gradient_pursuit"`` skips that solve
27and matches the update used in the paper.
29Two supports are exposed because the paper uses two inconsistent operationalizations. The
30``support`` is the numerically *active* set -- the selected atoms whose nonnegative
31coordinate materially contributes -- and ``coordinates`` are aligned with it; this is the
32paper's main-text sparse nonnegative reconstruction. The ``selected_support`` is every
33greedily selected atom, including any whose coordinate was driven to zero by the
34nonnegativity constraint; it defines the span for the paper's appendix projection. Hence
35``len(support) <= len(selected_support) <= k``.
37Two vector outputs correspondingly need not coincide: the ``reconstruction`` is the
38nonnegative combination over the active ``support``, while the J-space *component* is the
39orthogonal projection of the target onto the span of ``selected_support`` (the paper's
40appendix definition, and the residual its interventions use). For the exact NNLS re-solve
41the reconstruction equals the projection onto the *active* support (KKT stationarity), so the
42two vectors differ exactly when a selected atom has a zero coordinate.
44This module is model-free: it operates on a raw dictionary tensor, so it can be used and
45tested without loading a model.
47References:
48 - Gurnee et al. (2026), "Verbalizable Representations Form a Global Workspace in
49 Language Models," Transformer Circuits Thread -- the J-space decomposition.
50 - Pati, Rezaiifar & Krishnaprasad (1993), "Orthogonal Matching Pursuit: Recursive
51 Function Approximation with Applications to Wavelet Decomposition," 27th Asilomar
52 Conference on Signals, Systems and Computers -- the greedy atom selection shared
53 by both algorithms.
54 - Blumensath & Davies (2008), "Gradient Pursuits," IEEE Transactions on Signal
55 Processing 56(6):2370-2382 -- the ``gradient_pursuit`` coefficient update.
56 - Lawson & Hanson (1974), "Solving Least Squares Problems," Prentice-Hall -- the
57 active-set nonnegative least-squares re-solve used by the default algorithm.
58"""
60from __future__ import annotations
62import math
63import warnings
64from dataclasses import dataclass
65from typing import Callable, Dict, List, Tuple
67import torch
69#: The paper varies the sparsity level but "typically" chooses "no more than 25".
70DEFAULT_K = 25
72#: Active problems are normally small (``DEFAULT_K`` is 25), so solve them in float64 and use
73#: a conservative square-root-epsilon threshold for rank and KKT decisions. Unlike a threshold
74#: proportional to ``d_model``, this is invariant to appending zero rows to the same problem.
75_NNLS_RELATIVE_TOLERANCE = math.sqrt(torch.finfo(torch.float64).eps)
76_NNLS_RANK_RTOL = _NNLS_RELATIVE_TOLERANCE
78#: A selected atom counts as numerically *active* when its contribution ``c_i * ||v_i||`` is a
79#: materially nonzero fraction of the target scale ``||x||``. Using the contribution (not the
80#: raw coefficient) keeps the test scale-invariant across J-lens vectors of different native
81#: norms. This is a numerical-activity threshold, not an interpretability one. It matches the
82#: default NNLS coefficient-zeroing scale (:data:`_NNLS_RELATIVE_TOLERANCE`), so the active
83#: support equals the set of strictly-positive NNLS coordinates and the reconstruction over the
84#: active support is the projection onto it.
85_ACTIVE_RELATIVE_TOLERANCE = _NNLS_RELATIVE_TOLERANCE
87_GRADIENT_BACKTRACK_STEPS = 20
89#: Early-stopping correlation threshold. Once the best unselected normalized correlation drops
90#: to this fraction of ``||x||``, no atom can materially reduce the residual and selection
91#: stops. It sits at the float32 residual noise floor (the residual is computed in float32),
92#: so a full-rank target stops instead of selecting noise atoms -- this is what makes the
93#: selected support scale-invariant. It is deliberately coarser than the activity threshold:
94#: selection is a float32 correlation decision, activity a check against the float64 solve.
95_CORRELATION_RELATIVE_TOLERANCE = math.sqrt(torch.finfo(torch.float32).eps)
97#: Shared near-parallel policy for lens-space interventions. Two atoms whose absolute cosine
98#: reaches this threshold span an ill-conditioned pair: a coordinate swap between them is
99#: approximately a no-op. Both the anchored coordinate patch and ``JacobianLens.swap_hooks``
100#: read this one definition so their near-parallel diagnostics never drift apart.
101_SWAP_WARN_COSINE = 0.99
102#: Above this cosine ``swap_hooks`` refuses the intervention: inverting a two-atom basis whose
103#: columns are near-collinear is numerically hopeless. The anchored coordinate patch inverts
104#: nothing, so it only reads ``_SWAP_WARN_COSINE`` and never raises on this threshold.
105_SWAP_ERROR_COSINE = 0.999
108def _linalg_on_cpu_if_mps(op: Callable[..., torch.Tensor], *tensors: torch.Tensor) -> torch.Tensor:
109 """Run a linear-algebra op that MPS lacks a kernel for, restoring the input device.
111 Several LAPACK-backed ``torch.linalg`` routines (``pinv``, ``svdvals``) are unimplemented on
112 the MPS backend, so callers must round-trip through CPU. This centralizes the shared
113 detect-device -> ``.cpu()`` -> op -> restore idiom: off MPS ``op`` is called on the input
114 tensors unchanged, and on MPS it is called on their CPU copies with the result moved back to
115 the input device. All inputs must share one device.
116 """
117 device = tensors[0].device
118 if device.type == "mps": 118 ↛ 119line 118 didn't jump to line 119 because the condition on line 118 was never true
119 return op(*(tensor.cpu() for tensor in tensors)).to(device)
120 return op(*tensors)
123def _diagnose_intervention_pair(
124 unit_vectors: torch.Tensor, *, description: str, stacklevel: int
125) -> float:
126 """Apply the shared near-parallel warn/raise policy to two unit-normalized vectors.
128 Returns the signed cosine. Warns above :data:`_SWAP_WARN_COSINE` and raises above
129 :data:`_SWAP_ERROR_COSINE`. Used by interventions that invert a two-atom basis (e.g.
130 ``swap_hooks``); the anchored coordinate patch, which performs no inverse, warns without
131 raising and does not call this helper.
132 """
133 cosine = float((unit_vectors[0] @ unit_vectors[1]).item())
134 abs_cosine = abs(cosine)
135 if not math.isfinite(abs_cosine) or abs_cosine >= _SWAP_ERROR_COSINE:
136 raise ValueError(
137 f"{description} are numerically near-parallel "
138 f"(abs cosine={abs_cosine:.6f}); choose better-separated concepts"
139 )
140 if abs_cosine >= _SWAP_WARN_COSINE:
141 warnings.warn(
142 f"{description} are poorly conditioned "
143 f"(abs cosine={abs_cosine:.6f}); the intervention may be amplified",
144 UserWarning,
145 stacklevel=stacklevel,
146 )
147 return cosine
150@dataclass
151class JSpaceDecomposition:
152 """Result of a sparse J-space decomposition.
154 Attributes:
155 support: Indices of the numerically *active* selected atoms -- those whose
156 nonnegative coordinate materially contributes (token ids when the dictionary is
157 the vocabulary of J-lens vectors). A subset of ``selected_support``.
158 coordinates: Nonnegative pursuit coefficients aligned with ``support`` (the
159 "local J-space coordinates"); every entry is materially nonzero.
160 selected_support: Indices of every greedily selected atom, including any whose
161 coordinate was driven to zero by the nonnegativity constraint. Defines the span
162 for ``j_space_component``. Satisfies
163 ``support.numel() <= selected_support.numel() <= k``.
164 reconstruction: The nonnegative combination ``sum(coordinates * active atoms)`` over
165 ``support``.
166 j_space_component: The orthogonal projection of the target onto the span of
167 ``selected_support`` (the paper's "J-space component"). For the exact NNLS
168 re-solve it equals ``reconstruction`` unless a selected atom has a zero
169 coordinate, in which case the projection uses a larger span.
170 non_j_space_component: The residual ``target - j_space_component`` (the
171 "non-J-space component"), orthogonal to the selected span.
172 """
174 support: torch.Tensor
175 coordinates: torch.Tensor
176 selected_support: torch.Tensor
177 reconstruction: torch.Tensor
178 j_space_component: torch.Tensor
179 non_j_space_component: torch.Tensor
182def _nnls_tolerances(
183 active_atoms: torch.Tensor, target: torch.Tensor, coefficients: torch.Tensor
184) -> tuple[torch.Tensor, torch.Tensor]:
185 """Scale- and dtype-aware tolerances for the NNLS solver and its KKT check.
187 The dual score ``w_j = A[:, j]^T r`` (``r = target - A c``) scales as an atom norm times
188 the residual-computation scale. That scale includes ``|| |A| |c| ||`` so a nearly
189 cancelling fit receives enough absolute tolerance without making the relative threshold
190 depend on the number of rows. A numerically-zero coefficient has the corresponding inverse
191 atom-norm scale. Returns ``(dual_tol, coefficient_tol)``, each aligned with the columns of
192 ``active_atoms``.
193 """
194 tiny = torch.finfo(target.dtype).tiny
195 atom_norms = active_atoms.norm(dim=0).clamp_min(tiny)
196 scale = (target.norm() + (active_atoms.abs() @ coefficients.abs()).norm()).clamp_min(tiny)
197 dual_tol = _NNLS_RELATIVE_TOLERANCE * atom_norms * scale
198 coefficient_tol = _NNLS_RELATIVE_TOLERANCE * scale / atom_norms
199 return dual_tol, coefficient_tol
202def _validate_nnls_kkt(
203 active_atoms: torch.Tensor,
204 target: torch.Tensor,
205 coefficients: torch.Tensor,
206 dual_tol: torch.Tensor,
207 coefficient_tol: torch.Tensor,
208) -> None:
209 """Raise unless ``coefficients`` satisfy the NNLS KKT conditions within tolerance.
211 This is a postcondition check using the solver's numerical tolerances: primal
212 feasibility (``c >= 0``), dual feasibility on zeroed coordinates (``w_j <= 0``),
213 stationarity on positive coordinates (``w_j = 0``), and complementarity
214 (``c_i w_i = 0``).
215 """
216 dual = active_atoms.T @ (target - active_atoms @ coefficients)
217 if not bool(torch.isfinite(coefficients).all()) or not bool(torch.isfinite(dual).all()): 217 ↛ 218line 217 didn't jump to line 218 because the condition on line 217 was never true
218 raise RuntimeError("NNLS result contains non-finite values")
219 positive = coefficients > coefficient_tol
220 if bool((coefficients < -coefficient_tol).any()):
221 raise RuntimeError("NNLS result violates primal feasibility (c >= 0)")
222 if bool((dual[~positive] > dual_tol[~positive]).any()):
223 raise RuntimeError("NNLS result violates dual feasibility on zeroed coordinates")
224 if bool((dual[positive].abs() > dual_tol[positive]).any()):
225 raise RuntimeError("NNLS result violates stationarity on positive coordinates")
226 complementarity = coefficients.abs() * dual.abs()
227 complementarity_tol = (coefficients.abs() + coefficient_tol) * dual_tol
228 if bool((complementarity > complementarity_tol).any()):
229 raise RuntimeError("NNLS result violates complementarity (c_i * w_i = 0)")
232def _nonnegative_least_squares(active_atoms: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
233 """Minimize ``||active_atoms @ coefficients - target||`` over ``coefficients >= 0``.
235 Uses the Lawson-Hanson active-set method. A zeroed atom enters the passive (free) set when
236 it is most positively correlated with the residual; a passive atom leaves when its trial
237 coefficient is negative. Released atoms may later re-enter.
239 The active problem is solved in float64 with an explicit pseudoinverse rank threshold.
240 Coefficients are checked against the KKT conditions before and after they are cast back to
241 the target's dtype and device. Feasibility corrections use the ``3 * num_active`` budget;
242 a separate admission limit guards against numerical cycling. An exhausted limit, invalid
243 line-search step, stalled correction, or failed KKT check raises :class:`RuntimeError`.
244 Returns coefficients aligned with the columns of ``active_atoms``.
245 """
246 if active_atoms.ndim != 2 or target.ndim != 1: 246 ↛ 247line 246 didn't jump to line 247 because the condition on line 246 was never true
247 raise ValueError("active_atoms must be 2-D and target must be 1-D")
248 if active_atoms.shape[0] != target.shape[0]: 248 ↛ 249line 248 didn't jump to line 249 because the condition on line 248 was never true
249 raise ValueError("active_atoms and target have incompatible dimensions")
250 num_active = active_atoms.shape[1]
251 if num_active == 0:
252 return target.new_zeros(0)
254 result_dtype = target.dtype
255 result_device = target.device
256 work_device = torch.device("cpu") if target.device.type == "mps" else target.device
257 work_atoms = active_atoms.to(device=work_device).to(dtype=torch.float64)
258 work_target = target.to(device=work_device).to(dtype=torch.float64)
259 if not bool(torch.isfinite(work_atoms).all()) or not bool(torch.isfinite(work_target).all()):
260 raise ValueError("active_atoms and target must contain only finite values")
262 coefficients = work_target.new_zeros(num_active)
263 passive = work_target.new_zeros(num_active, dtype=torch.bool) # free coordinates (> 0)
264 max_corrections = 3 * num_active
265 # A correction can release several variables, each of which may later be re-admitted.
266 # This is only a cycle safeguard; the mathematical convergence budget is the correction count.
267 max_admissions = num_active * (max_corrections + 1)
268 corrections = 0
269 admissions = 0
270 # Main loop: free the zeroed atom most correlated with the residual until no zeroed atom
271 # has a materially positive correlation (KKT dual feasibility), or every atom is free.
272 while not bool(passive.all()):
273 dual_tol, coefficient_tol = _nnls_tolerances(work_atoms, work_target, coefficients)
274 dual = work_atoms.T @ (work_target - work_atoms @ coefficients)
275 dual = dual.masked_fill(passive, float("-inf"))
276 if bool((dual <= dual_tol).all()):
277 break
278 admissions += 1
279 if admissions > max_admissions:
280 raise RuntimeError("NNLS exceeded the active-set admission safeguard")
281 passive[int(torch.argmax(dual))] = True
282 # Feasibility loop: solve the unconstrained fit on the passive set; while it has a
283 # materially negative coordinate, step toward it until a passive coefficient reaches
284 # zero, then release the vanished coordinates. Each correction drops >= 1 passive atom
285 # (enforced by the progress guard), so the loop is bounded.
286 while True:
287 columns = passive.nonzero(as_tuple=True)[0]
288 before = int(columns.numel())
289 solution = torch.linalg.pinv(work_atoms[:, columns], rtol=_NNLS_RANK_RTOL) @ work_target
290 candidate = work_target.new_zeros(num_active)
291 candidate[columns] = solution
292 negative = passive & (candidate < -coefficient_tol)
293 if not bool(negative.any()):
294 passive = candidate > coefficient_tol
295 coefficients = torch.where(passive, candidate, torch.zeros_like(candidate))
296 break
297 corrections += 1
298 if corrections > max_corrections: 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true
299 raise RuntimeError("NNLS exceeded the Lawson-Hanson feasibility budget")
300 # Step toward the unconstrained fit, limited by the first coordinate to reach zero.
301 # Only materially-negative coordinates with a strictly positive denominator define
302 # the step; a zero-current/zero-candidate coordinate (e.g. a freshly admitted atom
303 # heading straight to <= 0) has a zero ratio, forcing a zero-length step -- it makes
304 # no move but is released by the cleanup below, so the passive set still shrinks.
305 denominator = coefficients - candidate
306 ratios = torch.where(
307 negative & (denominator > 0),
308 coefficients / denominator,
309 work_target.new_full((num_active,), float("inf")),
310 )
311 step = float(ratios.min())
312 if not math.isfinite(step) or not ( 312 ↛ 315line 312 didn't jump to line 315 because the condition on line 312 was never true
313 -_NNLS_RELATIVE_TOLERANCE <= step <= 1.0 + _NNLS_RELATIVE_TOLERANCE
314 ):
315 raise RuntimeError(f"NNLS line-search step {step} outside [0, 1]")
316 step = min(max(step, 0.0), 1.0)
317 updated = coefficients + step * (candidate - coefficients)
318 passive = updated > coefficient_tol
319 coefficients = torch.where(passive, updated, torch.zeros_like(updated))
320 if int(passive.sum()) >= before: 320 ↛ 321line 320 didn't jump to line 321 because the condition on line 320 was never true
321 raise RuntimeError("NNLS feasibility correction made no progress")
323 dual_tol, coefficient_tol = _nnls_tolerances(work_atoms, work_target, coefficients)
324 _validate_nnls_kkt(work_atoms, work_target, coefficients, dual_tol, coefficient_tol)
325 result = coefficients.to(device=result_device, dtype=result_dtype)
326 result_atoms = active_atoms.to(device=result_device, dtype=result_dtype)
327 result_target = target.to(device=result_device, dtype=result_dtype)
328 result_tiny = torch.finfo(result_dtype).tiny
329 result_relative_tolerance = math.sqrt(torch.finfo(result_dtype).eps)
330 result_atom_norms = result_atoms.norm(dim=0).clamp_min(result_tiny)
331 result_scale = (result_target.norm() + (result_atoms.abs() @ result.abs()).norm()).clamp_min(
332 result_tiny
333 )
334 result_dual_tol = result_relative_tolerance * result_atom_norms * result_scale
335 result_coefficient_tol = result_relative_tolerance * result_scale / result_atom_norms
336 _validate_nnls_kkt(result_atoms, result_target, result, result_dual_tol, result_coefficient_tol)
337 return result
340def _gradient_pursuit_step(
341 active_atoms: torch.Tensor, target: torch.Tensor, coefficients: torch.Tensor
342) -> torch.Tensor:
343 """One optimal-step-size projected-gradient update (Blumensath & Davies, 2008).
345 ``active_atoms`` is ``[d_model, num_active]`` (the selected atoms as columns) and
346 ``coefficients`` the current ``[num_active]`` coefficients (the newly added atom starts at
347 zero, so the incoming point is feasible and its residual is the previous residual). Moves
348 the coefficients along the steepest-descent direction ``active_atoms^T residual`` with the
349 exact line-search step, then projects onto the nonnegative orthant. The exact line search
350 happens before that projection, so the projected step can in principle increase the
351 residual. In that case the step is halved until the projected update no longer increases
352 the residual, with the feasible incoming point retained if the bounded search finds none.
353 """
354 residual = target - active_atoms @ coefficients
355 feasible = coefficients.clamp_min(0.0)
356 direction = active_atoms.T @ residual # [num_active]; gradient up to sign
357 projected = active_atoms @ direction # [d_model]
358 denominator = float((projected @ projected).detach())
359 if denominator <= 0.0: 359 ↛ 360line 359 didn't jump to line 360 because the condition on line 359 was never true
360 return feasible
361 step = float((projected @ residual).detach()) / denominator
362 current_residual_squared = float((residual @ residual).detach())
363 for _ in range(_GRADIENT_BACKTRACK_STEPS + 1): 363 ↛ 370line 363 didn't jump to line 370 because the loop on line 363 didn't complete
364 candidate = (coefficients + step * direction).clamp_min(0.0)
365 candidate_residual = target - active_atoms @ candidate
366 candidate_residual_squared = float((candidate_residual @ candidate_residual).detach())
367 if candidate_residual_squared <= current_residual_squared:
368 return candidate
369 step *= 0.5
370 return feasible
373def get_sparse_decomposition(
374 x: torch.Tensor,
375 dictionary: torch.Tensor,
376 k: int = DEFAULT_K,
377 *,
378 algorithm: str = "nonnegative_orthogonal_matching_pursuit",
379) -> JSpaceDecomposition:
380 """Greedily decompose ``x`` into a ``k``-sparse nonnegative combination of atoms.
382 Args:
383 x: Target vector, shape ``[d_model]``.
384 dictionary: Atom matrix, shape ``[num_atoms, d_model]`` (rows are atoms).
385 k: Upper bound on the number of atoms to select. Selection stops early once no
386 unselected atom is materially positively correlated with the residual, so fewer
387 than ``k`` atoms may be selected (and fewer still may be numerically active).
388 algorithm: Coefficient-update rule.
389 ``"nonnegative_orthogonal_matching_pursuit"`` (default) re-solves the selected-set
390 coefficients exactly as a nonnegative least-squares fit;
391 ``"gradient_pursuit"`` takes a single projected-gradient step per atom. See the
392 module docstring for the trade-off (they use the same selection rule, while the
393 exact re-solve is optimal on each selected set).
395 Returns:
396 A :class:`JSpaceDecomposition`. Its ``support`` holds only the numerically active
397 atoms and ``selected_support`` every selected atom, with
398 ``support.numel() <= selected_support.numel() <= k``.
400 Raises:
401 ValueError: On an unknown ``algorithm``, complex or non-finite inputs, a non-2-D
402 dictionary, a target whose length does not match ``d_model``, ``k`` outside
403 ``[1, num_atoms]``, or a dictionary with non-finite or zero-norm atoms.
404 RuntimeError: If ``algorithm="nonnegative_orthogonal_matching_pursuit"`` and the
405 nonnegative least-squares solve cannot be certified against its KKT conditions
406 within its numerical tolerance.
407 """
408 if algorithm not in ("nonnegative_orthogonal_matching_pursuit", "gradient_pursuit"):
409 raise ValueError(
410 "algorithm must be 'nonnegative_orthogonal_matching_pursuit' or "
411 f"'gradient_pursuit', got {algorithm!r}"
412 )
413 if dictionary.ndim != 2:
414 raise ValueError(
415 f"dictionary must be 2-D [num_atoms, d_model], got shape {tuple(dictionary.shape)}"
416 )
417 num_atoms, d_model = dictionary.shape
418 if x.ndim != 1 or x.shape[0] != d_model:
419 raise ValueError(f"x must be 1-D of length d_model={d_model}, got shape {tuple(x.shape)}")
420 if not 1 <= k <= num_atoms:
421 raise ValueError(f"k must be between 1 and num_atoms={num_atoms}, got {k}")
422 if torch.is_complex(x) or torch.is_complex(dictionary):
423 raise ValueError("x and dictionary must be real-valued")
425 target = x.float()
426 atoms = dictionary.float()
427 if not bool(torch.isfinite(target).all()):
428 raise ValueError("x contains non-finite entries")
429 if not bool(torch.isfinite(atoms).all()):
430 raise ValueError("dictionary contains non-finite entries")
431 atom_norms = torch.linalg.vector_norm(atoms, dim=1)
432 if not bool(torch.isfinite(atom_norms).all()) or bool((atom_norms == 0).any()):
433 raise ValueError("dictionary contains a non-finite or zero-norm atom")
435 x_norm = float(torch.linalg.vector_norm(target).detach())
436 correlation_tol = _CORRELATION_RELATIVE_TOLERANCE * x_norm
438 residual = target.clone()
439 selected: List[int] = []
440 coordinates = target.new_zeros(0)
441 for _ in range(k):
442 # Select the unselected atom most correlated with the current residual, using
443 # unit-norm atoms so high-norm atoms are not preferred for their scale alone.
444 correlation = (atoms @ residual) / atom_norms
445 for chosen in selected:
446 correlation[chosen] = float("-inf")
447 candidate = int(torch.argmax(correlation).item())
448 # Stop early once no unselected atom is materially positively correlated: under the
449 # nonnegativity constraint a non-positive correlation cannot reduce the residual, so
450 # ``k`` is an upper bound on the number of selected atoms, not a target.
451 if float(correlation[candidate].detach()) <= correlation_tol:
452 break
453 selected.append(candidate)
455 selected_atoms = atoms[selected].T # [d_model, len(selected)]
456 if algorithm == "nonnegative_orthogonal_matching_pursuit":
457 # Re-solve the coefficients jointly over the selected set as a nonnegative
458 # least-squares fit. Sequential per-atom updates are wrong once selected atoms
459 # are correlated.
460 coordinates = _nonnegative_least_squares(selected_atoms, target)
461 else:
462 # Carry the coefficients forward, initialising the new atom at zero, and take a
463 # single projected-gradient step over the selected set.
464 coordinates = _gradient_pursuit_step(
465 selected_atoms, target, torch.cat([coordinates, coordinates.new_zeros(1)])
466 )
467 residual = target - selected_atoms @ coordinates
469 # ``selected_support`` stays on CPU for token decoding; the vector-valued outputs stay on
470 # the computation device.
471 selected_support = torch.tensor(selected, dtype=torch.long)
472 selected_atoms = atoms[selected_support].T # [d_model, num_selected] on atoms.device
474 # Public active support: selected atoms whose contribution ``c_i * ||v_i||`` is a
475 # materially nonzero fraction of ``||x||``. For the NNLS re-solve this is exactly the
476 # strictly-positive coordinate set (its coefficient-zeroing uses the same scale); for
477 # gradient pursuit it prunes atoms left at (near-)zero by the final projected step.
478 contribution = coordinates * torch.linalg.vector_norm(selected_atoms, dim=0)
479 active = contribution > _ACTIVE_RELATIVE_TOLERANCE * x_norm # on the computation device
480 support = selected_support[active.cpu()] # CPU, aligned with the token-decoding tensors
481 active_coordinates = coordinates[active]
482 active_atoms = selected_atoms[:, active] # [d_model, num_active] on atoms.device
484 # Reconstruction is the nonnegative combination over the active support (empty -> zeros).
485 reconstruction = active_atoms @ active_coordinates
486 # The J-space component is the orthogonal projection of the target onto the span of every
487 # selected atom (paper appendix), computed with the same pseudoinverse construction as
488 # swap_hooks. That span can be larger than the active support when a selected coordinate is
489 # zero, so the projection differs from the nonnegative reconstruction in that case.
490 if selected:
491 j_space_component = _linalg_on_cpu_if_mps(
492 lambda atoms, vector: atoms @ (torch.linalg.pinv(atoms) @ vector),
493 selected_atoms,
494 target,
495 )
496 else:
497 j_space_component = target.new_zeros(d_model)
498 non_j_space_component = target - j_space_component
499 return JSpaceDecomposition(
500 support=support,
501 coordinates=active_coordinates,
502 selected_support=selected_support,
503 reconstruction=reconstruction,
504 j_space_component=j_space_component,
505 non_j_space_component=non_j_space_component,
506 )
509@dataclass
510class JSpaceOccupancy:
511 """Result of a J-space occupancy estimate.
513 Attributes:
514 occupancy: Estimated number of meaningfully-active atoms -- the step of maximum
515 separation between the real and random-control cumulative captured variance.
516 marginal_captured_variance: Per-step captured-variance gain of the real greedy selection,
517 shape ``[max_atoms]``.
518 control_captured_variance: Per-step captured-variance gain averaged over the random
519 control dictionaries, shape ``[max_atoms]``.
520 support: Greedily selected atom indices, shape ``[max_atoms]`` (token ids when the
521 dictionary is the vocabulary of J-lens vectors).
522 """
524 occupancy: int
525 marginal_captured_variance: torch.Tensor
526 control_captured_variance: torch.Tensor
527 support: torch.Tensor
530def _greedy_captured_variance_gains(
531 atoms: torch.Tensor, atom_norms: torch.Tensor, target: torch.Tensor, max_atoms: int
532) -> Tuple[torch.Tensor, torch.Tensor]:
533 """Greedily select exactly ``max_atoms`` atoms and return captured-variance gains.
535 At each step, add the unused atom with the greatest signed, norm-normalized correlation with
536 the current residual. Project ``target`` orthogonally onto the full selected span using a
537 pseudoinverse, then set the next residual to ``target - projection``. The captured variance is
538 ``||Pi_S target||^2 / ||target||^2``; the returned values are its per-step increments.
540 This shares the per-step correlation rule with :func:`get_sparse_decomposition`, but not its
541 residual recurrence: sparse decomposition uses a nonnegative coefficient-fit residual and may
542 stop early, while this recurrence does not stop early, so the selected supports can differ.
543 """
544 total_variance = float(target @ target)
545 support: List[int] = []
546 residual = target.clone()
547 captured_variance_gains: List[float] = []
548 previous_captured_variance = 0.0
549 for _ in range(max_atoms):
550 correlation = (atoms @ residual) / atom_norms
551 for chosen in support:
552 correlation[chosen] = float("-inf")
553 support.append(int(torch.argmax(correlation).item()))
554 active_atoms = atoms[support].T
555 projection = active_atoms @ (torch.linalg.pinv(active_atoms) @ target)
556 captured_variance = float((projection @ projection) / total_variance)
557 captured_variance_gains.append(captured_variance - previous_captured_variance)
558 previous_captured_variance = captured_variance
559 residual = target - projection
560 return torch.tensor(captured_variance_gains), torch.tensor(support, dtype=torch.long)
563def estimate_occupancy(
564 x: torch.Tensor,
565 dictionary: torch.Tensor,
566 *,
567 max_atoms: int = DEFAULT_K,
568 num_control_dictionaries: int = 32,
569 seed: int = 0,
570) -> JSpaceOccupancy:
571 """Estimate how many dictionary atoms are meaningfully active in ``x``.
573 Runs the projection-residual recurrence described in
574 :func:`_greedy_captured_variance_gains` for exactly ``max_atoms`` steps and compares the real
575 per-step captured-variance curve against the same recurrence on ``num_control_dictionaries``
576 random unit-norm dictionaries of the same size. This shares sparse decomposition's per-step
577 correlation rule, but uses an unconstrained span-projection residual rather than a nonnegative
578 coefficient-fit residual, so their supports need not match. The occupancy is the step of
579 maximum separation between the real and (averaged) control *cumulative* captured variance --
580 the point past which further atoms add no more than random directions would. Deterministic
581 given ``seed`` and needs no threshold. (Captured variance is a projection, hence scale-free,
582 so the random control atoms are simply unit-norm.)
584 Args:
585 x: Target vector, shape ``[d_model]``.
586 dictionary: Atom matrix, shape ``[num_atoms, d_model]`` (rows are atoms).
587 max_atoms: Number of atoms to select in the real and control recurrences.
588 num_control_dictionaries: Number of random control dictionaries to average over.
589 seed: Seed for the random control dictionaries (reproducibility).
591 Returns:
592 An :class:`JSpaceOccupancy`.
594 Raises:
595 ValueError: On complex inputs, a non-2-D dictionary, a target whose length does not match
596 ``d_model``, ``max_atoms`` outside ``[1, num_atoms]``,
597 ``num_control_dictionaries < 1``, a target with non-finite entries or a non-finite or
598 zero norm, or a dictionary with non-finite or zero-norm atoms.
599 """
600 if dictionary.ndim != 2: 600 ↛ 601line 600 didn't jump to line 601 because the condition on line 600 was never true
601 raise ValueError(
602 f"dictionary must be 2-D [num_atoms, d_model], got shape {tuple(dictionary.shape)}"
603 )
604 num_atoms, d_model = dictionary.shape
605 if x.ndim != 1 or x.shape[0] != d_model:
606 raise ValueError(f"x must be 1-D of length d_model={d_model}, got shape {tuple(x.shape)}")
607 if not 1 <= max_atoms <= num_atoms:
608 raise ValueError(f"max_atoms must be between 1 and num_atoms={num_atoms}, got {max_atoms}")
609 if num_control_dictionaries < 1:
610 raise ValueError(
611 f"num_control_dictionaries must be at least 1, got {num_control_dictionaries}"
612 )
613 if torch.is_complex(x) or torch.is_complex(dictionary):
614 raise ValueError("x and dictionary must be real-valued")
616 target = x.float()
617 atoms = dictionary.float()
618 if not bool(torch.isfinite(target).all()):
619 raise ValueError("x contains non-finite entries")
620 target_squared_norm = target @ target
621 if not bool(torch.isfinite(target_squared_norm)):
622 raise ValueError("x must have finite norm")
623 if float(target_squared_norm) <= 0.0:
624 raise ValueError("x must have non-zero norm")
625 if not bool(torch.isfinite(atoms).all()):
626 raise ValueError("dictionary contains non-finite entries")
627 atom_norms = (atoms * atoms).sum(dim=1).sqrt()
628 if not bool(torch.isfinite(atom_norms).all()) or bool((atom_norms == 0).any()):
629 raise ValueError("dictionary contains a non-finite or zero-norm atom")
631 real_captured_variance, support = _greedy_captured_variance_gains(
632 atoms, atom_norms, target, max_atoms
633 )
635 generator = torch.Generator(device=atoms.device).manual_seed(seed)
636 control_atom_norms = torch.ones(num_atoms, device=atoms.device)
637 control_variance_runs: List[torch.Tensor] = []
638 for _ in range(num_control_dictionaries):
639 random_atoms = torch.randn(
640 num_atoms, d_model, generator=generator, device=atoms.device, dtype=atoms.dtype
641 )
642 random_atoms = random_atoms / (random_atoms * random_atoms).sum(dim=1, keepdim=True).sqrt()
643 control_run_variance, _ = _greedy_captured_variance_gains(
644 random_atoms, control_atom_norms, target, max_atoms
645 )
646 control_variance_runs.append(control_run_variance)
647 control_captured_variance = torch.stack(control_variance_runs).mean(dim=0)
649 separation = real_captured_variance.cumsum(dim=0) - control_captured_variance.cumsum(dim=0)
650 occupancy = int(torch.argmax(separation).item()) + 1
651 return JSpaceOccupancy(
652 occupancy=occupancy,
653 marginal_captured_variance=real_captured_variance,
654 control_captured_variance=control_captured_variance,
655 support=support,
656 )
659@dataclass
660class JSpaceVarianceProfile:
661 """Per-layer J-space variance profile over a prompt corpus.
663 Produced by :meth:`JacobianLens.fraction_of_variance`.
665 Attributes:
666 layers: The source layers profiled, in order.
667 median: Per-layer median over positions of the J-space variance fraction
668 ``||j_space_component||^2 / ||activation||^2``.
669 pooled: Per-layer pooled ratio ``sum(||j_space_component||^2) / sum(||activation||^2)``
670 across the corpus (the paper's "fraction of total variance").
671 per_position: Per-layer 1-D tensor of the raw per-position variance fractions.
672 """
674 layers: List[int]
675 median: Dict[int, float]
676 pooled: Dict[int, float]
677 per_position: Dict[int, torch.Tensor]