Transformers Trained via Gradient Descent Can Provably Learn a Class of Teacher Models

Chenyang Zhang; Qingyue Zhao; Quanquan Gu; Yuan Cao

cs.LGMar 24, 2026

Transformers Trained via Gradient Descent Can Provably Learn a Class of Teacher Models

Authors: Chenyang Zhang, Qingyue Zhao, Quanquan Gu, Yuan Cao

Organizations: The University of Hong Kong · University of California, Los Angeles

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 include 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., 2025]. 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.

Figures & tables

Panel of Figure 1 — Figure 1: Visualization of parameter matrices for the transformer $\widetilde{\text{TF}}$ in ( 2.5 ), obtained after training to learn the teacher model $f^{*}$ and achieving loss convergence. The formal illustration of the loss function and training algorithm is provided in the next section.

Panel of Figure 2 — Figure 2: Excess training loss, excess OOD test loss (both in log-log scales), and cosine similarity between the value matrix $\mathbf{W}_{V}$ of one layer transformer ( 2.4 ), and ground truth value matrix $\mathbf{V}^{*}$ . These results are presented for six experimental sets, which originate from four distinct tasks.

Panel of Figure 3 — Figure 3: Heatmap of attention score matrix $\mathbf{S}^{(T)}$ when the training loss converges. The results are presented for six different experimental sets, indicated by the captions of sub-figures.

Panel of Figure 4 — Figure 4: Training loss and cosine similarity between the value matrix $\mathbf{W}_{V}$ of the one-layer transformer ( 2.4 ), and convolution kernel matrix $\mathbf{V}^{*}$ of the pre-trained teacher CNN.

Panel of Figure 5 — Figure 5: Heatmap of the ground truth softmax scores of average pooling, Heatmap of the attention scores $\mathbf{S}^{(T)}$ of trained one-layer transformer when loss converges, and an image example in MNIST.

Appendix figures & tables5 assets

Supplementary material from the paper’s appendix.

Appendix

Symbol	Meaning
$\mathbf{V}^{*}$	Ground truth value matrix in $f^{*}$ , a $M\times D$ matrix.
$\mathbf{S}^{*}$	Ground truth softmax score matrix in $f^{*}$ , a $D\times D$ column-stochastic matrix.
$D$	Sequence length (number of input tokens).
$d$	Feature dimension of each token.
$K$	Number of none zero entries in each column of $\mathbf{S}^{*}$ . It can represent:
	(i) the pooling size in CNN and pooling layer,

Appendix

Table 1: Key variables and their meanings.

Panel of Figure 6 — Figure 6: Excess training loss, excess OOD test loss (both in log-log scales), and cosine similarity between the value matrix $\mathbf{W}_{V}$ of one layer transformer ( 2.4 ), and ground truth value matrix $\mathbf{V}^{*}$ . These results are presented for experiments where training data is generated from Student-T distribution.

Panel of Figure 7 — Figure 7: Heatmap of attention score matrix $\mathbf{S}^{(T)}$ when the training loss converges. These results are presented for experiments where training data is generated from Student-T distribution.

Panel of Figure 8 — Figure 8: Excess training loss, excess OOD test loss (both in log-log scales), and cosine similarity between the value matrix $\mathbf{W}_{V}$ of one layer transformer ( 2.4 ), and ground truth value matrix $\mathbf{V}^{*}$ . These results are presented for experiments where training data is generated from mean-centered Gumbel distribution.

Panel of Figure 9 — Figure 9: Heatmap of attention score matrix $\mathbf{S}^{(T)}$ when the training loss converges. These results are presented for experiments where training data is generated from mean-centered Gumbel distribution.

Transformer

Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.

See all 540 papers

Generalization Analysis of Transformers in Distribution Regression
Jun 28, 2026Peilin Liu, Ding-Xuan ZhouNeural Network Generalization Transformer Attention
School of Mathematics and Statistics, University of Sydney, Sydney, NSW 2006, Australia
Transformers Learn Latent Mixture Models In-Context via Mirror Descent
Apr 12, 2026Francesco D'Angelo, Nicolas FlammarionIn-Context Learning Transformer Attention
TML Lab, EPFL, Switzerland
Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent
May 7, 2026Chenyang Zhang, Yuan CaoLogistic Regression Self-Attention
School of Computing & Data Science, The University of Hong Kong

Jun 28, 2026J/K move · Enter open · S save

Generalization Analysis of Transformers in Distribution Regression

Peilin Liu, Ding-Xuan Zhou

School of Mathematics and Statistics, University of Sydney, Sydney, NSW 2006, Australia

Neural Network Generalization Transformer Attention Distributional Regression Statistical Learning Theory

Open paper AskPDF arXiv

In recent years, models based on the Transformer architecture have seen widespread applications and have become one of the core tools in the field of deep learning. Numerous successful techniques, such as parameter-efficient fine-tuning and efficient scaling, have been proposed surrounding their applications to further enhance performance. However, the success of these strategies has always lacked the support of rigorous mathematical theory. To study the underlying mechanisms behind Transformers and related techniques, we first propose a Transformer learning framework motivated by distribution regression, with distributions being inputs, connect a two-stage sampling process with natural language processing, and present a mathematical formulation of the attention mechanism called attention operator. We demonstrate that by the attention operator, Transformers can compress distributions into function representations without loss of information. Moreover, with the advantages of our novel attention operator, Transformers exhibit a stronger capability to learn functionals with more complex structures than convolutional neural networks and fully connected networks. Finally, we obtain a generalization bound within the distribution regression framework. Through the aforementioned theoretical results, we further discuss some successful techniques emerging with large language models (LLMs), such as prompt tuning, parameter-efficient fine-tuning, and efficient scaling. We also provide theoretical insights behind these techniques within our novel analysis framework.

Abstract

Transformer

Explore similar work

Generalization Analysis of Transformers in Distribution Regression

Transformers Learn Latent Mixture Models In-Context via Mirror Descent

Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent