Research map / Large language models / Reasoning and chain of thought
Task Generalization with Autoregressive Compositional Structure: Can Learning from $D$ Tasks Generalize to $D^T$ 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.
Theoretically and empirically shows that models trained on Õ(D) tasks can generalize to D^T tasks when the task class has autoregressive compositional structure, demonstrated with chain-of-thought sparse parity from small GPT-2 models trained from scratch, plus arithmetic and translation.
Problem
When can learning from a small set of training tasks generalize to a much larger task family, and how many training tasks are needed to generalize to an exponentially large class of compositional tasks?
Approach
Defines the AutoRegressive Compositional (ARC) task class, where each task is a composition of T steps each chosen from D subtasks, and shows that identifying each step's subtask from per-component coverage needs only about D ln(DT) training tasks, enabling generalization to D^T tasks; chain-of-thought turns sparse parity with breadth C(d,k) into ARC(d,k).
Key results
- Training on 306 parity tasks with d=30, k=15 gives 98.10% accuracy on unseen tasks, a class of about 155 million tasks.
- With CoT, d=15, k=7 reaches 99.12% accuracy on unseen tasks after 121 training tasks, covering 6400 total tasks.
- Excluding all training tasks with s2=5 drops generalization to 45.6%, while excluding tasks containing both 4 and 6 keeps 96.2%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use CoT to decompose parity prediction into autoregressive subtasks | Transformer on sparse parity · Unseen parity tasks | accuracy | Helped |
| Train with 121 tasks using CoT | Transformer on sparse parity (d=15, k=7) · Unseen parity tasks | accuracy: 99.12% | Helped |
| Exclude tasks containing both 4 and 6 from training | Transformer with CoT on sparse parity · Missing Pair (d=10, k=3) | generalization accuracy: 96.2% | Hurt |
| Exclude training tasks with s2 = 5 | Transformer with CoT on sparse parity · Missing Coordinate (d=10, k=3) | generalization accuracy: 45.6% | Hurt |
| Exclude a missing coordinate instead of a missing pair | Transformer with CoT on sparse parity · Adversarial task selection (d=10, k=3) | generalization accuracy: 96.2% vs 45.6% | Mixed |
| Use i.i.d. task sampling rather than adversarial positional exclusion | Transformer with CoT on sparse parity · Unseen parity tasks | generalization accuracy | Helped |
| Use ICL without CoT (ablation) | Transformer on sparse parity · In-distribution parity generalization | accuracy | Hurt |
| Increase secret length k while holding d = 15 fixed | Transformer with CoT on sparse parity · Unseen parity tasks | accuracy | No clear effect |
| Scale training tasks approximately with d ln(d) while increasing d and k | Transformer with CoT on sparse parity · Unseen parity tasks | accuracy: 98.51; 99.12; 98.67; 98.60; 98.10 | No clear effect |
Surprising findings
- Adversarially chosen training task subsets (missing a positional coordinate) cause generalization to fail even when the excluded set is exponentially large.
- Translation chain experiments showed linear rather than logarithmic dependence on sequence length T, attributed to error accumulation.
Limitations
- Experiments are on small GPT-2-style models trained from scratch on synthetic tasks with a 3-layer, 1-head configuration.
- The theory assumes finite subtask families and uses a per-component coverage condition; the empirical translation T-scaling deviates from theory.
Experimental setup
Datasets: synthetic sparse parity sequences synthetic arithmetic sequences synthetic multi-language word translation chains
Benchmarks: held-out parity tasks held-out arithmetic tasks multi-step translation chains
Metrics: accuracy linear probe accuracy
Baselines: ICL without CoT
Builds on: chain-of-thought prompting sparse parity in-context learning study (Wies et al.)
Closest papers by meaning
- Compositional Capabilities of Autoregressive Transformers: A Study on Synthetic, Interpretable Tasks (ICML 2024 · similarity 0.90)
- Scaling can lead to compositional generalization (NeurIPS 2025 · similarity 0.89)
- Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasks (NeurIPS 2024 · similarity 0.89)
- Limits of Transformer Language Models on Learning to Compose Algorithms (NeurIPS 2024 · similarity 0.89)
- Extrapolation by Association: Length Generalization Transfer In Transformers (NeurIPS 2025 · similarity 0.88)
- Assign and Add: A Mechanistic Study of Compositional Arithmetic (NeurIPS 2026 · 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.