fix(jacobian_lens): ablate token directions jointly

tests/unit/tools/test_jacobian_lens.py

@@ -1059,6 +1059,28 @@ def test_intervention_hooks_use_bridge_names_and_run(
         assert not torch.allclose(intervened, baseline)
 
 
+def test_ablation_hooks_remove_every_direction_of_non_orthogonal_tokens(
+    toy_model: _ToyBridge, fitted_lens: JacobianLens
+) -> None:
+    tokens = [3, 5, 7]
+    layer = 1
+    vectors = fitted_lens.lens_vectors(toy_model, tokens, layer)
+    units = vectors / vectors.norm(dim=-1, keepdim=True)
+    off_diagonal = (units @ units.T)[~torch.eye(len(tokens), dtype=torch.bool)]
+    assert off_diagonal.abs().max() > 0.1  # one-at-a-time removal would leave a component
+
+    [(_, hook)] = fitted_lens.ablation_hooks(toy_model, tokens, layers=[layer])
+    activation = torch.randn(2, SEQ_LEN, D_MODEL, generator=torch.Generator().manual_seed(0))
+    ablated = hook(activation.clone(), SimpleNamespace(name=f"blocks.{layer}.hook_out"))
+
+    torch.testing.assert_close(
+        ablated @ vectors.T, torch.zeros(2, SEQ_LEN, len(tokens)), atol=1e-4, rtol=0
+    )
+    # Only the span of the vectors changes
+    projector = vectors.T @ torch.linalg.pinv(vectors.T)
+    torch.testing.assert_close(ablated, activation - activation @ projector, atol=1e-5, rtol=0)
+
+
 @pytest.mark.parametrize("intervention", ["steering", "ablation", "swap"])
 def test_intervention_negative_position_is_normalized(
     intervention: str, toy_model: _ToyBridge

transformer_lens/tools/analysis/jacobian_lens.py

@@ -99,6 +99,7 @@
     get_sparse_decomposition,
 )
 from transformer_lens.tools.analysis.logit_lens import logit_readout
+from transformer_lens.tools.analysis.projection_kernel import orthonormal_subspace
 from transformer_lens.utilities.hf_utils import call_hf_with_retry
 
 ConceptToken = Union[str, int]
@@ -1339,8 +1340,9 @@ def ablation_hooks(
     ) -> List[Tuple[str, Any]]:
         """Hooks that project token directions out of the residual stream.
 
-        For each token's unit lens vector ``v̂``: ``h <- h - (h·v̂) v̂``,
-        applied sequentially when several tokens are given.
+        With an orthonormal basis ``Q`` of the span of the tokens' lens vectors:
+        ``h <- h - Q Qᵀ h``, so ``h`` ends with zero projection onto every
+        requested vector. For a single token this is ``h - (h·v̂) v̂``.
 
         Args:
             model: The model the hooks will run on.
@@ -1357,20 +1359,20 @@ def ablation_hooks(
         for layer in [_normalize_layer(layer, model.cfg.n_layers) for layer in layers]:
             vectors = self.lens_vectors(model, tokens, layer)
             units = _unit_rows(vectors, layer=layer)
-            device_units: Dict[torch.device, torch.Tensor] = {}
+            # Lens vectors of related tokens are far from orthogonal, so removing them one
+            # at a time would reintroduce components along the earlier ones
+            basis = orthonormal_subspace(units.T).basis
+            device_bases: Dict[torch.device, torch.Tensor] = {}
 
             def transform(
                 selected: Float[torch.Tensor, "batch pos d_model"],
                 positions: List[int],
-                units: torch.Tensor = units,
-                device_units: Dict[torch.device, torch.Tensor] = device_units,
+                basis: torch.Tensor = basis,
+                device_bases: Dict[torch.device, torch.Tensor] = device_bases,
             ) -> Float[torch.Tensor, "batch pos d_model"]:
-                local_units = _cached_on_device(units, device_units, selected.device)
+                local_basis = _cached_on_device(basis, device_bases, selected.device)
                 result = selected.float()
-                for unit in local_units:
-                    coeff = result @ unit
-                    result = result - coeff.unsqueeze(-1) * unit
-                return result
+                return result - (result @ local_basis) @ local_basis.T
 
             hooks.append(
                 (