feat(jacobian_lens): J-space sparse decomposition - #1596
Open
janmenjayap wants to merge 5 commits into
Open
Conversation
- Add `get_sparse_decomposition` to decompose an activation into a k-sparse nonnegative combination of J-lens vectors (Gurnee et al., 2026). - Support `nonnegative_orthogonal_matching_pursuit` (default, exact NNLS re-solve) and `gradient_pursuit` algorithms. - Return both the nonnegative coordinates and the orthogonal-projection J-space component. - Distinguish the projection from the coefficient reconstruction; the projection residual matches `swap_hooks`. - Keep the implementation model-free by operating directly on the raw dictionary tensor. - Add tests covering both algorithms, exact-resolve NNLS correctness, a brute-force optimum oracle, and input validation. Part of TransformerLensOrg#1539 (Tier 2).
- Add `JacobianLens.lens_vector_dictionary(model, layer)` returning the `[d_vocab, d_model]` dictionary whose rows are the J-lens vectors `v_t = J[layer]^T W_U[:, t]`. - Cache the dictionary per (layer, device) and release it in `clear_device_cache`, so a sparse decomposition can reuse it; document its vocabulary-sized memory cost. - Add tests asserting the dictionary matches `lens_vectors` over every token, is cached and invalidated by `clear_device_cache`, and rejects an unfitted layer.
- Add `JacobianLens.decompose(model, activation_or_prompt, layer, *, position, k, algorithm)` decomposing either a raw activation vector or the `blocks.{layer}.hook_out` activation at a prompt position, validating inputs before building the dictionary.
- Build and cache the layer dictionary via `lens_vector_dictionary` and solve with `get_sparse_decomposition`.
- Export `JSpaceDecomposition` and `get_sparse_decomposition` from `transformer_lens.tools.analysis`.
- Add end-to-end tests for the raw-activation and prompt paths, the algorithm passthrough, and the input-validation error paths.
- Add a GPT-2 integration test (regular CI): `decompose` on a real `blocks.6.hook_out` activation returns k nonnegative atoms, the non-J-space residual is orthogonal to every selected J-lens vector, and the J-space component plus residual recover the activation. - Add a slow gemma-2-2b-it integration test validating `decompose` on the published lens artifact: support size, nonnegative coordinates, in-vocabulary token ids, and component-plus-residual reconstruction. - Document J-space sparse decomposition in `jacobian_lens_fitting.md`: the `decompose` API, local coordinates versus the orthogonal-projection J-space component, and the paper's variance facts with closed-model caveats.
Add a References section to the decomposition module docstring: Gurnee et al. (2026) for the J-space method, Pati et al. (1993) for the greedy orthogonal-matching-pursuit selection, Blumensath & Davies (2008) for the gradient-pursuit update, and Lawson & Hanson (1974) for the active-set nonnegative least-squares re-solve.
janmenjayap
changed the title
10 tasks
jlarson4
reviewed
Aug 3, 2026
jlarson4
left a comment
jlarson4
left a comment
Collaborator
There was a problem hiding this comment.
Hey @janmenjayap! Great work on this, it is an excellent piece of functionality. Just a couple adjustments I'd like you to make before I can merge. Please let me know if you have any questions!
| """Minimize ``||active_atoms @ coefficients - target||`` over ``coefficients >= 0``. | ||
|
|
||
| A small active-set solver: solve the unconstrained least squares, and while any | ||
| coefficient is negative, drop the most-negative atom and re-solve. Exact for the tiny |
Collaborator
There was a problem hiding this comment.
The active-set solver only drops atoms and never re-admits one, so it is not exact NNLS: 4 of 288 real GPT-2 decompositions returned coordinates violating the KKT dual condition, with relative violations up to 0.55. Can you complete the Lawson–Hanson re-entry step?
| residual = target.clone() | ||
| support: List[int] = [] | ||
| coordinates = target.new_zeros(0) | ||
| for _ in range(k): |
Collaborator
There was a problem hiding this comment.
Once a coefficient is clamped to zero the residual is unchanged, so the next pick uses a stale residual. On GPT-2 layer 10 only ~9 of 25 coordinates end up nonzero. Can you reconcile k with the docs, whose example prints all k support tokens as active concepts?
| # --------------------------------------------------------------------------- # | ||
| # Brute-force oracle (no external reference implementation exists) | ||
| # --------------------------------------------------------------------------- # | ||
| def _brute_force_best_support(x, dictionary, k): |
Collaborator
There was a problem hiding this comment.
The brute-force oracle calls _nonnegative_least_squares as its own inner solver, so it validates the greedy selection, but assumes the coefficient solver is correct. A wrong solver passes this test unchanged. Can the oracle use an independent reference, such as a KKT optimality check on the returned coordinates?
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Adds J-space sparse decomposition to the Jacobian lens: it writes an activation (or a steering / sparse-autoencoder direction) as a
k-sparse nonnegative combination of J-lens vectorsv_t = J_ℓ^T W_U[:, t], following Gurnee et al. (2026), "Verbalizable Representations Form a Global Workspace in Language Models" (Transformer Circuits Thread). This is the Tier-2 decomposition item tracked in #1539.New public surface (TransformerBridge only, matching the rest of the Jacobian lens)
get_sparse_decomposition(x, dictionary, k=25, *, algorithm=...)— a model-free greedy solver returning aJSpaceDecomposition(support, nonnegativecoordinates,reconstruction,j_space_component,non_j_space_component).JacobianLens.lens_vector_dictionary(model, layer)— the cached full-vocabulary J-lens dictionary for a layer.JacobianLens.decompose(model, activation_or_prompt, layer, *, position=None, k=25, algorithm=...)— decomposes either a raw[d_model]activation or theblocks.{layer}.hook_outresidual at a prompt position.JSpaceDecompositionandget_sparse_decompositionare exported fromtransformer_lens.tools.analysis.Two coefficient-update rules (
algorithm=)nonnegative_orthogonal_matching_pursuit(default) — an exact nonnegative least-squares re-solve on the active set (Lawson & Hanson, 1974); optimal on the selected support at the smallkused here.gradient_pursuit— the directional update of Blumensath & Davies (2008), for faithfulness to the paper and the large-active-set regime.Both share the same greedy orthogonal-matching-pursuit selection (Pati et al., 1993).
At vocabulary scale the per-step cost is dominated by the correlation over all atoms,
so the exact re-solve is effectively free while returning optimal coefficients on the
support — hence it is the default. The choice is disclosed here because the paper uses
gradient pursuit; both are provided.
Two outputs that need not coincide (paper appendix)
coordinates— the nonnegative pursuit coefficients (the "local J-space coordinates").j_space_component— the orthogonal projection of the activation onto the span of the selected vectors; its residual matchesswap_hooks.They diverge whenever a coordinate is clamped to zero by the nonnegativity constraint; both are returned.
Docs
A new "Sparse decomposition (J-space coordinates)" section in
jacobian_lens_fitting.md, including the paper's variance findings explicitly caveated as measured on closed Anthropic models (on open-weight models the shape may hold but the exact values will not necessarily transfer). A runnable demo notebook will follow on a separatedocs-named branch, per the Tier-2 scope.Part of #1539 (Tier 2).
Type of change
Checklist
Verification run locally
Full
make test-prequivalent (direct binaries, single Python) plus the static gates:check-format(pycln / isort / black), repo-wide — clean.mypy .— Success, no issues in 388 source files.unit(pytest tests/unit -m "not slow") — 4513 passed, 29 skipped, 10 xfailed.docstring(pytest transformer_lens/) — 18 passed.acceptance(pytest tests/acceptance -m "not slow") — passed.integration(pytest tests/integration -m "not slow") — 1145 passed. The only failure istests/integration/model_bridge/test_jamba_adapter.py::TestJambaGeneration::test_greedy_matches_hf, a pre-existing device-placement issue (index on cuda:0vs tensors oncpu) in the unrelated Jamba adapter — untouched by this PR, which changes onlytransformer_lens/tools/analysis/jacobian_lens*.build-docs(Sphinx) — build succeeded; the new docs page adds no warnings.decomposeand the slowgemma-2-2b-itdecomposecase both passed.CI re-runs the full suite across Python 3.10 / 3.11 / 3.12 (
compatibility-checks).cc: @jlarson4