Add distributed alignment search for interchange interventions

docs/source/content/analysis_tools.md

@@ -20,7 +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. |
+| 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 (fixed, or learned by DAS from pairs labelled by a high-level model), 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. |
 
@@ -43,7 +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. |
+| 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. DAS adds a forward and backward pass per batch per epoch; `micro_batch_size` bounds its memory. |
 | 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

@@ -33,25 +33,39 @@ from transformer_lens.tools.analysis import (
 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])
+
+def pairs(nouns):
+    """Bases and sources differ only in the subject's number; the target is the source's verb."""
+    singular = bridge.to_tokens([f"The {n} near the cabinet" for n in nouns])
+    plural = bridge.to_tokens([f"The {n}s near the cabinet" for n in nouns])
+    base, source = torch.cat([singular, plural]), torch.cat([plural, singular])
+    targets = (torch.arange(len(base)) < len(nouns)).long()  # " are" when the source is plural
+    subject = torch.full((len(base),), 2)  # the subject token, after the BOS token and "The"
+    return base, source, targets, subject
+
+
+fit_base, fit_source, fit_targets, fit_at = pairs(["key", "dog", "cat", "book", "car", "boy"])
+test_base, test_source, test_targets, test_at = pairs(["girl", "bird", "tree", "house"])
+
+# Fit the basis on the fit pairs' sources, labelled by subject number ...
+fit_acts = collect_activations(bridge, hook, fit_source, task_logits, positions=fit_at)
+basis = mean_difference_basis(fit_acts, labels=fit_targets.tolist())
+
+# ... and measure the interchange accuracy (IIA) on held-out pairs
+test_acts = collect_activations(bridge, hook, test_source, task_logits, positions=test_at)
 logits = intervened_logits(
-    bridge, hook, SubspaceInterchange(basis), base, source_acts, task_logits, positions=subject
+    bridge, hook, SubspaceInterchange(basis), test_base, test_acts, task_logits, positions=test_at
 )
-print(task_logits(base).argmax(-1), logits.argmax(-1))  # tensor([0, 1]) tensor([1, 0])
+print((logits.argmax(-1) == test_targets).float().mean())  # tensor(1.)
 ```
 
 `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.
+index per example. Fit a basis on one set of pairs and evaluate it on pairs it was not fit on, as
+here.
 
 `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
@@ -69,6 +83,36 @@ parameters.
 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.
 
+## Learn the subspace with DAS
+
+Distributed alignment search (DAS; Geiger et al., 2024) learns the basis instead of fixing it.
+`fit_das` trains an orthonormal basis so that the interchange makes the model produce, for each
+base/source pair, the class the high-level model predicts:
+
+```python
+from transformer_lens.tools.analysis import fit_das
+
+result = fit_das(
+    bridge, hook, fit_base, fit_source, fit_targets, task_logits, rank=1,
+    base_positions=fit_at, source_positions=fit_at, epochs=10,
+)
+logits = intervened_logits(
+    bridge, hook, SubspaceInterchange(result.subspace), test_base, test_acts, task_logits,
+    positions=test_at,
+)
+print((logits.argmax(-1) == test_targets).float().mean())  # tensor(1.)
+```
+
+`targets[i]` is the high-level model's output for base `i` with the variable taken from source
+`i`. Fit on one set of pairs and evaluate the IIA on pairs the basis was not trained on. Include
+training pairs where several variables differ between base and source, labelled by the high-level
+model's interchange: otherwise the basis can learn whatever the training pairs happen to move
+together. The model's parameters are frozen during training and restored afterwards;
+`micro_batch_size` accumulates gradients over smaller parts of each batch to save memory.
+
+DAS is expressive enough to fit arbitrary label mappings and to find effects in randomly
+initialised models (Arora et al., 2024), so its IIA needs controls trained the same way.
+
 ## Placebo
 
 `NormMatchedPlacebo(reference, basis)` moves each base exactly as far as `reference` does, but
@@ -86,5 +130,9 @@ IIA beside the placebo and the baselines.
 ## References
 
 - Geiger, Lu, Icard and Potts (2021). Causal Abstractions of Neural Networks. NeurIPS.
+- Arora, Jurafsky and Potts (2024). CausalGym: Benchmarking causal interpretability methods on
+  linguistic tasks. ACL.
+- Geiger, Wu, Potts, Icard and Goodman (2024). Finding Alignments Between Interpretable Causal
+  Variables and Distributed Neural Representations. CLeaR.
 - Makelov, Lange and Nanda (2024). Is This the Subspace You Are Looking for? An Interpretability
   Illusion for Subspace Activation Patching. ICLR.

tests/integration/model_bridge/test_interchange.py

@@ -6,6 +6,7 @@
     NormMatchedPlacebo,
     SubspaceInterchange,
     collect_activations,
+    fit_das,
     intervened_logits,
     random_basis,
 )
@@ -62,3 +63,59 @@ def test_the_placebo_matches_the_reference_displacement_on_real_activations(gpt2
     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)
+
+
+NOUNS = "key dog cat book car boy girl bird tree house road table door letter song plan".split()
+
+
+def _number_pairs(bridge, nouns):
+    """Base and source differ only in the subject's number; the target is the source's verb."""
+    singular = bridge.to_tokens([f"The {n} near the cabinet" for n in nouns])
+    plural = bridge.to_tokens([f"The {n}s near the cabinet" for n in nouns])
+    # BOS, "The", noun, "near", "the", "cabinet": every noun and its plural is one token
+    assert singular.shape[1] == plural.shape[1] == 6
+    base = torch.cat([singular, plural])
+    source = torch.cat([plural, singular])
+    targets = torch.cat([torch.ones(len(nouns)), torch.zeros(len(nouns))]).long()
+    return base, source, targets
+
+
+def test_das_finds_subject_number_on_held_out_nouns(gpt2_bridge):
+    answers = [gpt2_bridge.to_single_token(" is"), gpt2_bridge.to_single_token(" are")]
+    task_logits = lambda toks: gpt2_bridge(toks)[:, -1, answers]
+    hook = "blocks.1.hook_out"
+    base, source, targets = _number_pairs(gpt2_bridge, NOUNS[:10])
+    subject = torch.full((len(base),), 2)  # after the BOS token and "The"
+    result = fit_das(
+        gpt2_bridge,
+        hook,
+        base,
+        source,
+        targets,
+        task_logits,
+        rank=1,
+        base_positions=subject,
+        source_positions=subject,
+        epochs=10,
+        batch_size=20,
+        lr=1e-2,
+    )
+
+    base, source, targets = _number_pairs(gpt2_bridge, NOUNS[10:])
+    subject = torch.full((len(base),), 2)
+    source_acts = collect_activations(gpt2_bridge, hook, source, task_logits, positions=subject)
+
+    def accuracy(basis):
+        logits = intervened_logits(
+            gpt2_bridge,
+            hook,
+            SubspaceInterchange(basis),
+            base,
+            source_acts,
+            task_logits,
+            positions=subject,
+        )
+        return (logits.argmax(-1) == targets).float().mean().item()
+
+    assert accuracy(result.subspace) >= 0.9
+    assert accuracy(random_basis(gpt2_bridge.cfg.d_model, 1, seed=0)) <= 0.2