Coverage for transformer_lens/tools/analysis/jacobian_lens_decomposition.py: 95%
249 statements
« prev ^ index » next coverage.py v7.10.1, created at 2026-09-01 16:23 +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
63from dataclasses import dataclass
64from typing import Dict, List, Tuple
66import torch
68#: The paper varies the sparsity level but "typically" chooses "no more than 25".
69DEFAULT_K = 25
71#: Active problems are normally small (``DEFAULT_K`` is 25), so solve them in float64 and use
72#: a conservative square-root-epsilon threshold for rank and KKT decisions. Unlike a threshold
73#: proportional to ``d_model``, this is invariant to appending zero rows to the same problem.
74_NNLS_RELATIVE_TOLERANCE = math.sqrt(torch.finfo(torch.float64).eps)
75_NNLS_RANK_RTOL = _NNLS_RELATIVE_TOLERANCE
77#: A selected atom counts as numerically *active* when its contribution ``c_i * ||v_i||`` is a
78#: materially nonzero fraction of the target scale ``||x||``. Using the contribution (not the
79#: raw coefficient) keeps the test scale-invariant across J-lens vectors of different native
80#: norms. This is a numerical-activity threshold, not an interpretability one. It matches the
81#: default NNLS coefficient-zeroing scale (:data:`_NNLS_RELATIVE_TOLERANCE`), so the active
82#: support equals the set of strictly-positive NNLS coordinates and the reconstruction over the
83#: active support is the projection onto it.
84_ACTIVE_RELATIVE_TOLERANCE = _NNLS_RELATIVE_TOLERANCE
86_GRADIENT_BACKTRACK_STEPS = 20
88#: Early-stopping correlation threshold. Once the best unselected normalized correlation drops
89#: to this fraction of ``||x||``, no atom can materially reduce the residual and selection
90#: stops. It sits at the float32 residual noise floor (the residual is computed in float32),
91#: so a full-rank target stops instead of selecting noise atoms -- this is what makes the
92#: selected support scale-invariant. It is deliberately coarser than the activity threshold:
93#: selection is a float32 correlation decision, activity a check against the float64 solve.
94_CORRELATION_RELATIVE_TOLERANCE = math.sqrt(torch.finfo(torch.float32).eps)
97@dataclass
98class JSpaceDecomposition:
99 """Result of a sparse J-space decomposition.
101 Attributes:
102 support: Indices of the numerically *active* selected atoms -- those whose
103 nonnegative coordinate materially contributes (token ids when the dictionary is
104 the vocabulary of J-lens vectors). A subset of ``selected_support``.
105 coordinates: Nonnegative pursuit coefficients aligned with ``support`` (the
106 "local J-space coordinates"); every entry is materially nonzero.
107 selected_support: Indices of every greedily selected atom, including any whose
108 coordinate was driven to zero by the nonnegativity constraint. Defines the span
109 for ``j_space_component``. Satisfies
110 ``support.numel() <= selected_support.numel() <= k``.
111 reconstruction: The nonnegative combination ``sum(coordinates * active atoms)`` over
112 ``support``.
113 j_space_component: The orthogonal projection of the target onto the span of
114 ``selected_support`` (the paper's "J-space component"). For the exact NNLS
115 re-solve it equals ``reconstruction`` unless a selected atom has a zero
116 coordinate, in which case the projection uses a larger span.
117 non_j_space_component: The residual ``target - j_space_component`` (the
118 "non-J-space component"), orthogonal to the selected span.
119 """
121 support: torch.Tensor
122 coordinates: torch.Tensor
123 selected_support: torch.Tensor
124 reconstruction: torch.Tensor
125 j_space_component: torch.Tensor
126 non_j_space_component: torch.Tensor
129def _nnls_tolerances(
130 active_atoms: torch.Tensor, target: torch.Tensor, coefficients: torch.Tensor
131) -> tuple[torch.Tensor, torch.Tensor]:
132 """Scale- and dtype-aware tolerances for the NNLS solver and its KKT check.
134 The dual score ``w_j = A[:, j]^T r`` (``r = target - A c``) scales as an atom norm times
135 the residual-computation scale. That scale includes ``|| |A| |c| ||`` so a nearly
136 cancelling fit receives enough absolute tolerance without making the relative threshold
137 depend on the number of rows. A numerically-zero coefficient has the corresponding inverse
138 atom-norm scale. Returns ``(dual_tol, coefficient_tol)``, each aligned with the columns of
139 ``active_atoms``.
140 """
141 tiny = torch.finfo(target.dtype).tiny
142 atom_norms = active_atoms.norm(dim=0).clamp_min(tiny)
143 scale = (target.norm() + (active_atoms.abs() @ coefficients.abs()).norm()).clamp_min(tiny)
144 dual_tol = _NNLS_RELATIVE_TOLERANCE * atom_norms * scale
145 coefficient_tol = _NNLS_RELATIVE_TOLERANCE * scale / atom_norms
146 return dual_tol, coefficient_tol
149def _validate_nnls_kkt(
150 active_atoms: torch.Tensor,
151 target: torch.Tensor,
152 coefficients: torch.Tensor,
153 dual_tol: torch.Tensor,
154 coefficient_tol: torch.Tensor,
155) -> None:
156 """Raise unless ``coefficients`` satisfy the NNLS KKT conditions within tolerance.
158 This is a postcondition check using the solver's numerical tolerances: primal
159 feasibility (``c >= 0``), dual feasibility on zeroed coordinates (``w_j <= 0``),
160 stationarity on positive coordinates (``w_j = 0``), and complementarity
161 (``c_i w_i = 0``).
162 """
163 dual = active_atoms.T @ (target - active_atoms @ coefficients)
164 if not bool(torch.isfinite(coefficients).all()) or not bool(torch.isfinite(dual).all()): 164 ↛ 165line 164 didn't jump to line 165 because the condition on line 164 was never true
165 raise RuntimeError("NNLS result contains non-finite values")
166 positive = coefficients > coefficient_tol
167 if bool((coefficients < -coefficient_tol).any()):
168 raise RuntimeError("NNLS result violates primal feasibility (c >= 0)")
169 if bool((dual[~positive] > dual_tol[~positive]).any()):
170 raise RuntimeError("NNLS result violates dual feasibility on zeroed coordinates")
171 if bool((dual[positive].abs() > dual_tol[positive]).any()):
172 raise RuntimeError("NNLS result violates stationarity on positive coordinates")
173 complementarity = coefficients.abs() * dual.abs()
174 complementarity_tol = (coefficients.abs() + coefficient_tol) * dual_tol
175 if bool((complementarity > complementarity_tol).any()):
176 raise RuntimeError("NNLS result violates complementarity (c_i * w_i = 0)")
179def _nonnegative_least_squares(active_atoms: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
180 """Minimize ``||active_atoms @ coefficients - target||`` over ``coefficients >= 0``.
182 Uses the Lawson-Hanson active-set method. A zeroed atom enters the passive (free) set when
183 it is most positively correlated with the residual; a passive atom leaves when its trial
184 coefficient is negative. Released atoms may later re-enter.
186 The active problem is solved in float64 with an explicit pseudoinverse rank threshold.
187 Coefficients are checked against the KKT conditions before and after they are cast back to
188 the target's dtype and device. Feasibility corrections use the ``3 * num_active`` budget;
189 a separate admission limit guards against numerical cycling. An exhausted limit, invalid
190 line-search step, stalled correction, or failed KKT check raises :class:`RuntimeError`.
191 Returns coefficients aligned with the columns of ``active_atoms``.
192 """
193 if active_atoms.ndim != 2 or target.ndim != 1: 193 ↛ 194line 193 didn't jump to line 194 because the condition on line 193 was never true
194 raise ValueError("active_atoms must be 2-D and target must be 1-D")
195 if active_atoms.shape[0] != target.shape[0]: 195 ↛ 196line 195 didn't jump to line 196 because the condition on line 195 was never true
196 raise ValueError("active_atoms and target have incompatible dimensions")
197 num_active = active_atoms.shape[1]
198 if num_active == 0:
199 return target.new_zeros(0)
201 result_dtype = target.dtype
202 result_device = target.device
203 work_device = torch.device("cpu") if target.device.type == "mps" else target.device
204 work_atoms = active_atoms.to(device=work_device, dtype=torch.float64)
205 work_target = target.to(device=work_device, dtype=torch.float64)
206 if not bool(torch.isfinite(work_atoms).all()) or not bool(torch.isfinite(work_target).all()):
207 raise ValueError("active_atoms and target must contain only finite values")
209 coefficients = work_target.new_zeros(num_active)
210 passive = work_target.new_zeros(num_active, dtype=torch.bool) # free coordinates (> 0)
211 max_corrections = 3 * num_active
212 # A correction can release several variables, each of which may later be re-admitted.
213 # This is only a cycle safeguard; the mathematical convergence budget is the correction count.
214 max_admissions = num_active * (max_corrections + 1)
215 corrections = 0
216 admissions = 0
217 # Main loop: free the zeroed atom most correlated with the residual until no zeroed atom
218 # has a materially positive correlation (KKT dual feasibility), or every atom is free.
219 while not bool(passive.all()):
220 dual_tol, coefficient_tol = _nnls_tolerances(work_atoms, work_target, coefficients)
221 dual = work_atoms.T @ (work_target - work_atoms @ coefficients)
222 dual = dual.masked_fill(passive, float("-inf"))
223 if bool((dual <= dual_tol).all()):
224 break
225 admissions += 1
226 if admissions > max_admissions:
227 raise RuntimeError("NNLS exceeded the active-set admission safeguard")
228 passive[int(torch.argmax(dual))] = True
229 # Feasibility loop: solve the unconstrained fit on the passive set; while it has a
230 # materially negative coordinate, step toward it until a passive coefficient reaches
231 # zero, then release the vanished coordinates. Each correction drops >= 1 passive atom
232 # (enforced by the progress guard), so the loop is bounded.
233 while True:
234 columns = passive.nonzero(as_tuple=True)[0]
235 before = int(columns.numel())
236 solution = torch.linalg.pinv(work_atoms[:, columns], rtol=_NNLS_RANK_RTOL) @ work_target
237 candidate = work_target.new_zeros(num_active)
238 candidate[columns] = solution
239 negative = passive & (candidate < -coefficient_tol)
240 if not bool(negative.any()):
241 passive = candidate > coefficient_tol
242 coefficients = torch.where(passive, candidate, torch.zeros_like(candidate))
243 break
244 corrections += 1
245 if corrections > max_corrections: 245 ↛ 246line 245 didn't jump to line 246 because the condition on line 245 was never true
246 raise RuntimeError("NNLS exceeded the Lawson-Hanson feasibility budget")
247 # Step toward the unconstrained fit, limited by the first coordinate to reach zero.
248 # Only materially-negative coordinates with a strictly positive denominator define
249 # the step; a zero-current/zero-candidate coordinate (e.g. a freshly admitted atom
250 # heading straight to <= 0) has a zero ratio, forcing a zero-length step -- it makes
251 # no move but is released by the cleanup below, so the passive set still shrinks.
252 denominator = coefficients - candidate
253 ratios = torch.where(
254 negative & (denominator > 0),
255 coefficients / denominator,
256 work_target.new_full((num_active,), float("inf")),
257 )
258 step = float(ratios.min())
259 if not math.isfinite(step) or not ( 259 ↛ 262line 259 didn't jump to line 262 because the condition on line 259 was never true
260 -_NNLS_RELATIVE_TOLERANCE <= step <= 1.0 + _NNLS_RELATIVE_TOLERANCE
261 ):
262 raise RuntimeError(f"NNLS line-search step {step} outside [0, 1]")
263 step = min(max(step, 0.0), 1.0)
264 updated = coefficients + step * (candidate - coefficients)
265 passive = updated > coefficient_tol
266 coefficients = torch.where(passive, updated, torch.zeros_like(updated))
267 if int(passive.sum()) >= before: 267 ↛ 268line 267 didn't jump to line 268 because the condition on line 267 was never true
268 raise RuntimeError("NNLS feasibility correction made no progress")
270 dual_tol, coefficient_tol = _nnls_tolerances(work_atoms, work_target, coefficients)
271 _validate_nnls_kkt(work_atoms, work_target, coefficients, dual_tol, coefficient_tol)
272 result = coefficients.to(device=result_device, dtype=result_dtype)
273 result_atoms = active_atoms.to(device=result_device, dtype=result_dtype)
274 result_target = target.to(device=result_device, dtype=result_dtype)
275 result_tiny = torch.finfo(result_dtype).tiny
276 result_relative_tolerance = math.sqrt(torch.finfo(result_dtype).eps)
277 result_atom_norms = result_atoms.norm(dim=0).clamp_min(result_tiny)
278 result_scale = (result_target.norm() + (result_atoms.abs() @ result.abs()).norm()).clamp_min(
279 result_tiny
280 )
281 result_dual_tol = result_relative_tolerance * result_atom_norms * result_scale
282 result_coefficient_tol = result_relative_tolerance * result_scale / result_atom_norms
283 _validate_nnls_kkt(result_atoms, result_target, result, result_dual_tol, result_coefficient_tol)
284 return result
287def _gradient_pursuit_step(
288 active_atoms: torch.Tensor, target: torch.Tensor, coefficients: torch.Tensor
289) -> torch.Tensor:
290 """One optimal-step-size projected-gradient update (Blumensath & Davies, 2008).
292 ``active_atoms`` is ``[d_model, num_active]`` (the selected atoms as columns) and
293 ``coefficients`` the current ``[num_active]`` coefficients (the newly added atom starts at
294 zero, so the incoming point is feasible and its residual is the previous residual). Moves
295 the coefficients along the steepest-descent direction ``active_atoms^T residual`` with the
296 exact line-search step, then projects onto the nonnegative orthant. The exact line search
297 happens before that projection, so the projected step can in principle increase the
298 residual. In that case the step is halved until the projected update no longer increases
299 the residual, with the feasible incoming point retained if the bounded search finds none.
300 """
301 residual = target - active_atoms @ coefficients
302 feasible = coefficients.clamp_min(0.0)
303 direction = active_atoms.T @ residual # [num_active]; gradient up to sign
304 projected = active_atoms @ direction # [d_model]
305 denominator = float((projected @ projected).detach())
306 if denominator <= 0.0: 306 ↛ 307line 306 didn't jump to line 307 because the condition on line 306 was never true
307 return feasible
308 step = float((projected @ residual).detach()) / denominator
309 current_residual_squared = float((residual @ residual).detach())
310 for _ in range(_GRADIENT_BACKTRACK_STEPS + 1): 310 ↛ 317line 310 didn't jump to line 317 because the loop on line 310 didn't complete
311 candidate = (coefficients + step * direction).clamp_min(0.0)
312 candidate_residual = target - active_atoms @ candidate
313 candidate_residual_squared = float((candidate_residual @ candidate_residual).detach())
314 if candidate_residual_squared <= current_residual_squared:
315 return candidate
316 step *= 0.5
317 return feasible
320def get_sparse_decomposition(
321 x: torch.Tensor,
322 dictionary: torch.Tensor,
323 k: int = DEFAULT_K,
324 *,
325 algorithm: str = "nonnegative_orthogonal_matching_pursuit",
326) -> JSpaceDecomposition:
327 """Greedily decompose ``x`` into a ``k``-sparse nonnegative combination of atoms.
329 Args:
330 x: Target vector, shape ``[d_model]``.
331 dictionary: Atom matrix, shape ``[num_atoms, d_model]`` (rows are atoms).
332 k: Upper bound on the number of atoms to select. Selection stops early once no
333 unselected atom is materially positively correlated with the residual, so fewer
334 than ``k`` atoms may be selected (and fewer still may be numerically active).
335 algorithm: Coefficient-update rule.
336 ``"nonnegative_orthogonal_matching_pursuit"`` (default) re-solves the selected-set
337 coefficients exactly as a nonnegative least-squares fit;
338 ``"gradient_pursuit"`` takes a single projected-gradient step per atom. See the
339 module docstring for the trade-off (they use the same selection rule, while the
340 exact re-solve is optimal on each selected set).
342 Returns:
343 A :class:`JSpaceDecomposition`. Its ``support`` holds only the numerically active
344 atoms and ``selected_support`` every selected atom, with
345 ``support.numel() <= selected_support.numel() <= k``.
347 Raises:
348 ValueError: On an unknown ``algorithm``, complex or non-finite inputs, a non-2-D
349 dictionary, a target whose length does not match ``d_model``, ``k`` outside
350 ``[1, num_atoms]``, or a dictionary with non-finite or zero-norm atoms.
351 RuntimeError: If ``algorithm="nonnegative_orthogonal_matching_pursuit"`` and the
352 nonnegative least-squares solve cannot be certified against its KKT conditions
353 within its numerical tolerance.
354 """
355 if algorithm not in ("nonnegative_orthogonal_matching_pursuit", "gradient_pursuit"):
356 raise ValueError(
357 "algorithm must be 'nonnegative_orthogonal_matching_pursuit' or "
358 f"'gradient_pursuit', got {algorithm!r}"
359 )
360 if dictionary.ndim != 2:
361 raise ValueError(
362 f"dictionary must be 2-D [num_atoms, d_model], got shape {tuple(dictionary.shape)}"
363 )
364 num_atoms, d_model = dictionary.shape
365 if x.ndim != 1 or x.shape[0] != d_model:
366 raise ValueError(f"x must be 1-D of length d_model={d_model}, got shape {tuple(x.shape)}")
367 if not 1 <= k <= num_atoms:
368 raise ValueError(f"k must be between 1 and num_atoms={num_atoms}, got {k}")
369 if torch.is_complex(x) or torch.is_complex(dictionary):
370 raise ValueError("x and dictionary must be real-valued")
372 target = x.float()
373 atoms = dictionary.float()
374 if not bool(torch.isfinite(target).all()):
375 raise ValueError("x contains non-finite entries")
376 if not bool(torch.isfinite(atoms).all()):
377 raise ValueError("dictionary contains non-finite entries")
378 atom_norms = torch.linalg.vector_norm(atoms, dim=1)
379 if not bool(torch.isfinite(atom_norms).all()) or bool((atom_norms == 0).any()):
380 raise ValueError("dictionary contains a non-finite or zero-norm atom")
382 x_norm = float(torch.linalg.vector_norm(target).detach())
383 correlation_tol = _CORRELATION_RELATIVE_TOLERANCE * x_norm
385 residual = target.clone()
386 selected: List[int] = []
387 coordinates = target.new_zeros(0)
388 for _ in range(k):
389 # Select the unselected atom most correlated with the current residual, using
390 # unit-norm atoms so high-norm atoms are not preferred for their scale alone.
391 correlation = (atoms @ residual) / atom_norms
392 for chosen in selected:
393 correlation[chosen] = float("-inf")
394 candidate = int(torch.argmax(correlation).item())
395 # Stop early once no unselected atom is materially positively correlated: under the
396 # nonnegativity constraint a non-positive correlation cannot reduce the residual, so
397 # ``k`` is an upper bound on the number of selected atoms, not a target.
398 if float(correlation[candidate].detach()) <= correlation_tol:
399 break
400 selected.append(candidate)
402 selected_atoms = atoms[selected].T # [d_model, len(selected)]
403 if algorithm == "nonnegative_orthogonal_matching_pursuit":
404 # Re-solve the coefficients jointly over the selected set as a nonnegative
405 # least-squares fit. Sequential per-atom updates are wrong once selected atoms
406 # are correlated.
407 coordinates = _nonnegative_least_squares(selected_atoms, target)
408 else:
409 # Carry the coefficients forward, initialising the new atom at zero, and take a
410 # single projected-gradient step over the selected set.
411 coordinates = _gradient_pursuit_step(
412 selected_atoms, target, torch.cat([coordinates, coordinates.new_zeros(1)])
413 )
414 residual = target - selected_atoms @ coordinates
416 # ``selected_support`` stays on CPU for token decoding; the vector-valued outputs stay on
417 # the computation device.
418 selected_support = torch.tensor(selected, dtype=torch.long)
419 selected_atoms = atoms[selected_support].T # [d_model, num_selected] on atoms.device
421 # Public active support: selected atoms whose contribution ``c_i * ||v_i||`` is a
422 # materially nonzero fraction of ``||x||``. For the NNLS re-solve this is exactly the
423 # strictly-positive coordinate set (its coefficient-zeroing uses the same scale); for
424 # gradient pursuit it prunes atoms left at (near-)zero by the final projected step.
425 contribution = coordinates * torch.linalg.vector_norm(selected_atoms, dim=0)
426 active = contribution > _ACTIVE_RELATIVE_TOLERANCE * x_norm # on the computation device
427 support = selected_support[active.cpu()] # CPU, aligned with the token-decoding tensors
428 active_coordinates = coordinates[active]
429 active_atoms = selected_atoms[:, active] # [d_model, num_active] on atoms.device
431 # Reconstruction is the nonnegative combination over the active support (empty -> zeros).
432 reconstruction = active_atoms @ active_coordinates
433 # The J-space component is the orthogonal projection of the target onto the span of every
434 # selected atom (paper appendix), computed with the same pseudoinverse construction as
435 # swap_hooks. That span can be larger than the active support when a selected coordinate is
436 # zero, so the projection differs from the nonnegative reconstruction in that case.
437 if selected:
438 j_space_component = selected_atoms @ (torch.linalg.pinv(selected_atoms) @ target)
439 else:
440 j_space_component = target.new_zeros(d_model)
441 non_j_space_component = target - j_space_component
442 return JSpaceDecomposition(
443 support=support,
444 coordinates=active_coordinates,
445 selected_support=selected_support,
446 reconstruction=reconstruction,
447 j_space_component=j_space_component,
448 non_j_space_component=non_j_space_component,
449 )
452@dataclass
453class JSpaceOccupancy:
454 """Result of a J-space occupancy estimate.
456 Attributes:
457 occupancy: Estimated number of meaningfully-active atoms -- the step of maximum
458 separation between the real and random-control cumulative captured variance.
459 marginal_captured_variance: Per-step captured-variance gain of the real greedy selection,
460 shape ``[max_atoms]``.
461 control_captured_variance: Per-step captured-variance gain averaged over the random
462 control dictionaries, shape ``[max_atoms]``.
463 support: Greedily selected atom indices, shape ``[max_atoms]`` (token ids when the
464 dictionary is the vocabulary of J-lens vectors).
465 """
467 occupancy: int
468 marginal_captured_variance: torch.Tensor
469 control_captured_variance: torch.Tensor
470 support: torch.Tensor
473def _greedy_captured_variance_gains(
474 atoms: torch.Tensor, atom_norms: torch.Tensor, target: torch.Tensor, max_atoms: int
475) -> Tuple[torch.Tensor, torch.Tensor]:
476 """Greedily select exactly ``max_atoms`` atoms and return captured-variance gains.
478 At each step, add the unused atom with the greatest signed, norm-normalized correlation with
479 the current residual. Project ``target`` orthogonally onto the full selected span using a
480 pseudoinverse, then set the next residual to ``target - projection``. The captured variance is
481 ``||Pi_S target||^2 / ||target||^2``; the returned values are its per-step increments.
483 This shares the per-step correlation rule with :func:`get_sparse_decomposition`, but not its
484 residual recurrence: sparse decomposition uses a nonnegative coefficient-fit residual and may
485 stop early, while this recurrence does not stop early, so the selected supports can differ.
486 """
487 total_variance = float(target @ target)
488 support: List[int] = []
489 residual = target.clone()
490 captured_variance_gains: List[float] = []
491 previous_captured_variance = 0.0
492 for _ in range(max_atoms):
493 correlation = (atoms @ residual) / atom_norms
494 for chosen in support:
495 correlation[chosen] = float("-inf")
496 support.append(int(torch.argmax(correlation).item()))
497 active_atoms = atoms[support].T
498 projection = active_atoms @ (torch.linalg.pinv(active_atoms) @ target)
499 captured_variance = float((projection @ projection) / total_variance)
500 captured_variance_gains.append(captured_variance - previous_captured_variance)
501 previous_captured_variance = captured_variance
502 residual = target - projection
503 return torch.tensor(captured_variance_gains), torch.tensor(support, dtype=torch.long)
506def estimate_occupancy(
507 x: torch.Tensor,
508 dictionary: torch.Tensor,
509 *,
510 max_atoms: int = DEFAULT_K,
511 num_control_dictionaries: int = 32,
512 seed: int = 0,
513) -> JSpaceOccupancy:
514 """Estimate how many dictionary atoms are meaningfully active in ``x``.
516 Runs the projection-residual recurrence described in
517 :func:`_greedy_captured_variance_gains` for exactly ``max_atoms`` steps and compares the real
518 per-step captured-variance curve against the same recurrence on ``num_control_dictionaries``
519 random unit-norm dictionaries of the same size. This shares sparse decomposition's per-step
520 correlation rule, but uses an unconstrained span-projection residual rather than a nonnegative
521 coefficient-fit residual, so their supports need not match. The occupancy is the step of
522 maximum separation between the real and (averaged) control *cumulative* captured variance --
523 the point past which further atoms add no more than random directions would. Deterministic
524 given ``seed`` and needs no threshold. (Captured variance is a projection, hence scale-free,
525 so the random control atoms are simply unit-norm.)
527 Args:
528 x: Target vector, shape ``[d_model]``.
529 dictionary: Atom matrix, shape ``[num_atoms, d_model]`` (rows are atoms).
530 max_atoms: Number of atoms to select in the real and control recurrences.
531 num_control_dictionaries: Number of random control dictionaries to average over.
532 seed: Seed for the random control dictionaries (reproducibility).
534 Returns:
535 An :class:`JSpaceOccupancy`.
537 Raises:
538 ValueError: On complex inputs, a non-2-D dictionary, a target whose length does not match
539 ``d_model``, ``max_atoms`` outside ``[1, num_atoms]``,
540 ``num_control_dictionaries < 1``, a target with non-finite entries or a non-finite or
541 zero norm, or a dictionary with non-finite or zero-norm atoms.
542 """
543 if dictionary.ndim != 2: 543 ↛ 544line 543 didn't jump to line 544 because the condition on line 543 was never true
544 raise ValueError(
545 f"dictionary must be 2-D [num_atoms, d_model], got shape {tuple(dictionary.shape)}"
546 )
547 num_atoms, d_model = dictionary.shape
548 if x.ndim != 1 or x.shape[0] != d_model:
549 raise ValueError(f"x must be 1-D of length d_model={d_model}, got shape {tuple(x.shape)}")
550 if not 1 <= max_atoms <= num_atoms:
551 raise ValueError(f"max_atoms must be between 1 and num_atoms={num_atoms}, got {max_atoms}")
552 if num_control_dictionaries < 1:
553 raise ValueError(
554 f"num_control_dictionaries must be at least 1, got {num_control_dictionaries}"
555 )
556 if torch.is_complex(x) or torch.is_complex(dictionary):
557 raise ValueError("x and dictionary must be real-valued")
559 target = x.float()
560 atoms = dictionary.float()
561 if not bool(torch.isfinite(target).all()):
562 raise ValueError("x contains non-finite entries")
563 target_squared_norm = target @ target
564 if not bool(torch.isfinite(target_squared_norm)):
565 raise ValueError("x must have finite norm")
566 if float(target_squared_norm) <= 0.0:
567 raise ValueError("x must have non-zero norm")
568 if not bool(torch.isfinite(atoms).all()):
569 raise ValueError("dictionary contains non-finite entries")
570 atom_norms = (atoms * atoms).sum(dim=1).sqrt()
571 if not bool(torch.isfinite(atom_norms).all()) or bool((atom_norms == 0).any()):
572 raise ValueError("dictionary contains a non-finite or zero-norm atom")
574 real_captured_variance, support = _greedy_captured_variance_gains(
575 atoms, atom_norms, target, max_atoms
576 )
578 generator = torch.Generator(device=atoms.device).manual_seed(seed)
579 control_atom_norms = torch.ones(num_atoms, device=atoms.device)
580 control_variance_runs: List[torch.Tensor] = []
581 for _ in range(num_control_dictionaries):
582 random_atoms = torch.randn(
583 num_atoms, d_model, generator=generator, device=atoms.device, dtype=atoms.dtype
584 )
585 random_atoms = random_atoms / (random_atoms * random_atoms).sum(dim=1, keepdim=True).sqrt()
586 control_run_variance, _ = _greedy_captured_variance_gains(
587 random_atoms, control_atom_norms, target, max_atoms
588 )
589 control_variance_runs.append(control_run_variance)
590 control_captured_variance = torch.stack(control_variance_runs).mean(dim=0)
592 separation = real_captured_variance.cumsum(dim=0) - control_captured_variance.cumsum(dim=0)
593 occupancy = int(torch.argmax(separation).item()) + 1
594 return JSpaceOccupancy(
595 occupancy=occupancy,
596 marginal_captured_variance=real_captured_variance,
597 control_captured_variance=control_captured_variance,
598 support=support,
599 )
602@dataclass
603class JSpaceVarianceProfile:
604 """Per-layer J-space variance profile over a prompt corpus.
606 Produced by :meth:`JacobianLens.fraction_of_variance`.
608 Attributes:
609 layers: The source layers profiled, in order.
610 median: Per-layer median over positions of the J-space variance fraction
611 ``||j_space_component||^2 / ||activation||^2``.
612 pooled: Per-layer pooled ratio ``sum(||j_space_component||^2) / sum(||activation||^2)``
613 across the corpus (the paper's "fraction of total variance").
614 per_position: Per-layer 1-D tensor of the raw per-position variance fractions.
615 """
617 layers: List[int]
618 median: Dict[int, float]
619 pooled: Dict[int, float]
620 per_position: Dict[int, torch.Tensor]