Add subspace interchange interventions

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. |
 

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.

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

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)