Add trained controls for distributed alignment search by dipo101 · Pull Request #1894 · TransformerLensOrg/TransformerLens

dipo101 · 2026-10-13T00:34:52Z

Description

Part 3 of 3 for #1879. Stacked on #1893; only the last commit is new.

Adds the trained controls DAS needs: arbitrary_targets for an arbitrary-label control, and random_init_bridge / reinitialized for an untrained copy of the model, each trained with the same pairs, rank and budget. A test checks DAS beats the arbitrary-label control on held-out pairs; docs section added.

Type of change

New feature (non-breaking change which adds functionality)
This change requires a documentation update

Checklist:

I have commented my code, particularly in hard-to-understand areas
I have made corresponding changes to the documentation
My changes generate no new warnings
I have added tests that prove my fix is effective or that my feature works
New and existing unit tests pass locally with my changes
I have not rewritten tests relating to key interfaces which would affect backward compatibility

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 TransformerLensOrg#1879.

jlarson4

Just a couple small comments on this final phase @dipo101! Let me know if you have any questions

jlarson4 · 2026-10-14T20:53:06Z

+    torch.manual_seed(seed)
+    for module in clone.modules():
+        reset = getattr(module, "reset_parameters", None)
+        if module is not clone and callable(reset):


Parameters that don't live on a module with reset_parameters keep their trained values. That covers the root's own parameters (like scale on the test's _Planted) and layers like HF's Conv1D. The caller gets back a control that is in part the trained model, which this control exists to catch. Can we ensure reinitialized reinitializes every parameter, or refuse when it can't?

True, a partly trained control defeats the whole point. reinitialized now refuses, naming the uncovered parameters, whenever a parameter belongs to no module with reset_parameters (the root’s own, like scale, or layers like HF’s Conv1D), instead of guessing an initialisation. The error points HF models to random_init_bridge. It also no longer skips the root when the root itself has reset_parameters. Tests cover a fully reset model and both refusal cases.

jlarson4 · 2026-10-14T20:54:00Z

+    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)


Building the model here with AutoModelForCausalLM skips the parts of boot_transformers that pick the architecture's model class and prepare its config, causing failures across multiple architectures. Boot has a from_config path (load_weights=False). Can this helper build on that so the control always matches the architecture the trained bridge uses?

Good point. It now builds through boot_transformers(load_weights=False), so the control gets the trained bridge’s model class and prepared config, then casts to the requested dtype and switches to eval mode. The integration test checks the HF class matches the trained bridge’s for GPT-2 and TinyStories.

jlarson4 · 2026-10-14T20:54:45Z

