diff --git a/tests/unit/tools/test_jacobian_lens_span_projection.py b/tests/unit/tools/test_jacobian_lens_span_projection.py new file mode 100644 index 000000000..d336e0fbd --- /dev/null +++ b/tests/unit/tools/test_jacobian_lens_span_projection.py @@ -0,0 +1,113 @@ +"""Model-free span invariants for sparse Jacobian-lens projections.""" + +import pytest +import torch + +from tests.unit.tools.conftest import D_MODEL, _ToyBridge +from transformer_lens.tools.analysis import JacobianLens +from transformer_lens.tools.analysis.jacobian_lens_coordinate_patch import ( + solve_coordinate_patch, +) +from transformer_lens.tools.analysis.jacobian_lens_decomposition import ( + _project_onto_span, + estimate_occupancy, + get_sparse_decomposition, +) + + +def _orthogonal_problem(dim: int, scale: float): + dictionary = torch.zeros(3, dim) + dictionary[0, 0] = 1 + dictionary[1, 1] = scale + dictionary[2, 2] = 1 + target = torch.zeros(dim) + target[:2] = torch.tensor([2.0, 1.0]) + return target, dictionary + + +@pytest.mark.parametrize("dim", [3, 4096]) +@pytest.mark.parametrize("scale", [1.0, 3000.0]) +def test_decomposition_retains_orthogonal_selected_directions(dim: int, scale: float): + target, dictionary = _orthogonal_problem(dim, scale) + result = get_sparse_decomposition(target, dictionary, k=2) + assert result.support.tolist() == [0, 1] + torch.testing.assert_close(result.reconstruction, target) + torch.testing.assert_close(result.j_space_component, target) + torch.testing.assert_close(result.non_j_space_component, torch.zeros_like(target)) + torch.testing.assert_close(dictionary @ result.non_j_space_component, torch.zeros(3)) + + +@pytest.mark.parametrize("dim", [3, 4096]) +@pytest.mark.parametrize("scale", [1.0, 3000.0]) +def test_occupancy_gains_preserve_nested_spans(dim: int, scale: float): + target, dictionary = _orthogonal_problem(dim, scale) + result = estimate_occupancy(target, dictionary, max_atoms=2, num_control_dictionaries=1) + torch.testing.assert_close(result.marginal_captured_variance, torch.tensor([0.8, 0.2])) + + +def test_coordinate_reuse_accepts_correct_wide_span_projection(): + target, dictionary = _orthogonal_problem(4096, 3000.0) + decomposition = get_sparse_decomposition(target, dictionary, k=2) + # Supply the exact projection independently of the implementation under test. + decomposition.j_space_component = target.clone() + decomposition.non_j_space_component = torch.zeros_like(target) + patch = solve_coordinate_patch(target, dictionary, 0, 2, decomposition=decomposition) + expected = target.clone() + expected[0], expected[2] = 0, 2 + torch.testing.assert_close(patch.patched, expected) + decomposition.j_space_component[0] = 0 + decomposition.non_j_space_component = target - decomposition.j_space_component + with pytest.raises(ValueError, match="selected span"): + solve_coordinate_patch(target, dictionary, 0, 2, decomposition=decomposition) + + +@pytest.mark.parametrize("scale", [1e-12, 1.0, 1e12]) +def test_span_projection_is_invariant_to_atom_rescaling_and_zero_padding(scale: float): + atoms = torch.tensor([[1.0, 1.0], [0.0, 1.0], [0.0, 0.0]]) + target = torch.tensor([2.0, 1.0, 3.0]) + atoms[:, 1] *= scale + expected = torch.tensor([2.0, 1.0, 0.0]) + torch.testing.assert_close(_project_onto_span(atoms, target), expected) + padded_atoms = torch.nn.functional.pad(atoms, (0, 0, 0, 4093)) + padded_target = torch.nn.functional.pad(target, (0, 4093)) + projected = _project_onto_span(padded_atoms, padded_target) + torch.testing.assert_close(projected[:3], expected) + torch.testing.assert_close(projected[3:], torch.zeros(4093)) + + +@pytest.mark.parametrize("separation, retained", [(1e-5, True), (1e-10, False)]) +def test_span_projection_has_explicit_near_dependency_policy(separation: float, retained: bool): + atoms = torch.tensor([[1.0, 1.0], [0.0, separation], [0.0, 0.0]]) + target = torch.tensor([0.0, 1.0, 0.0]) + projected = _project_onto_span(atoms, target) + expected = target if retained else torch.zeros_like(target) + torch.testing.assert_close(projected, expected, atol=1e-8, rtol=1e-6) + + +def test_span_projection_handles_dependent_and_empty_columns(): + atoms = torch.tensor([[1.0, 2.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0], [0.0, 0.0, 0.0, 0.0]]) + target = torch.tensor([2.0, 1.0, 3.0]) + torch.testing.assert_close(_project_onto_span(atoms, target), torch.tensor([2.0, 1.0, 0.0])) + torch.testing.assert_close(_project_onto_span(atoms[:, :0], target), torch.zeros(3)) + + +def test_lens_variance_profile_uses_the_full_selected_span(): + model = _ToyBridge() + with torch.no_grad(): + model.embed.weight.zero_() + model.embed.weight[:, :2] = torch.tensor([2.0, 1.0]) + for block in model.blocks: + block.linear.weight.zero_() + model.unembed.weight.zero_() + model.unembed.weight[:, 2] = 1 + model.unembed.weight[0] = 0 + model.unembed.weight[0, 0] = 1 + model.unembed.weight[1] = 0 + model.unembed.weight[1, 1] = 3e6 + lens = JacobianLens({0: torch.eye(D_MODEL)}, n_prompts=1, d_model=D_MODEL) + result = lens.decompose(model, "test", layer=0, position=-1, k=2) + torch.testing.assert_close(result.non_j_space_component, torch.zeros(D_MODEL)) + profile = lens.fraction_of_variance(model, ["test"], layers=[0], k=2, skip_first=0) + assert profile.median[0] == pytest.approx(1.0) + assert profile.pooled[0] == pytest.approx(1.0) + torch.testing.assert_close(profile.per_position[0], torch.ones_like(profile.per_position[0])) diff --git a/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py b/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py index 5691620ac..61f986007 100644 --- a/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py +++ b/transformer_lens/tools/analysis/jacobian_lens_coordinate_patch.py @@ -18,6 +18,7 @@ _active_mask, _gradient_pursuit_extend, _linalg_on_cpu_if_mps, + _project_onto_span, _validate_nnls_kkt, get_sparse_decomposition, ) @@ -278,11 +279,7 @@ def _validate_decomposition( 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: - expected_component = _linalg_on_cpu_if_mps( - lambda atoms, vector: atoms @ (torch.linalg.pinv(atoms) @ vector), - selected_atoms, - target, - ) + expected_component = _project_onto_span(selected_atoms, target) else: expected_component = torch.zeros_like(target) if not _close(local_vectors[1], expected_component, span_condition): diff --git a/transformer_lens/tools/analysis/jacobian_lens_decomposition.py b/transformer_lens/tools/analysis/jacobian_lens_decomposition.py index 996d5b438..9a0f978e2 100644 --- a/transformer_lens/tools/analysis/jacobian_lens_decomposition.py +++ b/transformer_lens/tools/analysis/jacobian_lens_decomposition.py @@ -120,6 +120,26 @@ def _linalg_on_cpu_if_mps(op: Callable[..., torch.Tensor], *tensors: torch.Tenso return op(*tensors) +def _project_onto_span(atoms: torch.Tensor, target: torch.Tensor) -> torch.Tensor: + """Project onto atom columns using a scale- and ambient-dimension-independent rank policy. + + Unit columns are decomposed in float64, retaining singular values above the NNLS relative + rank tolerance. Exactly dependent and numerically indistinguishable directions are removed. + MPS performs this small selected-set decomposition on CPU, restoring dtype and device. + """ + if atoms.shape[1] == 0: + return torch.zeros_like(target) + + def project(local_atoms: torch.Tensor, local_target: torch.Tensor) -> torch.Tensor: + work = local_atoms.double() + norms = torch.linalg.vector_norm(work, dim=0).clamp_min(torch.finfo(torch.float64).tiny) + left, singular_values, _ = torch.linalg.svd(work / norms, full_matrices=False) + basis = left[:, singular_values > _NNLS_RANK_RTOL * singular_values.max()] + return (basis @ (basis.T @ local_target.double())).to(local_target.dtype) + + return _linalg_on_cpu_if_mps(project, atoms, target) + + def _diagnose_intervention_pair( unit_vectors: torch.Tensor, *, description: str, stacklevel: int ) -> float: @@ -167,6 +187,8 @@ class JSpaceDecomposition: ``selected_support`` (the paper's "J-space component"). For the exact NNLS re-solve it equals ``reconstruction`` unless a selected atom has a zero coordinate, in which case the projection uses a larger span. + Numerical rank is measured on unit atom columns in float64 with relative + tolerance ``sqrt(eps64)``, independently of ambient dimension and atom norms. non_j_space_component: The residual ``target - j_space_component`` (the "non-J-space component"), orthogonal to the selected span. """ @@ -498,15 +520,10 @@ def get_sparse_decomposition( # Reconstruction is the nonnegative combination over the active support (empty -> zeros). reconstruction = active_atoms @ active_coordinates # The J-space component is the orthogonal projection of the target onto the span of every - # selected atom (paper appendix), computed with the same pseudoinverse construction as - # swap_hooks. That span can be larger than the active support when a selected coordinate is + # selected atom (paper appendix). That span can exceed the active support when a coordinate is # zero, so the projection differs from the nonnegative reconstruction in that case. if selected: - j_space_component = _linalg_on_cpu_if_mps( - lambda atoms, vector: atoms @ (torch.linalg.pinv(atoms) @ vector), - selected_atoms, - target, - ) + j_space_component = _project_onto_span(selected_atoms, target) else: j_space_component = target.new_zeros(d_model) non_j_space_component = target - j_space_component @@ -548,8 +565,8 @@ def _greedy_captured_variance_gains( At each step, add the unused atom with the greatest signed, norm-normalized correlation with the current residual. Project ``target`` orthogonally onto the full selected span using a - pseudoinverse, then set the next residual to ``target - projection``. The captured variance is - ``||Pi_S target||^2 / ||target||^2``; the returned values are its per-step increments. + rank-revealing unit-column SVD, then set the next residual to ``target - projection``. + Captured variance is ``||Pi_S target||^2 / ||target||^2``; the returned values are its increments. This shares the per-step correlation rule with :func:`get_sparse_decomposition`, but not its residual recurrence: sparse decomposition uses a nonnegative coefficient-fit residual and may @@ -566,7 +583,7 @@ def _greedy_captured_variance_gains( correlation[chosen] = float("-inf") support.append(int(torch.argmax(correlation).item())) active_atoms = atoms[support].T - projection = active_atoms @ (torch.linalg.pinv(active_atoms) @ target) + projection = _project_onto_span(active_atoms, target) captured_variance = float((projection @ projection) / total_variance) captured_variance_gains.append(captured_variance - previous_captured_variance) previous_captured_variance = captured_variance