Add trained controls for distributed alignment search

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 (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. |
+| 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 a learned subspace with its trained controls (arbitrary labels, an untrained model), 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. |
 

docs/source/content/interchange_interventions.md

@@ -36,18 +36,19 @@ task_logits = lambda tokens: bridge(tokens)[:, -1, answers]
 hook = "blocks.1.hook_out"
 
 
-def pairs(nouns):
+def pairs(nouns, place):
     """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])
+    singular = bridge.to_tokens([f"The {n} {place}" for n in nouns])
+    plural = bridge.to_tokens([f"The {n}s {place}" 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"])
+nouns = ["key", "dog", "cat", "book", "car", "boy", "girl", "bird"]
+fit_base, fit_source, fit_targets, fit_at = pairs(nouns, "near the cabinet")
+test_base, test_source, test_targets, test_at = pairs(nouns, "behind the table")
 
 # 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)
@@ -65,7 +66,9 @@ print((logits.argmax(-1) == test_targets).float().mean())  # tensor(1.)
 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. Fit a basis on one set of pairs and evaluate it on pairs it was not fit on, as
-here.
+here: the held-out pairs put the same nouns in a new frame. Reusing the nouns gives the
+arbitrary-label control below keys it can learn; holding out nouns as well is a stronger test of
+the real alignment.
 
 `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
@@ -113,6 +116,59 @@ together. The model's parameters are frozen during training and restored afterwa
 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.
 
+## Trained controls
+
+Train each control with the same pairs, hook point, rank, budget and seeds as the real alignment,
+and evaluate it on the same held-out pairs:
+
+- **Arbitrary labels.** `arbitrary_targets(keys, n_classes, seed)` gives each key a fixed
+  pseudo-random class. DAS trained on these shows how much IIA the method's expressivity alone
+  produces. Key the labels by something that recurs between the fit and held-out pairs, such as
+  the source's subject word: on keys the fit pairs never contained, the labels are fresh coin
+  flips, every subspace scores chance, and the control can reject nothing.
+- **An untrained model.** `random_init_bridge(model_name, seed=...)` builds the architecture with
+  its HuggingFace initialisation and no training; `reinitialized(model, seed)` does the same for a
+  `HookedRootModule`. An effect DAS finds there does not come from anything the model learned.
+
+```python
+from transformer_lens.tools.analysis import arbitrary_targets, random_init_bridge
+
+# Arbitrary labels: a fixed pseudo-random class per source subject word, which recurs between the
+# fit and held-out pairs
+labels = lambda sources: arbitrary_targets([bridge.to_string(s[2]) for s in sources], 2, seed=0)
+arbitrary = fit_das(
+    bridge, hook, fit_base, fit_source, labels(fit_source), task_logits, rank=1,
+    base_positions=fit_at, source_positions=fit_at, epochs=10,
+)
+logits = intervened_logits(
+    bridge, hook, SubspaceInterchange(arbitrary.subspace), test_base, test_acts, task_logits,
+    positions=test_at,
+)
+print((logits.argmax(-1) == labels(test_source)).float().mean())  # tensor(0.6250)
+
+# An untrained model: the same pairs and targets, its own activations and readout
+untrained = random_init_bridge("gpt2", seed=0)
+untrained_logits = lambda tokens: untrained(tokens)[:, -1, answers]
+random_init = fit_das(
+    untrained, hook, fit_base, fit_source, fit_targets, untrained_logits, rank=1,
+    base_positions=fit_at, source_positions=fit_at, epochs=10,
+)
+untrained_acts = collect_activations(
+    untrained, hook, test_source, untrained_logits, positions=test_at
+)
+logits = intervened_logits(
+    untrained, hook, SubspaceInterchange(random_init.subspace), test_base, untrained_acts,
+    untrained_logits, positions=test_at,
+)
+print((logits.argmax(-1) == test_targets).float().mean())  # tensor(0.8125)
+```
+
+Here the arbitrary labels reach 0.62 and the untrained model 0.81, against 1.0 for the real
+alignment: DAS finds a subspace that moves even an untrained model's output, which is why the real
+alignment's IIA has to be read against both. Report the real alignment's IIA minus each control's
+(its selectivity), with uncertainty over the evaluation pairs. A subspace that does not clearly
+beat both controls is not evidence that the model carries the variable there.
+
 ## Placebo
 
 `NormMatchedPlacebo(reference, basis)` moves each base exactly as far as `reference` does, but

tests/integration/model_bridge/test_interchange.py

@@ -1,14 +1,18 @@
 """Real-Bridge integration checks for interchange interventions."""
 
+import pytest
 import torch
 
+from transformer_lens.model_bridge import TransformerBridge
+from transformer_lens.tools.analysis._model_state import require_eval_mode
 from transformer_lens.tools.analysis.interchange import (
     NormMatchedPlacebo,
     SubspaceInterchange,
     collect_activations,
     fit_das,
     intervened_logits,
     random_basis,
+    random_init_bridge,
 )
 
 PROMPTS = ["The key near the cabinet", "The keys near the cabinet"]
@@ -119,3 +123,39 @@ def accuracy(basis):
 
     assert accuracy(result.subspace) >= 0.9
     assert accuracy(random_basis(gpt2_bridge.cfg.d_model, 1, seed=0)) <= 0.2
+
+
+@pytest.mark.parametrize("name", ["gpt2", "tiny-stories-1M"])
+def test_random_init_bridge_matches_the_trained_architecture_without_the_training(name):
+    trained = TransformerBridge.boot_transformers(name, device="cpu")
+    before = torch.random.get_rng_state()
+    untrained = random_init_bridge(name, seed=0, device="cpu")
+    assert torch.equal(torch.random.get_rng_state(), before)  # the caller's stream is untouched
+    assert type(untrained.original_model) is type(trained.original_model)
+    assert untrained.cfg.n_layers == trained.cfg.n_layers
+    assert untrained.cfg.d_model == trained.cfg.d_model
+    assert untrained.W_E.dtype == torch.float32
+    assert not torch.allclose(untrained.W_E, trained.W_E)
+    assert torch.equal(untrained.W_E, random_init_bridge(name, seed=0, device="cpu").W_E)
+    require_eval_mode(untrained, operation="the random-init control")  # ready for fit_das
+    tokens = untrained.to_tokens(PROMPTS)
+    with torch.no_grad():
+        assert untrained(tokens).shape == trained(tokens).shape
+
+
+def test_random_init_bridge_casts_to_a_dtype_other_than_the_configs():
+    untrained = random_init_bridge(
+        "gpt2", seed=0, dtype=torch.bfloat16, device="cpu"
+    )  # config: float32
+    floating = [p for p in untrained.original_model.parameters() if p.is_floating_point()]
+    assert floating and all(p.dtype == torch.bfloat16 for p in floating)
+
+
+def test_random_init_bridge_puts_its_weights_on_the_device_it_records(monkeypatch):
+    # Stand in for a GPU machine, where boot's default device is not the CPU
+    import transformer_lens.model_bridge.sources.transformers.source as source
+
+    monkeypatch.setattr(source, "get_device", lambda: "meta")
+    untrained = random_init_bridge("gpt2", seed=0)
+    assert str(untrained.cfg.device) == "meta"
+    assert all(p.device.type == "meta" for p in untrained.original_model.parameters())

tests/unit/tools/test_interchange.py

@@ -12,12 +12,14 @@
     NormMatchedPlacebo,
     SubspaceInterchange,
     _model_device,
+    arbitrary_targets,
     collect_activations,
     fit_das,
     intervened_logits,
     mean_difference_basis,
     probe_basis,
     random_basis,
+    reinitialized,
     run_with_intervention,
 )
 
@@ -322,9 +324,9 @@ def test_das_leaves_the_callers_random_stream_alone():
     assert torch.equal(first.basis, again.basis)  # the seed still fixes the result
 
 
-def test_das_never_reseeds_a_device_generator(monkeypatch):
+def _forbid_reseeding(monkeypatch):
     def reseeded(*args, **kwargs):
-        raise AssertionError("fit_das reseeded a global or device generator")
+        raise AssertionError("a global or device generator was reseeded")
 
     for module, name in (
         (torch, "manual_seed"),
@@ -334,6 +336,10 @@ def reseeded(*args, **kwargs):
         (torch.mps, "manual_seed"),
     ):
         monkeypatch.setattr(module, name, reseeded)
+
+
+def test_das_never_reseeds_a_device_generator(monkeypatch):
+    _forbid_reseeding(monkeypatch)
     model, base, source, targets = _planted_pairs(n=32)
     fit_das(model, "hook_mid", base, source, targets, model, rank=1, epochs=1, seed=5)
 
@@ -366,3 +372,126 @@ def test_das_leaves_the_model_untouched():
     assert model.scale.requires_grad and torch.equal(model.scale.detach(), before)
     with pytest.raises(ValueError, match="same length"):
         fit_das(model, "hook_mid", base, source[:-1], targets, model, rank=1)
+
+
+def test_arbitrary_targets_are_fixed_per_key_and_seed():
+    keys = [f"source {i}" for i in range(400)]
+    targets = arbitrary_targets(keys, n_classes=3, seed=0)
+    assert targets.dtype == torch.long and set(targets.tolist()) == {0, 1, 2}
+    assert torch.equal(targets, arbitrary_targets(keys, n_classes=3, seed=0))
+    assert not torch.equal(targets, arbitrary_targets(keys, n_classes=3, seed=1))
+    assert torch.equal(arbitrary_targets(keys[:5] * 2, n_classes=3, seed=0)[5:], targets[:5])
+    assert all(abs(c - 400 / 3) < 50 for c in torch.bincount(targets).tolist())
+
+
+class _Layers(HookedRootModule):
+    """Only standard layers, each with reset_parameters, and a fixed buffer."""
+
+    def __init__(self):
+        super().__init__()
+        self.inp = nn.Linear(D_MODEL, D_MODEL)
+        self.norm = nn.LayerNorm(D_MODEL)
+        self.register_buffer("fixed", torch.randn(D_MODEL))
+        self.hook_mid = HookPoint()
+        self.setup()
+        self.eval()
+
+
+def test_reinitialized_resets_every_parameter_and_keeps_buffers():
+    model = _Layers()
+    with torch.no_grad():
+        model.norm.weight.add_(1.0)  # as training would, away from LayerNorm's ones and zeros
+        model.norm.bias.add_(1.0)
+    before = {name: p.detach().clone() for name, p in model.named_parameters()}
+    fresh = reinitialized(model, seed=0)
+    for name, parameter in fresh.named_parameters():
+        assert not torch.equal(parameter, before[name]), name
+        assert torch.equal(dict(model.named_parameters())[name], before[name])  # original untouched
+    assert torch.equal(fresh.fixed, model.fixed) and not fresh.training
+    assert torch.equal(reinitialized(model, seed=0).inp.weight, fresh.inp.weight)
+
+
+def test_reinitialized_resets_a_root_that_resets_its_own_parameters():
+    class _SelfResetting(_Layers):  # the root owns a parameter and knows how to reset it
+        def __init__(self):
+            super().__init__()
+            self.scale = nn.Parameter(torch.full((D_MODEL,), 3.0))
+
+        def reset_parameters(self):
+            nn.init.normal_(self.scale)
+
+    model = _SelfResetting()
+    fresh = reinitialized(model, seed=0)
+    assert not torch.equal(fresh.scale, model.scale)
+
+
+def test_reinitialized_refuses_parameters_it_cannot_reset():
+    with pytest.raises(ValueError, match="cannot reset scale"):
+        reinitialized(_Planted(torch.randn(D_MODEL)), seed=0)  # a parameter on the root itself
+
+    class _RawWeight(nn.Module):  # like HuggingFace's Conv1D: a parameter and no reset_parameters
+        def __init__(self):
+            super().__init__()
+            self.weight = nn.Parameter(torch.randn(D_MODEL, D_MODEL))
+
+    model = _Layers()
+    model.conv = _RawWeight()
+    with pytest.raises(ValueError, match="cannot reset conv.weight"):
+        reinitialized(model, seed=0)
+
+
+def test_reinitialized_leaves_every_random_stream_alone(monkeypatch):
+    model = _Layers()
+    torch.manual_seed(123)
+    expected = torch.rand(3)
+    torch.manual_seed(123)
+    _forbid_reseeding(monkeypatch)
+    reinitialized(model, seed=0)
+    assert torch.equal(torch.rand(3), expected)
+
+
+def _typed_pairs(prototypes: torch.Tensor, n: int, generator: torch.Generator):
+    """Random bases; sources drawn from a few recurring types, so the types act as keys."""
+    kinds = torch.randint(len(prototypes), (n,), generator=generator)
+    sources = prototypes[kinds] + 0.05 * torch.randn(n, D_MODEL, generator=generator)
+    return torch.randn(n, D_MODEL, generator=generator), sources, kinds
+
+
+def _held_out_iia(model, fit, test, fit_targets, test_targets) -> float:
+    result = fit_das(
+        model, "hook_mid", fit[0], fit[1], fit_targets, model, rank=1, epochs=40, lr=5e-2
+    )
+    acts = collect_activations(model, "hook_mid", test[1], model)
+    logits = intervened_logits(
+        model, "hook_mid", SubspaceInterchange(result.subspace), test[0], acts, model
+    )
+    return (logits.argmax(-1) == test_targets).float().mean().item()
+
+
+def test_das_beats_the_arbitrary_label_control_on_recurring_keys():
+    generator = torch.Generator().manual_seed(0)
+    model = _Planted(torch.randn(D_MODEL, generator=generator))
+    prototypes = torch.randn(64, D_MODEL, generator=generator)
+    fit, test = _typed_pairs(prototypes, 256, generator), _typed_pairs(prototypes, 256, generator)
+    # Held-out sources share their keys with training sources, so the control could be learned
+    assert set(test[2].tolist()) <= set(fit[2].tolist())
+    real = lambda sources: (sources @ model.direction > 0).long()
+    labels = lambda kinds: arbitrary_targets(kinds.tolist(), 2, seed=0)
+    assert _held_out_iia(model, fit, test, real(fit[1]), real(test[1])) > 0.9
+    assert _held_out_iia(model, fit, test, labels(fit[2]), labels(test[2])) < 0.7
+
+
+def test_the_arbitrary_label_control_can_reject_a_subspace():
+    # Types placed so that one direction fits their arbitrary labels: the control then does as
+    # well as the real alignment, and DAS does not beat it
+    generator = torch.Generator().manual_seed(0)
+    model = _Planted(torch.randn(D_MODEL, generator=generator))
+    prototypes = torch.randn(8, D_MODEL, generator=generator)
+    wanted = arbitrary_targets(list(range(8)), 2, seed=0)
+    prototypes[(prototypes @ model.direction > 0).long() != wanted] *= -1
+    fit, test = _typed_pairs(prototypes, 256, generator), _typed_pairs(prototypes, 256, generator)
+    real = lambda sources: (sources @ model.direction > 0).long()
+    labels = lambda kinds: arbitrary_targets(kinds.tolist(), 2, seed=0)
+    control = _held_out_iia(model, fit, test, labels(fit[2]), labels(test[2]))
+    assert control > 0.9
+    assert _held_out_iia(model, fit, test, real(fit[1]), real(test[1])) - control < 0.1

transformer_lens/tools/analysis/__init__.py

@@ -20,8 +20,10 @@
     - 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,
-      batched gradient-free evaluation, and distributed alignment search (DAS)
-      to learn the subspace.
+      difference-of-means, probe and random bases, a norm-matched placebo,
+      batched gradient-free evaluation, distributed alignment search (DAS) to
+      learn the subspace, and DAS's trained controls (arbitrary labels and a
+      randomly initialised model).
     - 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
@@ -83,12 +85,15 @@
     LearnedSubspace,
     NormMatchedPlacebo,
     SubspaceInterchange,
+    arbitrary_targets,
     collect_activations,
     fit_das,
     intervened_logits,
     mean_difference_basis,
     probe_basis,
     random_basis,
+    random_init_bridge,
+    reinitialized,
     run_with_intervention,
 )
 from transformer_lens.tools.analysis.jacobian_lens import (
@@ -209,6 +214,7 @@
     "SubspaceInterchange",
     "VocabularyRanking",
     "WeightLayout",
+    "arbitrary_targets",
     "attention_head_subspace_affinity",
     "attribution_patch",
     "collect_activations",
@@ -233,7 +239,9 @@
     "project_activations",
     "projection_kernel",
     "random_basis",
+    "random_init_bridge",
     "random_projection_kernel_moments",
+    "reinitialized",
     "run_causal_swap_benchmark",
     "run_with_intervention",
     "select_displacement_matched_control_token",