+        The untrained copy, in evaluation mode.
+    """
+    clone = copy.deepcopy(model)
+    torch.manual_seed(seed)


Same seed note here from the previous PR, we want to apply the same solution here we did on fit_das

Done, the same way as fit_das. both helpers build or reset on the CPU and seed only the CPU generator inside fork_rng(devices=[]), and reinitialized returns its copy on the model’s device. The no-reseeding test now covers reinitialized too, and the GPT-2/TinyStories test checks the caller’s CPU stream is unchanged.

jlarson4 · 2026-10-14T20:57:10Z

+            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


None of the held-out sources appear in training. Their arbitrary labels are fresh values and any subspace scores about 0.5. That means accuracy["arbitrary"] < 0.7 can't fail. The docs recipe has the same issue: the "near chance" it prints for the arbitrary control comes out that way for every model, meaning the control can never reject a subspace. Can we ensure the keys (in the docs and here) are something that recurs between training and evaluation pairs? Additionally would it be possible to adjust the test so the control side can fail?

Fair enough.

The test now draws sources from 64 recurring types, keyed by type and shared between fit and held-out pairs, and asserts every held-out type occurs in training. A second test places types so that a direction can fit their arbitrary labels, and checks the control then scores above 0.9 and DAS doesn’t beat it, so the control side can fail. The docs now hold out frames rather than nouns, key the labels by the source’s subject word, and explain why keys must recur. With that, the docs example reports 1.0 for the real alignment, 0.62 for the arbitrary-label control and 0.81 for the untrained model, so the text no longer claims “near chance”.

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

jlarson4

Thanks for wrapping up those for comments! I found a couple additional items on this pass, let me know if you have any questions. Should be good to merge after these.

jlarson4 · 2026-10-15T03:38:56Z

+    # 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()


With device=None, boot records its default device (CUDA when available) in cfg.device, but the load_weights=False model stays on the CPU. When using random_init_bridge(name, seed=0), as in the docs example, fails at the first forward on any GPU machine. Can the default use the bridge's config and weights on the same device?

Ok, thanks. random_init_bridge now passes device through to boot and then moves the bridge to bridge.cfg.device, so the weights will be where the config says, which with no device given is boot’s default. CI has no GPU, so the test makes boot’s default device meta and checks the weights end up there.

jlarson4 · 2026-10-15T03:40:24Z

    assert accuracy(random_basis(gpt2_bridge.cfg.d_model, 1, seed=0)) <= 0.2
+
+
+@pytest.mark.parametrize("name", ["gpt2", "tiny-stories-1M"])


Both models for this test have float32 configs, both pass if the maybe_cast_floating_params call is removed. Can the test cover a case where the requested dtype differs from the config's?

Fair enough. A new test asks for bf16 from GPT-2, whose config is f32, and checks every floating param comes back in bf16. It fails without the maybe_cast_floating_params call.

jlarson4 · 2026-10-15T03:41:02Z

+    assert torch.equal(reinitialized(model, seed=0).inp.weight, fresh.inp.weight)
+
+
+def test_reinitialized_refuses_parameters_it_cannot_reset():


Putting back the module is not clone skip still passes every test, because no test model's root defines reset_parameters. If it came back, a root that owns parameters and resets them itself would keep its trained values. Can we add a test to cover that case?

cool, added. A root that owns a param and defines its own reset_parameters now gets reset. the test also fails if the module is not clone skip comes back.

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

jlarson4 · 2026-10-15T15:16:19Z

Looks great, merging as is and closing the issue, thanks for putting this together.

jlarson4 mentioned this pull request Oct 13, 2026

Add distributed alignment search for interchange interventions #1893

Merged

8 tasks done

jlarson4 linked an issue Oct 13, 2026 that may be closed by this pull request

[Proposal] Distributed Alignment Search (DAS) with trained controls #1879

Closed

1 task done

dipo101 force-pushed the das-trained-controls branch 2 times, most recently from b8831e2 to e7497b8 Compare October 13, 2026 20:48

dipo101 mentioned this pull request Oct 13, 2026

Add subspace interchange interventions #1892

Merged

8 tasks done

dipo101 force-pushed the das-trained-controls branch 3 times, most recently from a218a89 to 9e0922c Compare October 14, 2026 15:35

dipo101 force-pushed the das-trained-controls branch from 9e0922c to 9695f86 Compare October 14, 2026 17:07

jlarson4 reviewed Oct 14, 2026

View reviewed changes

dipo101 requested a review from jlarson4 October 14, 2026 22:51

jlarson4 reviewed Oct 15, 2026

View reviewed changes

dipo101 requested a review from jlarson4 October 15, 2026 07:13

jlarson4 merged commit edaf8a7 into TransformerLensOrg:dev Oct 15, 2026
28 checks passed

		assert accuracy(random_basis(gpt2_bridge.cfg.d_model, 1, seed=0)) <= 0.2


		@pytest.mark.parametrize("name", ["gpt2", "tiny-stories-1M"])

		assert torch.equal(reinitialized(model, seed=0).inp.weight, fresh.inp.weight)


		def test_reinitialized_refuses_parameters_it_cannot_reset():

Conversation

dipo101 commented Oct 13, 2026

Description

Type of change

Checklist:

Uh oh!

jlarson4 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jlarson4 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jlarson4 commented Oct 15, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants