fix: validate reused coordinate decompositions by pursuit algorithm

tests/unit/tools/test_jacobian_lens_coordinate_patch.py

@@ -284,6 +284,124 @@ def test_rejects_incompatible_precomputed_decomposition() -> None:
         solve_coordinate_patch(activation, changed_dictionary, 0, 2, decomposition=decomposition)
 
 
+@pytest.mark.parametrize(
+    "algorithm", ["nonnegative_orthogonal_matching_pursuit", "gradient_pursuit"]
+)
+def test_correlated_decomposition_reuse_and_position_cache(
+    algorithm: str, monkeypatch: pytest.MonkeyPatch
+) -> None:
+    import transformer_lens.tools.analysis.jacobian_lens_coordinate_patch as core_module
+
+    dictionary = torch.tensor([[1.0, 0.0, 0.0], [0.7, 0.7, 0.0], [0.0, 0.0, 1.0]])
+    activation = torch.tensor([1.0, 0.3, 0.0])
+    decomposition = get_sparse_decomposition(activation, dictionary, k=2, algorithm=algorithm)
+    fresh = solve_coordinate_patch(activation, dictionary, 0, 2, k=2, algorithm=algorithm)
+    reused = solve_coordinate_patch(
+        activation, dictionary, 0, 2, decomposition=decomposition, algorithm=algorithm
+    )
+    torch.testing.assert_close(reused.patched, fresh.patched)
+    torch.testing.assert_close(
+        reused.delta,
+        torch.tensor([-decomposition.coordinates[0], 0.0, decomposition.coordinates[0]]),
+    )
+    calls = []
+    original = core_module.get_sparse_decomposition
+
+    def spy(*args, **kwargs):
+        calls.append(1)
+        return original(*args, **kwargs)
+
+    monkeypatch.setattr(core_module, "get_sparse_decomposition", spy)
+    cache: dict = {}
+    for supplied_cache in (None, cache, cache):
+        patched, patches = solve_coordinate_patch_positions(
+            activation.view(1, 1, 3),
+            dictionary,
+            [7],
+            0,
+            2,
+            layer=0,
+            decomposition_cache=supplied_cache,
+            k=2,
+            algorithm=algorithm,
+        )
+        torch.testing.assert_close(patched[0, 0], fresh.patched)
+        torch.testing.assert_close(patches[(0, 0, 7)].residual, fresh.residual)
+    assert len(calls) == 2
+
+
+@pytest.mark.parametrize(
+    "algorithm", ["nonnegative_orthogonal_matching_pursuit", "gradient_pursuit"]
+)
+def test_correlated_reuse_rejects_scaled_coordinates(algorithm: str) -> None:
+    dictionary = torch.tensor([[1.0, 0.0, 0.0], [0.7, 0.7, 0.0], [0.0, 0.0, 1.0]])
+    activation = torch.tensor([1.0, 0.3, 0.0])
+    decomposition = get_sparse_decomposition(activation, dictionary, k=2, algorithm=algorithm)
+    scaled = replace(
+        decomposition,
+        coordinates=decomposition.coordinates * 2,
+        reconstruction=decomposition.reconstruction * 2,
+    )
+    with pytest.raises(ValueError, match="activation"):
+        solve_coordinate_patch(
+            activation, dictionary, 0, 2, decomposition=scaled, algorithm=algorithm
+        )
+
+
+@pytest.mark.parametrize("scale", [1e-3, 1.0, 1e3])
+def test_gradient_pursuit_reuse_with_multiple_correlated_atoms(scale: float) -> None:
+    generator = torch.Generator().manual_seed(17)
+    dictionary = torch.randn(8, 5, generator=generator)
+    activation = torch.randn(5, generator=generator) * scale
+    decomposition = get_sparse_decomposition(
+        activation, dictionary, k=4, algorithm="gradient_pursuit"
+    )
+    source = int(decomposition.support[0])
+    target = (source + 1) % len(dictionary)
+    fresh = solve_coordinate_patch(
+        activation, dictionary, source, target, k=4, algorithm="gradient_pursuit"
+    )
+    reused = solve_coordinate_patch(
+        activation,
+        dictionary,
+        source,
+        target,
+        decomposition=decomposition,
+        algorithm="gradient_pursuit",
+    )
+    torch.testing.assert_close(reused.patched, fresh.patched)
+    with pytest.raises(ValueError, match="activation"):
+        solve_coordinate_patch(
+            activation * 2,
+            dictionary,
+            source,
+            target,
+            decomposition=decomposition,
+            algorithm="gradient_pursuit",
+        )
+
+
+def test_gradient_pursuit_empty_support_reports_missing_source() -> None:
+    dictionary = torch.eye(3)
+    activation = torch.zeros(3)
+    decomposition = get_sparse_decomposition(
+        activation, dictionary, k=2, algorithm="gradient_pursuit"
+    )
+    with pytest.raises(ValueError, match="active support"):
+        solve_coordinate_patch(
+            activation, dictionary, 0, 1, decomposition=decomposition, algorithm="gradient_pursuit"
+        )
+
+
+def test_precomputed_decomposition_rejects_unknown_algorithm() -> None:
+    activation, dictionary = _problem()
+    decomposition = get_sparse_decomposition(activation, dictionary, k=2)
+    with pytest.raises(ValueError, match="algorithm"):
+        solve_coordinate_patch(
+            activation, dictionary, 0, 1, decomposition=decomposition, algorithm="unknown"
+        )
+
+
 def test_rejects_scaled_but_self_consistent_decomposition() -> None:
     activation, dictionary = _problem()
     decomposition = get_sparse_decomposition(activation, dictionary, k=2)

tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py

@@ -15,12 +15,19 @@
 
 
 def _active_source_and_distinct_target(
-    lens: JacobianLens, model: _ToyBridge, prompt: str, layer: int, position: int
+    lens: JacobianLens,
+    model: _ToyBridge,
+    prompt: str,
+    layer: int,
+    position: int,
+    algorithm: str = "nonnegative_orthogonal_matching_pursuit",
 ) -> tuple[int, int]:
     """Discover a real active source at (layer, position) instead of guessing a token id --
     coordinate_patch_hooks requires the source to be active, and a toy model's real activations
     are not hand-computable in advance."""
-    decomposition = lens.decompose(model, prompt, layer=layer, position=position, k=SOLVE_K)
+    decomposition = lens.decompose(
+        model, prompt, layer=layer, position=position, k=SOLVE_K, algorithm=algorithm
+    )
     source_id = int(decomposition.support[0])
     target_id = (source_id + 1) % D_VOCAB
     return source_id, target_id
@@ -83,19 +90,31 @@ def test_coordinate_patch_hooks_changes_only_requested_positions(toy_model: _Toy
     torch.testing.assert_close(delta[:, :-1], torch.zeros_like(delta[:, :-1]))
 
 
+@pytest.mark.parametrize(
+    "algorithm", ["nonnegative_orthogonal_matching_pursuit", "gradient_pursuit"]
+)
 def test_coordinate_patch_hooks_oracle_parity_with_offline_coordinate_patch(
     toy_model: _ToyBridge,
+    algorithm: str,
 ) -> None:
     lens = _lens()
     prompt = "a toy prompt"
     tokens = toy_model.to_tokens(prompt)
     _, baseline = toy_model.run_with_cache(tokens)
     pre_hook_activation = baseline["blocks.0.hook_out"][0, -1, :].float()
-    source_id, target_id = _active_source_and_distinct_target(lens, toy_model, prompt, 0, -1)
+    source_id, target_id = _active_source_and_distinct_target(
+        lens, toy_model, prompt, 0, -1, algorithm=algorithm
+    )
 
     with pytest.warns(UserWarning):
         hooks = lens.coordinate_patch_hooks(
-            toy_model, source_id, target_id, layers=[0], positions=[-1], k=SOLVE_K
+            toy_model,
+            source_id,
+            target_id,
+            layers=[0],
+            positions=[-1],
+            k=SOLVE_K,
+            algorithm=algorithm,
         )
     with toy_model.hooks(fwd_hooks=hooks):
         _, patched = toy_model.run_with_cache(tokens)
@@ -107,19 +126,27 @@ def test_coordinate_patch_hooks_oracle_parity_with_offline_coordinate_patch(
         source_token=source_id,
         target_token=target_id,
         k=SOLVE_K,
+        algorithm=algorithm,
     )
     torch.testing.assert_close(patched["blocks.0.hook_out"][0, -1, :].float(), expected.patched)
 
 
+@pytest.mark.parametrize(
+    "algorithm", ["nonnegative_orthogonal_matching_pursuit", "gradient_pursuit"]
+)
 def test_decomposition_cache_hit_skips_resolve_across_hook_firings(
-    toy_model: _ToyBridge, monkeypatch: pytest.MonkeyPatch
+    toy_model: _ToyBridge,
+    monkeypatch: pytest.MonkeyPatch,
+    algorithm: str,
 ) -> None:
     import transformer_lens.tools.analysis.jacobian_lens_coordinate_patch as core_module
 
     lens = _lens()
     prompt = "a toy prompt"
     tokens = toy_model.to_tokens(prompt)
-    source_id, target_id = _active_source_and_distinct_target(lens, toy_model, prompt, 0, -1)
+    source_id, target_id = _active_source_and_distinct_target(
+        lens, toy_model, prompt, 0, -1, algorithm=algorithm
+    )
     cache: dict = {}
     calls = []
     original = core_module.get_sparse_decomposition
@@ -138,6 +165,7 @@ def spy(*args: Any, **kwargs: Any) -> Any:
             positions=[-1],
             decomposition_cache=cache,
             k=SOLVE_K,
+            algorithm=algorithm,
         )
     with toy_model.hooks(fwd_hooks=hooks):
         toy_model(tokens)

