diff --git a/tests/unit/tools/test_jacobian_lens.py b/tests/unit/tools/test_jacobian_lens.py index e62ece269..04baef613 100644 --- a/tests/unit/tools/test_jacobian_lens.py +++ b/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 diff --git a/transformer_lens/tools/analysis/jacobian_lens.py b/transformer_lens/tools/analysis/jacobian_lens.py index 6d8b6d8a1..0cfb5984c 100644 --- a/transformer_lens/tools/analysis/jacobian_lens.py +++ b/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( (