Task Generalization with Autoregressive Compositional Structure: Can Learning from $d$ Tasks Generalize to $D^T$ Tasks?
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 demonstrate this exponential generalization in arithmetic and language translation, extending beyond parity functions.
Cite
Text
Abedsoltan et al. "Task Generalization with Autoregressive Compositional Structure: Can Learning from $d$ Tasks Generalize to $D^T$ Tasks?." Proceedings of the 42nd International Conference on Machine Learning, 2025.Markdown
[Abedsoltan et al. "Task Generalization with Autoregressive Compositional Structure: Can Learning from $d$ Tasks Generalize to $D^T$ Tasks?." Proceedings of the 42nd International Conference on Machine Learning, 2025.](https://mlanthology.org/icml/2025/abedsoltan2025icml-task/)BibTeX
@inproceedings{abedsoltan2025icml-task,
title = {{Task Generalization with Autoregressive Compositional Structure: Can Learning from $d$ Tasks Generalize to $D^T$ Tasks?}},
author = {Abedsoltan, Amirhesam and Zhang, Huaqing and Wen, Kaiyue and Lin, Hongzhou and Zhang, Jingzhao and Belkin, Mikhail},
booktitle = {Proceedings of the 42nd International Conference on Machine Learning},
year = {2025},
pages = {154-173},
volume = {267},
url = {https://mlanthology.org/icml/2025/abedsoltan2025icml-task/}
}