From 9695f86db3efe2f4be2be88f3ca7446574725626 Mon Sep 17 00:00:00 2001
From: dipo101 <21123483+dipo101@users.noreply.github.com>
Date: Tue, 13 Oct 2026 02:02:36 +0200
Subject: [PATCH 1/3] Add trained controls for distributed alignment search

Add arbitrary_targets, reinitialized and random_init_bridge to
tools.analysis.interchange: an arbitrary-label control and an untrained
copy of the model, to train DAS on with the same pairs, rank and budget.
A learned subspace should clearly beat both before its interchange accuracy
is read as evidence. Includes unit tests, a GPT-2 test of the untrained
bridge, and a docs section.

Part of #1879.
---
 docs/source/content/analysis_tools.md         |   2 +-
 .../content/interchange_interventions.md      |  48 ++++++++
 .../model_bridge/test_interchange.py          |  16 +++
 tests/unit/tools/test_interchange.py          |  58 ++++++++++
 transformer_lens/tools/analysis/__init__.py   |  12 +-
 .../tools/analysis/interchange.py             | 104 +++++++++++++++++-
 6 files changed, 235 insertions(+), 5 deletions(-)

diff --git a/docs/source/content/analysis_tools.md b/docs/source/content/analysis_tools.md
index b9683ecf5..bb0c8c1e5 100644
--- a/docs/source/content/analysis_tools.md
+++ b/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. |
 
