Coverage for transformer_lens/tools/analysis/jacobian_lens_decomposition.py: 94%
280 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-10-23 20:45 +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 _project_onto_span(atoms: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
124 """Project onto atom columns using a scale- and ambient-dimension-independent rank policy.
126 Unit columns are decomposed in float64, retaining singular values above the NNLS relative
127 rank tolerance. Exactly dependent and numerically indistinguishable directions are removed.
128 MPS performs this small selected-set decomposition on CPU, restoring dtype and device.
129 """
130 if atoms.shape[1] == 0: 130 ↛ 133line 130 didn't jump to line 133 because the condition on line 130 was always true
131 return torch.zeros_like(target)
133 def project(local_atoms: torch.Tensor, local_target: torch.Tensor) -> torch.Tensor:
134 work = local_atoms.double()
135 norms = torch.linalg.vector_norm(work, dim=0).clamp_min(torch.finfo(torch.float64).tiny)
136 left, singular_values, _ = torch.linalg.svd(work / norms, full_matrices=False)
137 basis = left[:, singular_values > _NNLS_RANK_RTOL * singular_values.max()]
138 return (basis @ (basis.T @ local_target.double())).to(local_target.dtype)
140 return _linalg_on_cpu_if_mps(project, atoms, target)
143def _diagnose_intervention_pair(
144 unit_vectors: torch.Tensor, *, description: str, stacklevel: int
145) -> float:
146 """Apply the shared near-parallel warn/raise policy to two unit-normalized vectors.
148 Returns the signed cosine. Warns above :data:`_SWAP_WARN_COSINE` and raises above
149 :data:`_SWAP_ERROR_COSINE`. Used by interventions that invert a two-atom basis (e.g.
150 ``swap_hooks``); the anchored coordinate patch, which performs no inverse, warns without
151 raising and does not call this helper.
152 """
153 cosine = float((unit_vectors[0] @ unit_vectors[1]).item())
154 abs_cosine = abs(cosine)
155 if not math.isfinite(abs_cosine) or abs_cosine >= _SWAP_ERROR_COSINE:
156 raise ValueError(
157 f"{description} are numerically near-parallel "
158 f"(abs cosine={abs_cosine:.6f}); choose better-separated concepts"
159 )
160 if abs_cosine >= _SWAP_WARN_COSINE:
161 warnings.warn(
162 f"{description} are poorly conditioned "
163 f"(abs cosine={abs_cosine:.6f}); the intervention may be amplified",
164 UserWarning,
165 stacklevel=stacklevel,
166 )
167 return cosine
170@dataclass
171class JSpaceDecomposition:
172 """Result of a sparse J-space decomposition.
174 Attributes:
175 support: Indices of the numerically *active* selected atoms -- those whose
176 nonnegative coordinate materially contributes (token ids when the dictionary is
177 the vocabulary of J-lens vectors). A subset of ``selected_support``.
178 coordinates: Nonnegative pursuit coefficients aligned with ``support`` (the
179 "local J-space coordinates"); every entry is materially nonzero.
180 selected_support: Indices of every greedily selected atom, including any whose
181 coordinate was driven to zero by the nonnegativity constraint. Defines the span
182 for ``j_space_component``. Satisfies
183 ``support.numel() <= selected_support.numel() <= k``.
184 reconstruction: The nonnegative combination ``sum(coordinates * active atoms)`` over
185 ``support``.
186 j_space_component: The orthogonal projection of the target onto the span of
187 ``selected_support`` (the paper's "J-space component"). For the exact NNLS
188 re-solve it equals ``reconstruction`` unless a selected atom has a zero
189 coordinate, in which case the projection uses a larger span.
190 Numerical rank is measured on unit atom columns in float64 with relative
191 tolerance ``sqrt(eps64)``, independently of ambient dimension and atom norms.
192 non_j_space_component: The residual ``target - j_space_component`` (the
193 "non-J-space component"), orthogonal to the selected span.
194 """
196 support: torch.Tensor
197 coordinates: torch.Tensor
198 selected_support: torch.Tensor
199 reconstruction: torch.Tensor
200 j_space_component: torch.Tensor
201 non_j_space_component: torch.Tensor
204def _nnls_tolerances(
205 active_atoms: torch.Tensor, target: torch.Tensor, coefficients: torch.Tensor
206) -> tuple[torch.Tensor, torch.Tensor]:
207 """Scale- and dtype-aware tolerances for the NNLS solver and its KKT check.
209 The dual score ``w_j = A[:, j]^T r`` (``r = target - A c``) scales as an atom norm times
210 the residual-computation scale. That scale includes ``|| |A| |c| ||`` so a nearly
211 cancelling fit receives enough absolute tolerance without making the relative threshold
212 depend on the number of rows. A numerically-zero coefficient has the corresponding inverse
213 atom-norm scale. Returns ``(dual_tol, coefficient_tol)``, each aligned with the columns of
214 ``active_atoms``.
215 """
216 tiny = torch.finfo(target.dtype).tiny
217 atom_norms = active_atoms.norm(dim=0).clamp_min(tiny)
218 scale = (target.norm() + (active_atoms.abs() @ coefficients.abs()).norm()).clamp_min(tiny)
219 dual_tol = _NNLS_RELATIVE_TOLERANCE * atom_norms * scale
220 coefficient_tol = _NNLS_RELATIVE_TOLERANCE * scale / atom_norms
221 return dual_tol, coefficient_tol
224def _validate_nnls_kkt(
225 active_atoms: torch.Tensor,
226 target: torch.Tensor,
227 coefficients: torch.Tensor,
228 dual_tol: torch.Tensor,
229 coefficient_tol: torch.Tensor,
230) -> None:
231 """Raise unless ``coefficients`` satisfy the NNLS KKT conditions within tolerance.
233 This is a postcondition check using the solver's numerical tolerances: primal
234 feasibility (``c >= 0``), dual feasibility on zeroed coordinates (``w_j <= 0``),
235 stationarity on positive coordinates (``w_j = 0``), and complementarity
236 (``c_i w_i = 0``).
237 """
238 dual = active_atoms.T @ (target - active_atoms @ coefficients)
239 if not bool(torch.isfinite(coefficients).all()) or not bool(torch.isfinite(dual).all()): 239 ↛ 240line 239 didn't jump to line 240 because the condition on line 239 was never true
240 raise RuntimeError("NNLS result contains non-finite values")
241 positive = coefficients > coefficient_tol
242 if bool((coefficients < -coefficient_tol).any()):
243 raise RuntimeError("NNLS result violates primal feasibility (c >= 0)")
244 if bool((dual[~positive] > dual_tol[~positive]).any()):
245 raise RuntimeError("NNLS result violates dual feasibility on zeroed coordinates")
246 if bool((dual[positive].abs() > dual_tol[positive]).any()):
247 raise RuntimeError("NNLS result violates stationarity on positive coordinates")
248 complementarity = coefficients.abs() * dual.abs()
249 complementarity_tol = (coefficients.abs() + coefficient_tol) * dual_tol
250 if bool((complementarity > complementarity_tol).any()):
251 raise RuntimeError("NNLS result violates complementarity (c_i * w_i = 0)")
254def _nonnegative_least_squares(active_atoms: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
255 """Minimize ``||active_atoms @ coefficients - target||`` over ``coefficients >= 0``.
257 Uses the Lawson-Hanson active-set method. A zeroed atom enters the passive (free) set when
258 it is most positively correlated with the residual; a passive atom leaves when its trial
259 coefficient is negative. Released atoms may later re-enter.
261 The active problem is solved in float64 with an explicit pseudoinverse rank threshold.
262 Coefficients are checked against the KKT conditions before and after they are cast back to
263 the target's dtype and device. Feasibility corrections use the ``3 * num_active`` budget;
264 a separate admission limit guards against numerical cycling. An exhausted limit, invalid
265 line-search step, stalled correction, or failed KKT check raises :class:`RuntimeError`.
266 Returns coefficients aligned with the columns of ``active_atoms``.
267 """
268 if active_atoms.ndim != 2 or target.ndim != 1: 268 ↛ 269line 268 didn't jump to line 269 because the condition on line 268 was never true
269 raise ValueError("active_atoms must be 2-D and target must be 1-D")
270 if active_atoms.shape[0] != target.shape[0]: 270 ↛ 271line 270 didn't jump to line 271 because the condition on line 270 was never true
271 raise ValueError("active_atoms and target have incompatible dimensions")
272 num_active = active_atoms.shape[1]
273 if num_active == 0:
274 return target.new_zeros(0)
276 result_dtype = target.dtype
277 result_device = target.device
278 work_device = torch.device("cpu") if target.device.type == "mps" else target.device
279 work_atoms = active_atoms.to(device=work_device).to(dtype=torch.float64)
280 work_target = target.to(device=work_device).to(dtype=torch.float64)
281 if not bool(torch.isfinite(work_atoms).all()) or not bool(torch.isfinite(work_target).all()):
282 raise ValueError("active_atoms and target must contain only finite values")
284 coefficients = work_target.new_zeros(num_active)
285 passive = work_target.new_zeros(num_active, dtype=torch.bool) # free coordinates (> 0)
286 max_corrections = 3 * num_active
287 # A correction can release several variables, each of which may later be re-admitted.
288 # This is only a cycle safeguard; the mathematical convergence budget is the correction count.
289 max_admissions = num_active * (max_corrections + 1)
290 corrections = 0
291 admissions = 0
292 # Main loop: free the zeroed atom most correlated with the residual until no zeroed atom
293 # has a materially positive correlation (KKT dual feasibility), or every atom is free.
294 while not bool(passive.all()):
295 dual_tol, coefficient_tol = _nnls_tolerances(work_atoms, work_target, coefficients)
296 dual = work_atoms.T @ (work_target - work_atoms @ coefficients)
297 dual = dual.masked_fill(passive, float("-inf"))
298 if bool((dual <= dual_tol).all()):
299 break
300 admissions += 1
301 if admissions > max_admissions:
302 raise RuntimeError("NNLS exceeded the active-set admission safeguard")
303 passive[int(torch.argmax(dual))] = True
304 # Feasibility loop: solve the unconstrained fit on the passive set; while it has a
305 # materially negative coordinate, step toward it until a passive coefficient reaches
306 # zero, then release the vanished coordinates. Each correction drops >= 1 passive atom
307 # (enforced by the progress guard), so the loop is bounded.
308 while True:
309 columns = passive.nonzero(as_tuple=True)[0]
310 before = int(columns.numel())
311 solution = torch.linalg.pinv(work_atoms[:, columns], rtol=_NNLS_RANK_RTOL) @ work_target
312 candidate = work_target.new_zeros(num_active)
313 candidate[columns] = solution
314 negative = passive & (candidate < -coefficient_tol)
315 if not bool(negative.any()):
316 passive = candidate > coefficient_tol
317 coefficients = torch.where(passive, candidate, torch.zeros_like(candidate))
318 break
319 corrections += 1
320 if corrections > max_corrections: 320 ↛ 321line 320 didn't jump to line 321 because the condition on line 320 was never true
321 raise RuntimeError("NNLS exceeded the Lawson-Hanson feasibility budget")
322 # Step toward the unconstrained fit, limited by the first coordinate to reach zero.
323 # Only materially-negative coordinates with a strictly positive denominator define
324 # the step; a zero-current/zero-candidate coordinate (e.g. a freshly admitted atom
325 # heading straight to <= 0) has a zero ratio, forcing a zero-length step -- it makes
326 # no move but is released by the cleanup below, so the passive set still shrinks.
327 denominator = coefficients - candidate
328 ratios = torch.where(
329 negative & (denominator > 0),
330 coefficients / denominator,
331 work_target.new_full((num_active,), float("inf")),
332 )
333 step = float(ratios.min())
334 if not math.isfinite(step) or not ( 334 ↛ 337line 334 didn't jump to line 337 because the condition on line 334 was never true
335 -_NNLS_RELATIVE_TOLERANCE <= step <= 1.0 + _NNLS_RELATIVE_TOLERANCE
336 ):
337 raise RuntimeError(f"NNLS line-search step {step} outside [0, 1]")
338 step = min(max(step, 0.0), 1.0)
339 updated = coefficients + step * (candidate - coefficients)
340 passive = updated > coefficient_tol
341 coefficients = torch.where(passive, updated, torch.zeros_like(updated))
342 if int(passive.sum()) >= before: 342 ↛ 343line 342 didn't jump to line 343 because the condition on line 342 was never true
343 raise RuntimeError("NNLS feasibility correction made no progress")
345 dual_tol, coefficient_tol = _nnls_tolerances(work_atoms, work_target, coefficients)
346 _validate_nnls_kkt(work_atoms, work_target, coefficients, dual_tol, coefficient_tol)
347 result = coefficients.to(device=result_device, dtype=result_dtype)
348 result_atoms = active_atoms.to(device=result_device, dtype=result_dtype)
349 result_target = target.to(device=result_device, dtype=result_dtype)
350 result_tiny = torch.finfo(result_dtype).tiny
351 result_relative_tolerance = math.sqrt(torch.finfo(result_dtype).eps)
352 result_atom_norms = result_atoms.norm(dim=0).clamp_min(result_tiny)
353 result_scale = (result_target.norm() + (result_atoms.abs() @ result.abs()).norm()).clamp_min(
354 result_tiny
355 )
356 result_dual_tol = result_relative_tolerance * result_atom_norms * result_scale
357 result_coefficient_tol = result_relative_tolerance * result_scale / result_atom_norms
358 _validate_nnls_kkt(result_atoms, result_target, result, result_dual_tol, result_coefficient_tol)
359 return result
362def _gradient_pursuit_step(
363 active_atoms: torch.Tensor, target: torch.Tensor, coefficients: torch.Tensor
364) -> torch.Tensor:
365 """One optimal-step-size projected-gradient update (Blumensath & Davies, 2008).
367 ``active_atoms`` is ``[d_model, num_active]`` (the selected atoms as columns) and
368 ``coefficients`` the current ``[num_active]`` coefficients (the newly added atom starts at
369 zero, so the incoming point is feasible and its residual is the previous residual). Moves
370 the coefficients along the steepest-descent direction ``active_atoms^T residual`` with the
371 exact line-search step, then projects onto the nonnegative orthant. The exact line search
372 happens before that projection, so the projected step can in principle increase the
373 residual. In that case the step is halved until the projected update no longer increases
374 the residual, with the feasible incoming point retained if the bounded search finds none.
375 """
376 residual = target - active_atoms @ coefficients
377 feasible = coefficients.clamp_min(0.0)
378 direction = active_atoms.T @ residual # [num_active]; gradient up to sign
379 projected = active_atoms @ direction # [d_model]
380 denominator = float((projected @ projected).detach())
381 if denominator <= 0.0: 381 ↛ 382line 381 didn't jump to line 382 because the condition on line 381 was never true
382 return feasible
383 step = float((projected @ residual).detach()) / denominator
384 current_residual_squared = float((residual @ residual).detach())
385 for _ in range(_GRADIENT_BACKTRACK_STEPS + 1): 385 ↛ 392line 385 didn't jump to line 392 because the loop on line 385 didn't complete
386 candidate = (coefficients + step * direction).clamp_min(0.0)
387 candidate_residual = target - active_atoms @ candidate
388 candidate_residual_squared = float((candidate_residual @ candidate_residual).detach())
389 if candidate_residual_squared <= current_residual_squared:
390 return candidate
391 step *= 0.5
392 return feasible
395def _gradient_pursuit_extend(
396 selected_atoms: torch.Tensor, target: torch.Tensor, coordinates: torch.Tensor
397) -> torch.Tensor:
398 """Gradient-pursuit update after the newest column of ``selected_atoms`` joins at zero."""
399 return _gradient_pursuit_step(
400 selected_atoms, target, torch.cat([coordinates, coordinates.new_zeros(1)])
401 )
404def _active_mask(
405 coordinates: torch.Tensor, selected_atoms: torch.Tensor, target_norm: float
406) -> torch.Tensor:
407 """Selected atoms whose contribution ``c_i * ||v_i||`` is a material fraction of ``||x||``."""
408 contribution = coordinates * torch.linalg.vector_norm(selected_atoms, dim=0)
409 return contribution > _ACTIVE_RELATIVE_TOLERANCE * target_norm
412def get_sparse_decomposition(
413 x: torch.Tensor,
414 dictionary: torch.Tensor,
415 k: int = DEFAULT_K,
416 *,
417 algorithm: str = "nonnegative_orthogonal_matching_pursuit",
418) -> JSpaceDecomposition:
419 """Greedily decompose ``x`` into a ``k``-sparse nonnegative combination of atoms.
421 Args:
422 x: Target vector, shape ``[d_model]``.
423 dictionary: Atom matrix, shape ``[num_atoms, d_model]`` (rows are atoms).
424 k: Upper bound on the number of atoms to select. Selection stops early once no
425 unselected atom is materially positively correlated with the residual, so fewer
426 than ``k`` atoms may be selected (and fewer still may be numerically active).
427 algorithm: Coefficient-update rule.
428 ``"nonnegative_orthogonal_matching_pursuit"`` (default) re-solves the selected-set
429 coefficients exactly as a nonnegative least-squares fit;
430 ``"gradient_pursuit"`` takes a single projected-gradient step per atom. See the
431 module docstring for the trade-off (they use the same selection rule, while the
432 exact re-solve is optimal on each selected set).
434 Returns:
435 A :class:`JSpaceDecomposition`. Its ``support`` holds only the numerically active
436 atoms and ``selected_support`` every selected atom, with
437 ``support.numel() <= selected_support.numel() <= k``.
439 Raises:
440 ValueError: On an unknown ``algorithm``, complex or non-finite inputs, a non-2-D
441 dictionary, a target whose length does not match ``d_model``, ``k`` outside
442 ``[1, num_atoms]``, or a dictionary with non-finite or zero-norm atoms.
443 RuntimeError: If ``algorithm="nonnegative_orthogonal_matching_pursuit"`` and the
444 nonnegative least-squares solve cannot be certified against its KKT conditions
445 within its numerical tolerance.
446 """
447 if algorithm not in ("nonnegative_orthogonal_matching_pursuit", "gradient_pursuit"):
448 raise ValueError(
449 "algorithm must be 'nonnegative_orthogonal_matching_pursuit' or "
450 f"'gradient_pursuit', got {algorithm!r}"
451 )
452 if dictionary.ndim != 2:
453 raise ValueError(
454 f"dictionary must be 2-D [num_atoms, d_model], got shape {tuple(dictionary.shape)}"
455 )
456 num_atoms, d_model = dictionary.shape
457 if x.ndim != 1 or x.shape[0] != d_model:
458 raise ValueError(f"x must be 1-D of length d_model={d_model}, got shape {tuple(x.shape)}")
459 if not 1 <= k <= num_atoms:
460 raise ValueError(f"k must be between 1 and num_atoms={num_atoms}, got {k}")
461 if torch.is_complex(x) or torch.is_complex(dictionary):
462 raise ValueError("x and dictionary must be real-valued")
464 target = x.float()
465 atoms = dictionary.float()
466 if not bool(torch.isfinite(target).all()):
467 raise ValueError("x contains non-finite entries")
468 if not bool(torch.isfinite(atoms).all()):
469 raise ValueError("dictionary contains non-finite entries")
470 atom_norms = torch.linalg.vector_norm(atoms, dim=1)
471 if not bool(torch.isfinite(atom_norms).all()) or bool((atom_norms == 0).any()):
472 raise ValueError("dictionary contains a non-finite or zero-norm atom")
474 x_norm = float(torch.linalg.vector_norm(target).detach())
475 correlation_tol = _CORRELATION_RELATIVE_TOLERANCE * x_norm
477 residual = target.clone()
478 selected: List[int] = []
479 coordinates = target.new_zeros(0)
480 for _ in range(k):
481 # Select the unselected atom most correlated with the current residual, using
482 # unit-norm atoms so high-norm atoms are not preferred for their scale alone.
483 correlation = (atoms @ residual) / atom_norms
484 for chosen in selected:
485 correlation[chosen] = float("-inf")
486 candidate = int(torch.argmax(correlation).item())
487 # Stop early once no unselected atom is materially positively correlated: under the
488 # nonnegativity constraint a non-positive correlation cannot reduce the residual, so
489 # ``k`` is an upper bound on the number of selected atoms, not a target.
490 if float(correlation[candidate].detach()) <= correlation_tol:
491 break
492 selected.append(candidate)
494 selected_atoms = atoms[selected].T # [d_model, len(selected)]
495 if algorithm == "nonnegative_orthogonal_matching_pursuit":
496 # Re-solve the coefficients jointly over the selected set as a nonnegative
497 # least-squares fit. Sequential per-atom updates are wrong once selected atoms
498 # are correlated.
499 coordinates = _nonnegative_least_squares(selected_atoms, target)
500 else:
501 # Carry the coefficients forward, initialising the new atom at zero, and take a
502 # single projected-gradient step over the selected set.
503 coordinates = _gradient_pursuit_extend(selected_atoms, target, coordinates)
504 residual = target - selected_atoms @ coordinates
506 # ``selected_support`` stays on CPU for token decoding; the vector-valued outputs stay on
507 # the computation device.
508 selected_support = torch.tensor(selected, dtype=torch.long)
509 selected_atoms = atoms[selected_support].T # [d_model, num_selected] on atoms.device
511 # Public active support: selected atoms whose contribution ``c_i * ||v_i||`` is a
512 # materially nonzero fraction of ``||x||``. For the NNLS re-solve this is exactly the
513 # strictly-positive coordinate set (its coefficient-zeroing uses the same scale); for
514 # gradient pursuit it prunes atoms left at (near-)zero by the final projected step.
515 active = _active_mask(coordinates, selected_atoms, x_norm) # on the computation device
516 support = selected_support[active.cpu()] # CPU, aligned with the token-decoding tensors
517 active_coordinates = coordinates[active]
518 active_atoms = selected_atoms[:, active] # [d_model, num_active] on atoms.device
520 # Reconstruction is the nonnegative combination over the active support (empty -> zeros).
521 reconstruction = active_atoms @ active_coordinates
522 # The J-space component is the orthogonal projection of the target onto the span of every
523 # selected atom (paper appendix). That span can exceed the active support when a coordinate is
524 # zero, so the projection differs from the nonnegative reconstruction in that case.
525 if selected:
526 j_space_component = _project_onto_span(selected_atoms, target)
527 else:
528 j_space_component = target.new_zeros(d_model)
529 non_j_space_component = target - j_space_component
530 return JSpaceDecomposition(
531 support=support,
532 coordinates=active_coordinates,
533 selected_support=selected_support,
534 reconstruction=reconstruction,
535 j_space_component=j_space_component,
536 non_j_space_component=non_j_space_component,
537 )
540@dataclass
541class JSpaceOccupancy:
542 """Result of a J-space occupancy estimate.
544 Attributes:
545 occupancy: Estimated number of meaningfully-active atoms -- the step of maximum
546 separation between the real and random-control cumulative captured variance.
547 marginal_captured_variance: Per-step captured-variance gain of the real greedy selection,
548 shape ``[max_atoms]``.
549 control_captured_variance: Per-step captured-variance gain averaged over the random
550 control dictionaries, shape ``[max_atoms]``.
551 support: Greedily selected atom indices, shape ``[max_atoms]`` (token ids when the
552 dictionary is the vocabulary of J-lens vectors).
553 """
555 occupancy: int
556 marginal_captured_variance: torch.Tensor
557 control_captured_variance: torch.Tensor
558 support: torch.Tensor
561def _greedy_captured_variance_gains(
562 atoms: torch.Tensor, atom_norms: torch.Tensor, target: torch.Tensor, max_atoms: int
563) -> Tuple[torch.Tensor, torch.Tensor]:
564 """Greedily select exactly ``max_atoms`` atoms and return captured-variance gains.
566 At each step, add the unused atom with the greatest signed, norm-normalized correlation with
567 the current residual. Project ``target`` orthogonally onto the full selected span using a
568 rank-revealing unit-column SVD, then set the next residual to ``target - projection``.
569 Captured variance is ``||Pi_S target||^2 / ||target||^2``; the returned values are its increments.
571 This shares the per-step correlation rule with :func:`get_sparse_decomposition`, but not its
572 residual recurrence: sparse decomposition uses a nonnegative coefficient-fit residual and may
573 stop early, while this recurrence does not stop early, so the selected supports can differ.
574 """
575 total_variance = float(target @ target)
576 support: List[int] = []
577 residual = target.clone()
578 captured_variance_gains: List[float] = []
579 previous_captured_variance = 0.0
580 for _ in range(max_atoms):
581 correlation = (atoms @ residual) / atom_norms
582 for chosen in support:
583 correlation[chosen] = float("-inf")
584 support.append(int(torch.argmax(correlation).item()))
585 active_atoms = atoms[support].T
586 projection = _project_onto_span(active_atoms, target)
587 captured_variance = float((projection @ projection) / total_variance)
588 captured_variance_gains.append(captured_variance - previous_captured_variance)
589 previous_captured_variance = captured_variance
590 residual = target - projection
591 return torch.tensor(captured_variance_gains), torch.tensor(support, dtype=torch.long)
594def estimate_occupancy(
595 x: torch.Tensor,
596 dictionary: torch.Tensor,
597 *,
598 max_atoms: int = DEFAULT_K,
599 num_control_dictionaries: int = 32,
600 seed: int = 0,
601) -> JSpaceOccupancy:
602 """Estimate how many dictionary atoms are meaningfully active in ``x``.
604 Runs the projection-residual recurrence described in
605 :func:`_greedy_captured_variance_gains` for exactly ``max_atoms`` steps and compares the real
606 per-step captured-variance curve against the same recurrence on ``num_control_dictionaries``
607 random unit-norm dictionaries of the same size. This shares sparse decomposition's per-step
608 correlation rule, but uses an unconstrained span-projection residual rather than a nonnegative
609 coefficient-fit residual, so their supports need not match. The occupancy is the step of
610 maximum separation between the real and (averaged) control *cumulative* captured variance --
611 the point past which further atoms add no more than random directions would. Deterministic
612 given ``seed`` on every device and needs no threshold. (Captured variance is a projection,
613 hence scale-free, so the random control atoms are simply unit-norm.)
615 Control dictionaries are drawn on the CPU (for cross-device reproducibility), so for a large
616 dictionary on an accelerator that draw dominates run time; lower ``num_control_dictionaries``
617 to trade control precision for speed.
619 Args:
620 x: Target vector, shape ``[d_model]``.
621 dictionary: Atom matrix, shape ``[num_atoms, d_model]`` (rows are atoms).
622 max_atoms: Number of atoms to select in the real and control recurrences.
623 num_control_dictionaries: Number of random control dictionaries to average over; each
624 is a CPU draw the size of ``dictionary``.
625 seed: Seed for the random control dictionaries (reproducibility).
627 Returns:
628 An :class:`JSpaceOccupancy`.
630 Raises:
631 ValueError: On complex inputs, a non-2-D dictionary, a target whose length does not match
632 ``d_model``, ``max_atoms`` outside ``[1, num_atoms]``,
633 ``num_control_dictionaries < 1``, a target with non-finite entries or a non-finite or
634 zero norm, or a dictionary with non-finite or zero-norm atoms.
635 """
636 if dictionary.ndim != 2: 636 ↛ 637line 636 didn't jump to line 637 because the condition on line 636 was never true
637 raise ValueError(
638 f"dictionary must be 2-D [num_atoms, d_model], got shape {tuple(dictionary.shape)}"
639 )
640 num_atoms, d_model = dictionary.shape
641 if x.ndim != 1 or x.shape[0] != d_model:
642 raise ValueError(f"x must be 1-D of length d_model={d_model}, got shape {tuple(x.shape)}")
643 if not 1 <= max_atoms <= num_atoms:
644 raise ValueError(f"max_atoms must be between 1 and num_atoms={num_atoms}, got {max_atoms}")
645 if num_control_dictionaries < 1:
646 raise ValueError(
647 f"num_control_dictionaries must be at least 1, got {num_control_dictionaries}"
648 )
649 if torch.is_complex(x) or torch.is_complex(dictionary):
650 raise ValueError("x and dictionary must be real-valued")
652 target = x.float()
653 atoms = dictionary.float()
654 if not bool(torch.isfinite(target).all()):
655 raise ValueError("x contains non-finite entries")
656 target_squared_norm = target @ target
657 if not bool(torch.isfinite(target_squared_norm)):
658 raise ValueError("x must have finite norm")
659 if float(target_squared_norm) <= 0.0:
660 raise ValueError("x must have non-zero norm")
661 if not bool(torch.isfinite(atoms).all()):
662 raise ValueError("dictionary contains non-finite entries")
663 atom_norms = (atoms * atoms).sum(dim=1).sqrt()
664 if not bool(torch.isfinite(atom_norms).all()) or bool((atom_norms == 0).any()):
665 raise ValueError("dictionary contains a non-finite or zero-norm atom")
667 real_captured_variance, support = _greedy_captured_variance_gains(
668 atoms, atom_norms, target, max_atoms
669 )
671 # Drawn on CPU so a seed gives the same control dictionaries on every device.
672 generator = torch.Generator(device="cpu").manual_seed(seed)
673 control_atom_norms = torch.ones(num_atoms, device=atoms.device)
674 control_variance_runs: List[torch.Tensor] = []
675 for _ in range(num_control_dictionaries):
676 random_atoms = torch.randn(num_atoms, d_model, generator=generator, dtype=atoms.dtype)
677 random_atoms = random_atoms.to(atoms.device)
678 random_atoms = random_atoms / (random_atoms * random_atoms).sum(dim=1, keepdim=True).sqrt()
679 control_run_variance, _ = _greedy_captured_variance_gains(
680 random_atoms, control_atom_norms, target, max_atoms
681 )
682 control_variance_runs.append(control_run_variance)
683 control_captured_variance = torch.stack(control_variance_runs).mean(dim=0)
685 separation = real_captured_variance.cumsum(dim=0) - control_captured_variance.cumsum(dim=0)
686 occupancy = int(torch.argmax(separation).item()) + 1
687 return JSpaceOccupancy(
688 occupancy=occupancy,
689 marginal_captured_variance=real_captured_variance,
690 control_captured_variance=control_captured_variance,
691 support=support,
692 )
695@dataclass
696class JSpaceVarianceProfile:
697 """Per-layer J-space variance profile over a prompt corpus.
699 Produced by :meth:`JacobianLens.fraction_of_variance`.
701 Attributes:
702 layers: The source layers profiled, in order.
703 median: Per-layer median over positions of the J-space variance fraction
704 ``||j_space_component||^2 / ||activation||^2``.
705 pooled: Per-layer pooled ratio ``sum(||j_space_component||^2) / sum(||activation||^2)``
706 across the corpus (the paper's "fraction of total variance").
707 per_position: Per-layer 1-D tensor of the raw per-position variance fractions.
708 """
710 layers: List[int]
711 median: Dict[int, float]
712 pooled: Dict[int, float]
713 per_position: Dict[int, torch.Tensor]