transformer_lens.tools.analysis.jacobian_lens_decomposition module¶

J-space sparse decomposition for the Jacobian lens.

Decomposes an activation (or a steering / sparse-autoencoder direction) into a sparse nonnegative combination of J-lens vectors, following Gurnee et al. (2026), “Verbalizable Representations Form a Global Workspace in Language Models” (Transformer Circuits Thread).

The decomposition is greedy: at each of at most k steps the atom most correlated with the current residual (using unit-normalised atoms, so high-norm atoms are not preferred for their scale alone) is added to the selected set, and the selected-set coefficients are updated under a nonnegativity constraint. k is an upper bound, not a target: selection stops early once no unselected atom has a materially positive residual correlation (under nonnegativity a negatively-correlated atom cannot reduce the residual), so fewer than k atoms may be selected. Two coefficient-update rules are provided via algorithm:

"nonnegative_orthogonal_matching_pursuit" (default) – the active-set coefficients are
re-solved with a float64 active-set nonnegative least-squares fit at each step, then checked against the KKT conditions.
"gradient_pursuit" – the directional update of Blumensath & Davies (2008): a single optimal-step-size gradient step on the active coefficients, projected onto the nonnegative orthant.

Both use the same greedy selection rule and both are nonnegative; they differ in the coefficient update and therefore can choose different atoms at later steps. At vocabulary scale, correlation over all atoms costs O(num_atoms * d_model) per step, while the exact re-solve adds float64 linear algebra on the small selected set. The NNLS update is the default because it optimizes all selected coefficients jointly. "gradient_pursuit" skips that solve and matches the update used in the paper.

Two supports are exposed because the paper uses two inconsistent operationalizations. The support is the numerically active set – the selected atoms whose nonnegative coordinate materially contributes – and coordinates are aligned with it; this is the paper’s main-text sparse nonnegative reconstruction. The selected_support is every greedily selected atom, including any whose coordinate was driven to zero by the nonnegativity constraint; it defines the span for the paper’s appendix projection. Hence len(support) <= len(selected_support) <= k.

Two vector outputs correspondingly need not coincide: the reconstruction is the nonnegative combination over the active support, while the J-space component is the orthogonal projection of the target onto the span of selected_support (the paper’s appendix definition, and the residual its interventions use). For the exact NNLS re-solve the reconstruction equals the projection onto the active support (KKT stationarity), so the two vectors differ exactly when a selected atom has a zero coordinate.

This module is model-free: it operates on a raw dictionary tensor, so it can be used and tested without loading a model.

References

Gurnee et al. (2026), “Verbalizable Representations Form a Global Workspace in Language Models,” Transformer Circuits Thread – the J-space decomposition.
Pati, Rezaiifar & Krishnaprasad (1993), “Orthogonal Matching Pursuit: Recursive Function Approximation with Applications to Wavelet Decomposition,” 27th Asilomar Conference on Signals, Systems and Computers – the greedy atom selection shared by both algorithms.
Blumensath & Davies (2008), “Gradient Pursuits,” IEEE Transactions on Signal Processing 56(6):2370-2382 – the gradient_pursuit coefficient update.
Lawson & Hanson (1974), “Solving Least Squares Problems,” Prentice-Hall – the active-set nonnegative least-squares re-solve used by the default algorithm.

transformer_lens.tools.analysis.jacobian_lens_decomposition.DEFAULT_K = 25¶: The paper varies the sparsity level but “typically” chooses “no more than 25”.

class transformer_lens.tools.analysis.jacobian_lens_decomposition.JSpaceDecomposition(support: Tensor, coordinates: Tensor, selected_support: Tensor, reconstruction: Tensor, j_space_component: Tensor, non_j_space_component: Tensor)¶

Bases: object

Result of a sparse J-space decomposition.

support¶

Indices of the numerically active selected atoms – those whose nonnegative coordinate materially contributes (token ids when the dictionary is the vocabulary of J-lens vectors). A subset of selected_support.

Type:: torch.Tensor

coordinates¶

Nonnegative pursuit coefficients aligned with support (the “local J-space coordinates”); every entry is materially nonzero.

Type:: torch.Tensor

selected_support¶

Indices of every greedily selected atom, including any whose coordinate was driven to zero by the nonnegativity constraint. Defines the span for j_space_component. Satisfies support.numel() <= selected_support.numel() <= k.

Type:: torch.Tensor

reconstruction¶

The nonnegative combination sum(coordinates * active atoms) over support.

Type:: torch.Tensor

j_space_component¶

The orthogonal projection of the target onto the span of selected_support (the paper’s “J-space component”). For the exact NNLS re-solve it equals reconstruction unless a selected atom has a zero coordinate, in which case the projection uses a larger span.

Type:: torch.Tensor

non_j_space_component¶

The residual target - j_space_component (the “non-J-space component”), orthogonal to the selected span.

Type:: torch.Tensor

coordinates: Tensor¶

j_space_component: Tensor¶

non_j_space_component: Tensor¶

reconstruction: Tensor¶

selected_support: Tensor¶

support: Tensor¶

