Research map / Architectures and efficiency / Attention and transformers
Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic 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 GPT-style transformers on linear modular functions z = ax + by mod p with a subset of task vectors, showing out-of-distribution in-context generalization emerges with enough pre-training tasks and analyzing the learned circuits.
Problem
How in-context learning and skill composition emerge in autoregressive transformers, studied on a controlled algorithmic dataset of modular linear functions where test tasks are unseen during pre-training.
Approach
Defines a family of p^2 modular linear tasks labeled by (a,b), pre-trains on a subset of tasks with structured rectangular task sampling and aligned batches, and shows a memorization-to-OOD-generalization transition as the number of tasks grows, with mechanistic analysis revealing circle-of-circles representations and an algorithmic shift from ratio matching to modular regression in deeper models.
Key results
- Out-of-distribution generalization emerges only when the number of pre-training tasks n_i.d. is at least 256 for a six-layer model at p=29.
- Two-layer models are the smallest capable of OOD generalization, but their OOD phase is transient for deeper models requiring early stopping.
- The d=4 model shows an algorithmic shift from ratio matching to modular regression as in-context examples increase to about 10.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| increase the number of pre-training tasks | GPT-style transformer · modular arithmetic out-of-distribution tasks | accuracy | Helped |
| increase the number of pre-training tasks | GPT-style transformer · modular arithmetic in-context learning | few-shot sample efficiency | Helped |
| increase the number of pre-training tasks | GPT-style transformer · modular arithmetic in-distribution pre-training | accuracy | Hurt |
| use a two-block transformer instead of a deeper model | GPT-style transformer · modular arithmetic out-of-distribution generalization | accuracy | Mixed |
| use a four-block rather than a two-block transformer | GPT-style transformer · modular arithmetic with many in-context examples | out-of-distribution accuracy | Helped |
| use a four-block rather than a two-block transformer | GPT-style transformer · modular arithmetic with many in-context examples | label-corruption robustness | Helped |
| use a six-block transformer | GPT-style transformer · modular arithmetic out-of-distribution generalization during training | accuracy | Mixed |
| early-stop deeper models | GPT-style transformer · modular arithmetic out-of-distribution generalization | accuracy | Helped |
| sample pre-training task vectors randomly instead of using the rectangular rule (ablation) | six-block GPT-style transformer · modular arithmetic out-of-distribution generalization | accuracy | Hurt |
| remove balanced task representation and aligned input sequences from batches (ablation) | GPT-style transformer · modular arithmetic pre-training | training loss | Hurt |
| corrupt one in-context label (ablation) | four-block transformer · modular arithmetic with longer context sequences | prediction performance | No clear effect |
| corrupt multiple in-context labels at random locations (ablation) | four-block transformer · modular arithmetic with longer context sequences | prediction performance | Mixed |
Surprising findings
- The two-layer model skips the in-distribution generalization phase and jumps directly from memorization to OOD generalization.
- Random task selection prevented learning; structured rectangular task sampling and aligned batches were needed for training to succeed.
Limitations
- Restricted to algorithmic synthetic datasets; transfer to realistic LLMs unclear.
- Interpretability of deeper models (d=4,6) is much harder; components not fully understood.
- LayerNorm role left for future work.
Experimental setup
Datasets: synthetic modular arithmetic sequences (p=29)
Benchmarks: in-distribution and out-of-distribution test sets of modular linear tasks
Metrics: accuracy loss vs number of shots phase diagram classification (75% threshold)
Baselines: ratio matching modular regression
Builds on: grokking in-context linear regression induction heads RoPE AdamW
Closest papers by meaning
- Assign and Add: A Mechanistic Study of Compositional Arithmetic (NeurIPS 2026 · similarity 0.91)
- Task Generalization with Autoregressive Compositional Structure: Can Learning from $D$ Tasks Generalize to $D^T$ Tasks? (ICML 2025 · similarity 0.89)
- Emergence in non-neural models: grokking modular arithmetic via average gradient outer product (ICML 2025 · similarity 0.88)
- Can In-context Learning Really Generalize to Out-of-distribution Tasks? (ICLR 2025 · similarity 0.88)
- Teaching Arithmetic to Small Transformers (ICLR 2024 · 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.