CoChat
Sign Up
PreprintCited by 2Machine Learning and AlgorithmsNeural Networks and ApplicationsStochastic Gradient Optimization Techniques

Learning Compositional Functions with Transformers from Easy-to-Hard Data

Open access
Wang, Zixuan, Nichani, Eshaan, Bietti, Alberto, Damian, Alex, Hsu, Daniel, Lee, Jason D., Wu, Denny·2025·arXiv
DOI 10.48550/arxiv.2505.23683arXiv 2505.23683
Open this paper in CoChat — free accountOpen PDF

Abstract

Transformer-based language models have demonstrated impressive capabilities across a range of complex reasoning tasks. Prior theoretical work exploring the expressive power of transformers has shown that they can efficiently perform multi-step reasoning tasks involving parallelizable computations. However, the learnability of such constructions, particularly the conditions on the data distribution that enable efficient learning via gradient-based optimization, remains an open question. Towards answering this question, in this work we study the learnability of the $k$-fold composition task, which requires computing an interleaved composition of $k$ input permutations and $k$ hidden permutations, and can be expressed by a transformer with $O(\log k)$ layers. On the negative front, we prove a Statistical Query (SQ) lower bound showing that any SQ learner that makes only polynomially-many queries to an SQ oracle for the $k$-fold composition task distribution must have sample size exponential in $k$, thus establishing a statistical-computational gap. On the other hand, we show that this function class can be efficiently learned, with runtime and sample complexity polynomial in $k$, by gradient descent on an $O(\log k)$-depth transformer via two different curriculum learning strategies: one in which data consists of $k'$-fold composition functions with $k' \le k$ presented in increasing difficulty, and another in which all such data is presented simultaneously. Our work sheds light on the necessity and sufficiency of having both easy and hard examples in the data distribution for transformers to learn complex compositional tasks.

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Integrity signals

Registry

DataCite · 10.48550/arxiv.2505.23683

Cites

Statistics

Claims the paper makes

The claims and statistics of this paper haven’t been checked yet.

Run the check — free, no account needed. It reads the open-access text and checks the sources it cites.

A few free checks a day · papers already verified open instantly · results are public

Claims other papers make about this one

Each is checked against this paper’s own text: does it say what the citing paper says it does?

  • Wang et al. (2025) proved and empirically demonstrated the benefit of training with a good mixture of problems where the number of functions to be composed was appropriately controlled.

    “Wang et al. (2025) who, in the context of learning to compose functions, proved and also empirically demonstrated the benefit of training with a 'good mixture' of problems where the number of functions to be composed was appropriately controlled.”

    From Beyond Statistical Learning: Exact Learning Is Essential for General Intelligence (2025) · cites this paper as Wang et al. (2025)

Self-check

Percentages, sums and reported differences re-computed from the paper’s own numbers.

The self-check needs the paper’s full text.

Cites

No references could be resolved yet.