class transformer_lens.tools.analysis.jacobian_lens_decomposition.JSpaceOccupancy(occupancy: int, marginal_captured_variance: Tensor, control_captured_variance: Tensor, support: Tensor)¶

Bases: object

Result of a J-space occupancy estimate.

occupancy¶

Estimated number of meaningfully-active atoms – the step of maximum separation between the real and random-control cumulative captured variance.

Type:: int

marginal_captured_variance¶

Per-step captured-variance gain of the real greedy selection, shape [max_atoms].

Type:: torch.Tensor

control_captured_variance¶

Per-step captured-variance gain averaged over the random control dictionaries, shape [max_atoms].

Type:: torch.Tensor

support¶

Greedily selected atom indices, shape [max_atoms] (token ids when the dictionary is the vocabulary of J-lens vectors).

Type:: torch.Tensor

control_captured_variance: Tensor¶

marginal_captured_variance: Tensor¶

occupancy: int¶

support: Tensor¶

class transformer_lens.tools.analysis.jacobian_lens_decomposition.JSpaceVarianceProfile(layers: List[int], median: Dict[int, float], pooled: Dict[int, float], per_position: Dict[int, Tensor])¶

Bases: object

Per-layer J-space variance profile over a prompt corpus.

Produced by JacobianLens.fraction_of_variance().

layers¶

The source layers profiled, in order.

Type:: List[int]

median¶

Per-layer median over positions of the J-space variance fraction ||j_space_component||^2 / ||activation||^2.

Type:: Dict[int, float]

pooled¶

Per-layer pooled ratio sum(||j_space_component||^2) / sum(||activation||^2) across the corpus (the paper’s “fraction of total variance”).

Type:: Dict[int, float]

per_position¶

Per-layer 1-D tensor of the raw per-position variance fractions.

Type:: Dict[int, torch.Tensor]

layers: List[int]¶

median: Dict[int, float]¶

per_position: Dict[int, Tensor]¶

pooled: Dict[int, float]¶

transformer_lens.tools.analysis.jacobian_lens_decomposition.estimate_occupancy(x: Tensor, dictionary: Tensor, *, max_atoms: int = 25, num_control_dictionaries: int = 32, seed: int = 0) → JSpaceOccupancy¶

Estimate how many dictionary atoms are meaningfully active in x.

Runs the projection-residual recurrence described in _greedy_captured_variance_gains() for exactly max_atoms steps and compares the real per-step captured-variance curve against the same recurrence on num_control_dictionaries random unit-norm dictionaries of the same size. This shares sparse decomposition’s per-step correlation rule, but uses an unconstrained span-projection residual rather than a nonnegative coefficient-fit residual, so their supports need not match. The occupancy is the step of maximum separation between the real and (averaged) control cumulative captured variance – the point past which further atoms add no more than random directions would. Deterministic given seed and needs no threshold. (Captured variance is a projection, hence scale-free, so the random control atoms are simply unit-norm.)

Parameters:

x – Target vector, shape [d_model].
dictionary – Atom matrix, shape [num_atoms, d_model] (rows are atoms).
max_atoms – Number of atoms to select in the real and control recurrences.
num_control_dictionaries – Number of random control dictionaries to average over.
seed – Seed for the random control dictionaries (reproducibility).

Returns:

An JSpaceOccupancy.

Raises:

ValueError – On complex inputs, a non-2-D dictionary, a target whose length does not match d_model, max_atoms outside [1, num_atoms], num_control_dictionaries < 1, a target with non-finite entries or a non-finite or zero norm, or a dictionary with non-finite or zero-norm atoms.

transformer_lens.tools.analysis.jacobian_lens_decomposition.get_sparse_decomposition(x: Tensor, dictionary: Tensor, k: int = 25, *, algorithm: str = 'nonnegative_orthogonal_matching_pursuit') → JSpaceDecomposition¶

Greedily decompose x into a k-sparse nonnegative combination of atoms.

Parameters:

x – Target vector, shape [d_model].
dictionary – Atom matrix, shape [num_atoms, d_model] (rows are atoms).
k – Upper bound on the number of atoms to select. Selection stops early once no unselected atom is materially positively correlated with the residual, so fewer than k atoms may be selected (and fewer still may be numerically active).
algorithm – Coefficient-update rule. "nonnegative_orthogonal_matching_pursuit" (default) re-solves the selected-set coefficients exactly as a nonnegative least-squares fit; "gradient_pursuit" takes a single projected-gradient step per atom. See the module docstring for the trade-off (they use the same selection rule, while the exact re-solve is optimal on each selected set).

Returns:

A JSpaceDecomposition. Its support holds only the numerically active atoms and selected_support every selected atom, with support.numel() <= selected_support.numel() <= k.

Raises:

ValueError – On an unknown algorithm, complex or non-finite inputs, a non-2-D dictionary, a target whose length does not match d_model, k outside [1, num_atoms], or a dictionary with non-finite or zero-norm atoms.
RuntimeError – If algorithm="nonnegative_orthogonal_matching_pursuit" and the nonnegative least-squares solve cannot be certified against its KKT conditions within its numerical tolerance.