Research map / Architectures and efficiency / Attention and transformers
Compositional Capabilities of Autoregressive Transformers: A Study on Synthetic, Interpretable Tasks
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.
Trains small nanoGPT Transformers on synthetic compositions of bijection/permutation functions marked by task tokens, showing they generalize to exponentially/combinatorially many unseen compositions with step-by-step outputs and identifying attention-selects/MLP-executes mechanism.
Problem
Whether autoregressive Transformers trained on compositional data can learn to compose their atomic capabilities and generalize to exponentially or combinatorially many unseen function compositions, and what factors drive this.
Approach
Uses a minimal synthetic setup with task tokens specifying function compositions over data tokens, comparing in-order vs out-of-order training data, step-by-step vs direct prompt formats, and probing mechanisms via linear probes, attention maps, and a theoretical construction.
Key results
- A Transformer trained on 30-100 function compositions generalizes almost perfectly to 3125 unseen in-order compositions of 5 bijections (step-by-step format).
- Direct prompting fails to generalize to new in-order compositions of bijections even when trained on 2000 of 3125 compositions.
- LSTMs fail to compositionally generalize in the same synthetic setup while Transformers do.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Generate intermediate outputs step by step instead of predicting only the final composition | Autoregressive Transformer · Unseen in-order compositions of five bijections | accuracy: 100 compositions -> all 3125 possible in-order compositions | Helped |
| Train with atomic functions only, rather than including compositions (ablation) | Autoregressive Transformer · Compositional generalization | accuracy | Hurt |
| Include in-order compositions in the training data (ablation) | Autoregressive Transformer · Unseen in-order compositions | accuracy: 10 | Helped |
| Include out-of-order compositions in the training data (ablation) | Autoregressive Transformer · Unseen out-of-order compositions | accuracy: 100 | Helped |
| Use direct prompting instead of step-by-step prompting (ablation) | Autoregressive Transformer · Unseen in-order compositions of five bijections | accuracy: 2000 of the 3125 possible in-order compositions | Hurt |
| Train on 250 observed combinations of a permutation and a bijection (ablation) | Autoregressive Transformer with direct prompting · Two-function compositions of permutations and bijections | accuracy: 250 -> all 625 possible compositions | Helped |
| Use direct prompting for compositions of two bijections (ablation) | Autoregressive Transformer · Unseen in-order compositions of two bijections | accuracy: 500 of 625 compositions | Hurt |
| Replace the Transformer with an LSTM | Autoregressive model · Unseen compositions in direct and step-by-step formats | accuracy | Hurt |
| Increase Transformer depth in the step-by-step setup (ablation) | Autoregressive Transformer · Unseen in-order compositions of five bijections | accuracy | Hurt |
| Increase embedding dimension (ablation) | Autoregressive Transformer · Unseen in-order compositions of five bijections | accuracy | Helped |
| Vary the number of attention heads in a 1-layer Transformer (ablation) | Autoregressive Transformer · Compositional generalization | accuracy | No clear effect |
| Train on random compositions rather than only atomic functions (ablation) | Autoregressive Transformer · Compositions with few versus many non-identity functions | accuracy: 25; 50 | Mixed |
Surprising findings
- Direct prompt format fails to compositionally generalize even with high data diversity, while step-by-step format succeeds.
- Accuracy decreases for very deep models (many layers), attributed to optimization difficulty.
Limitations
- Synthetic setup with small vocabulary and fixed function families; may not transfer to natural language pretraining.
- Direct prompt format does not generalize even with large training diversity, leaving the failure mode unexplained beyond optimization arguments.
- Out-of-order generalization only partial, requiring more data diversity.
Experimental setup
Datasets: synthetic compositions of bijections and permutations over 6 data tokens (vocabulary size 10)
Benchmarks: held-out in-order and out-of-order compositions of synthetic functions
Metrics: accuracy over last six tokens accuracy over 1000 data token inputs
Baselines: LSTM BASE training subset RANDOM training subset
Builds on: nanoGPT chain-of-thought induction heads
Closest papers by meaning
- Task Generalization with Autoregressive Compositional Structure: Can Learning from $D$ Tasks Generalize to $D^T$ Tasks? (ICML 2025 · similarity 0.90)
- LeaPformer: Enabling Linear Transformers for Autoregressive and Simultaneous Tasks via Learned Proportions (ICML 2024 · similarity 0.89)
- How do Transformers Perform In-Context Autoregressive Learning ? (ICML 2024 · similarity 0.89)
- Algorithmic Capabilities of Random Transformers (NeurIPS 2024 · similarity 0.88)
- Flexible Language Modeling in Continuous Space with Transformer-based Autoregressive Flows (NeurIPS 2025 · similarity 0.88)
- Limits of Transformer Language Models on Learning to Compose Algorithms (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.