Research works printed but not linked: 67

  1. Emmanuel Abbe, Enric Boix-Adsera, Matthew S Brennan, Guy Bresler, and Dheeraj Nagaraj. The staircase property: How hierarchical structure can guide deep learning. Advances in Neural Information Processing Systems, 34:26989-27002, 2021.
  2. Emmanuel Abbe, Enric Boix Adsera, and Theodor Misiakiewicz. The merged-staircase property: a necessary and nearly sufficient condition for sgd learning of sparse functions on two-layer neural networks. In Conference on Learning Theory, pages 4782-4887. PMLR, 2022.
  3. Emmanuel Abbe, Elisabetta Cornacchia, and Aryo Lotfi. Provable advantage of curriculum learning on parity targets with mixed inputs. Advances in Neural Information Processing Systems, 36:24291-24321, 2023.
  4. Emmanuel Abbe, Samy Bengio, Aryo Lotfi, Colin Sandon, and Omid Saremi. How far can transformers reason? The locality barrier and inductive scratchpad. In Advances in Neural Information Processing Systems, 2024.
  5. Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn to implement preconditioned gradient descent for in-context learning. In Advances in Neural Information Processing Systems, volume 36, 2024.
  6. Sepehr Assadi and Vishvajeet N. Graph streaming lower bounds for parameter estimation and property testing via a streaming xor lemma. In Proceedings of the 53rd Annual ACM SIGACT Symposium on Theory of Computing, pages 612-625, 2021.
  7. Gregor Bachmann and Vaishnavh Nagarajan. The pitfalls of next-token prediction. In International Conference on Machine Learning (ICML), 2024.
  8. Yoshua Bengio, Jérôme Louradour, Ronan Collobert, and Jason Weston. Curriculum learning. In Proceedings of the 26th Annual International Conference on Machine Learning, ICML '09, page 41-48, New York, NY, USA, 2009. Association for Computing Machinery. ISBN 9781605585161. doi: 10.1145/1553374.1553380. URL https://doi.org/10.1145/ 1553374.1553380. · DOI · Paper page
  9. Satwik Bhattamishra, Michael Hahn, Phil Blunsom, and Varun Kanade. Separations in the representational capabilities of transformers and recurrent architectures. arXiv preprint arXiv:2406.09347, 2024. · arXiv
  10. Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: A memory viewpoint. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
  11. Stephanie Chan, Adam Santoro, Andrew Lampinen, Jane Wang, Aaditya Singh, Pierre Richemond, James McClelland, and Felix Hill. Data distributional properties drive emergent in-context learning in transformers. In Advances in neural information processing systems, 2022.
  12. Lijie Chen, Binghui Peng, and Hongxun Wu. Theoretical limitations of multi-layer transformer. arXiv preprint arXiv:2412.02975, 2024. · arXiv
  13. Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi- head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402.19442, 2024. · arXiv
  14. Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. arXiv preprint arXiv:2409.10559, 2024. · arXiv
  15. Sitao Cheng, Liangming Pan, Xunjian Yin, Xinyi Wang, and William Yang Wang. Understand- ing the interplay between parametric and contextual knowledge for large language models. arXiv preprint arXiv:2410.08414, 2024. · arXiv
  16. Alexandru Damian, Jason Lee, and Mahdi Soltanolkotabi. Neural networks can learn represen- tations with gradient descent. In Conference on Learning Theory, pages 5413-5452. PMLR, 2022.
  17. Yuntian Deng, Yejin Choi, and Stuart Shieber. From explicit cot to implicit cot: Learning to internalize cot step by step. arXiv preprint arXiv:2405.14838, 2024. · arXiv
  18. Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.21783, 2024. · arXiv
  19. Nouha Dziri, Ximing Lu, Melanie Sclar, Xiang Lorraine Li, Liwei Jiang, Bill Yuchen Lin, Peter West, Chandra Bhagavatula, Ronan Le Bras, Jena D. Hwang, Soumya Sanyal, Sean Welleck, Xiang Ren, Allyson Ettinger, Zaid Harchaoui, and Yejin Choi. Faith and fate: Limits of transformers on compositionality (2023). arXiv preprint arXiv:2305.18654, 2023. · arXiv
  20. Jeffrey L Elman. Learning and development in neural networks: The importance of starting small. Cognition, 48(1):71-99, 1993.
  21. Guhao Feng, Bohang Zhang, Yuntian Gu, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. Advances in Neural Information Processing Systems, 36, 2024.
  22. Cheng Gao, Yuan Cao, Zihao Li, Yihan He, Mengdi Wang, Han Liu, Jason Matthew Klusowski, and Jianqing Fan. Global convergence in training large-scale transformers. arXiv preprint arXiv:2410.23610, 2024. · arXiv
  23. Fabian Gloeckle, Badr Youbi Idrissi, Baptiste Rozière, David Lopez-Paz, and Gabriel Syn- naeve. Better & faster large language models via multi-token prediction. arXiv preprint arXiv:2404.19737, 2024. · arXiv
  24. William Timothy Gowers and Emanuele Viola. Interleaved group products. SIAM Journal on Computing, 48(2):554-580, 2019.
  25. Jianhao Huang, Zixuan Wang, and Jason D Lee. Transformers learn to implement multi-step gradient descent with chain of thought. arXiv preprint arXiv:2502.21212, 2025. · arXiv
  26. Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint arXiv:2310.05249, 2023. · arXiv
  27. Sungjin Im, Ravi Kumar, Silvio Lattanzi, Benjamin Moseley, and Sergei Vassilvitskii. Mas- sively parallel computation: Algorithms and applications. Foundations and Trends® in Optimization, 5(4):340-417, 2023.
  28. Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
  29. Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Transformers are better than state space models at copying. arXiv preprint arXiv:2402.01032, 2024. · arXiv
  30. Michael Kearns. Efficient noise-tolerant learning from statistical queries. J. ACM, 45(6): 983-1006, 1998.
  31. Juno Kim and Taiji Suzuki. Transformers provably solve parity efficiently with chain of thought. arXiv preprint arXiv:2410.08633, 2024. · arXiv
  32. Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint arXiv:1412.6980, 2014. · arXiv
  33. Jack Lanchantin, Shubham Toshniwal, Jason Weston, and Sainbayar Sukhbaatar. Learning to reason and memorize with self-notes. In Advances in Neural Information Processing Systems, volume 36, 2024.
  34. Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, Yuhuai Wu, Behnam Neyshabur, Guy Gur-Ari, and Vedant Misra. Solving quantitative reasoning problems with language models. In Advances in Neural Information Processing Systems, volume 35, 2022.
  35. Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In ICML, 2023.
  36. Zhiyuan Li, Hong Liu, Denny Zhou, and Tengyu Ma. Chain of thought empowers transformers to solve inherently serial problems. arXiv preprint arXiv:2402.12875, 2024. · arXiv
  37. Hunter Lightman, Vineet Kosaraju, Yura Burda, Harri Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe. Let's verify step by step. arXiv preprint arXiv:2305.20050, 2023. · arXiv
  38. Aixin Liu, Bei Feng, Bing Xue, Bingxuan Wang, Bochao Wu, Chengda Lu, Chenggang Zhao, Chengqi Deng, Chenyu Zhang, Chong Ruan, et al. Deepseek-v3 technical report. arXiv preprint arXiv:2412.19437, 2024. · arXiv
  39. Bingbin Liu, Jordan T. Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transform- ers learn shortcuts to automata. In The Eleventh International Conference on Learning Repre- sentations, 2023. URL https://openreview.net/forum?id=De4FYqjFueZ.
  40. Qian Liu, Xiaosen Zheng, Niklas Muennighoff, Guangtao Zeng, Longxu Dou, Tianyu Pang, Jing Jiang, and Min Lin. Regmix: Data mixture as regression for language model pre-training. arXiv preprint arXiv:2407.01492, 2024. · arXiv
  41. William Merrill and Ashish Sabharwal. The expresssive power of transformers with chain of thought. arXiv preprint arXiv:2310.07923, 2023. · arXiv
  42. William Merrill and Ashish Sabharwal. The parallelism tradeoff: Limitations of log-precision transformers. Transactions of the Association for Computational Linguistics, 11:531-545, 2023.
  43. Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. In International Conference on Machine Learning (ICML), 2024.
  44. Eshaan Nichani, Jason D Lee, and Alberto Bietti. Understanding factual recall in transformers via associative memories. arXiv preprint arXiv:2412.06538, 2024. · arXiv
  45. Noam Nisan and Avi Wigderson. Rounds in communication complexity revisited. SIAM Journal on Computing, 22(1):211-219, 1993. doi: 10.1137/0222016. URL https://doi. org/10.1137/0222016. · DOI
  46. Maxwell Nye, Anders Johan Andreassen, Guy Gur-Ari, Henryk Michalewski, Jacob Austin, David Bieber, David Dohan, Aitor Lewkowycz, Maarten Bosma, David Luan, Charles Sutton, and Augustus Odena. Show your work: scratchpads for intermediate computation with language models. arXiv preprint arXiv:2112.00114, 2021. · arXiv
  47. Abhishek Panigrahi, Bingbin Liu, Sadhika Malladi, Andrej Risteski, and Surbhi Goel. Pro- gressive distillation induces an implicit curriculum. arXiv preprint arXiv:2410.05464, 2024. · arXiv
  48. Christos H. Papadimitriou and Michael Sipser. Communication complexity. Journal of Computer and System Sciences, 28(2):260-269, 1984. ISSN 0022-0000. doi: https://doi.org/10. 1016/0022-0000(84)90069-2. URL https://www.sciencedirect.com/science/ article/pii/0022000084900692. · DOI
  49. Binghui Peng, Srini Narayanan, and Christos Papadimitriou. On limitations of the trans- former architecture. In First Conference on Language Modeling, 2024. URL https: //openreview.net/forum?id=KidynPuLNW.
  50. Yunwei Ren, Zixuan Wang, and Jason D Lee. Learning and transferring sparse contextual bi- grams with linear transformers. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
  51. Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. In Advances in Neural Information Processing Systems 36, 2023.
  52. Clayton Sanford, Bahare Fatemi, Ethan Hall, Anton Tsitsulin, Mehran Kazemi, Jonathan Hal- crow, Bryan Perozzi, and Vahab Mirrokni. Understanding transformer reasoning capabilities via graph algorithms. arXiv preprint arXiv:2405.18512, 2024. · arXiv
  53. Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Transformers, parallel computation, and logarithmic depth. In Forty-First International Conference on Machine Learning, 2024.
  54. Sainbayar Sukhbaatar, Jason Weston, and Rob Fergus. End-to-end memory networks. In Advances in Neural Information Processing Systems, volume 28, 2015.
  55. Balázs Szörényi. Characterizing statistical query learning: simplified notions and proofs. In International Conference on Algorithmic Learning Theory, pages 186-200. Springer, 2009.
  56. Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. In Advances in Neural Information Processing Systems (NeurIPS), 2023.
  57. Jonathan Uesato, Nate Kushman, Ramana Kumar, Francis Song, Noah Siegel, Lisa Wang, Antonia Creswell, Geoffrey Irving, and Irina Higgins. Solving math word problems with process-and outcome-based feedback. arXiv preprint arXiv:2211.14275, 2022. · arXiv
  58. Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D. Lee. Transformers provably learn sparse token selection while fully-connected nets cannot. In ICML, 2024.
  59. Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, and Denny Zhou. Chain-of-thought prompting elicits reasoning in large language models. In Advances in Neural Information Processing Systems, volume 35, 2022.
  60. Kaiyue Wen, Huaqing Zhang, Hongzhou Lin, and Jingzhao Zhang. From sparse dependence to sparse attention: unveiling how chain-of-thought enhances transformer sample efficiency. arXiv preprint arXiv:2410.05459, 2024. · arXiv
  61. Jason Weston, Sumit Chopra, and Antoine Bordes. Memory networks. arXiv preprint arXiv:1410.3916, 2014. · arXiv
  62. Jason Weston, Antoine Bordes, Sumit Chopra, Alexander M Rush, Bart Van Merriënboer, Armand Joulin, and Tomas Mikolov. Towards ai-complete question answering: A set of prerequisite toy tasks. arXiv preprint arXiv:1502.05698, 2015. · arXiv
  63. Sang Michael Xie, Hieu Pham, Xuanyi Dong, Nan Du, Hanxiao Liu, Yifeng Lu, Percy S Liang, Quoc V Le, Tengyu Ma, and Adams Wei Yu. Doremi: Optimizing data mixtures speeds up language model pretraining. In Advances in Neural Information Processing Systems, volume 36, 2024.
  64. Sohee Yang, Elena Gribovskaya, Nora Kassner, Mor Geva, and Sebastian Riedel. Do large language models latently perform multi-hop reasoning? arXiv preprint arXiv:2402.16837, 2024. · arXiv
  65. Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Tom Griffiths, Yuan Cao, and Karthik Narasimhan. Tree of thoughts: deliberate problem solving with large language models. Advances in Neural Information Processing Systems, 36, 2024.
  66. Amir Yehudayoff. Pointer chasing via triangular discrimination. Combinatorics, Probability and Computing, 29(4):485-494, 2020.
  67. Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. arXiv preprint arXiv:2306.09927, 2023. · arXiv

