Research map / Architectures and efficiency / Attention and transformers
Selective Induction Heads: How Transformers Select Causal Structures in Context
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Introduces interleaved Markov chains with varying lags but fixed transition matrix to study how 3-layer attention-only transformers select causal structure in-context via a "selective induction head".
Problem
How transformers dynamically select the correct causal structure (lag) in context when the same token pairs can have different causal relationships depending on context, which fixed-structure Markov chain studies do not capture.
Approach
A synthetic task of interleaved Markov chains with a shared transition matrix but different lags per sequence; a 3-layer disentangled/attention-only transformer construction aggregates normalized transition probabilities per lag and the final layer copies the token at the selected lag (selective induction head), asymptotically approximating maximum likelihood lag selection.
Key results
- The 3-layer constructed disentangled transformer matches the maximum likelihood performance on lags {1,2,3}, and trained disentangled and standard transformers match the construction.
- A 2-layer attention-only transformer fails to solve the task for any head configuration.
- With a single head in the second layer and two lags, the single-head construction attains the optimal sample complexity, matching the two-head version.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use the three-layer constructed transformer with the optimal number of second-layer heads | Interleaved Markov Chain next-token prediction · Contiguous lags | KL divergence | No clear effect |
| Train a three-layer disentangled transformer from scratch | Interleaved Markov Chain next-token prediction · Contiguous and non-contiguous lags | KL divergence | No clear effect |
| Train a three-layer standard transformer from scratch | Interleaved Markov Chain next-token prediction · Contiguous and non-contiguous lags | KL divergence | No clear effect |
| Use one second-layer head instead of the construction's optimal number (ablation) | Constructed transformer · Interleaved Markov Chain next-token prediction | KL divergence | Hurt |
| Use one second-layer head instead of the optimal number (ablation) | Trained transformers · Interleaved Markov Chain next-token prediction | KL divergence | Helped |
| Use the alternative single-head construction for two lags (ablation) | Constructed transformer · Two-lag Interleaved Markov Chain | KL divergence | No clear effect |
| Reduce the transformer from three layers to two layers (ablation) | Standard transformer · Interleaved Markov Chain next-token prediction | KL divergence | Hurt |
| Add layers beyond three (ablation) | Standard transformer · Interleaved Markov Chain next-token prediction | KL divergence | No clear effect |
| Increase second-layer head count up to the number of lags (ablation) | Standard transformer · Five-lag Interleaved Markov Chain | KL divergence | Helped |
| Add second-layer heads beyond the number of lags (ablation) | Standard transformer · Five-lag Interleaved Markov Chain | KL divergence | No clear effect |
Surprising findings
- Trained transformers with one head in the second layer can outperform the interpretable construction, suggesting more efficient non-interpretable aggregation mechanisms.
- For two lags, trained transformers can achieve performance closer to BMA than ML at small sequence lengths.
Limitations
- The key claim about expected normalized transition probabilities is proven only for two lags, independent lags, and no-normalization cases; the general case is left as a conjecture supported empirically.
- Analysis is restricted to attention-only transformers and a synthetic Markov chain task with a fixed alphabet of size 5.
Experimental setup
Datasets: synthetic interleaved Markov chain sequences
Benchmarks: synthetic interleaved Markov chain task
Metrics: KL divergence between true and predicted next-token distribu
Baselines: maximum likelihood estimator Bayesian model average 2-layer attention-only transformer n-gram estimation
Builds on: induction heads disentangled transformers in-context Markov chain bigram learning
Closest papers by meaning
- How Transformers Learn Causal Structure with Gradient Descent (ICML 2024 · similarity 0.92)
- What needs to go right for an induction head? A mechanistic study of in-context learning circuits and their formation (ICML 2024 · similarity 0.91)
- Towards Understanding the Universality of Transformers for Next-Token Prediction (ICLR 2025 · similarity 0.90)
- Causal Head Gating: A Framework for Interpreting Roles of Attention Heads in Transformers (NeurIPS 2025 · similarity 0.89)
- What One Cannot, Two Can: Two-Layer Transformers Provably Represent Induction Heads on Any-Order Markov Chains (NeurIPS 2025 · similarity 0.89)
- Linking In-context Learning in Transformers to Human Episodic Memory (NeurIPS 2024 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.