diff --git a/tests/unit/tools/test_jacobian_lens_coordinate_patch.py b/tests/unit/tools/test_jacobian_lens_coordinate_patch.py index 4b847727b..9081d3eca 100644 --- a/tests/unit/tools/test_jacobian_lens_coordinate_patch.py +++ b/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) diff --git a/tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py b/tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py index 9a2ce793b..4960ab421 100644 --- a/tests/unit/tools/test_jacobian_lens_coordinate_patch_hooks.py +++ b/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) diff --git a/transformer_lens/tools/analysis/jacobian_lens.py b/transformer_lens/tools/analysis/jacobian_lens.py index f0034d1a0..6d8b6d8a1 100644 --- a/transformer_lens/tools/analysis/jacobian_lens.py +++ b/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=...)``. diff --git a/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py b/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py index 807ae50f1..7be1b759f 100644 --- a/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py +++ b/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