Transformers Trained via Gradient Descent Can Provably Learn a Class of Teacher Models
Chenyang Zhang, Qingyue Zhao, Quanquan Gu, Yuan Cao
Abstract
Transformers have achieved great success across a wide range of applications, yet the theoretical foundations underlying their success remain largely unexplored. To demystify the strong capacities of transformers applied to versatile scenarios and tasks, we theoretically investigate utilizing transformers as students to learn from a class of teacher models. Specifically, the teacher models covered in our analysis encompass convolution layers with average pooling, graph convolution layers, and various classic statistical learning models, including a variant of sparse token selection models (Sanford et al., 2023; Wang et al., 2024) and group-sparse linear predictors (Zhang et al., 2025c). When learning from this class of teacher models, we prove that one-layer transformers with simplified "position-only" attention can successfully recover all parameter blocks of the teacher models, thus achieving the optimal population loss. Building upon the efficient mimicry of trained transformers towards teacher models, we further demonstrate that they can generalize well to a broad class of out-of-distribution data under mild assumptions. The key in our analysis is to identify a fundamental bilinear structure shared by various learning tasks, which enables us to establish unified learning guarantees for these tasks when treating them as teachers for transformers.
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 1354e002-ce05-4e6f-ac77-e844288f3298Cited by top-tier papers1
Ask how each one uses itBuilds on38
- Swin Transformer: Hierarchical Vision Transformer using Shifted WindowsZe Liu, Yutong Lin, Yue Cao, Han Hu et al.ICCV 2021 · 31,683 citations
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- An Image is Worth 16x16 Words: Transformers for Image Recognition at ScaleAlexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn et al.ICLR 2021 · 21,477 citations
- Decision Transformer: Reinforcement Learning via Sequence ModelingLili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee et al.NeurIPS 2021 · 2,557 citations
- Tokens-to-Token ViT: Training Vision Transformers from Scratch on ImageNetLi Yuan, Yunpeng Chen, Tao Wang, Weihao Yu et al.ICCV 2021 · 2,462 citations
Related papers
- Transformer Learns Optimal Variable Selection in Group-Sparse ClassificationChenyang Zhang, Xuran Meng, Yuan CaoICLR 2025
- Transformers Provably Learn Sparse Token Selection While Fully-Connected Nets CannotZixuan Wang, Stanley Wei, Daniel Hsu, Jason D. LeeICML 2024 · 22 citations
- Towards Understanding Transformers in Learning Random WalksWei Shi, Yuan CaoNeurIPS 2025 · 1 citation
- Are Transformers universal approximators of sequence-to-sequence functions?Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi et al.ICLR 2020 · 481 citations
- A Theoretical Understanding of Shallow Vision Transformers: Learning, Generalization, and Sample ComplexityHongkang Li, Meng Wang, Sijia Liu, Pin-Yu ChenICLR 2023 · 1 citation
