diff --git a/docs/source/content/analysis_tools.md b/docs/source/content/analysis_tools.md index cd9e940a5..e5346ee5c 100644 --- a/docs/source/content/analysis_tools.md +++ b/docs/source/content/analysis_tools.md @@ -20,6 +20,7 @@ the [4.0 migration guide](migrating_to_v4.md) for existing `HookedTransformer` c | What does each layer "predict", and where does a token's rank change? | **Logit Lens** | A prompt or activation cache, optional target tokens, top-k or a vocabulary subset → per-layer logits, log-probabilities, ranks and entropy. | What the residual stream reads as through the real final norm and unembedding; the final entry reproduces the model's own logits. | | What vocabulary directions appear in an MLP's gradient factors? | **Backward Lens** | A prompt, one target token, and selected layers → gradient factors and vocabulary rankings. | A diagnostic of the forward inputs and backward signals composing a weight gradient. | | How can I read or edit residuals through a fitted transport map? | **Jacobian Lens** | A matching lens artifact and model, plus prompts or activations → vocabulary readouts, decompositions, or interventions. | Readouts under the fitted map; causal effects require running and measuring an intervention. | +| Does a subspace at one hook point carry a variable the output depends on? | **Interchange Interventions** | A hook point, base and source inputs, a basis, and a task readout → task logits with the base's subspace coordinates set to the source's. | The effect of that intervention on the chosen pairs; compare with a norm-matched placebo and random bases before reading it as use. | | How much do two subspaces overlap? | **Projection Kernel** | Two subspace bases, or Bridge attention-head weight spaces → overlap scores and principal-angle information. | Shared geometric support, independent of the choice of basis within each subspace. | | What are a head's QK/OV singular directions? | **SVD Circuits: head decomposition** | Model weights and a layer/head index → singular values, vectors, and degeneracy reports. | The linear structure of the head's weight maps. | @@ -42,6 +43,7 @@ available hooks, and the meaning of the selected tensor axes. | Logit Lens | Targets TransformerBridge, raw or compatibility mode. Needs a decoder with an `ln_final` component (or no final norm) and `W_U`; encoder-decoder and encoder-only bridges are rejected, as are quantized unembeddings. Hybrid SSM stacks work with `incl_mid=False`. Run the model in evaluation mode. | One cached forward pass over the residual hooks, or reuse of a cache. Logits are produced in chunks of `chunk_size` rows; full-vocabulary output is materialized only when no `targets`, `top_k` or `vocab` selector is given. | | Backward Lens | Requires a raw decoder-only TransformerBridge with dense, non-gated MLPs and without compatibility mode, such as GPT-2 or Pythia/GPT-NeoX. Use one target token and check the restrictions in the [tool guide](backward_lens.md). | Gradient computation plus vocabulary projections for selected layers and positions; retaining full logits increases memory. | | Jacobian Lens | Requires a fresh, causal decoder-only Bridge with raw HF weights, without compatibility mode or weight processing. Validate the lens against the model. Fitting additionally requires all modules in evaluation mode. | Loading an existing artifact avoids fitting. The ordinary fitting estimator uses one forward and `ceil(d_model / dim_batch)` backwards per prompt; larger batches increase memory. | +| Interchange Interventions | Any model with the hook API (TransformerBridge or `HookedRootModule`). Activations at the hook must be `[batch, d_model]`, or `[batch, pos, d_model]` with one position index per example; base and source rows are paired. | One forward pass per `batch_size` sources to collect activations, then one per `batch_size` bases per intervention with `intervened_logits`; both run without gradients, so memory is bounded by `batch_size`. `run_with_intervention` is a single differentiable pass over all the bases it is given. | | Projection Kernel | The numerical API accepts finite, real floating-point matrices via orthonormal bases. The attention-head wrapper requires Bridge weights with compatible dimensions and ranks. | Basis extraction uses SVD. All-head comparisons allocate basis stacks and a pairwise score grid; they can be large despite requiring no forward pass. | | SVD head decomposition | Uses Bridge per-block weight accessors; requires accessible, compatible `W_Q`, `W_K`, `W_V`, and `W_O` for the selected head. No compatibility mode or activation cache is needed. | Factored QK/OV SVD, with rank bounded by `d_head`; the dense `d_model × d_model` product is not materialized. | diff --git a/docs/source/content/interchange_interventions.md b/docs/source/content/interchange_interventions.md new file mode 100644 index 000000000..635aa86ec --- /dev/null +++ b/docs/source/content/interchange_interventions.md @@ -0,0 +1,90 @@ +# Interchange Interventions + +An interchange intervention runs the model on a *base* input while setting the base's +activation at one hook point, inside a linear subspace, to its value on a *source* input: + +$$ +h' = h_{\mathrm{base}} + s \, (h_{\mathrm{source}} - h_{\mathrm{base}}) B B^\top, +$$ + +where $B$ is an orthonormal `[d_model, rank]` basis (the convention of +[`orthonormal_subspace`](projection_kernel.md)) and $s$ is the strength, 1 by default. The base +keeps its own coordinates outside the subspace. + +The question it answers is causal. Suppose a high-level model says the output depends on a +variable, such as a subject's grammatical number. If interchanging a subspace makes the model +produce the output that the high-level model predicts for the base with that variable taken from +the source, the subspace carries the variable at that hook point (Geiger et al., 2021). The share +of base/source pairs where this happens is the *interchange intervention accuracy* (IIA). + +## Run an intervention + +```python +import torch + +from transformer_lens.model_bridge import TransformerBridge +from transformer_lens.tools.analysis import ( + SubspaceInterchange, + collect_activations, + intervened_logits, + mean_difference_basis, +) + +bridge = TransformerBridge.boot_transformers("gpt2") +answers = [bridge.to_single_token(" is"), bridge.to_single_token(" are")] +task_logits = lambda tokens: bridge(tokens)[:, -1, answers] + +base = bridge.to_tokens(["The key near the cabinet", "The keys near the cabinet"]) +source = base.flip(0) # each base takes the other prompt's subject number +subject = torch.full((2,), 2) # the subject token, after the BOS token +hook = "blocks.1.hook_out" + +source_acts = collect_activations(bridge, hook, source, task_logits, positions=subject) +basis = mean_difference_basis(source_acts, labels=[1, 0]) +logits = intervened_logits( + bridge, hook, SubspaceInterchange(basis), base, source_acts, task_logits, positions=subject +) +print(task_logits(base).argmax(-1), logits.argmax(-1)) # tensor([0, 1]) tensor([1, 0]) +``` + +`task_logits` runs the model in one forward pass and reads the task out of its output, so the +tools do not assume next-token prediction; the hook must fire exactly once in that pass. +Activations at the hook must be `[batch, d_model]`, or `[batch, pos, d_model]` with one position +index per example. A real experiment fits the basis on held-out pairs, not on the pairs it +evaluates. + +`collect_activations` and `intervened_logits` run in batches of `batch_size` without gradients and +return CPU tensors, so memory does not grow with the number of pairs. `run_with_intervention` is +the same intervention as one differentiable forward pass, for training an intervention's +parameters. + +## Bases and baselines + +- `mean_difference_basis(activations, labels)`: the span of the centred class means; for two + classes, the difference-in-means direction. +- `probe_basis(weights)`: the directions of a linear probe's weights. Undo any input scaling + before passing them. +- `random_basis(d_model, rank, seed)`: a seeded random orthonormal basis. + +Run every basis on the same pairs. A basis whose IIA does not clearly exceed a random basis of +the same rank is not evidence that the subspace carries the variable. + +## Placebo + +`NormMatchedPlacebo(reference, basis)` moves each base exactly as far as `reference` does, but +inside another subspace, usually `random_basis` with the same rank. If the placebo also changes the +output, part of the effect comes from displacing the activation by that much, not from the +aligned subspace. + +## Interpretation + +An IIA measures what an intervention does under the chosen pairs, hook point, positions and +basis. It does not by itself show that the model uses the variable: subspace interventions can +activate pathways the model does not use on its own inputs (Makelov et al., 2024). Report the +IIA beside the placebo and the baselines. + +## References + +- Geiger, Lu, Icard and Potts (2021). Causal Abstractions of Neural Networks. NeurIPS. +- Makelov, Lange and Nanda (2024). Is This the Subspace You Are Looking for? An Interpretability + Illusion for Subspace Activation Patching. ICLR. diff --git a/docs/source/index.md b/docs/source/index.md index 3a513e97e..30b9c97b4 100644 --- a/docs/source/index.md +++ b/docs/source/index.md @@ -67,6 +67,7 @@ generated/demos/Jacobian_Lens_Coordinate_Patch_Benchmark_Demo content/backward_lens content/debugging_numerical_divergence content/sparse_probing +content/interchange_interventions content/svd_circuits generated/demos/SVD_Circuits_Demo generated/demos/Main_Demo diff --git a/tests/integration/model_bridge/test_interchange.py b/tests/integration/model_bridge/test_interchange.py new file mode 100644 index 000000000..5c4b79d81 --- /dev/null +++ b/tests/integration/model_bridge/test_interchange.py @@ -0,0 +1,64 @@ +"""Real-Bridge integration checks for interchange interventions.""" + +import torch + +from transformer_lens.tools.analysis.interchange import ( + NormMatchedPlacebo, + SubspaceInterchange, + collect_activations, + intervened_logits, + random_basis, +) + +PROMPTS = ["The key near the cabinet", "The keys near the cabinet"] + + +def _setup(bridge): + tokens = bridge.to_tokens(PROMPTS) + answers = [bridge.to_single_token(" is"), bridge.to_single_token(" are")] + last = torch.full((len(PROMPTS),), tokens.shape[1] - 1) + return tokens, answers, last, (lambda toks: bridge(toks)[:, -1, answers]) + + +def test_collected_activations_match_the_cache(gpt2_bridge): + tokens, _, last, task_logits = _setup(gpt2_bridge) + acts = collect_activations( + gpt2_bridge, "blocks.5.hook_out", tokens, task_logits, positions=last + ) + _, cache = gpt2_bridge.run_with_cache(tokens) + assert torch.allclose(acts, cache["blocks.5.hook_out"][:, -1].cpu(), atol=1e-5) + + +def test_a_full_interchange_at_the_last_block_reproduces_the_source_logits(gpt2_bridge): + tokens, _, last, task_logits = _setup(gpt2_bridge) + hook = f"blocks.{gpt2_bridge.cfg.n_layers - 1}.hook_out" + swapped = tokens.flip(0) # each prompt takes the other's activation + source_acts = collect_activations(gpt2_bridge, hook, swapped, task_logits, positions=last) + out = intervened_logits( + gpt2_bridge, + hook, + SubspaceInterchange(torch.eye(gpt2_bridge.cfg.d_model)), + tokens, + source_acts, + task_logits, + positions=last, + batch_size=1, + ) + with torch.no_grad(): + expected = task_logits(swapped) + assert torch.allclose(out, expected, atol=1e-4) + # Singular and plural prompts prefer different verbs, so the swap flips both predictions + assert torch.equal(out.argmax(-1), task_logits(tokens).argmax(-1).flip(0)) + + +def test_the_placebo_matches_the_reference_displacement_on_real_activations(gpt2_bridge): + tokens, _, last, task_logits = _setup(gpt2_bridge) + acts = collect_activations( + gpt2_bridge, "blocks.3.hook_out", tokens, task_logits, positions=last + ) + d_model = gpt2_bridge.cfg.d_model + reference = SubspaceInterchange(random_basis(d_model, 4, seed=0)) + placebo = NormMatchedPlacebo(reference, random_basis(d_model, 4, seed=1)) + base, source = acts, acts.flip(0) + moved = (placebo(base, source) - base).norm(dim=-1) + assert torch.allclose(moved, (reference(base, source) - base).norm(dim=-1), rtol=1e-4) diff --git a/tests/unit/tools/test_interchange.py b/tests/unit/tools/test_interchange.py new file mode 100644 index 000000000..6b016c122 --- /dev/null +++ b/tests/unit/tools/test_interchange.py @@ -0,0 +1,229 @@ +"""Unit tests for interchange interventions on linear subspaces.""" + +import pytest +import torch +from torch import nn + +from transformer_lens.hook_points import HookPoint +from transformer_lens.HookedRootModule import HookedRootModule +from transformer_lens.tools.analysis import SubspaceInterchange as exported_interchange +from transformer_lens.tools.analysis.interchange import ( + NormMatchedPlacebo, + SubspaceInterchange, + collect_activations, + intervened_logits, + mean_difference_basis, + probe_basis, + random_basis, + run_with_intervention, +) + +D_MODEL = 6 + + +class _Toy(HookedRootModule): + """Linear in, hook, linear out; with ``seq``, the output reads only the last position.""" + + def __init__(self, seq: bool = False): + super().__init__() + torch.manual_seed(0) + self.inp = nn.Linear(D_MODEL, D_MODEL) + self.out = nn.Linear(D_MODEL, 2) + self.hook_mid = HookPoint() + self.seq = seq + self.setup() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.hook_mid(self.inp(x)) + return self.out(h[:, -1] if self.seq else h) + + +def _orthonormal(basis: torch.Tensor) -> bool: + return torch.allclose(basis.T @ basis, torch.eye(basis.shape[1]), atol=1e-5) + + +def test_public_analysis_exports(): + assert exported_interchange is SubspaceInterchange + + +def test_interchange_takes_the_source_inside_the_subspace_and_keeps_the_base_outside(): + base, source = torch.randn(4, D_MODEL), torch.randn(4, D_MODEL) + B = random_basis(D_MODEL, 2, seed=0) + out = SubspaceInterchange(B)(base, source) + assert torch.allclose(out @ B, source @ B, atol=1e-5) + complement = torch.eye(D_MODEL) - B @ B.T + assert torch.allclose(out @ complement, base @ complement, atol=1e-5) + assert torch.allclose(SubspaceInterchange(torch.eye(D_MODEL))(base, source), source, atol=1e-5) + assert torch.allclose(SubspaceInterchange(B, strength=0.0)(base, source), base) + + +def test_random_basis_is_orthonormal_and_seeded(): + assert _orthonormal(random_basis(D_MODEL, 3, seed=1)) + assert torch.equal(random_basis(D_MODEL, 3, seed=1), random_basis(D_MODEL, 3, seed=1)) + assert not torch.equal(random_basis(D_MODEL, 3, seed=1), random_basis(D_MODEL, 3, seed=2)) + + +def test_mean_difference_basis_spans_the_centred_class_means(): + acts = torch.randn(200, D_MODEL) + labels = [i % 2 for i in range(200)] + acts[torch.tensor(labels) == 1, 0] += 5.0 + direction = mean_difference_basis(acts, labels) + assert direction.shape == (D_MODEL, 1) + expected = acts[1::2].mean(0) - acts[0::2].mean(0) + cosine = torch.nn.functional.cosine_similarity(direction[:, 0], expected, dim=0) + assert cosine.abs() > 1 - 1e-5 + + three = mean_difference_basis(torch.randn(90, D_MODEL), [i % 3 for i in range(90)]) + assert three.shape == (D_MODEL, 2) and _orthonormal(three) + with pytest.raises(ValueError, match="two classes"): + mean_difference_basis(torch.randn(10, D_MODEL), [0] * 10) + + +def test_probe_basis_follows_the_weights(): + w = torch.randn(1, D_MODEL) + one = probe_basis(w) + assert torch.allclose(one[:, 0].abs(), (w[0] / w[0].norm()).abs(), atol=1e-5) + rows = torch.randn(3, D_MODEL) + centred = probe_basis(rows) + assert centred.shape == (D_MODEL, 2) and _orthonormal(centred) + # The centred rows lie in the returned span + residual = (rows - rows.mean(0)).T - centred @ (centred.T @ (rows - rows.mean(0)).T) + assert residual.abs().max() < 1e-4 + + +def test_placebo_moves_as_far_as_the_reference_inside_its_own_subspace(): + base, source = torch.randn(8, D_MODEL), torch.randn(8, D_MODEL) + reference = SubspaceInterchange(random_basis(D_MODEL, 2, seed=0)) + placebo_basis = random_basis(D_MODEL, 2, seed=1) + moved = NormMatchedPlacebo(reference, placebo_basis)(base, source) + expected = (reference(base, source) - base).norm(dim=-1) + assert torch.allclose((moved - base).norm(dim=-1), expected, atol=1e-5) + outside = (moved - base) - (moved - base) @ placebo_basis @ placebo_basis.T + assert outside.abs().max() < 1e-5 + + +def test_a_full_interchange_reproduces_the_source_output(): + model = _Toy() + base, source = torch.randn(5, D_MODEL), torch.randn(5, D_MODEL) + source_acts = collect_activations(model, "hook_mid", source, model) + out = run_with_intervention( + model, "hook_mid", SubspaceInterchange(torch.eye(D_MODEL)), base, source_acts, model + ) + assert torch.allclose(out, model(source), atol=1e-5) + + +class _SumToy(_Toy): + """Reads every position, so each row's write shows in its output.""" + + def __init__(self): + super().__init__(seq=True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.out(self.hook_mid(self.inp(x)).sum(1)) + + +def test_positions_select_one_vector_per_example(): + rows = torch.tensor([0, 2, 1, 2]) + base, source = torch.randn(4, 3, D_MODEL), torch.randn(4, 3, D_MODEL) + full = SubspaceInterchange(torch.eye(D_MODEL)) + + model = _Toy(seq=True) # reads only the last position + _, cache = model.run_with_cache(source) + acts = collect_activations(model, "hook_mid", source, model, positions=rows, batch_size=3) + assert torch.allclose(acts, cache["hook_mid"][torch.arange(4), rows]) + out = run_with_intervention(model, "hook_mid", full, base, acts, model, positions=rows) + at_last = rows == 2 + assert torch.allclose(out[at_last], model(source)[at_last], atol=1e-5) + assert torch.allclose(out[~at_last], model(base)[~at_last], atol=1e-6) + + summed = _SumToy() + source_acts = collect_activations(summed, "hook_mid", source, summed, positions=rows) + out = run_with_intervention(summed, "hook_mid", full, base, source_acts, summed, positions=rows) + with torch.no_grad(): + expected = summed.inp(base) + expected[torch.arange(4), rows] = summed.inp(source)[torch.arange(4), rows] + assert torch.allclose(out, summed.out(expected.sum(1)), atol=1e-5) + + with pytest.raises(ValueError, match="one position index per example"): + collect_activations(model, "hook_mid", source, model) + + +class _Pattern(HookedRootModule): + """A 4-D hook, shaped like an attention pattern: [batch, head, query, key].""" + + def __init__(self): + super().__init__() + self.hook_pattern = HookPoint() + self.setup() + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.hook_pattern(x).flatten(1)[:, :2] + + +def test_positions_must_match_the_activation(): + pattern = _Pattern() + x, at = torch.randn(4, 2, 3, 3), torch.zeros(4, dtype=torch.long) + with pytest.raises(ValueError, match="must be \\[batch, pos, d_model\\]"): + collect_activations(pattern, "hook_pattern", x, pattern, positions=at) + with pytest.raises(ValueError, match="must be \\[batch, pos, d_model\\]"): + run_with_intervention( + pattern, + "hook_pattern", + SubspaceInterchange(torch.eye(3)), + x, + torch.randn(4, 3), + pattern, + positions=at, + ) + model = _Toy(seq=True) + with pytest.raises(ValueError, match="one index per example"): + collect_activations(model, "hook_mid", torch.randn(4, 3, D_MODEL), model, positions=at[:3]) + + +def test_the_intervention_is_differentiable(): + model = _Toy() + basis = random_basis(D_MODEL, 2, seed=0).requires_grad_(True) + base, source = torch.randn(4, D_MODEL), torch.randn(4, D_MODEL) + source_acts = collect_activations(model, "hook_mid", source, model) + out = run_with_intervention( + model, "hook_mid", SubspaceInterchange(basis), base, source_acts, model + ) + out.sum().backward() + assert basis.grad is not None and basis.grad.abs().sum() > 0 + + +def test_a_hook_must_fire_exactly_once(): + class _Twice(_Toy): + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.hook_mid(self.inp(x)) + return self.out(self.hook_mid(h)) + + class _Unused(_Toy): + def __init__(self): + super().__init__() + self.hook_unused = HookPoint() + self.setup() + + x, source, full = ( + torch.randn(2, D_MODEL), + torch.randn(2, D_MODEL), + SubspaceInterchange(torch.eye(D_MODEL)), + ) + for model, hook, times in ((_Twice(), "hook_mid", 2), (_Unused(), "hook_unused", 0)): + with pytest.raises(RuntimeError, match=f"fired {times} times"): + collect_activations(model, hook, x, model) + with pytest.raises(RuntimeError, match=f"fired {times} times"): + run_with_intervention(model, hook, full, x, source, model) + + +def test_intervened_logits_matches_one_pass_in_bounded_batches(): + model = _Toy() + base, source = torch.randn(10, D_MODEL), torch.randn(10, D_MODEL) + acts = collect_activations(model, "hook_mid", source, model) + intervention = SubspaceInterchange(random_basis(D_MODEL, 2, seed=0).requires_grad_(True)) + one = run_with_intervention(model, "hook_mid", intervention, base, acts, model) + batched = intervened_logits(model, "hook_mid", intervention, base, acts, model, batch_size=3) + assert torch.allclose(batched, one.detach(), atol=1e-6) + assert not batched.requires_grad and batched.device.type == "cpu" + with pytest.raises(ValueError, match="one row per base"): + intervened_logits(model, "hook_mid", intervention, base, acts[:-1], model) diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py index cedf22f57..c624c13b5 100644 --- a/transformer_lens/tools/analysis/__init__.py +++ b/transformer_lens/tools/analysis/__init__.py @@ -17,6 +17,10 @@ layers, or attention heads. - direct_path_patching: Direct path patching for head-to-head circuit analysis. + - interchange: Interchange interventions on linear subspaces: set a base's + activation inside a subspace to its value on a source input, with + difference-of-means, probe and random bases, a norm-matched placebo, and + batched, gradient-free evaluation. - jacobian_lens: The Jacobian lens (J-lens) -- per-layer causal transport to the output vocabulary basis, with loading of published lens artifacts, native fitting, readouts, interventions, J-space sparse decomposition, and @@ -73,6 +77,16 @@ get_act_patch_direct_path, get_act_patch_direct_path_all_sources, ) +from transformer_lens.tools.analysis.interchange import ( + NormMatchedPlacebo, + SubspaceInterchange, + collect_activations, + intervened_logits, + mean_difference_basis, + probe_basis, + random_basis, + run_with_intervention, +) from transformer_lens.tools.analysis.jacobian_lens import ( JacobianLens, JacobianLensReadout, @@ -172,6 +186,7 @@ "LogitReadout", "LogitSignature", "Node", + "NormMatchedPlacebo", "PatchResult", "ProjectedFactor", "ProjectionKernelResult", @@ -185,10 +200,12 @@ "SparseProbeResult", "SparseProbeSweep", "SubspaceBasis", + "SubspaceInterchange", "VocabularyRanking", "WeightLayout", "attention_head_subspace_affinity", "attribution_patch", + "collect_activations", "decompose_head", "direct_logit_attribution", "estimate_occupancy", @@ -197,16 +214,21 @@ "get_act_patch_direct_path", "get_act_patch_direct_path_all_sources", "get_sparse_decomposition", + "intervened_logits", "load_artifact", "logit_lens", "logit_readout", "logit_signature", + "mean_difference_basis", "orthonormal_subspace", "patch_along_directions", + "probe_basis", "project_activations", "projection_kernel", + "random_basis", "random_projection_kernel_moments", "run_causal_swap_benchmark", + "run_with_intervention", "select_displacement_matched_control_token", "solve_coordinate_patch", "solve_coordinate_patch_positions", diff --git a/transformer_lens/tools/analysis/interchange.py b/transformer_lens/tools/analysis/interchange.py new file mode 100644 index 000000000..bd7cc2bbb --- /dev/null +++ b/transformer_lens/tools/analysis/interchange.py @@ -0,0 +1,320 @@ +"""Interchange interventions on linear subspaces. + +An interchange intervention runs the model on a *base* input while setting the base's activation +at one hook point, inside a subspace, to its value on a *source* input:: + + h' = h_base + strength * (h_source - h_base) @ B @ B.T + +``B`` is an orthonormal ``[d_model, rank]`` basis, the convention of +:func:`~transformer_lens.tools.analysis.projection_kernel.orthonormal_subspace`. The base keeps its +own coordinates outside the subspace. If the model then produces the output a high-level causal +model predicts for the base with one variable taken from the source, the subspace carries that +variable at that hook point (Geiger et al., 2021). + +The module provides the intervention, fixed-direction bases to intervene on (the difference of +class means, a linear probe's weights, a random basis), and a norm-matched random-subspace placebo. +An interchange accuracy is not evidence on its own: compare it with the placebo and the baselines. + +Works with any model exposing the hook API (``TransformerBridge`` or a ``HookedRootModule``). +Activations at the hook must be ``[batch, d_model]``, or ``[batch, pos, d_model]`` with one +position index per example. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Callable, Hashable, Iterator, Optional, Sequence, Union + +import numpy as np +import torch +from jaxtyping import Float, Int + +from transformer_lens.tools.analysis.projection_kernel import orthonormal_subspace + +TaskLogits = Callable[[torch.Tensor], Float[torch.Tensor, "batch n_classes"]] +"""Runs the model on a batch of inputs in one forward pass and returns task logits, e.g. +``lambda tokens: model(tokens)[:, -1, answer_token_ids]``.""" + +Intervention = Callable[[torch.Tensor, torch.Tensor], torch.Tensor] +"""Maps (base activations, source activations), both ``[batch, d_model]``, to new base activations.""" + + +def random_basis(d_model: int, rank: int, seed: int) -> Float[torch.Tensor, "d_model rank"]: + """A seeded random orthonormal basis: the baseline any aligned basis should beat. + + Args: + d_model: Dimension of the activation. + rank: Number of basis directions. + seed: Seed for the Gaussian matrix the basis is drawn from. + + Returns: + Orthonormal basis, ``[d_model, rank]``. + """ + generator = torch.Generator().manual_seed(seed) + return orthonormal_subspace(torch.randn(d_model, rank, generator=generator), rank=rank).basis + + +def mean_difference_basis( + activations: Float[torch.Tensor, "batch d_model"], labels: Sequence[Hashable] +) -> Float[torch.Tensor, "d_model rank"]: + """The span of the centred class means: for two classes, the difference-in-means direction. + + Args: + activations: Activations, one row per example. + labels: The variable's value for each example; at least two distinct values. + + Returns: + Orthonormal basis of rank ``n_classes - 1``, ``[d_model, n_classes - 1]``. + + Raises: + ValueError: If fewer than two classes are present. + """ + values = np.asarray(labels) + classes = np.unique(values) + if len(classes) < 2: + raise ValueError("mean_difference_basis needs at least two classes") + acts = activations.float().cpu() + means = torch.stack([acts[torch.as_tensor(values == c)].mean(0) for c in classes]) + return orthonormal_subspace((means - means.mean(0)).T, rank=len(classes) - 1).basis + + +def probe_basis( + weights: Float[torch.Tensor, "n_rows d_model"], +) -> Float[torch.Tensor, "d_model rank"]: + """The directions of a linear probe's weights, in activation space. + + Undo any input scaling first: for a probe fit on standardised inputs, divide each weight by its + feature's standard deviation. + + Args: + weights: Probe weights; one row for a binary probe, one row per class otherwise. + + Returns: + Orthonormal basis: the weight direction for one row, the span of the centred rows (rank + ``n_rows - 1``) for several. + """ + w = torch.as_tensor(weights, dtype=torch.float32) + if w.shape[0] == 1: + return orthonormal_subspace(w.T, rank=1).basis + return orthonormal_subspace((w - w.mean(0)).T, rank=w.shape[0] - 1).basis + + +@dataclass(frozen=True) +class SubspaceInterchange: + """Set the base's coordinates inside ``basis`` to the source's. + + Attributes: + basis: Orthonormal basis, ``[d_model, rank]``. + strength: 1 interchanges the coordinates; other values scale the move, 0 leaves the base. + """ + + basis: torch.Tensor + strength: float = 1.0 + + def __call__(self, base: torch.Tensor, source: torch.Tensor) -> torch.Tensor: + B = self.basis.to(base.device, base.dtype) + return base + self.strength * ((source - base) @ B) @ B.T + + +@dataclass(frozen=True) +class NormMatchedPlacebo: + """Move each base exactly as far as ``reference`` does, but inside a random subspace. + + An effect that the placebo also produces comes from displacing the activation by that much, + not from the aligned subspace. + + Attributes: + reference: The intervention whose displacement is matched. + basis: Orthonormal basis of the placebo subspace, usually from :func:`random_basis` with + the reference's rank. + """ + + reference: SubspaceInterchange + basis: torch.Tensor + + def __call__(self, base: torch.Tensor, source: torch.Tensor) -> torch.Tensor: + distance = (self.reference(base, source) - base).norm(dim=-1, keepdim=True) + B = self.basis.to(base.device, base.dtype) + move = ((source - base) @ B) @ B.T + return base + move * (distance / move.norm(dim=-1, keepdim=True).clamp_min(1e-8)) + + +def _select(act: torch.Tensor, positions: Optional[torch.Tensor]) -> torch.Tensor: + if positions is None: + if act.ndim != 2: + raise ValueError( + f"activation has shape {tuple(act.shape)}; pass one position index per example" + ) + return act + if act.ndim != 3: + raise ValueError( + f"with positions, the activation must be [batch, pos, d_model]; " + f"got shape {tuple(act.shape)}" + ) + if positions.shape != (act.shape[0],): + raise ValueError( + f"positions must hold one index per example, shape ({act.shape[0]},); " + f"got {tuple(positions.shape)}" + ) + return act[torch.arange(act.shape[0], device=act.device), positions.to(act.device)] + + +def _write(act: torch.Tensor, positions: Optional[torch.Tensor], new: torch.Tensor) -> torch.Tensor: + if positions is None: + return new + out = act.clone() + out[torch.arange(act.shape[0], device=act.device), positions.to(act.device)] = new + return out + + +def _chunks(n: int, size: int) -> Iterator[slice]: + for start in range(0, n, size): + yield slice(start, min(start + size, n)) + + +def _at( + positions: Optional[torch.Tensor], idx: Union[slice, torch.Tensor] +) -> Optional[torch.Tensor]: + return None if positions is None else positions[idx] + + +@torch.no_grad() +def collect_activations( + model: Any, + hook_name: str, + inputs: torch.Tensor, + task_logits: TaskLogits, + *, + positions: Optional[Int[torch.Tensor, "batch"]] = None, + batch_size: int = 256, +) -> Float[torch.Tensor, "batch d_model"]: + """Activations at ``hook_name``, at one position per example if given, returned on the CPU. + + Args: + model: A model exposing the hook API. + hook_name: The hook point to read. + inputs: Model inputs, one row per example. + task_logits: Runs the model on a batch of inputs. + positions: For ``[batch, pos, d_model]`` activations, the position to read per example. + batch_size: Rows per forward pass. + + Returns: + Activations, ``[batch, d_model]``. + + Raises: + ValueError: If the activation does not match the positions given (or not given). + RuntimeError: If the hook does not fire exactly once per forward pass. + """ + out = [] + for idx in _chunks(len(inputs), batch_size): + store: list[torch.Tensor] = [] + + def grab( + act: torch.Tensor, hook: Any, _positions: Optional[torch.Tensor] = _at(positions, idx) + ) -> None: + store.append(_select(act, _positions).detach().cpu()) + + with model.hooks(fwd_hooks=[(hook_name, grab)]): + task_logits(inputs[idx]) + if len(store) != 1: + raise RuntimeError(f"{hook_name} fired {len(store)} times in one forward pass") + out.append(store[0]) + return torch.cat(out) + + +def run_with_intervention( + model: Any, + hook_name: str, + intervention: Intervention, + base_inputs: torch.Tensor, + source_activations: Float[torch.Tensor, "batch d_model"], + task_logits: TaskLogits, + *, + positions: Optional[Int[torch.Tensor, "batch"]] = None, +) -> Float[torch.Tensor, "batch n_classes"]: + """Task logits for ``base_inputs`` with ``intervention`` applied at ``hook_name``, in one pass. + + Differentiable with respect to any parameters of ``intervention``, for training an + intervention. To evaluate one over many pairs in bounded memory, use :func:`intervened_logits`. + + Args: + model: A model exposing the hook API. + hook_name: The hook point to intervene on. + intervention: Maps base and source activations to new base activations, e.g. a + :class:`SubspaceInterchange`. + base_inputs: Model inputs to run. + source_activations: Activations from the source inputs at the same hook point, one row per + base, e.g. from :func:`collect_activations`. + task_logits: Runs the model on a batch of inputs. + positions: For ``[batch, pos, d_model]`` activations, the position to intervene on per base. + + Returns: + Task logits under the intervention. + + Raises: + ValueError: If the activation does not match the positions given (or not given). + RuntimeError: If the hook does not fire exactly once in the forward pass. + """ + calls = 0 + + def hook(act: torch.Tensor, hook: Any) -> torch.Tensor: + nonlocal calls + calls += 1 + selected = _select(act, positions) + source = source_activations.to(act.device, act.dtype) + return _write(act, positions, intervention(selected, source)) + + with model.hooks(fwd_hooks=[(hook_name, hook)]): + logits = task_logits(base_inputs) + # A second firing would reuse source activations collected for the first + if calls != 1: + raise RuntimeError(f"{hook_name} fired {calls} times in one forward pass") + return logits + + +@torch.no_grad() +def intervened_logits( + model: Any, + hook_name: str, + intervention: Intervention, + base_inputs: torch.Tensor, + source_activations: Float[torch.Tensor, "batch d_model"], + task_logits: TaskLogits, + *, + positions: Optional[Int[torch.Tensor, "batch"]] = None, + batch_size: int = 256, +) -> Float[torch.Tensor, "batch n_classes"]: + """Task logits under ``intervention``, in batches without gradients, returned on the CPU. + + The evaluation counterpart of :func:`run_with_intervention`, with the same arguments and + checks: memory is bounded by ``batch_size`` rather than by the number of pairs. + + Args: + model: A model exposing the hook API. + hook_name: The hook point to intervene on. + intervention: Maps base and source activations to new base activations. + base_inputs: Model inputs to run. + source_activations: Activations from the source inputs at the same hook point, one row per + base. + task_logits: Runs the model on a batch of inputs. + positions: For ``[batch, pos, d_model]`` activations, the position to intervene on per base. + batch_size: Rows per forward pass. + + Returns: + Task logits under the intervention, ``[batch, n_classes]``. + """ + if len(source_activations) != len(base_inputs): + raise ValueError("source_activations must have one row per base input") + out = [] + for idx in _chunks(len(base_inputs), batch_size): + logits = run_with_intervention( + model, + hook_name, + intervention, + base_inputs[idx], + source_activations[idx], + task_logits, + positions=_at(positions, idx), + ) + out.append(logits.detach().cpu()) + return torch.cat(out)