Cited by

Sorted by citations

Authors
Wang, Z., Nichani, E., Bietti, A., Damian, A., Hsu, D., Lee, J. D., & Wu, D.
Affiliations
  • Princeton University (US)
  • Flatiron Institute (US)
  • Columbia University (US)
ORCID
Venue
arXiv
Publisher
arXiv
Published
2025
ISSN
2331-8422 (ISSN-L)
DOI
10.48550/arxiv.2505.23683
arXiv
2505.23683
OpenAlex
W4416613006
Cited by
0
References
0
Type
posted content
Language
en
Open access
Licence
arxiv.org/licenses/nonexclusive-distrib/1.0/
Locations
  • arXiv (Cornell University) (submittedVersion)
  • arXiv (Cornell University)
Funders
  • National Science Foundation (2212262, 2144994, 2002272, CAREER, 2107304)
  • Simons Institute for the Theory of Computing, University of California Berkeley
  • Office of Naval Research (N00014)
Topics
  • Physical Sciences › Computer Science › Artificial Intelligence › Machine Learning and Algorithms
  • Physical Sciences › Computer Science › Artificial Intelligence › Neural Networks and Applications
  • Physical Sciences › Computer Science › Artificial Intelligence › Stochastic Gradient Optimization Techniques
Keywords
compositional functions, transformer models, learnability, statistical-computational gap, curriculum learning, gradient descent, sample complexity

About these checks

Registry, retraction, Cites and statistics results are computed, not judged, with one exception: a printed reference that no index lists is read once by a model before it counts as a citation error. Claims are extracted by a model and kept only when their quote appears verbatim in the text; each cited source is then checked, with a model’s help, for whether it says what the claim relies on. Those verdicts can be wrong — read the claim’s quote and the cited source before relying on one.

Statistics and arithmetic are recomputed, not judged.

Confidence bands are shown only for fields calibrated on our annotated test set. None is calibrated yet, so no band is shown.

What counts as a claim ›How we measure confidence ›How we check this ›

Are you an author, or do you see a problem with this page? Request a review