Skip to content

Add distributed alignment search for interchange interventions - #1893

Open
dipo101 wants to merge 3 commits into
TransformerLensOrg:devfrom
dipo101:distributed-alignment-search
Open

dipo101 wants to merge 3 commits into
TransformerLensOrg:devfrom
dipo101:distributed-alignment-search

Conversation

@dipo101

@dipo101 dipo101 commented Oct 13, 2026

Copy link
Copy Markdown

Description

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

Adds fit_das and LearnedSubspace: distributed alignment search trains an orthonormal basis so that interchanging it makes the model produce the outputs a high-level causal model predicts for each base/source pair. Tested on a planted direction and on GPT-2 subject number with held-out nouns; 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 transformer_lens.tools.analysis.interchange: an intervention that sets a
base input's activation inside an orthonormal subspace to its value on a
source input, at one hook point and optionally one position per example.
It comes with difference-of-means, probe and random bases to intervene on
and a norm-matched random-subspace placebo, plus unit tests, GPT-2
integration tests and a docs page.

Part of TransformerLensOrg#1879.
@jlarson4

Copy link
Copy Markdown
Collaborator

I will wait to review #1893 & #1894 until #1892 is resolved and merged.

@jlarson4 jlarson4 linked an issue Oct 13, 2026 that may be closed by this pull request
1 task done
run_with_intervention now raises unless the hook fires exactly once in the
forward pass, as collect_activations already did. With positions, the
activation must be [batch, pos, d_model] and positions must hold one index
per example; other shapes raise instead of slicing the wrong axis. Add
intervened_logits, which evaluates an intervention in batches without
gradients and returns CPU logits, so memory is bounded by batch_size. The
position tests now use a different position per row.
Add fit_das and LearnedSubspace to tools.analysis.interchange: DAS trains an
orthonormal basis so that interchanging it makes the model produce the
outputs a high-level causal model predicts for each base/source pair. The
model is frozen during training and restored afterwards, and batches can be
split to save memory. Includes unit tests on a planted direction, a GPT-2
test on held-out nouns, and a docs section.

Part of TransformerLensOrg#1879.
@dipo101
dipo101 force-pushed the distributed-alignment-search branch 2 times, most recently from a625238 to 15dddc6 Compare October 13, 2026 20:48
@dipo101 dipo101 mentioned this pull request Oct 13, 2026
8 tasks done
@jlarson4

Copy link
Copy Markdown
Collaborator

Hi @dipo101, please rebase this onto dev when you have a moment so I can review. Thanks!

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

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

2 participants