feat(jacobian_lens): J-space sparse decomposition

docs/source/content/jacobian_lens_fitting.md

@@ -232,3 +232,67 @@ lens = JacobianLens.from_pretrained(
 To propose a short-name entry in TransformerLens, open a pull request that adds the
 published file to `transformer_lens/tools/analysis/jacobian_lens_registry.json` and
 include the fitting provenance and validation results.
+
+## Sparse decomposition (J-space coordinates)
+
+A fitted lens also decomposes an activation into the concepts it is *disposed to say*.
+`JacobianLens.decompose` writes an activation `x` at layer ℓ as a `k`-sparse **nonnegative**
+combination of J-lens vectors `v_t = J_ℓ^T W_U[:, t]` (one direction per vocabulary token),
+selected greedily:
+
+```python
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.tools.analysis import JacobianLens
+
+model = TransformerBridge.boot_transformers("gpt2", device="cpu")
+lens = JacobianLens.from_pretrained("gpt2-small", model=model)
+
+# decompose the activation at a prompt position ...
+result = lens.decompose(model, "The Eiffel Tower is in the city of", layer=6, position=-1, k=8)
+# ... or a raw [d_model] activation you already have (leave position=None):
+#   result = lens.decompose(model, activation, layer=6, k=8)
+
+tokens = [model.to_string(int(t)) for t in result.support]  # the k selected J-lens vectors
+coordinates = result.coordinates                            # their nonnegative coefficients
+```
+
+The result carries two things that need not coincide (per the paper's appendix):
+
+- `coordinates` (the *local J-space coordinates*) -- the nonnegative pursuit coefficients on the
+  selected J-lens vectors, with `reconstruction = sum(coordinates * v_t)`.
+- `j_space_component` (the *J-space component*) -- the orthogonal projection of the activation
+  onto the span of the selected vectors -- and `non_j_space_component = x - j_space_component`,
+  the residual the interventions leave unchanged.
+
+```
+x  in R^d_model
+   |-- decompose(x, layer, k)
+        |-- coordinates            a_t >= 0      (local J-space coordinates)
+        |-- j_space_component      Pi_S x        (orthogonal projection onto span of selected v_t)
+        \-- non_j_space_component  x - Pi_S x     (orthogonal to the selected vectors)
+```
+
+Two algorithms are available via `algorithm=`: `"nonnegative_orthogonal_matching_pursuit"`
+(default; an exact nonnegative least-squares re-solve on the selected set, optimal at the small
+sparsity used here) and `"gradient_pursuit"` (the directional update of Blumensath & Davies
+(2008) that the paper uses, which only pays off for large active sets).
+
+### Interpreting the numbers honestly
+
+The quantitative findings below are from Gurnee et al. (2026) and were measured on **closed
+Anthropic models** (Sonnet / Haiku / Opus); on open-weight models the *shape* may hold but the
+exact values will not necessarily transfer.
+
+- The decomposition is **not** a top-k logit-lens readout: because the J-lens vectors are
+  overcomplete and non-orthogonal, it gives "a different (and typically less redundant) set of
+  active concepts than simply taking the top-k by inner product."
+- The J-space is a **small fraction** of the activation: its component "never [exceeds] more than
+  10%" of total activation variance, and for concept vectors carries "a median of only 6-7% ...
+  the remaining ~93% lying outside the J-space." Do not expect `j_space_component` to capture
+  most of `x`.
+- `k` defaults to 25 because the paper "typically choose[s] it to be no more than 25, which we
+  empirically observed to be the number of J-lens vectors that are meaningfully active at a given
+  time."
+
+The full-vocabulary dictionary is cached on the model's device and is vocabulary-sized
+(gigabytes for large models); release it with `lens.clear_device_cache()`.

tests/integration/test_jacobian_lens.py

@@ -426,3 +426,59 @@ def assert_valid_topk(
     )
     torch.testing.assert_close(result.model_logits[0], expected_final_logits, atol=1e-5, rtol=1e-5)
     torch.testing.assert_close(result.lens_logits[final_layer], result.model_logits)
+
+
+def test_decompose_gpt2_activation_reconstructs_and_is_orthogonal(published_gpt2_lens, gpt2_bridge):
+    """decompose on a real GPT-2 activation: k nonnegative atoms, the non-J-space residual is
+    orthogonal to the selected J-lens vectors, and component + residual recover the activation."""
+    layer, k = 6, 8
+    result = published_gpt2_lens.decompose(gpt2_bridge, PROMPT, layer=layer, position=-1, k=k)
+
+    assert result.support.numel() == k
+    assert (result.coordinates >= 0).all()
+    assert (result.support >= 0).all() and (result.support < gpt2_bridge.cfg.d_vocab).all()
+
+    # the decomposed activation is the model's blocks.{layer}.hook_out at the last position
+    tokens = gpt2_bridge.to_tokens(PROMPT)
+    hook = f"blocks.{layer}.hook_out"
+    _, cache = gpt2_bridge.run_with_cache(tokens, names_filter=lambda name: name == hook)
+    activation = cache[hook][0, -1, :].float()
+    assert torch.allclose(
+        result.j_space_component + result.non_j_space_component, activation, atol=1e-3
+    )
+
+    # the non-J-space residual is orthogonal to every selected J-lens vector (cosine ~ 0)
+    dictionary = published_gpt2_lens.lens_vector_dictionary(gpt2_bridge, layer)
+    residual = result.non_j_space_component
+    for atom_id in result.support.tolist():
+        atom = dictionary[atom_id]
+        cosine = torch.dot(residual, atom) / (residual.norm() * atom.norm())
+        assert cosine.abs().item() < 1e-3
+
+
+@pytest.mark.slow
+def test_decompose_gemma_activation_is_valid():
+    """Decompose a real gemma-2-2b-it activation via its published lens (slow: real download)."""
+    from transformer_lens.model_bridge import TransformerBridge
+    from transformer_lens.tools.analysis import JacobianLens
+
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    model = TransformerBridge.boot_transformers(GEMMA_MODEL, dtype=torch.bfloat16, device=device)
+    lens = JacobianLens.from_pretrained(
+        LENS_REPO, filename=GEMMA_LENS_FILE, revision=LENS_REVISION, model=model
+    )
+    layer = lens.source_layers[len(lens.source_layers) // 2]
+    k = 16
+    result = lens.decompose(model, PROMPT, layer=layer, position=-1, k=k)
+
+    assert result.support.numel() == k
+    assert (result.coordinates >= 0).all()
+    assert (result.support >= 0).all() and (result.support < model.cfg.d_vocab).all()
+
+    tokens = model.to_tokens(PROMPT)
+    hook = f"blocks.{layer}.hook_out"
+    _, cache = model.run_with_cache(tokens, names_filter=lambda name: name == hook)
+    activation = cache[hook][0, -1, :].float()
+    assert torch.allclose(
+        result.j_space_component + result.non_j_space_component, activation, atol=1e-2
+    )

tests/unit/tools/test_jacobian_lens.py

@@ -18,7 +18,11 @@
 from transformer_lens.model_bridge.supported_architectures.deepseek_v4 import (
     DeepseekV4BlockBridge,
 )
-from transformer_lens.tools.analysis import JacobianLens
+from transformer_lens.tools.analysis import (
+    JacobianLens,
+    JSpaceDecomposition,
+    get_sparse_decomposition,
+)
 from transformer_lens.utilities.activation_functions import apply_softcap
 
 D_MODEL = 6
@@ -1154,3 +1158,107 @@ def fake_retry(function: Any, **kwargs: Any) -> str:
 
         assert len(set(filenames)) == 1, "All three aliases should resolve to the same filename"
         assert filenames[0].endswith("gpt2_jacobian_lens.pt")
+
+
+def test_lens_vector_dictionary_matches_lens_vectors_and_caches(toy_model: _ToyBridge) -> None:
+    """The full-vocabulary dictionary equals lens_vectors over every token, is cached per
+    (layer, device), and is released by clear_device_cache."""
+    torch.manual_seed(0)
+    d_model = toy_model.cfg.d_model
+    layer = 1
+    lens = JacobianLens({layer: torch.randn(d_model, d_model)}, n_prompts=1, d_model=d_model)
+    d_vocab = toy_model.W_U.shape[1]
+
+    dictionary = lens.lens_vector_dictionary(toy_model, layer)
+    assert dictionary.shape == (d_vocab, d_model)
+
+    all_vectors = lens.lens_vectors(toy_model, list(range(d_vocab)), layer)
+    assert torch.allclose(dictionary, all_vectors, atol=1e-5)
+
+    # cached: the same object is returned on a repeat call, and clearing releases it
+    assert lens.lens_vector_dictionary(toy_model, layer) is dictionary
+    lens.clear_device_cache()
+    assert lens.lens_vector_dictionary(toy_model, layer) is not dictionary
+
+
+def test_lens_vector_dictionary_rejects_unfitted_layer(toy_model: _ToyBridge) -> None:
+    """Requesting a layer the lens was not fitted at raises (delegated to _matrix_on)."""
+    d_model = toy_model.cfg.d_model
+    lens = JacobianLens({1: torch.randn(d_model, d_model)}, n_prompts=1, d_model=d_model)
+    with pytest.raises(ValueError):
+        lens.lens_vector_dictionary(toy_model, 0)  # layer 0 was not fitted
+
+
+def test_decompose_raw_activation_returns_jspace_decomposition(
+    toy_model: _ToyBridge, fitted_lens: JacobianLens
+) -> None:
+    """decompose on a raw activation vector runs the solver against the layer's dictionary."""
+    torch.manual_seed(0)
+    activation = torch.randn(toy_model.cfg.d_model)
+    result = fitted_lens.decompose(toy_model, activation, layer=0, k=3)
+    assert isinstance(result, JSpaceDecomposition)
+    assert result.support.numel() == 3
+    assert (result.coordinates >= 0).all()
+    assert result.j_space_component.shape == (toy_model.cfg.d_model,)
+
+
+def test_decompose_prompt_matches_manual_activation(
+    toy_model: _ToyBridge, fitted_lens: JacobianLens
+) -> None:
+    """decompose(prompt, position) decomposes the blocks.{layer}.hook_out activation at that
+    position -- identical to fetching it manually and decomposing directly."""
+    layer, position, k = 0, -1, 3
+    result = fitted_lens.decompose(toy_model, "a toy prompt", layer=layer, position=position, k=k)
+
+    tokens = toy_model.to_tokens("a toy prompt")
+    hook = f"blocks.{layer}.hook_out"
+    _, cache = toy_model.run_with_cache(tokens, names_filter=lambda name: name == hook)
+    activation = cache[hook][0, position, :]
+    dictionary = fitted_lens.lens_vector_dictionary(toy_model, layer)
+    expected = get_sparse_decomposition(activation.float(), dictionary, k)
+
+    assert torch.equal(result.support, expected.support)
+    assert torch.allclose(result.coordinates, expected.coordinates, atol=1e-5)
+
+
+def test_decompose_rejects_bad_inputs(toy_model: _ToyBridge, fitted_lens: JacobianLens) -> None:
+    d_model = toy_model.cfg.d_model
+    # a string with no position is neither a raw activation nor a positioned prompt
+    with pytest.raises(ValueError):
+        fitted_lens.decompose(toy_model, "a toy prompt", layer=0, k=3)
+    # raw activation of the wrong width
+    with pytest.raises(ValueError):
+        fitted_lens.decompose(toy_model, torch.randn(d_model + 1), layer=0, k=3)
+    # a batched prompt
+    with pytest.raises(ValueError):
+        fitted_lens.decompose(
+            toy_model, torch.zeros(2, 3, dtype=torch.long), layer=0, position=1, k=3
+        )
+    # a raw activation paired with a position is ambiguous
+    with pytest.raises(ValueError):
+        fitted_lens.decompose(toy_model, torch.randn(d_model), layer=0, position=0, k=3)
+
+
+def test_decompose_passes_algorithm_through(
+    toy_model: _ToyBridge, fitted_lens: JacobianLens
+) -> None:
+    """The wrapper forwards ``algorithm`` to the solver."""
+    torch.manual_seed(0)
+    activation = torch.randn(toy_model.cfg.d_model)
+    result = fitted_lens.decompose(
+        toy_model, activation, layer=0, k=3, algorithm="gradient_pursuit"
+    )
+    dictionary = fitted_lens.lens_vector_dictionary(toy_model, 0)
+    expected = get_sparse_decomposition(
+        activation.float(), dictionary, 3, algorithm="gradient_pursuit"
+    )
+    assert torch.equal(result.support, expected.support)
+    assert torch.allclose(result.coordinates, expected.coordinates, atol=1e-5)
+
+
+def test_decompose_rejects_unfitted_layer(toy_model: _ToyBridge, fitted_lens: JacobianLens) -> None:
+    """Decomposing at a layer the lens was not fitted at raises (the final layer is never fit)."""
+    with pytest.raises(ValueError):
+        fitted_lens.decompose(
+            toy_model, torch.randn(toy_model.cfg.d_model), layer=N_LAYERS - 1, k=3
+        )