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¶

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.