Task Generalization with Autoregressive Compositional Structure: Can Learning from D Tasks Generalize to DT Tasks?
Amirhesam Abedsoltan, Huaqing Zhang, Kaiyue Wen, Hongzhou Lin, Jingzhao Zhang, Mikhail Belkin
Abstract
Large language models (LLMs) exhibit remarkable task generalization, solving tasks they were never explicitly trained on with only a few demonstrations. This raises a fundamental question: When can learning from a small set of tasks generalize to a large task family? In this paper, we investigate task generalization through the lens of autoregressive compositional structure, where each task is a composition of T operations, and each operation is among a finite family of D subtasks. This yields a total class of size D T . We first show that generalization to all D T tasks is theoretically achievable by training on only Õ(D) tasks. Empirically, we demonstrate that Transformers achieve such exponential task generalization on sparse parity functions via In-context Learning (ICL) and chain-of-thought (CoT) reasoning. We further show generalization in arithmetic and translation, beyond parity functions.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 1aa50b5c-3353-4b3d-b912-0f89b5b67935Cited by top-tier papers4
- Efficient and Minimax Optimal In-context Nonparametric Regression with TransformersMichelle Ching, Ioana Popescu, Nico Smith, Tianyi Ma et al.ICML 2026 · 6 citations
- Provable Sample Efficiency of Curriculum Post-Training for Transformer ReasoningDake Bu, Wei Huang, Andi Han, Atsushi Nitanda et al.ICML 2026
- Generalization in LLM Problem Solving: The Case of the Shortest PathYao Tong, Jiayuan Ye, Anastasia Borovykh, Reza ShokriICLR 2026
- Learning to Extrapolate to New Tasks: A Relational Approach to Task ExtrapolationAdam Ousherovitch, Yixin WangICML 2026
Builds on18
- Language Models are Few-Shot LearnersTom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah et al.NeurIPS 2020 · 64,255 citations
- What Can Transformers Learn In-Context? A Case Study of Simple Function ClassesShivam Garg, Dimitris Tsipras, Percy Liang, Gregory ValiantNeurIPS 2022 · 883 citations
- Faith and Fate: Limits of Transformers on CompositionalityNouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li et al.NeurIPS 2023 · 728 citations
- Measuring Compositional Generalization: A Comprehensive Method on Realistic DataDaniel Keysers, Nathanael Schärli, Nathan Scales, Hylke Buisman et al.ICLR 2020 · 401 citations
- Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm SelectionYu Bai, Fan Chen, Huan Wang, Caiming Xiong et al.NeurIPS 2023 · 356 citations
Related papers
- Compositional Capabilities of Autoregressive Transformers: A Study on Synthetic, Interpretable TasksRahul Ramesh, Ekdeep Singh Lubana, Mikail Khona, Robert P. Dick et al.ICML 2024 · 17 citations
- Chain-of-Thought Provably Enables Learning the (Otherwise) UnlearnableChenxiao Yang, Zhiyuan Li, David WipfICLR 2025
- Learning to grok: Emergence of in-context learning and skill composition in modular arithmetic tasksTianyu He, Darshil Doshi, Aritra Das, Andrey GromovNeurIPS 2024 · 52 citations
- Everything Everywhere All at Once: LLMs can In-Context Learn Multiple Tasks in SuperpositionZheyang Xiong, Ziyang Cai, John Cooper, Albert Ge et al.ICML 2025
- What Algorithms can Transformers Learn? A Study in Length GeneralizationHattie Zhou, Arwen Bradley, Etai Littwin, Noam Razin et al.ICLR 2024 · 189 citations
