TransformerLensOrg
diff --git a/‎docs/source/content/jacobian_lens_fitting.md‎
Lines changed: 88 additions & 0 deletions b/‎docs/source/content/jacobian_lens_fitting.md‎
Lines changed: 88 additions & 0 deletions
diff --git a/‎tests/integration/test_jacobian_lens.py‎
Lines changed: 67 additions & 0 deletions b/‎tests/integration/test_jacobian_lens.py‎
Lines changed: 67 additions & 0 deletions
diff --git a/‎tests/unit/tools/test_jacobian_lens.py‎
Lines changed: 114 additions & 1 deletion b/‎tests/unit/tools/test_jacobian_lens.py‎
Lines changed: 114 additions & 1 deletion
@@ -233,6 +233,94 @@ To propose a short-name entry in TransformerLens, open a pull request that adds
 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 sparse **nonnegative**
+combination of J-lens vectors `v_t = J_ℓ^T W_U[:, t]` (one direction per vocabulary token),
+selected greedily. `k` is an **upper bound**, not a target: selection stops early once no
+unselected vector is materially positively correlated with the residual (under nonnegativity a
+negatively-correlated vector cannot reduce it), so fewer than `k` vectors may be selected and
+fewer still may be numerically active.
+
+```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 (up to k) *active* J-lens vectors
+coordinates = result.coordinates                            # their nonnegative coefficients
+```
+
+The result exposes **two supports**, because the paper uses two inconsistent operationalizations
+(a main-text sparse nonnegative reconstruction and an appendix projection onto a selected span):
+
+- `support` -- the numerically **active** vectors: the selected vectors whose contribution
+  `coordinates[i] * ||v_t||` is a materially nonzero fraction of `||x||`. `coordinates` is aligned
+  with `support`, and `reconstruction = sum(coordinates * v_t)` over `support`.
+- `selected_support` -- **every** greedily selected vector, including any whose coordinate the
+  nonnegativity constraint drove to zero. It defines the span for `j_space_component`. Hence
+  `len(support) <= len(selected_support) <= k`.
+
+So two vector outputs also need not coincide:
+
+- `reconstruction` -- the nonnegative combination over the active `support`.
+- `j_space_component` (the *J-space component*) -- the orthogonal projection of the activation onto
+  the span of `selected_support` -- with `non_j_space_component = x - j_space_component`, the
+  residual the interventions leave unchanged.
+
+For the default exact NNLS re-solve the `reconstruction` equals the projection onto the *active*
+support (KKT stationarity), so it differs from `j_space_component` exactly when a selected vector
+has a zero coordinate (the projection then uses a strictly larger span).
+
+```
+x  in R^d_model
+   |-- decompose(x, layer, k)
+        |-- support / coordinates  a_t >= 0      (active vectors; reconstruction = sum a_t v_t)
+        |-- selected_support       S             (all selected vectors; defines the span below)
+        |-- 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=`. The default,
+`"nonnegative_orthogonal_matching_pursuit"`, solves a nonnegative least-squares (NNLS) problem
+over the selected atoms in float64 after each step. It checks the result against the KKT
+conditions and raises `RuntimeError` if the check fails. `"gradient_pursuit"` skips that solve
+and uses the directional update from Blumensath & Davies (2008), matching the update used in
+the paper; its projected step is accepted only when it does not increase the residual. The two
+algorithms share the same greedy selection *rule* but, because their coefficient residuals
+differ, may select different vectors at later steps and so return a different `support` and
+`reconstruction`.
+
+### 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: the paper's span projection (the
+  `selected_support` operationalization here) "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." Those figures are the paper's own measurements on its models; do not read
+  them off this implementation's `j_space_component` without matching the operationalization.
+- `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." Here `k` is an **upper bound**: `support` returns *at most* `k` active vectors (often
+  fewer), never `k` padded with zero-coefficient slots.
+
+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()`.
+
 ## Importing an existing lens
 
 `JacobianLens.load()` accepts two file schemas: the standard artifact format (four
 
@@ -426,3 +426,70 @@ 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: up to k active 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)
+
+    # ``k`` is an upper bound: ``support`` holds only numerically active atoms, a subset of the
+    # selected set. Do not assert a closed-model active count -- just record it on failure.
+    assert (
+        result.support.numel() <= result.selected_support.numel() <= k
+    ), f"active={result.support.numel()} selected={result.selected_support.numel()} k={k}"
+    assert set(result.support.tolist()).issubset(set(result.selected_support.tolist()))
+    assert result.coordinates.numel() == result.support.numel()
+    assert (result.coordinates > 0).all()  # every returned coordinate is active
+    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 (the projection
+    # span is the whole selected support, not only the active atoms); cosine ~ 0
+    dictionary = published_gpt2_lens.lens_vector_dictionary(gpt2_bridge, layer)
+    residual = result.non_j_space_component
+    for atom_id in result.selected_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() <= result.selected_support.numel() <= k
+    ), f"active={result.support.numel()} selected={result.selected_support.numel()} k={k}"
+    assert set(result.support.tolist()).issubset(set(result.selected_support.tolist()))
+    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
+    )
@@ -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
@@ -1278,3 +1282,112 @@ 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)
+    # ``k`` is an upper bound: only active atoms are returned, as a subset of the selected set.
+    assert result.support.numel() <= result.selected_support.numel() <= 3
+    assert set(result.support.tolist()).issubset(set(result.selected_support.tolist()))
+    assert result.coordinates.numel() == result.support.numel()
+    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.equal(result.selected_support, expected.selected_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.equal(result.selected_support, expected.selected_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
+        )