diff --git a/docs/source/content/interchange_interventions.md b/docs/source/content/interchange_interventions.md
index 0564378f8..a24b74624 100644
--- a/docs/source/content/interchange_interventions.md
+++ b/docs/source/content/interchange_interventions.md
@@ -113,6 +113,54 @@ 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 source a fixed
+  pseudo-random class. DAS trained on these shows how much IIA the method's expressivity alone
+  produces.
+- **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, for fit and held-out pairs alike
+labels = lambda sources: arbitrary_targets([bridge.to_string(s) 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())  # near chance, 0.5
+
+# 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())  # near chance, 0.5
+```
+
+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
diff --git a/tests/integration/model_bridge/test_interchange.py b/tests/integration/model_bridge/test_interchange.py
index 74ad81abc..62939c17b 100644
--- a/tests/integration/model_bridge/test_interchange.py
+++ b/tests/integration/model_bridge/test_interchange.py
@@ -2,6 +2,7 @@
 
 import torch
 
+from transformer_lens.tools.analysis._model_state import require_eval_mode
 from transformer_lens.tools.analysis.interchange import (
     NormMatchedPlacebo,
     SubspaceInterchange,
@@ -9,6 +10,7 @@
     fit_das,
     intervened_logits,
     random_basis,
+    random_init_bridge,
 )
 
 PROMPTS = ["The key near the cabinet", "The keys near the cabinet"]
@@ -119,3 +121,17 @@ def accuracy(basis):
 
     assert accuracy(result.subspace) >= 0.9
     assert accuracy(random_basis(gpt2_bridge.cfg.d_model, 1, seed=0)) <= 0.2
+
+
+def test_random_init_bridge_keeps_the_architecture_without_the_training(gpt2_bridge):
+    untrained = random_init_bridge("gpt2", seed=0, device="cpu")
+    again = random_init_bridge("gpt2", seed=0, device="cpu")
+    assert untrained.cfg.n_layers == gpt2_bridge.cfg.n_layers
+    assert untrained.cfg.d_model == gpt2_bridge.cfg.d_model
+    assert untrained.W_E.dtype == torch.float32
+    assert not torch.allclose(untrained.W_E, gpt2_bridge.W_E)
+    assert torch.equal(untrained.W_E, again.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 == gpt2_bridge(tokens).shape
diff --git a/tests/unit/tools/test_interchange.py b/tests/unit/tools/test_interchange.py
index f037d556f..620d9f766 100644
--- a/tests/unit/tools/test_interchange.py
+++ b/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,
 )
 
@@ -366,3 +368,59 @@ 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())
+
+
+def test_reinitialized_resets_parameters_and_keeps_buffers():
+    model = _Planted(torch.randn(D_MODEL))
+    model.inner = nn.Linear(D_MODEL, D_MODEL)
+    before = model.inner.weight.detach().clone()
+    fresh = reinitialized(model, seed=0)
+    assert not torch.equal(fresh.inner.weight, before)
+    assert torch.equal(model.inner.weight, before)  # the original is untouched
+    assert torch.equal(fresh.direction, model.direction)
+    assert not fresh.training
+    assert torch.equal(reinitialized(model, seed=0).inner.weight, fresh.inner.weight)
+
+
+def test_das_beats_the_arbitrary_label_control_on_held_out_pairs():
+    model, base, source, targets = _planted_pairs()
+    keys = lambda rows: [tuple(r.tolist()) for r in rows]
+    fits = {
+        "real": fit_das(
+            model, "hook_mid", base, source, targets, model, rank=1, epochs=30, lr=5e-2
+        ),
+        "arbitrary": fit_das(
+            model,
+            "hook_mid",
+            base,
+            source,
+            arbitrary_targets(keys(source), 2, seed=0),
+            model,
+            rank=1,
+            epochs=30,
+            lr=5e-2,
+        ),
+    }
+    _, test_base, test_source, _ = _planted_pairs(seed=1)
+    expected = {
+        "real": (test_source @ model.direction > 0).long(),
+        "arbitrary": arbitrary_targets(keys(test_source), 2, seed=0),
+    }
+    source_acts = collect_activations(model, "hook_mid", test_source, model)
+    accuracy = {}
+    for name, result in fits.items():
+        logits = intervened_logits(
+            model, "hook_mid", SubspaceInterchange(result.subspace), test_base, source_acts, model
+        )
+        accuracy[name] = (logits.argmax(-1) == expected[name]).float().mean().item()
+    assert accuracy["real"] > 0.9 and accuracy["arbitrary"] < 0.7
diff --git a/transformer_lens/tools/analysis/__init__.py b/transformer_lens/tools/analysis/__init__.py
index db9b450bf..82f94960f 100644
--- a/transformer_lens/tools/analysis/__init__.py
+++ b/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",
diff --git a/transformer_lens/tools/analysis/interchange.py b/transformer_lens/tools/analysis/interchange.py
index 2182d3093..a01f936af 100644
--- a/transformer_lens/tools/analysis/interchange.py
+++ b/transformer_lens/tools/analysis/interchange.py
@@ -15,8 +15,13 @@
 class means, a linear probe's weights, a random basis), and a norm-matched random-subspace placebo.
 :func:`fit_das` learns the basis instead, by distributed alignment search (DAS; Geiger et al.,
 2024): it trains an orthonormal basis so that the interchange produces the high-level model's
-predicted outputs. An interchange accuracy is not evidence on its own: compare it with the placebo
-and the baselines.
+predicted outputs.
+
+An interchange accuracy is not evidence on its own. DAS fits arbitrary label mappings and finds
+effects in randomly initialised models (Arora et al., 2024), so the module also provides the
+controls to train the same way: :func:`arbitrary_targets` for an arbitrary-label control, and
+:func:`random_init_bridge` or :func:`reinitialized` for an untrained copy of the model. Compare a
+learned subspace with both, and 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
@@ -25,6 +30,8 @@ class means, a linear probe's weights, a random basis), and a norm-matched rando
 
 from __future__ import annotations
 
+import copy
+import hashlib
 from dataclasses import dataclass
 from typing import Any, Callable, Hashable, Iterator, Optional, Sequence, Union, cast
 
@@ -494,3 +501,96 @@ def fit_das(
         for p, flag in zip(model.parameters(), flags):
             p.requires_grad_(flag)
     return DASResult(subspace.eval(), tuple(losses), hook_name, rank)
+
+
+def arbitrary_targets(
+    keys: Sequence[Hashable], n_classes: int, seed: int
+) -> Int[torch.Tensor, "batch"]:
+    """A fixed pseudo-random class per key: targets for the arbitrary-label control.
+
+    Train :func:`fit_das` on these instead of the high-level model's targets, with the same pairs,
+    rank and budget. The interchange accuracy it reaches shows how much accuracy the method's
+    expressivity alone produces; a learned subspace should clearly exceed it.
+
+    Args:
+        keys: One key per pair identifying its source, e.g. the source text.
+        n_classes: Number of output classes.
+        seed: Seed for the mapping.
+
+    Returns:
+        A class in ``[0, n_classes)`` per key; equal keys always get equal classes.
+    """
+
+    def draw(key: Hashable) -> int:
+        digest = hashlib.sha256(f"{seed}:{key!r}".encode()).digest()
+        return int.from_bytes(digest[:8], "big") % n_classes
+
+    return torch.tensor([draw(k) for k in keys], dtype=torch.long)
+
+
+def reinitialized(model: nn.Module, seed: int) -> nn.Module:
+    """A copy of ``model`` with freshly initialised parameters: the random-init control.
+
+    For models built from standard ``nn`` layers, such as a ``HookedRootModule``: every submodule
+    with ``reset_parameters`` is reset. Buffers, such as fixed architectural matrices, are kept. For
+    HuggingFace models use :func:`random_init_bridge`, which applies the architecture's own
+    initialisation scheme.
+
+    Args:
+        model: The trained model; it is not modified.
+        seed: Seed for the new parameters.
+
+    Returns:
+        The untrained copy, in evaluation mode.
+    """
+    clone = copy.deepcopy(model)
+    torch.manual_seed(seed)
+    for module in clone.modules():
+        reset = getattr(module, "reset_parameters", None)
+        if module is not clone and callable(reset):
+            reset()
+    return clone.eval()
+
+
+def random_init_bridge(
+    model_name: str,
+    *,
+    seed: int,
+    model_class: Any = None,
+    revision: Optional[str] = None,
+    dtype: torch.dtype = torch.float32,
+    device: Optional[Union[str, torch.device]] = None,
+    **boot_kwargs: Any,
+) -> Any:
+    """A ``TransformerBridge`` with the architecture's HuggingFace initialisation and no training.
+
+    The random-init control: train :func:`fit_das` on it with the same pairs, rank and budget as
+    on the trained model. Effects DAS finds here come from the method and the architecture, not
+    from anything the model learned.
+
+    Args:
+        model_name: HuggingFace model name, as for ``TransformerBridge.boot_transformers``.
+        seed: Seed for the initialisation.
+        model_class: HuggingFace model class to build; ``AutoModelForCausalLM`` by default.
+        revision: Model revision whose configuration to use.
+        dtype: Parameter dtype.
+        device: Device to move the bridge to.
+        **boot_kwargs: Passed on to ``TransformerBridge.boot_transformers``.
+
+    Returns:
+        The untrained ``TransformerBridge``, in evaluation mode.
+    """
+    from transformers import AutoConfig, AutoModelForCausalLM
+
+    from transformer_lens.model_bridge import TransformerBridge
+
+    torch.manual_seed(seed)
+    config = AutoConfig.from_pretrained(model_name, revision=revision)
+    # from_config follows the config's dtype, and a supplied hf_model keeps its own dtype and
+    # placement, so both are set explicitly
+    hf_model = (model_class or AutoModelForCausalLM).from_config(config).to(dtype)
+    bridge = TransformerBridge.boot_transformers(
+        model_name, revision=revision, hf_model=hf_model, dtype=dtype, **boot_kwargs
+    )
+    # HuggingFace builds models from a config in training mode; the analyses require eval mode
+    return (bridge if device is None else bridge.to(device)).eval()

From 3403db19e868dc3d88d538a228034307bff4bd16 Mon Sep 17 00:00:00 2001
From: dipo101 <21123483+dipo101@users.noreply.github.com>
Date: Wed, 14 Oct 2026 23:31:56 +0100
Subject: [PATCH 2/3] Make the trained controls reset everything, match the
 architecture, and key on recurring inputs

reinitialized now refuses a model with any parameter that no module's
reset_parameters covers, instead of returning a control that keeps some
trained values, and resets a CPU copy with only the CPU generator seeded
inside a fork. random_init_bridge builds through boot_transformers with
load_weights=False, so the control gets the trained bridge's model class
and prepared config, and seeds the same way. The arbitrary-label tests and
docs now key labels by inputs that recur between fit and held-out pairs,
and a new test shows the control rejecting a subspace when a direction can
fit its labels.
---
 .../content/interchange_interventions.md      |  36 +++--
 .../model_bridge/test_interchange.py          |  22 ++-
 tests/unit/tools/test_interchange.py          | 143 ++++++++++++------
 .../tools/analysis/interchange.py             |  81 ++++++----
 4 files changed, 186 insertions(+), 96 deletions(-)

diff --git a/docs/source/content/interchange_interventions.md b/docs/source/content/interchange_interventions.md
index a24b74624..aba48f40b 100644
--- a/docs/source/content/interchange_interventions.md
+++ b/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
@@ -118,9 +121,11 @@ initialised models (Arora et al., 2024), so its IIA needs controls trained the s
 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 source a fixed
+- **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.
+  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.
@@ -128,8 +133,9 @@ and evaluate it on the same held-out pairs:
 ```python
 from transformer_lens.tools.analysis import arbitrary_targets, random_init_bridge
 
-# Arbitrary labels: a fixed pseudo-random class per source, for fit and held-out pairs alike
-labels = lambda sources: arbitrary_targets([bridge.to_string(s) for s in sources], 2, seed=0)
+# 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,
@@ -138,7 +144,7 @@ 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())  # near chance, 0.5
+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)
@@ -154,11 +160,13 @@ 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())  # near chance, 0.5
+print((logits.argmax(-1) == test_targets).float().mean())  # tensor(0.8125)
 ```
 
-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
+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
diff --git a/tests/integration/model_bridge/test_interchange.py b/tests/integration/model_bridge/test_interchange.py
index 62939c17b..b378601f0 100644
--- a/tests/integration/model_bridge/test_interchange.py
+++ b/tests/integration/model_bridge/test_interchange.py
@@ -1,7 +1,9 @@
 """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,
@@ -123,15 +125,19 @@ def accuracy(basis):
     assert accuracy(random_basis(gpt2_bridge.cfg.d_model, 1, seed=0)) <= 0.2
 
 
-def test_random_init_bridge_keeps_the_architecture_without_the_training(gpt2_bridge):
-    untrained = random_init_bridge("gpt2", seed=0, device="cpu")
-    again = random_init_bridge("gpt2", seed=0, device="cpu")
-    assert untrained.cfg.n_layers == gpt2_bridge.cfg.n_layers
-    assert untrained.cfg.d_model == gpt2_bridge.cfg.d_model
+@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, gpt2_bridge.W_E)
-    assert torch.equal(untrained.W_E, again.W_E)
+    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 == gpt2_bridge(tokens).shape
+        assert untrained(tokens).shape == trained(tokens).shape
diff --git a/tests/unit/tools/test_interchange.py b/tests/unit/tools/test_interchange.py
index 620d9f766..043949325 100644
--- a/tests/unit/tools/test_interchange.py
+++ b/tests/unit/tools/test_interchange.py
@@ -324,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"),
@@ -336,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)
 
@@ -380,47 +384,100 @@ def test_arbitrary_targets_are_fixed_per_key_and_seed():
     assert all(abs(c - 400 / 3) < 50 for c in torch.bincount(targets).tolist())
 
 
-def test_reinitialized_resets_parameters_and_keeps_buffers():
-    model = _Planted(torch.randn(D_MODEL))
-    model.inner = nn.Linear(D_MODEL, D_MODEL)
-    before = model.inner.weight.detach().clone()
+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)
-    assert not torch.equal(fresh.inner.weight, before)
-    assert torch.equal(model.inner.weight, before)  # the original is untouched
-    assert torch.equal(fresh.direction, model.direction)
-    assert not fresh.training
-    assert torch.equal(reinitialized(model, seed=0).inner.weight, fresh.inner.weight)
+    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_das_beats_the_arbitrary_label_control_on_held_out_pairs():
-    model, base, source, targets = _planted_pairs()
-    keys = lambda rows: [tuple(r.tolist()) for r in rows]
-    fits = {
-        "real": fit_das(
-            model, "hook_mid", base, source, targets, model, rank=1, epochs=30, lr=5e-2
-        ),
-        "arbitrary": fit_das(
-            model,
-            "hook_mid",
-            base,
-            source,
-            arbitrary_targets(keys(source), 2, seed=0),
-            model,
-            rank=1,
-            epochs=30,
-            lr=5e-2,
-        ),
-    }
-    _, test_base, test_source, _ = _planted_pairs(seed=1)
-    expected = {
-        "real": (test_source @ model.direction > 0).long(),
-        "arbitrary": arbitrary_targets(keys(test_source), 2, seed=0),
-    }
-    source_acts = collect_activations(model, "hook_mid", test_source, model)
-    accuracy = {}
-    for name, result in fits.items():
-        logits = intervened_logits(
-            model, "hook_mid", SubspaceInterchange(result.subspace), test_base, source_acts, model
-        )
-        accuracy[name] = (logits.argmax(-1) == expected[name]).float().mean().item()
-    assert accuracy["real"] > 0.9 and accuracy["arbitrary"] < 0.7
+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
diff --git a/transformer_lens/tools/analysis/interchange.py b/transformer_lens/tools/analysis/interchange.py
index a01f936af..d14c9bbcc 100644
--- a/transformer_lens/tools/analysis/interchange.py
+++ b/transformer_lens/tools/analysis/interchange.py
@@ -531,33 +531,53 @@ def draw(key: Hashable) -> int:
 def reinitialized(model: nn.Module, seed: int) -> nn.Module:
     """A copy of ``model`` with freshly initialised parameters: the random-init control.
 
-    For models built from standard ``nn`` layers, such as a ``HookedRootModule``: every submodule
-    with ``reset_parameters`` is reset. Buffers, such as fixed architectural matrices, are kept. For
-    HuggingFace models use :func:`random_init_bridge`, which applies the architecture's own
-    initialisation scheme.
+    For models built from standard ``nn`` layers, such as a ``HookedRootModule``: every module with
+    ``reset_parameters`` is reset, so every parameter must belong to one. Buffers, such as fixed
+    architectural matrices, are kept. For HuggingFace models use :func:`random_init_bridge`, which
+    applies the architecture's own initialisation scheme.
 
     Args:
         model: The trained model; it is not modified.
         seed: Seed for the new parameters.
 
     Returns:
-        The untrained copy, in evaluation mode.
+        The untrained copy, on the device of the model's parameters, in evaluation mode.
+
+    Raises:
+        ValueError: If a parameter belongs to no module with ``reset_parameters``, so it would keep
+            its trained value.
     """
-    clone = copy.deepcopy(model)
-    torch.manual_seed(seed)
-    for module in clone.modules():
-        reset = getattr(module, "reset_parameters", None)
-        if module is not clone and callable(reset):
-            reset()
-    return clone.eval()
+    covered = {
+        id(parameter)
+        for module in model.modules()
+        if callable(getattr(module, "reset_parameters", None))
+        for parameter in module.parameters(recurse=False)
+    }
+    missing = [name for name, parameter in model.named_parameters() if id(parameter) not in covered]
+    if missing:
+        more = f" and {len(missing) - 5} more" if len(missing) > 5 else ""
+        raise ValueError(
+            f"reinitialized cannot reset {', '.join(missing[:5])}{more}: no module's "
+            "reset_parameters covers them, so they would keep their trained values. Reset them "
+            "yourself, or use random_init_bridge for a HuggingFace model."
+        )
+    device = _model_device(model)
+    clone = copy.deepcopy(model).cpu()
+    # Reset on the CPU, seeding only the CPU generator inside a fork that restores it;
+    # torch.manual_seed would also reseed every CUDA, MPS and XPU generator
+    with torch.random.fork_rng(devices=[]):
+        torch.default_generator.manual_seed(seed)
+        for module in clone.modules():
+            reset = getattr(module, "reset_parameters", None)
+            if callable(reset):
+                reset()
+    return clone.to(device).eval()
 
 
 def random_init_bridge(
     model_name: str,
     *,
     seed: int,
-    model_class: Any = None,
-    revision: Optional[str] = None,
     dtype: torch.dtype = torch.float32,
     device: Optional[Union[str, torch.device]] = None,
     **boot_kwargs: Any,
@@ -566,31 +586,30 @@ def random_init_bridge(
 
     The random-init control: train :func:`fit_das` on it with the same pairs, rank and budget as
     on the trained model. Effects DAS finds here come from the method and the architecture, not
-    from anything the model learned.
+    from anything the model learned. The model is built by ``boot_transformers`` with
+    ``load_weights=False``, so it has the same model class and prepared configuration as the
+    trained bridge.
 
     Args:
-        model_name: HuggingFace model name, as for ``TransformerBridge.boot_transformers``.
+        model_name: Model name or alias, as for ``TransformerBridge.boot_transformers``.
         seed: Seed for the initialisation.
-        model_class: HuggingFace model class to build; ``AutoModelForCausalLM`` by default.
-        revision: Model revision whose configuration to use.
         dtype: Parameter dtype.
-        device: Device to move the bridge to.
-        **boot_kwargs: Passed on to ``TransformerBridge.boot_transformers``.
+        device: Device to move the bridge to; the CPU by default.
+        **boot_kwargs: Passed on to ``TransformerBridge.boot_transformers``, e.g. ``revision``.
 
     Returns:
         The untrained ``TransformerBridge``, in evaluation mode.
     """
-    from transformers import AutoConfig, AutoModelForCausalLM
-
     from transformer_lens.model_bridge import TransformerBridge
+    from transformer_lens.utilities.multi_gpu import maybe_cast_floating_params
 
-    torch.manual_seed(seed)
-    config = AutoConfig.from_pretrained(model_name, revision=revision)
-    # from_config follows the config's dtype, and a supplied hf_model keeps its own dtype and
-    # placement, so both are set explicitly
-    hf_model = (model_class or AutoModelForCausalLM).from_config(config).to(dtype)
-    bridge = TransformerBridge.boot_transformers(
-        model_name, revision=revision, hf_model=hf_model, dtype=dtype, **boot_kwargs
-    )
-    # HuggingFace builds models from a config in training mode; the analyses require eval mode
+    # The model is built on the CPU, so seed only the CPU generator, inside a fork that restores it
+    with torch.random.fork_rng(devices=[]):
+        torch.default_generator.manual_seed(seed)
+        bridge = TransformerBridge.boot_transformers(
+            model_name, load_weights=False, dtype=dtype, **boot_kwargs
+        )
+    # Building from a config follows the config's dtype, and HuggingFace leaves the model in
+    # training mode; the analyses require eval mode
+    maybe_cast_floating_params(bridge.original_model, dtype)
     return (bridge if device is None else bridge.to(device)).eval()

From 0b653cde7c53b46b133e866a1c48c8285e2afd38 Mon Sep 17 00:00:00 2001
From: dipo101 <21123483+dipo101@users.noreply.github.com>
Date: Thu, 15 Oct 2026 08:08:33 +0100
Subject: [PATCH 3/3] Put the untrained bridge on its recorded device, and test
 dtype and root resets

random_init_bridge now moves the bridge to the device boot recorded, which
by default is boot's own choice (CUDA when available), since a model built
from a config stays on the CPU. New tests cover a requested dtype that
differs from the config's, a non-CPU default device, and a root module that
resets its own parameters.
---
 .../content/interchange_interventions.md       |  4 ++--
 .../model_bridge/test_interchange.py           | 18 ++++++++++++++++++
 tests/unit/tools/test_interchange.py           | 14 ++++++++++++++
 transformer_lens/tools/analysis/interchange.py | 10 +++++-----
 4 files changed, 39 insertions(+), 7 deletions(-)

diff --git a/docs/source/content/interchange_interventions.md b/docs/source/content/interchange_interventions.md
index aba48f40b..f0abf118f 100644
--- a/docs/source/content/interchange_interventions.md
+++ b/docs/source/content/interchange_interventions.md
@@ -166,8 +166,8 @@ 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.
+(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
 
diff --git a/tests/integration/model_bridge/test_interchange.py b/tests/integration/model_bridge/test_interchange.py
index b378601f0..215cc20ec 100644
--- a/tests/integration/model_bridge/test_interchange.py
+++ b/tests/integration/model_bridge/test_interchange.py
@@ -141,3 +141,21 @@ def test_random_init_bridge_matches_the_trained_architecture_without_the_trainin
     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())
diff --git a/tests/unit/tools/test_interchange.py b/tests/unit/tools/test_interchange.py
index 043949325..c1c6f81af 100644
--- a/tests/unit/tools/test_interchange.py
+++ b/tests/unit/tools/test_interchange.py
@@ -411,6 +411,20 @@ def test_reinitialized_resets_every_parameter_and_keeps_buffers():
     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
diff --git a/transformer_lens/tools/analysis/interchange.py b/transformer_lens/tools/analysis/interchange.py
index d14c9bbcc..54c234637 100644
--- a/transformer_lens/tools/analysis/interchange.py
+++ b/transformer_lens/tools/analysis/interchange.py
@@ -594,7 +594,7 @@ def random_init_bridge(
         model_name: Model name or alias, as for ``TransformerBridge.boot_transformers``.
         seed: Seed for the initialisation.
         dtype: Parameter dtype.
-        device: Device to move the bridge to; the CPU by default.
+        device: Device for the bridge; by default the device ``boot_transformers`` would choose.
         **boot_kwargs: Passed on to ``TransformerBridge.boot_transformers``, e.g. ``revision``.
 
     Returns:
@@ -607,9 +607,9 @@ def random_init_bridge(
     with torch.random.fork_rng(devices=[]):
         torch.default_generator.manual_seed(seed)
         bridge = TransformerBridge.boot_transformers(
-            model_name, load_weights=False, dtype=dtype, **boot_kwargs
+            model_name, load_weights=False, dtype=dtype, device=device, **boot_kwargs
         )
-    # Building from a config follows the config's dtype, and HuggingFace leaves the model in
-    # training mode; the analyses require eval mode
+    # Building from a config follows the config's dtype and leaves the weights on the CPU,
+    # whatever device boot recorded; HuggingFace also leaves the model in training mode
     maybe_cast_floating_params(bridge.original_model, dtype)
-    return (bridge if device is None else bridge.to(device)).eval()
+    return bridge.to(bridge.cfg.device).eval()