# Out-of-distribution generalization via composition: a lens through induction heads in Transformers

arXiv: https://arxiv.org/abs/2408.09503 · HTML: https://mlsanity.com/paper/2408.09503/

## Key facts

- Cited by 28 papers in the map (data through 14 Oct 2026).
- Builds on 92 of its 121 references that are in the map.
- First posted on arXiv 18 Aug 2024.

## Builds on (most cited first, 30 of 92)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [The Llama 3 Herd of Models](https://mlsanity.com/paper/2407.21783/) (2024)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Llama 2: Open Foundation and Fine-Tuned Chat Models](https://mlsanity.com/paper/2307.09288/) (2023)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Efficient Estimation of Word Representations in Vector Space](https://mlsanity.com/paper/1301.3781/) (2013)
- [Large Language Models are Zero-Shot Reasoners](https://mlsanity.com/paper/2205.11916/) (2022)
- [Mistral 7B](https://mlsanity.com/paper/2310.06825/) (2023)
- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [Emergent Abilities of Large Language Models](https://mlsanity.com/paper/2206.07682/) (2022)
- [Gemma 2: Improving Open Language Models at a Practical Size](https://mlsanity.com/paper/2408.00118/) (2024)
- [Pythia: A Suite for Analyzing Large Language Models Across Training and Scaling](https://mlsanity.com/paper/2304.01373/) (2023)
- [Rethinking the Role of Demonstrations: What Makes In-Context Learning Work?](https://mlsanity.com/paper/2202.12837/) (2022)
- [Calibrate Before Use: Improving Few-Shot Performance of Language Models](https://mlsanity.com/paper/2102.09690/) (2021)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Inference-Time Intervention: Eliciting Truthful Answers from a Language Model](https://mlsanity.com/paper/2306.03341/) (2023)
- [Interpretability in the Wild: a Circuit for Indirect Object Identification in GPT-2 small](https://mlsanity.com/paper/2211.00593/) (2022)
- [Fantastically Ordered Prompts and Where to Find Them: Overcoming Few-Shot Prompt Order Sensitivity](https://mlsanity.com/paper/2104.08786/) (2021)
- [Editing Models with Task Arithmetic](https://mlsanity.com/paper/2212.04089/) (2022)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [An Explanation of In-context Learning as Implicit Bayesian Inference](https://mlsanity.com/paper/2111.02080/) (2021)

## Built on by (most cited first, 28 of 28)

- [Is Chain-of-Thought Reasoning of LLMs a Mirage? A Data Distribution Lens](https://mlsanity.com/paper/2508.01191/) (2025)
- [Can In-context Learning Really Generalize to Out-of-distribution Tasks?](https://mlsanity.com/paper/2410.09695/) (2024)
- [Learning to Inference Adaptively for Multimodal Large Language Models](https://mlsanity.com/paper/2503.10905/) (2025)
- [Task Generalization With AutoRegressive Compositional Structure: Can Learning From D Tasks Generalize to DT Tasks?](https://mlsanity.com/paper/2502.08991/) (2025)
- [Unifying Attention Heads and Task Vectors via Hidden State Geometry in In-Context Learning](https://mlsanity.com/paper/2505.18752/) (2025)
- [How Transformers Get Rich: Approximation and Dynamics Analysis](https://mlsanity.com/paper/2410.11474/) (2024)
- [KV Shifting Attention Enhances Language Modeling](https://mlsanity.com/paper/2411.19574/) (2024)
- [Born a Transformer -- Always a Transformer? On the Effect of Pretraining on Architectural Abilities](https://mlsanity.com/paper/2505.21785/) (2025)
- [Task Diversity Shortens the ICL Plateau](https://mlsanity.com/paper/2410.05448/) (2024)
- [Can GRPO Help LLMs Transcend Their Pretraining Origin?](https://mlsanity.com/paper/2510.15990/) (2025)
- [ICL CIPHERS: Quantifying "Learning" in In-Context Learning via Substitution Ciphers](https://mlsanity.com/paper/2504.19395/) (2025)
- [Decomposing Complex Visual Comprehension into Atomic Visual Skills for Vision Language Models](https://mlsanity.com/paper/2505.20021/) (2025)
- [Task Vectors, Learned Not Extracted: Performance Gains and Mechanistic Insight](https://mlsanity.com/paper/2509.24169/) (2025)
- [Shattered Compositionality: Counterintuitive Learning Dynamics of Transformers for Arithmetic](https://mlsanity.com/paper/2601.22510/) (2026)
- [Localizing Task Recognition and Task Learning in In-Context Learning via Attention Head Analysis](https://mlsanity.com/paper/2509.24164/) (2025)
- [Generalization vs. Specialization under Concept Shift](https://mlsanity.com/paper/2409.15582/) (2024)
- [Understanding the Staged Dynamics of Transformers in Learning Latent Structure](https://mlsanity.com/paper/2511.19328/) (2025)
- [How Does Unfaithful Reasoning Emerge from Autoregressive Training? A Study of Synthetic Experiments](https://mlsanity.com/paper/2602.01017/) (2026)
- [Critical Windows of Complexity Control: When Transformers Decide to Reason or Memorize](https://mlsanity.com/paper/2605.04396/) (2026)
- [ResidualDroppath: Enhancing Feature Reuse over Residual Connections](https://mlsanity.com/paper/2411.09475/) (2024)
- [Can Language Models Compose Skills In-Context?](https://mlsanity.com/paper/2510.22993/) (2025)
- [Multi-Agent Cooperative Learning for Robust Vision-Language Alignment under OOD Concepts](https://mlsanity.com/paper/2601.09746/) (2026)
- [Bridging the Semantic Chasm: Synergistic Conceptual Anchoring for Generalized Few-Shot and Zero-Shot OOD Perception](https://mlsanity.com/paper/2602.00340/) (2026)
- [Finding Sense in Nonsense with Generated Contexts: Perspectives from Humans and Language Models](https://mlsanity.com/paper/2602.11699/) (2026)
- [Distance generalization in transformers: why bother with positional encoding?](https://mlsanity.com/paper/2609.11913/) (2026)
- [Metacognitive Steering: Learning the Structure of Scientific Judgment](https://mlsanity.com/paper/2609.16245/) (2026)
- [A Survey on the Linear Representation Hypothesis](https://mlsanity.com/paper/2609.22695/) (2026)
- [Theory on Attention Dynamics for Out-of-Distribution In-Context Learning](https://mlsanity.com/paper/2609.36448/) (2026)

## Abstract (arXiv)

Large language models (LLMs) such as GPT-4 sometimes appear to be creative, solving novel tasks often with a few demonstrations in the prompt. These tasks require the models to generalize on distributions different from those from training data -- which is known as out-of-distribution (OOD) generalization. Despite the tremendous success of LLMs, how they approach OOD generalization remains an open and underexplored question. We examine OOD generalization in settings where instances are generated according to hidden rules, including in-context learning with symbolic reasoning. Models are required to infer the hidden rules behind input prompts without any fine-tuning. We empirically examined the training dynamics of Transformers on a synthetic example and conducted extensive experiments on a variety of pretrained LLMs, focusing on a type of components known as induction heads. We found that OOD generalization and composition are tied together -- models can learn rules by composing two self-attention layers, thereby achieving OOD generalization. Furthermore, a shared latent subspace in the embedding (or feature) space acts as a bridge for composition by aligning early layers and later layers, which we refer to as the common bridge representation hypothesis.