transformer_lens/tools/analysis/jacobian_lens.py

@@ -1013,7 +1013,8 @@ def coordinate_patch(
             k: Sparse-solver upper bound when ``decomposition`` is not supplied.
             mode: ``"substitute"`` or ``"swap"``.
             alpha: Finite interpolation strength; zero is an exact no-op.
-            algorithm: Sparse coefficient-update rule when solving a fresh decomposition.
+            algorithm: Sparse coefficient-update rule, also used to validate a supplied
+                decomposition. Use the algorithm that produced its coordinates.
 
         Returns:
             A :class:`CoordinatePatch` containing the edited frame, diagnostics, and activation.
@@ -1558,7 +1559,8 @@ def coordinate_patch_hooks(
             k: Sparse-solver upper bound on a cache miss.
             mode: ``"substitute"`` or ``"swap"``.
             alpha: Finite interpolation strength; zero is an exact no-op.
-            algorithm: Sparse coefficient-update rule on a cache miss.
+            algorithm: Sparse coefficient-update rule on a cache miss and validation rule on a
+                hit. Cached decompositions must use this algorithm.
 
         Returns:
             ``[(hook_name, fn), ...]`` for ``model.hooks(fwd_hooks=...)``.

transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py

@@ -12,9 +12,11 @@
 
 from transformer_lens.tools.analysis._model_state import require_int
 from transformer_lens.tools.analysis.jacobian_lens_decomposition import (
+    _ACTIVE_RELATIVE_TOLERANCE,
     _SWAP_WARN_COSINE,
     DEFAULT_K,
     JSpaceDecomposition,
+    _gradient_pursuit_step,
     _linalg_on_cpu_if_mps,
     _validate_nnls_kkt,
     get_sparse_decomposition,
@@ -169,6 +171,7 @@ def _validate_decomposition(
     target: torch.Tensor,
     dictionary: torch.Tensor,
     target_idx: int,
+    algorithm: str,
 ) -> None:
     num_atoms, d_model = dictionary.shape
     support = decomposition.support
@@ -223,27 +226,36 @@ def _validate_decomposition(
     expected_reconstruction = support_atoms.T @ coordinates
     if not _close(local_vectors[0], expected_reconstruction, span_condition):
         raise ValueError("precomputed decomposition reconstruction is incompatible with dictionary")
-    # Couple the coordinates to ``target`` (the activation), not just to the dictionary: the
-    # reconstruction/dictionary check above passes for any self-consistently scaled pair, so a
-    # wrongly-scaled decomposition would slip through and silently anchor the edit to a bogus
-    # residual. Reuse the solver's NNLS KKT stationarity test -- the active coordinates are all
-    # strictly positive (validated above), so stationarity ``A^T (target - A c) == 0`` is the
-    # exact condition for ``c`` being the nonnegative fit of ``target`` over the active support.
-    # Tolerances mirror the solver's result-dtype certification (sqrt(eps) of the float32
-    # coordinates) so genuine float32-rounded decompositions are not falsely rejected.
-    active_atoms = support_atoms.T
-    tiny = torch.finfo(target.dtype).tiny
-    relative_tolerance = math.sqrt(torch.finfo(target.dtype).eps)
-    atom_norms = active_atoms.norm(dim=0).clamp_min(tiny)
-    scale = (target.norm() + (active_atoms.abs() @ coordinates.abs()).norm()).clamp_min(tiny)
-    dual_tol = relative_tolerance * atom_norms * scale
-    coefficient_tol = relative_tolerance * scale / atom_norms
-    try:
-        _validate_nnls_kkt(active_atoms, target, coordinates, dual_tol, coefficient_tol)
-    except RuntimeError as error:
-        raise ValueError(
-            f"precomputed decomposition coordinates do not fit the activation ({error})"
-        ) from error
+    if algorithm == "gradient_pursuit":
+        # Replay only the selected-set updates; approximate coefficients need not satisfy NNLS KKT.
+        replay = target.new_zeros(0)
+        selected_rows = dictionary[selected_cpu].float()
+        for count in range(1, selected_cpu.numel() + 1):
+            replay = _gradient_pursuit_step(
+                selected_rows[:count].T, target, torch.cat([replay, replay.new_zeros(1)])
+            )
+        active = replay * selected_rows.norm(dim=1) > (
+            _ACTIVE_RELATIVE_TOLERANCE * float(target.norm().detach())
+        )
+        if not torch.equal(support_cpu, selected_cpu[active.cpu()]) or (
+            coordinates.numel() and not _close(coordinates, replay[active])
+        ):
+            raise ValueError("precomputed decomposition coordinates do not fit the activation")
+    else:
+        # Couple the coordinates to the activation, including self-consistently rescaled inputs.
+        active_atoms = support_atoms.T
+        tiny = torch.finfo(target.dtype).tiny
+        relative_tolerance = math.sqrt(torch.finfo(target.dtype).eps)
+        atom_norms = active_atoms.norm(dim=0).clamp_min(tiny)
+        scale = (target.norm() + (active_atoms.abs() @ coordinates.abs()).norm()).clamp_min(tiny)
+        dual_tol = relative_tolerance * atom_norms * scale
+        coefficient_tol = relative_tolerance * scale / atom_norms
+        try:
+            _validate_nnls_kkt(active_atoms, target, coordinates, dual_tol, coefficient_tol)
+        except RuntimeError as error:
+            raise ValueError(
+                f"precomputed decomposition coordinates do not fit the activation ({error})"
+            ) from error
     if not _close(local_vectors[1] + local_vectors[2], target, span_condition):
         raise ValueError("precomputed decomposition is incompatible with activation")
     if selected_atoms is not None:
@@ -362,7 +374,8 @@ def solve_coordinate_patch(
         k: Sparse-solver upper bound when no decomposition is supplied.
         mode: ``"substitute"`` to overwrite the target, or ``"swap"`` to exchange coordinates.
         alpha: Finite interpolation strength. Zero returns an unchanged float32 clone.
-        algorithm: Sparse coefficient-update rule when no decomposition is supplied.
+        algorithm: Sparse coefficient-update rule, also used to validate supplied coordinates.
+            Set this to the algorithm that produced a supplied decomposition.
 
     Returns:
         A :class:`CoordinatePatch` with the anchored edit and numerical diagnostics.
@@ -388,13 +401,18 @@ def solve_coordinate_patch(
         raise ValueError(f"mode must be 'substitute' or 'swap', got {mode!r}")
     if isinstance(alpha, bool) or not isinstance(alpha, Real) or not math.isfinite(float(alpha)):
         raise ValueError("alpha must be a finite real number")
+    if algorithm not in ("nonnegative_orthogonal_matching_pursuit", "gradient_pursuit"):
+        raise ValueError(
+            "algorithm must be 'nonnegative_orthogonal_matching_pursuit' or "
+            f"'gradient_pursuit', got {algorithm!r}"
+        )
 
     if decomposition is None:
         decomposition = get_sparse_decomposition(target, dictionary, k, algorithm=algorithm)
     elif not isinstance(decomposition, JSpaceDecomposition):
         raise ValueError("decomposition must be a JSpaceDecomposition")
     else:
-        _validate_decomposition(decomposition, target, dictionary, target_idx)
+        _validate_decomposition(decomposition, target, dictionary, target_idx, algorithm)
 
     support_before = decomposition.support.detach().cpu().clone()
     support_ids = support_before.tolist()
@@ -524,7 +542,8 @@ def solve_coordinate_patch_positions(
         k: Sparse-solver upper bound on a cache miss.
         mode: ``"substitute"`` or ``"swap"``, forwarded to :func:`solve_coordinate_patch`.
         alpha: Finite interpolation strength, forwarded to :func:`solve_coordinate_patch`.
-        algorithm: Sparse coefficient-update rule on a cache miss.
+        algorithm: Sparse coefficient-update rule on a cache miss and validation rule on a hit.
+            Cached decompositions must use this algorithm.
 
     Returns:
         A tuple ``(patched, patches)``: ``patched`` has the same shape as ``activations``, with