Transformers Trained via Gradient Descent Can Provably Learn a Class of Teacher Models
Authors: Chenyang Zhang
Qingyue Zhao
Quanquan Gu
Yuan Cao
Venue: ICLR 2026 Poster (2026)
Keywords: transformer
gradient descent
teacher-student setting
Primary area: learning theory

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.

Published as a conference paper at ICLR 2026
TRANSFORMERS TRAINED VIA GRADIENT DESCENT CAN PROVABLY LEARN A CLASS OF TEACHER MOD- ELS
Chenyang Zhang', Qingyue Zhao', Quanquan Gu Yuan Cao f The University of Hong Kong 1 Universsiyy of California, Los Angeles chyzhang@connece hku. hoku.hi hk, zhaogy2 2ucclllll .euu qguêcs. ucla. edu, yuancaochku.hk
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 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 analy- sis 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 teacher mod- els, we prove that one-layer transformers with simplified "position-only" attention can successfully recover all parameter blocks of the teacher models, thus achiev- ing the optimal population loss. Building upon the efficient mimicry of trained transformers towards teacher models, we further demonstrate that they can gen- eralize well 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.
INTRODUCTION
Transformers have rapidly become cornerstone in the field of modern machine learning, demon- strating exceptional performance and versatility across diverse applications, including natural lan- guage processing Vaswani al., 2017 Radford et al. 2019 |penAl| 2023 Devlin 2018 Achiam
al. 2023 Vig Belinkov 2019 Touvron et al. [2023] Ouyang et al.) 2022). computer vision
(Dosovitskiy et al. 2020 Rao et al. [222] Liu et al. [202 Yuan et al. 202 |
al. 2025b a)
and reinforcement learning (Jumper et al. 2021 Chen et al. 2021] Janner et al. 2021 Reed et al.
2022 Acting critical component of transformers, self-attention layers assign varying weights to features hased on their relevance and embedded positional context. This design principle intuitively endows transformers with a remarkable ability efficiently process both structural and positional information, as empirically validated in numerous applications mentioned above. However, despite their profound impact, the theoretical foundations of transformers, especially the mechanisms of how self-attention layers work, remain largely unexplored due to their intricate architecture.
Some recent theoretical studies aimed understand transformers by analyzing their capability in solving specific tasks (Zhang et al.) (0.4b) & Vardi 2025]Jelasst Jelassi et al. 2022)) Wang et al.) 2024]] Zhang et al., 2025c). Specifically, Zhhan et al. (2024b) considered in-context linear regression, and demonstrated that for Gaussian data, one-layer transformer with linear attention can perform linear regression based on the context, and then apply the obtained linear model to make predictions on query data. Later, Frei & Vardi (2025) further extended the setting to in-context linear classification, and studied the in-context benign overfitting phenomena when learning from Gaussian mixture data. Jelassi et al. (2022) investigated specific data model based the 'patch association' assumption, where an image is divided into disjoint partitions, and patches within the same partition share similar characteristics. They theoretically demonstrate that one-layer vision transformer (ViT) can extract

---

Published as a conference paper at ICLR 2026
the spatial structure among patches when trained on this data model. Wang et al. (2024) studied problem termed ssparse token selections', , where the objective is to find the average of several tokens from specific positions, and they proved that one-layer transformer
successfully solve this task on Gaussian data when the positional information of the target positions is embedded into the query token. [Zhang al. (2025c) considered group sparse linear model, where the input's label is determined by features from only one of several input feature groups (the 'label-relevant group'), ), and prove that for Gaussian data, a trained one-layer transformer can achieve correct classification by identifying features from this group and learning the ground truth linear classifier. Although these works have offered valuable insights into the underlying mechanisms of transformers, their focus on very specific learning tasks limits the generality of their theoretical findings, prompting us to seek unified theoretical framework accounting for a broader range of examples.
Despite the distinctions among the model simplifications and technical assumptions, we observe that for some learning tasks discussed above, including a variant of the sparse token selection (Sanford et al. 2023 Wang et al.) 2024), the group sparse linear predictors (Zhang et al. (025c), and patch association (Jelassi et al. 2022), their true responses are essentially given by bilinear functions. In addition, the linear attention studied in Zhang et al. (2024a); ; Frei & Vardi (2025) inherently constitutes a bilinear structure with respect to its parameter matrices. Motivated by this observation, define a general class of "teacher models" that employ bilinear structure, and investigate the setting where one-layer transformers are trained as "student" models under the supervision from these teacher models. Our framework not only encompasses the learning tasks from prior works but also covers popular, previously unexplored models such as convolution layers with average pooling and graph convolution layers on regular graphs. The purpose of our analysis is to establish unified theoretical guarantees for one-layer transformer models trained with gradient descent in learning this class of teacher models.
The major contributions of this work are as follows.
We theoretically demonstrate that one-layer transformers trained via gradient descent can effec- tively recover a general class of teacher models. To support this claim, we establish a tight con- vergence guarantee for the population loss, with matching upper and lower bounds at the rate of (+)), where T is the iteration number of gradient descent. We also establish out-of cut-of-diisributionn generalization bounds for the obtained transformer model and demonstrate that is competitive with the teacher model over wide rage of learning tasks. This illustrates the effectiveness and robustness of transformer models in learning from diverse teacher models.
Our theory covers a wide range of learning tasks, including some settings closely related those studied in (Wang et al. 2024| Zhang et al. 2025c). Specifically, Wang et al. (2024) study a type of "sparse token selection" task where the goal is to select number of target input tokens specified by a query column, and then output their average. Assuming that the positions of the target tokens
are randomly generated for each data point, the authors establish an O (lo((2) convergence rate. comparison, our setting covers a slightly different task where the target positions are fixed but are not explicitly fed to model, and our theoretical results demonstrate tight θ ( convergence rate with matching upper and lower bounds. Compared with ZZang et al. (2025c) which focuses on group sparse linear classification, that transformers can also perform efficient group sparse linear regression.
work provides complementary results and demonstrates
Experiments on both synthetic and real-world data are conducted to verify our theory through the examples of learning convolution layer with average pooling, learning graph convolution layer with regular graphs, learning sparse token selection, and group sparse linear regression. In all experiments, we can observe clear loss convergence and parameter convergence that match our theory. The experiments setup does not exactly match our theory assumptions, indicating that our theory conclusions can also hold in more practical training setups and real-data learning tasks.
2 PROBLEM SETUP
In this section, we introduce the definition of the teacher models we study this paper, and give various examples covered in our definition.
We consider a teacher model with an input matrix X E Rd of the following form:
= (
(V'ss,,,
(2.1)
2

---

Published as a conference paper at ICLR 2026
where V IR Mxd is the ground truth value matrix of the teacher model, and S RD is the ground truth softmax scores. Each column of S* has K non-zero entries equivalent to addition, o--) denotes either an identity map, ReLU, or Leaky ReLU activation function.
1 K- In
The teacher models defined in (2.1) can cover a general class functions (models). Notably, when K = and all the non-zero entries of S appear on its diagonal, S* equals the identity matrix Ip. In this scenario, the teacher model (2.1) reduces to (X)
(V) and can be seen as a single-layer neural network. Besides this naive example where S Ip, the teacher model (2.1) also includes some other common architectures and models. We discuss these examples in the following.
Example 2.1 (Single convolutional layer with average pooling). We consider a convolution layer consisting of convolution operation, average pooling, and then the activation function. The convo- lution operation essentially performed by taking inner products between each convolution ker- nel with each patch of the imput. We consider a convolution layer with M (vectorized) kernels vi:..., Viz. and consider an input consisting of D (vectorized) patches XI:...; ;.....,,p xp- In average pooling, we take averages according to partition of the 1) patches. Let G = 12.......... ..2] be disjoint partition of [D], forming pooling groups with [g,| = K, j E [J]. Then the final output of this convolution layer corresponding the j-th pooling group and the m-th kernel is given
(K EV'vin (vi xi) ovv Xl,;/K), E [M) [
iEgs
where a is the activation function, X = X1, Xg,X,, xpl R and 1g, E R is a vector whose entries are 1 for indices in 9j. and otherwise. Then, we can summarize all outputs into matrix:
FCNN(X
(V' 1/..///......./// E MR.
where V Ivi...., V E I Mxd Here, the j-th column of (X) corresponds to the output of j-th pooling group gý, and m-th row of FeNN(X corresponds to the output of m-th kernel v
To formulate the convolution layer above as teacher for transformers, we further specify the cor- respondence between each input patch and the output. The teacher model can then be given as (X) oVV* where the i-th column of S* is 1 K, with gj being the group containing
Example 2.2 (Single graph convolution layer matrix of a degree- K 1) regular graph with D nodes, and X = [x1, .2,....,pJ E IR D be the feature matrix of this graph, with each column x, (for all in [D]) representing the d-dimensional feature vector of the i-th node. A typical single graph convolution layer (Kipf & Welling 2017). with weight matrix V E Mxd is defined as
regular graph). Let A DR be an adjacency
FGcN(X) = =DVV'Xx-12AA-111
(2.2)
where Ā = A + ID is the adjacency matrix with self-connections added, and D is the diagonal degree matrix of A. For a degree- K - 1) regular graph, each node has K - neighbors, and hence each column of A contains K ones and D D K zeroes, and D K.Ip Therefore, the GCN defined in (.2) is equivalent to (X) oVV-voVV XS*) with V and S D 1-2 AD 1/2 A/K.
Example 2.3 (Sparse token selection model (Sanford et al. 2023 Wang et al. 2024 Let X [x1, X2, -......-p)) xD] R<x be a sequence of d-dimensional tokens. Given K felement index set C [D], the goal of sparse token selection to (i) select the tokens x,, i E g, and (ii) take an average over the selected tokens. Hence, can define
FSTS(X)=
Σ Xi. K
kg
- Then it is clear that f'(X) (VV*XS*) with V y Ip. S K 1g 1. E DxD and ((-) being identity map is equivalent to I'STS (X), except that [x( duplicates the output D times to match the output dimensions of a self-attention layer.
Remark 2.4. The "sparse token selection" task defined in Example 2.3 is slightly different from
that studied in Wang et al. 2024). In our setting, the index set y is specified as part of the learning objective and therefore remains fixed across all inputs. In contrast, Wang et al. (2024) considers a setting in which g is provided as part of the input, allowing target positions vary between different inputs. We remark that despite the difference, our learning task and that studied in Wang et al. (2024) essentially lead to very similar learning dynamics. We provide detailed discussion in Appendix
3

---

Published as a conference paper at ICLR 2026
Example 2.5 (Group sparse linear predictors (Zhang al. 2025c)). Let X Rd* be a sequence of d-dimensional feature groups. For a given ground truth vector v E R" and label-relevant group index i", the group sparse linear predictor will first search for the variable group x; corresponding to the label-relevant index i", and then calculate its inner product with the ground truth vector v². Hence, we define
[x1, X2, ........ XD
FGSLP (v",
Consider a teacher model f*(X) = (VV with V v by reducing M to 1, Sr = e -1p and (()) being identity map. Then similar to Example 2.3 for D times, and is essentially equivalent to FGSLP(X)
2.3 f (XX duplicates the output of FGSLP(X)
One-layer transformer. A one-layer transformer model Vaswani et al. (0017) 2020 can be defined
Dosovitskiy et al.
W K
WKW VD
TF(Z: Wo: Wk) =a
W Zs ZS
(2.3)
In this formulation, Z represents the input matrix of the transformers, obtained by concatenating the original feature matrix X with its positional encoding matrix P. Specifically, for each column x; (for all E (D) of the original feature matrix X, we concatenate it with the position encoding vector Ps. which contains the positional information of this specific index, to generate a column of as z; = xx IT. The complete positional encoding matrix is denoted and we employ an orthogonal design for P. meaning that P is an D x D) orthogonal matrix. For analytical convenience, the practice of concatenating feature and positional encoding matrices has been widely adopted in recent theoretical studies (Nichani et al. [2024] Bai et al., 2024, Wang al.) 2024 Zhang et al., 2025c). Furthermore, S(.) RXX + 1XxD denotes the softmax operator, which implements the softmax function column-wisely and Wv. We. Wx represent the value matrix, query matrix, and key matrix typical self-attention structure, respectively. Instead of studying the typical structure (2.3), we consider moderately simplified "position-only" softmax self-attention this paper, which is defined as
P P1. .PI.P..... P2..... PD),
P WKOQ)) PWOOP))= o (W ) MXD (2.4)
TF(Z; W W Wk)) =
W WrXX
VD
In comparison with the typical single-head self-attention architecture 22.3 our model (2.4) is sim-
plified from the following two aspects: (i). We re-parameterize the original key matrix W K and
query matrix We into one trainable key-query matrix W ko. which has been adopted in almost theoretical studies regarding the optimization of transformers (Tan
2023 Zhang et al. 2024b Wang et al.) 2224: Huang et al.] 2024| Frei & Vardi 2025) Zhang et al., | He et al. 2025a).
(ii). We employ an architecture such that only the positional encoding matrix P involved when calculating the softmax attention score, and the value matrix Wy only interacts with the feature
matrix X. To illustrate a rationale for this design, consider the following one-layer transformers:
)
Z W WKoZ) VD
)
TF(Z: W W WKQ)
W Wv ZS
(2.5)
where the entire input matrix Z is involved both the calculation of attention score and interac- tions with the value matrix. Empirical observations (illustrated in Figure I reveal that when the transformer model TF in (25) is used to learn a teacher model f in 2..I. substantial training pre- dominantly occurs in the left block of Wy and the 'bottom- nrght' block of W KQQ These actively trained blocks map to Wy v and W KQ respectively in model (2.4) while other parameter blocks of TF exhibit negligible changes from their initial values. Consequently, our model (.4 can be considered essentially equivalent to the transformer model TF if these rarely updated blocks within Wv and WKQ are fixed to zero. This strategy of fixing certain transformer parameters during training is widely adopted in the theoretical studies on the optimization of transformers (Wu et al.,
2023
Tarzanagh et al. 2023a; Huang et al., 2024 Saaaamoto & Sato 2024 Frei & Vardi 2025 He 2025a), and analogous "position-only attention structures are also adopted in Jelassi et al.
ct al.
2022
Wang et al. (2024).
4

---

Published as a conference paper at ICLR 2026
X
(a) Heatmap of Wv
(b) Heatmap of W KQ
Figure 1: Visualization of parameter matrices for the transformer TF in 2.5 obtained after training 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.
MAIN RESULTS
In this section, we demonstrate our theoretical conclusions of utilizing a one-layer transformer 2.4 to learn a given teacher model f in [.I).. For teacher model f parameterized with the ground truth value matrix V and ground truth softmax scores S the observed label Y for an input matrix X is assumed to be generated as:
Y (x
+ RR
(3.1)
where & E MxD is noise matrix independent of X and following a zero-mean distribution. To train a one-layer transformer (2.4), we consider the population mean squared error as the objective loss function. Specifically, given input-label pair (X, Y). the loss function is defined
1 Ex [|Y TF Wv: W [
C(Wv: W wxQ)
(3.2)
2
Here, each column of X is assumed to independently follow the standard Gaussian distribution during the training stage of (2.4. i.e. i..d by the noise component E, even the loss of the ground truth model f has an irreducible term, and we denote this term as the optimal loss, i.e.
~ NNO,, Ia) for all E (D). Due to the variance introduced
Lopt - = - [[181]]
To evaluate the performance of one-layer transformer with different Wv and W KQ. consider the excess loss defined as: L(Wv; WkQ) - Copt. While the choice population loss implicitly suggests infinite training data set-a scenario not feasible in practice-it significantly simplifies the technical challenges of conducting a rigorous optimization analysis for transformer models. This approach enables to focus the global optimization trajectories, and has been adopted in most of the recent theoretical studies regarding the optimization of transformer models (Zhang et al. 2024b
Huang et al. 2024 Wang al. 2024 Jelassi al. 2022 Frei & Vardi 2025 Zhang al. 2025c).
For the training objective loss 3.2p. we utilize the gradient descent to derive the optimal solutions for the value matrix Wv, and key-query matrix WKQ. The iterative rule for Wy and W KQ during the learning process can be expressed as
((I+1) V
w(O) - VV,, C(W KQ):
(3.3)
((+11) KQ
w - 7/ V
(3.4)
KQ
KQ
where 1/ is the learning rate, and the initializations set as W liminaries, the following theorem characterizes the convergence of gradient descent (33) and 33.
W (0)
KQ 0. Based these pre-
5

---

Published as a conference paper at ICLR 2026
Theorem 3.1. Suppose that D 2 ())) OMM-ID Under these conditions, there exists T θ
k V-Il
such that for all T 7", the following results hold.
1. The attention scores achieved by the one-layer transformer (2.4. match the ground truth softmax scores of the teacher model: ( at the T-th iteration satisfies that
D
S - s =
||
2. The value matrix Wy of the one-layer transformer (2.4) of the teacher model:
aligns with the ground truth value matrix
W v+ o D2
K
n/T
3. The excess loss is minimized with matching lower and upper bounds:
cKD4
r/T S < <<(w(7)
(W w Copt S
eKD
1/T"
where e and are two positive constants satisfying S
The proof of Theorem 3.1 is given in Appendix D Theorem 3.1 demonstrates that a one-layer transformer can Iearn the teacher model f formulated in (2.1 from two aspects. The first and second results show that the one-layer transformer's value matrix w and attention scores S converge (in the Frobenius norm) to the teacher model's ground truth value matrix V and softmax scores S respectively. This reveals that a one-layer transformer trained via gradient descent can correctly recover the teacher model by accurately learning all its core components. The third result in Theorem 3.1 shows that the training loss will eventually converge to the optimal loss at a rate of O K The third result characterizes the convergence of the training loss. It shows that the
excess loss decreases at the rate O KDD) ). with matching upper and lower bounds. We note that the x/T
factor D indicates that the convergence takes large number of iterations when the sequence length D is large. However, the matching lower bound in Theorem [.1 confirms that this rate is already optimal and cannot be improved under our current setting, In fact, this polynomial dependence on D originates from two intrinsic aspects of the learning task: (i) Since the loss is the squared Frobenius distance between two M x D matrices, it necessarily aggregates errors over all D columns, and thus scales proportionally with the sequence length; (ii) The 1/vD factor appears in the gradients of W KQ and requires W "Q scale larger achieve sufficient convergence, thereby introducing additional factors of D into the convergence rate.
2.5 our teacher model fr encompasses settings that are closely related to the learning tasks studied in Wang et al. (2024) and /Zhang et al.] (2025c). For the "sparse token selection" problem, Theorem 3.1 establishes learning guarantee for the setting in which the target index set is fixed by the learning objective and not provided as part of the input. This offers a complementary perspective to the settings Wang al. (2024), where the target index set is given
As illustrated in Examples 2.3 and
part of input, and may vary across different data points. Under our setting, Theo- rem 3.1 yields tight (+)) convergence rate with matching upper and lower bounds, sharper than the O [og(T") guarantee obtained under the different problem formulation of Wang al. (024)) A detailed comparison between the convergence rate is provided in Appendix C Regarding group- sparse linear prediction, Zhang et al. (2025c) focus primarily on the classification setting, while Theorem 3.1 delivers complementary result by addressing the regression setting.
The learning guarantee in Theorem established under the assumption that the data input matrix X is Gaussian, and the target response matrix Y is provided by the teacher with noises. Here, we can also study the dfitt---ditiiiuution (OOD) generalization guarantee of the obtained transformer model on data without such assumptions. Specifically, we consider any feature and response matrices X Rd Y E MXD with bounded second moments, and establish bounds the OOD loss
LOOD (Wo; WxQ) = Ex | (2 Wy: W k

---

Published as a conference paper at ICLR 2026
by comparing it with the loss achieved by the teacher model. We have the following theorem.
Theorem 3.2. Suppose that D 2 ((poly( (M, K)) and 1j S (MM-)D))//)) In addition, the OOD input pairs (X, Y) satisfy the condition that each column X, and y, has finite second moments, i.e. there exists a constant >> such that |[| E I for all E [D]. Then for any e > 0,
there exists T, = 0 & > IIv 11
such that for any T Te, the OOD loss satisfies that:
LOOD
W
W R2)
S 2 E - f (X) +
Theorem 3.2 requires only the mild assumption that X and Y have bounded second moments. Notably, the response matrix Y need not be generated by or correlated with the output of the teacher model (X) Therefore, the term [ loss, analogous to the role of Copt in Theorem 3.I This shows that the trained transformer's O.O.D. loss exceeds that of the teacher model by at most e, demonstrating its robustness to distribution shift. In addition, although
[IP)] measures the teacher model's 's O.O.D. test
challenging to establishing a matching lower bound for all pairs (X, Y) like Theorem 3.1 a worst-case can be constructed to demonstrate that this upper bound attainable, thereby validating the tightness of Theorem 3.2 The complete proof of Theorem 3.2 and the worst-case example are provided in Section
EXPERIMENTS
In this section, we present our experimental results. As detailed in Section 2 the teacher model can cover various models, including (i). convolution layer with average pooling, (ii). graph convolution layer on a regular graph, (iii). sparse token selection model, and (iv). group sparse linear predictor. Our experiments also focus on these four cases.
We conduct experiments on both synthetic data and real-world data sets, respectively. For experi- ments on synthetic data, follow the exact definitions in Section to build up teacher models f" For experiments on real-world datasets, we pre-train a teacher CNN on the MNIST dataset, whose first convolution layer then served as the teacher model train the student transformer.
4.1 SYNTHETIC DATA EXPERIMENTS
We begin by detailing the common experimental setups on synthetic data. Given parameters d and D, fixed orthogonal matrix P E IRDxD serves the positional encoding matrix We adopt an online gradient descent algorithm to simulate training over the population loss. At each iteration, we sample a new batch of N 100 standard d D Gaussian matrices, i.c. {X,) C DdxD For each X,, with E [N]. its corresponding label Y, ff (X,.) + En, where E,, E dependently sampled Gaussian matrix. We concatenate each X,, with the fixed positional encoding matrix P to form Z,, as the inputs to the transformer Subsequently. a gradient descent update is per- formed using this batch of N 100 data pairs [ZZ, Y, ) Furthermore, we also generate an- other batch of N = 100 data pairs {ZZ.. ) following the almost identical procedure, except that each X,, is generated from the exponential distribution. This batch of data pairs {(Z,, Y nn) )y prepared for caleulating the excess OOD loss, defined as LOOD Z E. YY, (',(x,)
E RRRMD is another in-
In the next, we introduce the distinct settings for different tasks, specifically the ground-truth soft- max score matrices S*. For the task of learning convolution layer with average pooling, we set D = 36 and K = 4, where the pooling groups are partitioned by aggregating the K neighbor patches into group. Given this partition teacher model can be formulated into diagonal block matrix as S' = [ 1...,,ee 1xxx), with totally D/ K blocks. For the task learning a graph convolution layer, we consider a 'cycle- graph' with D = 20 nodes, where each node is connected to exactly two other nodes, i.e. the i-th node is connected to its adjacent nodes (i 1) and (i + 1). Under this setup, the ground-truth softmax score S* is constructed as follows: for each column i, the entries at (i+1) 1) are set to 1/K with K = 3, while all other entries are zero. For both the tasks of learning the sparse token selection model and the group sparse linear predictor, we the total number of
pooling groups, the ground truth softmax score of the
(e- 1), and
7

---

Published as a conference paper at ICLR 2026
tokens/feature groups D 20, and randomly generate K indices from [D] indices of target to- kens/ label-relevant group, where K = 4 and respectively, In these two sets of tasks, the rows representing the target tokens/ label-relevant group equal to 1/K, while other rows are filled with O.
top eg.
Ong log)
Crears a milrrtt behwnen Wy and
PILLA CAR
Loaky PALLL
(TALL GCM
03
KeLld CNN
67C
07S
(a) Excess training loss (log-log) (b) Excess OOD test loss (log-log)
(c) Cosine similarity
Figure 2: Excess training loss, excess OOD test loss (both in log-log scales), and cosine similarity between the value matrix Wy of one layer transformer (2.4). and ground truth value matrix V*. These results are presented for six experimental sets, which originate from four distinct tasks.
For the task of learning CNN and GCN, we conduct two sets for each with ReLU and Leaky ReLU respectively. Experiment results are given in Figures 2 and 3 Figure 2(a) and Figure 2(b) demon- strate the convergence curves for the excess training loss and the excess OOD test loss (both in log-log scales). We can clearly observe that both the excess training loss and the OOD test loss converge to a small value on all six sets of experiments. After initial iterations, the curves for excess training loss appear almost straight with slopes equal -1, and excess OOD loss curves have ap- proximate -0.5 slopes. These observations validate the (1/TT) convergence rate in Theorem 3.1 the cosine similarity curve be- tween the value matrix w and the ground truth value matrix V*. It shows that ((() directionally aligns with the ground truth value matrix V* in all six experiments since the very beginning.
and ((1/VT) convergence rate in Theorem 3.2 Figure 2(c)
Furthermore, providess the heatmaps of the attention scores when the loss converges. Specif- ically, Figure 3(a) and Figure respectivvy display the attention scores when learning a convo- lution layer with ReLU and Leaky ReLU. In both figures, the attention scores exhibit a diagonal block matrix pattern, where each diagonal block has approximately equal values 1/4. Figure (3() and Figure Figure[Xdd)sshoo show the attention scores when learning a graph convolution layer on a cycle graph. Specifically, the attention scores show pattern of cyelic tridiagonal matrix, with all the significant entries having approximately equal values 1/3. Figure (re) and Figure 3( show the attention scores when learning a sparse token selection task and group sparse linear predictor. We can observe that only the rows corresponding to the target positions are assigned significant values in both tasks. In summary, all these patterns match the ground truth softmax scores, which are described previously.
4.2 REAL DATA EXPERIMENTS
We also conduct experiments on the MNIST dataset. Each image is normalized and resized to 27 x 27 pixels. We train a two-layer CNN with M 16 convolution kernels, each having a 3 x kernel size. Given the 27 x 27 image dimensions, each image is divided into D = 81 patches. An average pooling layer with 3 x pooling receptive field (i.e K = 9) is additive to the first convolution layer, and then cascaded with activation and a linear layer for classification. This two-layer CNN is trained by minimizing the cross-entropy loss, achieving a moderate test accuracy of about 71% on the test set after 20 epochs. After training of this teacher CNN, its first convolution layer with average pooling is extracted as the teacher model f", with its hidden-layer outputs supervising a one-layer transformer (2.4). The training of the layer transformer is still conducted on the MNIST dataset, and the mean-squared loss is employed for optimization.
The experiment results are given in Figure and Figure Figure 4(a) displays the training loss curves. We can observe that for both ReL.U and Leaky ReI.U, the training loss very quickly con- verges to small value. Figure 4(b) demonstrates the cosine similarity curve between the value matrix ((t) of the transformer and the convolution kernel matrix V of the teacher convolution layer. The similarity rises above 0.9, indicating that the transformer successfully learns the ground-
V
8

---

Published as a conference paper at ICLR 2026
soore
00
02
02
01
01
-0.
-02
(a) ReLU CNN
(b) Leaky ReLU CNN
(c) ReLU GCN
Attention score sm
51ll
0.75
0.2
0.0
0:3
025
0.00
0.00
0.1
-.00
-035
0050
007
-0.3
-0.9
(d) Leaky ReLU GCN
(c) Sparse token selection
(f) Group sparse linear predictor
Figure 3: Heatmap of attention score matrix ( when the training loss converges. The results are presented for six different experimental sets, indicated by the captions of sub-figures.
MNIST: Cosine similarity
0.5
0.8
0.2
Datccees
100K
10k
100k
JM
(a) Training loss on MNIST dataset
(b) Cosine similarity between parameter matrices
Figure 4: Training loss and cosine similarity between the value matrix Wy of the one-layer trans-
former (.4). and convolution kernel matrix V of the pre-trained teacher CNN.
truth value matrix of the teacher model. Furthermore, Figure 5(a) provides the heatmap of the ground truth softmax score derived from the teacher CNN's ' average pooling layer. Figure 5(b) and Figure 5(c) respectively present heatmaps of attention scores convergence for the transformers with ReL.U and Leaky ReLUU activations. We can observe that both the attention scores achieved capture the pattern of the ground truth softmax scores, with notable exceptions in the first and last nine rows in the softmax heatmap. We remark that the failure in learning these rows of ground-truth softmax scores is due the fact that they correspond MNIST image patches that are mostly all background (all zero). Figure 5(d) highlights the image regions corresponding dille-tooleeam softmax scores, marked by yellow rectangles. We can see that they are indeed boundary regions and are mostly pure background. Consequently, they offer minimal informative content to model, explaining why transformers can not attend to these positions. Overall, clear that the real-world data experiments corroborate our theory.
by transformers
5 PROOF SKETCH OF THEOREM 3.1
In this section, outline the major steps in the proof of Theorem 3.1 For simplicity, here focus the case where o-)) is the identity map. More details, including more general choices of (-)), are formally proved in Appendix D The proof consists of three main steps:

---

Published as a conference paper at ICLR 2026
S
(a) teacher CNN
(b) ReLU result
(c) Leaky ReLU result
(d) Example image
Figure 5: Heatmap of the ground truth sofimax scores of average pooling, Heatmap of the attention
scores ((T) of trained one-layer transformer when loss converges, and an image example in MNIST.
Step Structures Wv and WKQ during training. A critical step in our proof is to show that throughout training, the parameter matrices W and W KQ preserve the following decompositions:
w Ci(tV*;; KQQ C2(() Σ ΣE P.P. C3(t)
ΣEPPP. P.P.
where G denotes the index set of entries of value 1/K in i-th column of S The details of this
conclusion are given in Lemma D.2 Based on the decompositions, we can express ((1) as: S
if $ Gi. Comparing these results with the definition of the teacher model (*), can further observe that
K+(D-K) (=p(=(()(()+ if i' E G:: s
esp(-(C
K+(D-K) (xx
( V + Ci(t) 1;
((t) S + C2(() + Ca(t) + oc.
In this way, the original optimization analysis regarding full matrices Wv and W KQ is simplified into studying the updates of three scalars (I), C2((), Ca(
Step 2. Accurate characterization of convergence that Ci(t) -1 and C2(() + Cs(() + oc. The decompositions obtained in Step implies that the coefficients (t), C2(t), Ca(t) essentially follow gradient descent starting from zero initialization minimizing the loss
2
D-K K
KC1
K
E(C,, C2. C3)
1
C7
C2+Cs VD
C2
(D - K)E
K+ K)e
VD
We remark that this expression of E(Ci, C22 C2) corresponds to the special case where (.)) is the identity map. The general formulation for (-) is activation deferred to Lemma D.2 carefully analyzing the training dynamics, can show that for sufficiently large T.
Then by
DVK
VD
1 K
-
C2(T) C.(T)
log +
T + +eXx)
|V
2
The details are provided in Lemmas D.2 D.5
D.15
D.18 and F.12
Step 3. Final convergence results. Combining the convergence rates obtained in Step 2. and the formulations S(T) and w!(") in Step 1., we can further obtain that ||S(T) St [w(") V*Il= e ) Under mean mean-squared loss, the O /T convergence of the matrices ((t) and ( directly suggests that loss will decay the rate of O (+), which finishes the proof.
CONCLUSIONS AND LIMITATIONS
In this paper, we provide the theoretical guarantee that a one-layer transformer can learn a class of teacher models, covering a wide range of common models in machine learning. Specifically, we establish tight convergence bound at the rate of O() (z) for the population loss. We also establish out-of-distribution generalization bounds for the obtained transformer model, demonstrating its ro- bustness. To empirically support our findings, we conduct experiments both synthetic data and real data, and all results align with our theoretical conclusion. Our current theory focuses on one- layer models, and make certain simplifications and assumptions on the model and data, which present a limitation. We believe establishing teacher-student learning guarantees for more complex models and under milder assumptions is interesting and promising further work direction.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGMENTS
We would like to thank the anonymous reviewers and area chairs for their helpful comments. Yuan Cao is supported in part by NSFC 12301657, Hong Kong RGC ECS 27308624, and Hong Kong RGC GRF 17301825.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303.08777 2023.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
Alon Brutzkus and Amir Globerson. Globally optimal gradient descent for convnet with gaussian inputs. In International conference on machine learning, pp. 605-614. PMLR, 2017.
Yuan Cao, Difan Zou, Yuanzhi Li, and Quanquan Gu. The implicit bias of batch normalization in linear models and two-layer linear convolutional neural networks. In The Thirty Sixth Anmual Conference on Learning Theory, pp. 5699-5753. PMLR, 2023.
Yuan Cao, Yihan He, Dennis Wu, Hong- Yu Chen, Jianqing Fan, and Han Liu. Transformers simulate mle for sequence generation in bayesian networks. arXiv preprint arXiv:2501.02547 2025.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learninngg via sequence modeling. Advances in neural information processing systems, 34:1 5084-15097, 2021.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Unveiling induction heads: Provable training dynamics and feature learning in transformers. In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024a.
Xingwu Chen, Lei Zhao, and Difan Zou. How transformers utilize multi-head attention in-context learning? case study on sparse linear regression. In ICML 2024 Workshop on Theoretical Foundations of Foundation Models, 2024b.
Xingwu Chen, Miao Lu, Beining Wu, and Difan Zou. Towards thcoretical understanding of transformer test-time computing: Investigation on in-context linear regression. arXiv preprint wwXiv:2508.07571, 2025.
Jacob Devlin. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 1111.880000000) 2018
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010. 11929, 2020.
Simon Du, Jason Lee, Yuandong Tian, Aarti Singh, and Barnabas Poczos. Gradient descent learns une-hidden-layee enn: Don't be afraid of spurious local minima. International Conference on Machine Learning, pp. 1339-1348. PMLR, 2018a.
Simon S Du, Jason D Lee, and Yuandong Tian. When is convolutional filter easy to learn? In International Conference on Learning Representations, 2018b.
Spencer Frei and Gal Vardi. Trained transformer classifiers generalize and exhibit benign overfitting in-context. In The Thirteenth International Conference Learning Representations, 2025.
Cheng Gao, Yuan Cao, Zihao Li. Yihan He, Mengdi Wang, Han Liu, Jason Klusowski, and Jianqing Fan. Global convergence in training large-scale transformers. Advances in Neural Information Processing Systems, 37:29213-29284 2024.

---

Published as a conference paper at ICLR 2026
Surbhi Goel, Adam Klivans, and Raghu Meka. Learning convolutional layer with overlapping patches. In International conference on machine learning, pp. 1783-1791. PMLR, 2018.
Jianliang He, Xintian Pan, Siyu Chen, and Zhuoran Yang. In-context linear regression demysti- fied: Training dynamics and mechanistic interpretability of multi-head softmax attention. arXiv preprint Stti:2503.12733: 2025a.
Yihan He, Yuan Cao, Hong-Yu Yu Chen, Dennis Wu, Jianqing Fan, and Han Liu. Learning spectral methods by transformers. arXiv preprint arXiv:2501.03112 2025b.
Yihan He, Hong-Yo Yu Chen, Yuan Cao, Jianqing Fan, and Han Liu. Transformers versus the em algorithm in multi-class clustering. arXiv preprint arXiv:2502.06007 2025c.
Jianhao Huang, Zixuan Wang, and Jason D Lee. Transformers learn to implement multi-step gra- dient descent with chain thought. In The Thirteenth International Conference on Learning Representations, 2025.
Yu Huang. Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In Forty-first International Conference on Machine Learning, 2024.
Muhammed Emrullah Ildiz, Yixiao HUANG, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self-attention to markov models: Unveiling the dynamics of generative transformers. In Forty-first International Conference on Machine Learning. 2024.
Michael Janner, Qiyang Li, and Sergey Levine. Offline reinforcement learning as one big sequence modeling problem. Advances in neural information processing systems, 14:1273-1286, 2021.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-378366 2022.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger,
Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-588) 2021.
Thomas N Kipf and Max Welling. Semi-supervised classification with graph convolutional net- works. In International Conference on Learning Representations, 2017.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise not the main factor behind the gap between sgd and adam on transformers, but sign descent might be, In The Eleventh International Conference on Learning Representations, 2023.
Bingrui Li, Wei Huang, Andi Han, Zhampeng Zhou, Taiji Suzuki, Jun Zhu, and Jianfei Chen. On the optimization and generalization of two-layer transformers with sign gradient descent. arXiv preprint arXiv:2410.044 04870, 2024a.
Hongkang Li, Meng Wang, Sijia Liu, and Pin- Yu Chen. A theoretical understanding of shallow vision transformers; Learning, generalization, and sample complexity. In The Eleventh Interna- tional Conference on Learning Representations, 2023a.
Tianle Li, Chenyang Zhang, Xingwu Chen, Yuan Cao, and Difan Zou. On the robustness of trans- formers against context hijacking for linear classification. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards mechanistic understanding. In International Conference on Machine Learning, pp. 19689-19729. PMLR, 2023b.
Zihao Li, Yuan Cao, Cheng Gao, Yihan He, Han Liu, Jason Klusowski, Jianqing Fan, and Mengdi Wang. One-layer transformer provably learns nne-nearest neighbor in context, Advances in Neu- ral Information Processing Systems, 77:82166-82200 2024b.
Tianyi Liu, Minshuo Chen, Mo Zhou, Simon S Du, Enlu Zhou, and Tuo Zhao. Towards understand- ing the importance of shortcut connections in residual networks. Advances in neural information processing systems, 32, 2019.
12

---

Published as a conference paper at ICLR 2026
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In Proceedings of the IEEE/CVF international conference computer vision, pp. 10012-10022, 2021.
Xuran Meng, Difan Zou, and Yuan Cao. Benign overfitting in two-layer relu convolutional neu- ral networks for xor data. In The Forty-first International Conference on Machine Learning (21007/2024-2770722222 Vienna), 2024.
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. In Forty-first International Conference on Machine Learning. 2024.
OpenAl. Gpt-4 technical report. arxiv 2303.08774. View in Article, 2(5), 2023.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to fol- low instructions with human feedback. Advances in neural information processing systems, 35; 27730-21 27744, 2022.
Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than sgd for transform- ers. arxiv preprint wrXiv:2366.00200 2023.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
Yongming Rao, Wenliang Zhao, Benlin Liu, Jiwen Lu, Jie Zhou, and Cho-Jui Hsich. Dynamicvit: Efficient vision transformers with dynamic token sparsification. Advances in neural information processing systems, 34:1 13937-13949, 2021.
Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06175. 2022.
Keitaro Sakamoto and Issei Sato. Benign not-benign overfitting token selection of attention mechanism. arXiv preprint arXiv:2409. 17625, 2024.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. Advances in Neural Information Processing Systems, 36:36677-36777, 2023.
Wei Shi and Yuan Cao. Towards understanding transformers in learning random walks. In The Thirty-ninth Annual Conference on Neural Information Processing Systems, 2025.
Chen Siyu, Sheen Heejune, Wang Tianhao, and Yang Zhuoran. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. In The Thirty Seventh Annual Conference on Learning Theory, 4573-4573, PMLR, 2024.
Mahdi Soltanolkotabi. Learning relus via gradient descent. Advances in neural information pro-
cessing systems, 30, 2017.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. In NeurlPS 2023 Workshop on Mathematics of Modern Machine Learning. 2023a.
Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token se- lection in attention mechanism. Advances in Neural Information Processing Systems, 36:48314- 48362, 2023b.
Yuandong Tian. An analytical formula of population gradient for two-layered relu network and its applications in convergence and critical point analysis. In International conference on machine learning, pp. 3404-3413, PMLR, 2017.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding train- ing dynamics and token composition in 1-layer transformer. Advances Neural Information Processing Systems, 36:7191 1-71199 2023.
13

---

Published as a conference paper at ICLR 2026
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Shaolei Du. Joma: Demys- tifying multilayer transformers via joint dynamics of mlp and attention. The Tweifth Interna- tional Conference on Learning Representations, 2024.
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, al. Llama: Open and efficient foundation language models. arXiv preprint arXiv:2302. 13971, 2023.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Eukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural informa- tion processing systems, 30, 2017.
Jesse Vig and Yonatan Belinkov. Analyzing the structure of attention in transformer language model. In Proceedings of the 2019 ACL Workshop BlackboxNLP: Analyzing and Interpreting Neural Networks for NLP, pp. 63-76, 2019.
Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D Lee. Transformers provably learn sparse token selection while fully-connected nets cannot. In Forty-first International Conference Machine Learning. 2024.
Jingfeng Wu, Difan Zou, Zixiang Chen, Vladimir Braverman, Quanquan Gu, and Peter L Bartlett. How many pretraining tasks are needed for in-context learning of linear regression? arXiv preprint (3t::3310008391 2023.
Weihang Xu and Simon Du Over- ave-parramttriization exponentially slows down gradient descent for learning a single neuron. In The Thirty Sixth Annual Conference on Learning Theory, pp. 1155- 1198. PMLR, 2023.
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zi-Hang Jiang, Francis EH Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token vit: Training vision transformers from scratch on imagenet. In Proceedings of the IEEE/CVF international conference on computer vision, pp- 558-567, 2021.
Chang-Bin Zhang, Jinhong Ni, Yujie Zhong, and Kai Han. v-clr. View-consistent learning for open- world instance segmentation. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp. 20307-20317. 2025a.
Chang-Bin Zhang, Yujie Zhong, and Kai Han. Mr. detr: Instructive multi-route training for detection transformers. In Proceedings of the Computer Vision and Pattern Recognition Conference, pp- 9933-9943, 2025b.
Chenyang Zhang, Gao Peifeng, Difan Zou, and Yuan Cao. Gradient descent robustly learns the in- trinsic dimension of data in training convolutional neural networks. In High-dimensional Learn- ing Dynamics 2024: The Emergence of Structure and Reasoning, 2024a.
Chenyang Zhang, Xuran Meng, and Yuan Cao. Transformer leamn optimal variable selection in group-sparse classification. In The Thirteenth International Conference on Learning Representa- tions, 2025c.
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 13:15-3315,333, 2020.
Ruiqi Zhang, Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49):1-55, 2024b.
Ruiqi Zhang, Jingfeng Wu, and Peter L Bartlett. In-context learning of linear transformer block: benefits of the mip component and one-step gd initialization. arXiv preprint arXiv:2402.1495 14951, 2024c.
Qingyue Zhao and Banghua Zhu. Towards the fundamental limits of knowledge transfer over finite domains. In The Twelfth International Conference on Learning Representations, 2024.
Mo Zhou, Tianyi Liu, Yan Li, Dachao Lin, Enlu Zhou, and Tuo Zhao. Toward understanding the importance of noise in training neural networks. In International Conference on Machine Learning, pp. 7594-7602. PMLR, 2019.
14

---

Published as a conference paper at ICLR 2026
A NOTATION
In this section, we introduce the key notations we use throughout paper. We first introduce the following mathematical notations.
Mathematical notations. Given two sequences {ra} and {yn. we denote In = -On)) if there exist some absolute constant C > 0 and N >> 0 such that |2.,] S Cl for all 2 N.. Similarly, we denote In = (yyn if there exist C2 > 0 and N > 0 such that Fnl 2 C2|ynl for all n > N. We say In = (v,)) if In (yn) and In, (y,,) both holds. We use O(-) (2(-). and 6(.) to hide logarithmic factors in these notations respectively. Moreover, we denote In = poly(yn) Tn Ovv))) for some positive constant D. and x,, polylog(yn) if 2,,
(loy()o(()))) (v.)). For two scalars a and b, we denote a V b = max(a, b) and a A 6 = min a. b). For any n E N+, we use [n] In addition, we use 1,, to denote a ne-dimensional vector with all entries. For an index set g. 1 denotes a vector whose entries are 1 for indices in g. and otherwise. Let A1, A..... A, be matrices with the same dimensionality d1 x d2, then Diag( [A A,y..., An) is ndj x ndz diagonal block matrix, with A , A, being the block entries.
to denote the set {1, 2, .)..
In addition, we also provide a summary table of the key variables in our study in Table
Table I: Key variables and their meanings.
Symbol
Meaning
V
Ground truth value matrix in f", a M x D matrix.
S
Ground truth softmax
matrix in f", a D x 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 S*. It can represent: (i) the pooling size in CNN and pooling layer, (ii) the number of neighbors of GCN layer,
(iii) the number of target tokens in sparse token selection, (iv) it equals to in group-sparse linear models.
G
Target index set of i-th input token, namely Sp.z = K if E G', and 0 otherwise. Number of gradient descent iterations. Learning rate.
t
n Wv. WKQ C
Parameter matrices of the transformer. Population loss (objective function). Out of distribution loss.
LO.DD.
C(()
Coefficients of the decompositions of Wy and W KQ during the training.
B ADDITIONAL RELATED WORKS
Optimization of transformers. There exist multiple recent works studying the optimizations of transformers, most of which focus on the single-layer architecture. Zhang et al. (2020); Kunstner et al. (203);; Pan & Li (2023): Li al. (2024a) investigate performance comparison between the adaptive methods and SGD under different settings from both theoretical and empirical perspec- tives. et al. (2023b) investigates the optimal parameters of transformers applied a masked topic structure model similar to the Bert framework through a two-stage training regime [ldiz et al. (2024); Chen al. (2024a); Shi & Cao (2025) explain the mechanism of attention from the perspective of Markov chans. Tian et al. (2023 2024) study the training dynamics of transformers, jointly with decoder layer and a fully-connected layer, respectively. [Li et al. (2024b) analyzes transformer train- ing behavior in the context of one-nearest neighbor selection. Gat et al. (2024) addresses the global convergence of transformers given certain prerequisites. Tarzanagh et al. (2023a b) demonstrates that single-layer attention mechanisms can converge directionally towards the hard margin solution typical of Support Vector Machines (SVMs). Furthermore, Li et al. (2023a) presents a generaliza- tion error bound for vision transformers optimized using stochastic gradient descent. Some works consider the transformers to perform certain algorithms. He et al. (202cc) theoretically characterizes Softmax attention approximating the Expectation and Maximization updates in EM for Gaussian
15

---

Published as a conference paper at ICLR 2026
mixture models. He et al.
(2025b) show that multi-layer Transformers can provably learn and im- plement spectral methods for Gaussian mixture models via vre-training. Furthermore, many other existing works investigate the optimization of transformers under the so-called "in-context learn- ing" settings (Chen et al.] 2024b Huang et al.,
2024
Zhang et al., 2024b e Nichani et al. 2024]
Huang al.; 2025|Chen et al. 2025 Li al.
2025
Cao et al. (025). Based on the framework (024)) extends this result to one-layer softmax at-
proposed in (Zhang ct al., 22244). Huang et al.
tention transformers. Siyu al. (2024) investigates the multi-head self-attention under this setting, and summarizes two distinct patterns among all heads. (iich.ii et al. (2024) demonstrates that when solving in-context learning tasks with latent causal structure, transformers can encode the latent causal graph. Huang et al. (2025) demonstrates that Chain of Thought (CoT) prompting enables Transformer models to learn to perform multi-step gradient descent and effectively recover true weights CCen et al. (2025) focuses on the test time computing on the in-context linear regression. Lict al. (225)) studies the context hijacking phenomenon by investigating an optimization procedure with different learning rates. Cao et al. (2025) proves that Transformers
implement in-context maximum likelihood estimation and autoregressive sampling for Bayesian networks, establishing their capability to simulate MLE-based sequence generation.
Teacher-student framework for training neural networks. We also introduce some related theo- retical works regarding the training of "student" neural network under the guidance of "teacher
model" (Brutzkus& & Globerson| [017 Tian 2017:Soltanolkotale [017| Gool et al. [0118] Du et al.)
2018b at Zhou et al. [219 Liu al. 2019 Xu Du,
Dul 2023). Several studies establish convergence guarantees for gradient descent in specific ReLU network settings: [Brutkkus & Globerson (2017)
demonstrated polynomial-time global convergence for mee-hiden-llayr non-overlapping convolu- tional ReLU networks with Gaussian inputs: Tian (2017) characterized critical points and proved gradient descent convergence for two-layer ReLU student-teacher networks under Gaussian inputs; and Du et al. (2018b)a) a) provided polynomial-time recovery guarantees for learning convolutional ReLU filters and networks, respectively, using (stochastic) gradient descent, even with potential spurious minimizers and for general or Gaussian inputs. Furthermore, Zhou et al. (2019) and Liu ctt al. (2019) showed that methods like perturbed gradient descent with noise annealing or specific normalizations and initializations can achieve polynomial-time global convergence in convolutional neural networks (including ResNets) despite the presence of spurious local optima. Research fo- cusing single ReLU scenarios includes [Soltanolkotabi (2017)'s 's analysis of linear convergence for a single ReLU in a high-dimensional Gaussian model with structured weights, and Xu & Du 2023)'s 's finding that over-parame parameterrizin a student network to learn single target ReLU neuron under Gaussian inputs can surprisingly slow convergence. Goel et al. (2018) introduced Convotron, provably efficient algorithm for see-hidden-layer convolutional networks with general patches, achieving global convergence through noise-tolerant stochastic updates without requiring special initialization or learning rate tuning. Zhhao & Zhu (2024) studies the statistical limits of knowledge transfer over finite domains, characterizing minimax rates under different levels of teacher supervi- sion.
C COMPARISON WITH WANG ET AL. (0024)
In this section, we compare the essential optimization dynamics in Wang et al. (024) and our works. Wang et al.] (2024) and our work both rely on the symmetry of Gaussian data and the uniform distribution among the target tokens expected to be selected. A critical technical step shared by both analyses is to simplify the optimization regarding the full parameter matrices to investigate the evolutions of several specific scalars, as demonstrated in Lemma 3.2 in Wang et al. (2024) and in our Lemma D.2 Specifically, the analysis in Wang et al. (2024) tracks the evolution of two scalars, (t) and C(t), for which the coefficients are essentially minimizing the loss
d Ī(a, C) 2(D K)
K(D K)
a
1 K
+a² +
K
2
K (D k)--c
K - K) K)e-c (C.1)
as demonstrated top of Page 31 in Wang al. (2024).
As demonstrated in Lemma D.2, our analysis focus on the scalars (t), Ct(,), (t). When the teacher model is reduced to the "sparse token selection" task defined in Example 2.3 with V
Ia
16

---

Published as a conference paper at ICLR 2026
and without activation function, the coefficients C1 (itt), C2(t)) essentially minimize the loss
dD
C1
2
C2, 2(D - K)
K(D K)
Cat w2
K
K + (D
K)e
c9
K
(C.2)
Cotes VD
K+ + (D K)e
Comparing these two loss functions in (C.2) and C.2 we can observe that they essentially share the same function structur.. Specifically, if we regard Ca in (C.1) as one term, playing the role
as C' in (C.I). then these two functions only differ by factor D. Therefore, while the setting of the "sparse token selection" task in our work is different from that considered in Wang et al. (2024), they can be formulated into essentially identical optimization problem. Notably, the loss in (C22 is only the special case in our setting with V = Ia and without activation function, while the general complicated and provided in Lemma D.2 Therefore, the setting considered in
case is much our work is more general compared with that in Wang et al. (2024), from a technical perspective. This also highlights that establishing tight convergence rate with a matching lower bound indeed constitutes a technical advantage of our work.
D PROOF OF THEOREM 3.1
In this section, we provide a detailed proof for Theorem 3.1 We first introduce several notations used in the following proof. For each i E [D], we use G to denote the index set to which the if E G and 0 otherwise. With this notation, we can express that (X)] m,i = (vv Xl()) = K R PeEes (v xx))). In addition we let
entries of i-th column of S* k i.. S
V [vi, ...... v . and Wv = [Wv.1, Wv:....., wy.ar] MR Based on this notation, it is equivalent to consider the gradient descent updating regarding each wvm for all m E [M], expressed as
Vyn
(e+1)
w (K)
V.m V/V L(W
KQ).
(D.1)
In the following proof, we will consider the gradient descent updating details for each Wym and derive the conclusion for w!(h) based on the result of w./m for all m E [M]. For simplicity of presentation, we assume that each v is normalized in the remaining sections, i.e. ||v for all m [M]. without loss of generality (W.L.O.G.). However, our theoretical findings and proofs can be directly extended to the case where Vm is not normalized. For each v let I,,
[ven, vi, is not normalized, the first column of Im will be Ty
E Rd be an orthogonal matrix with Vm being its first column. (Actually, V
Furthermore, we introduce several definitions regarding the expectations of Gaussian random vari- ables. Let ~ (000, a). 22 ~ (0, b), and x3 ~ NO, c) be three independent Gaussian random variables. In addition, (-)) can be the identity map, the ReLU activation function, and the Leaky ReLU activation function, with denoting the coefficient of the Leaky ReLU activation function when the input is negative. Specifically, when (-) indicates the Leaky ReLU activation function, (x)) x1 (o20) + Kxl (x<c0- Then, based on these notations, define that
(i(a) [[ (21));
(D.2)
F2(a, b) = E[zio + z2))o'(n + 22));;
(D.3)
Es(a, b) = E[(z, + 22) (mn + x2))::
(D.4)
F, = E[r + 22)0 o (z) + 22 + rs)]:
(D.5)
b,(a,b,c) = ( + x2 + rz)]
(D.6)
We provide the detailed calculations for these expectations in Section E1
D.1 DETAILED GRADIENT DESCENT UPDATING RULES
In this subsection, we introduce and prove several lemmas regarding the calculation details regarding the gradient descent iterative rule (D.)) and 3.4
17

---

Published as a conference paper at ICLR 2026
Lemma D.I. The gradient descent updating regarding Wmm for all m E [M] and w
(0)
KQ. which
have been defined in D.I
and
(.4). can be rewritten
D
D
W.m
(++1)
w. (e) Vi +11 17
2
E
Y m,a
(4) Vi
!s!!
0'
(w (2)
V. x /5
x; S
(D.7)
KQ
= witt)
KQ
+
M Σ
Σ
m,s - a Σ (w x1, ) (
) o'
(wmm w
Vim X1, ss
m=1 1=1
D
ΣΣ
w (1) Vima X
/S (P.
Pig Jp!
(1).8
1==1i==1
Proof of Lemmal D.I
By the chain rule of derivatives, we have L(W (1) w KQ)
w (1+1)
w. (1)
V/ wwmm
= (e)
V
E I
(Z Wy: W wkq)l
Vmm
Vim
Vmm 2
W. Vim
(t)
Wv.
E
(Y
m aW (x XS
m' [s]
V.m
(t)
2
Σ
M
D
V
Wy.m
E YYmri
Σ
D
w
Vimi Xin ) ss!
$1.5
)}
m m'=1 1=1
D
]
D (t)
(t)
V.m
17
E
Y m,²
o
(w
Vim
x
(w(t)
Vim= x !!!
Xi
((1) 12."
where the last equality holds simply by the chain rule of differentiation. This proves (.7]. Next for WKQ, we have 0
(1+1) w KQ
C(W
V
(t)
KQ)
w KQ= 2
! VWwxoBlIy B TF(Z: Wr; WxQ) [+]
M
D
-wi!) KQ
"
Σ
Σ
V B (Y m, o (W XS(1)) m)
m
M
D
D Y Ym,i o
Vim X;, |i!!! (1.s )
D
-W KQ + 1
Σ
E
Σ
(w
d'
Σ
(w/h
v.m x,,
15
is,?
m=1
11=1
(W
VWKO (W (W)) TXS
KQP:
(D.9)
VD
For the derivative calculation of I, we have
D
(w
(t)
D
PW!!! (1)
==
V
WkQ
(w
(t) V.m
X
S
KOP:
(w Xi,) Vwxe
S
KQPi
VD
vD
d
PW(!)
KoPi VD
D
D
PW (
KQP
V
PW(
VW
VD
1==1
12
d
VD
l2
D
PW(O
(1)
D = Σ 61 187
V.n
S'
DD
KQP:
P.P.
Σ
(wo Von x,, (((1) (t) 11 22. (P., -p..D.T.
i1.i2
(D.10)
'Here we slightly abuse the notation of S(-). If the input a D-dimensional vector, S(-) denotes the softmax function from RD + RD. If the input is a Dj x Dz-dimensional matrix, S(-) represents the softmax operator which implements the softmax normalization defined above column-wisely.
18

---

Published as a conference paper at ICLR 2026
The last equality holds as (a) diag(a) S(a)S(a) R for any vector R and consequently,
8
Pw
KQp.
-
s
Pw
Ro
)1,)
S (1-
(
if 12;
PW KOP:
vD
vD
VD
61,42
[[s
Pw
KoP.
s
PW vD
KqP
--S s
otherwise.
VD
1. 62.1
ig
By substituting the result of I from (D100) into D.9
we complete the proof of (D.8)
The next lemma demonstrates that the training dynamics of w Vim for all m E [M] and W (b)
KQ exhibit
specific patterns. Analyzing the training processes described in (D.1) and (D.21 can be reframed as an investigation into the coefficients of these patterns.
Lemma D.2. Under the same conditions of Theorem 3.1 there exist a time dependent non-negative scalar (t), and non-negative, monotonically increasing scalars (2()) and Ca(t). such that
(1) Vm
(t())
Vm for all m E [M];
( KQ
Ct())
Σ P.P. Cs(() Σ Σ P.P.
= EGG
GGG
Due to the specific pattern of Wil)
kQ demonstrated above, there exist time dependent scalar
p(t)
C2(0+Ca(0) V
K + (D K)e
such that (00) p(t) for all i [D] and i G¹. Otherwise, S
1- Kptt). D-K Additionally, 1 $ ptt S K and p(t) is monotonically increasing. Based on the definition of p(t), Ct),, C2(t), D
and (t()) have the following iterative rules:
G(O0))) ci((+ + DDF (1- Ci(t) Ky(() cT())
- Ci( DT
Kp(()
-)(+1)== = C2(1) + Ci())
1
r3)) - -9:(0 [22 +0(+()+)))))))))))))))))))))))))
VD
K
P(()
Cs(t+1) =)(()) (1 -11(1-Kp)) Ky(t))
(D - (D-K)F5
-Ci(t) (7(1) 1- Lpp((t) )
(D-K)F(2
Ci(t)
VD(D - K)
Kp(()) Kp(t)(1 Kp(t))
where F(1) = Fi (Kp(t) + (1-Kp(0)²)
Fill]
F2(p(t) 1)p(t)2 + (I-Kp(0))))
D-K
F F/ (1--Kp(0)
[W-KF= Kp(1)² + (D-K ---K))(0KK()))
135 Fs (F(((()) (1- (I-Kp(t))² p(())
F, (p(t)),(() ,K (K- 1))(()) (1- (1-Kp(0)²) F( Fs Kp(t)², (1-Kp(0)2 (D-K)²
and
(D-K)²
F(())
(t)
K. In addition, based on all these definitions, the coefficients (t)). C2), and
Cs(t) are essentially minimizing the following loss function by gradient descent
Z(C1, C2. C3) co D/VV*1
2(D K(D K) -Cip) + + (I(1-Kp) kp) KD) 1
where Co is an absolute constant such that Co ((((.) is identity map)
+ } +
1t*" 1 (a()) Leaky In addition, Fo(Ci,p) is defined as
00;
If ffo(-))s identity map
pC?
Kp arctan
Py K(D-K) I-Kp
-1 1
(1-Kp)vK #VD-K
If (ro(-)ss ReLU activation
Fs
(1 - k==))ppp
Kp (xp(farees arctan
PV/KID-K) 1-Kp
+ (1-KD)KK FVD-K
If (fo(-))s Leaky ReLU activation
19

---

Published as a conference paper at ICLR 2026
We establish these conclusions by induction. It can be easily verified that all these conclusions hold 0, since the parameters are initialized as w Oxxx and / KQ = ODXD- However, for the sake of conciseness and coherence the presentation, we rearrange the contents of Lemma D.2 into Lemma D.4 and Lemma D.8 including the relevant details regarding Wv.m and WKQ re- spectively. To prevent the proof of a single Lemma D.2 from becoming overly lengthy, we prove Lemmas D.4 and D.8 separately.
As we use induction, we assume that the conclusions of both Lemma[D.4]and Lemma the current iteration. We then demonstrate that the conclusion of either Lemma D.4oor Lemma D.Shholds at the next iteration, depending on which lemma we are proving. It is important to clarify that this is not circular reasoning: all these contents can indeed be organized into single Lemma D.2 It is reasonable to assume that all conclusions hold for each iteration and to verify that these conclusions remain valid for the next iteration, long rigorously demonstrate their validity at the outset.
In the following, we introduce and prove Lemma D.4 and Lemma D respectively. Besides, the notations defined in Lemma D.2 containing p()), P1 consistent unless stated otherwise.
F())
[())
(t) p(t) and p.()) will remain P
2,1 2,2
We first introduce and prove lemma regarding the ratio between F and Fs) which will be utilized in the proof of Lemma D.4
Lemma D.3. Under the same conditions of Theorem 3.1 for F10 and p defined in Lemma D.2 it holds that
(()
Kp(t)= S vDK((t).
F
Proof of Lemma D.3
By Lemma F. and Lemma
E5
we can derive that
. .If (() is the identity map, then
F(
(D - K))2(1?
D-K K
DKp(t) - 2Kp(t)+ +1 = Kp(()
2 Kp(t));
F
DKp(1) + p()
-2K
F
(D (0-K)pp(t)2
D-K
DKp(1)2 - pKp(1) +1 Kv(1)
S VKP(t).
DKp(t) + p(t)
I - -2K
The last incquality is derived by 2/DK 2KK S DKp(t)+ 1 -2K S D
K as 1 D
S sp(t) < K.
p(t)
If (-) ) is ReLU activation function, is also straightforward that Kp(t)2
F(O 2(D- -0)-K K)
D- K
F 2 - 2Kp(t) +1 Kp(t) DKp(1) + p()) 2K
Z
2 Kp(t).
On the other hand, by Lemma
F55
it can be derived that
2(D K)
Kp(t)²
+
2m V p(t)(1 Kp(t)))
1
DKp(1)2 - 2Kv(()
+1
D-K DKp(t) + M()
D-K
1- Kp(t) K ((DKp())-2KK(()) - +1)
-Kp(1)
+ +1
- 2K
1 L/D 2V K
D 1
D D-K K
1
<Kp(1)
+
VV
K
2
< <DDKK().
where the penultimate inequality holds since DKp(t) p() 2K > 2/DK 2K, and 1-Kp(t) D DKp is a decreasing function w.r.t. p(t) as the numerator is decreasing w.r.t. p(() while denominator is increasing w.r.t. p((). Therefore, takes the maximum value when p(t) 3. and consequently
1-Kp(t) 22K00)
D-K
R
20

---

Published as a conference paper at ICLR 2026
If (+) is Leaky ReLU activation function, by utilizing similar calculation, it holds that
F()
2(D - K) (
Kp(() (D K)t(t)
DKp(t)² 2Kp(t) +1 Ky(t;;
(1+ (DKp(()) - 2Kp(K p(t) +1)
F(1) 2(D -K) (LtN) + --=aa"
K /D=RD(U(1-K(())
S
2
S vDKp(t).
Filt)
(1 + [++))((KD(C))) 2Kp(()+1)
This completes the proof.
Lemma D.4 (Restatement of Lemma D.2 the first part). Under the same conditions of Theorem 3.1 there exist time dependent non-neg negative scalars Ci(t), such that
(i()) Vin. for all m E [M],
(D.11)
where (t) has the following iterative rule:
Ky(() ((()) Ci(t)+ nDF (t) Ky(()
F
(1 Ci(t)
+ Dm
(D.12)
(i())
pultt
where (t) K
Proof of Lemma D.4 First at the initialization t = 0. we have w V = Ox satisfying (D.11 Next, we assume that at f-th iteration, the conclusion of (D.IT) still holds, and we will prove that continues to hold at the t + 1-th iteration. Actually, it suffices to show that
V www. L(W w(()) kQ) = (t) vmn for all m E [M],
(D.13)
where (t) is a time-dependent scalar. By Lemma D.I
we have
D
V ww. L(W
KQ)
Σ
ΣΣ
E
Y
(w
V. x., is 1
(t)
0' Σ
w
V.m x,, )s!
i=1
11
i1=1
1
Σ
E
Y
Σ
(w
x )s! xi, ((t)
Vu
1==1
E
W
V. X >s
(4)
0'
V.
(t)
x,, (s!0
(1.1 xi, !!!
=14=1
1,=1
12
(D.14)
For 11, we have
D
D
h = Σ
M
E
Y
oo
[
(mmm r I'm T !!!!
1=
m=
D
(wlom X., (8!"0 ) (Vin xi, ys ] v
[F))]]
Σ
mit o'
1==1
E [F(X)
(w (t)
x3 (5/? (Em,ki 1800 si
Em,
2=
D
D
EE
[x))] m.t o'
Σ
(w Xi, X /5! (vi xx, ( vin
21

---

Published as a conference paper at ICLR 2026
The first quality holds as E is mean-zero and independent with X, and the last equality holds as the orthogonality between v and Em,k implies that (Vi X1a is independent (Em, Xi,) for all i2 E [D]. Notice that [f(()) m.d = Ro Evec: (Vin: xx)) and X., (8
o' d' (c) (S
( (v d'(Ci(t)
o' o' X (s!
Consequently, (Sm, x1, is mean zero Gaussian random variable, and independent with both [F*(x)] m,s and o'( (Ei= 2 Mim x()))
simultaneously, implying that
E [M(X)
m. o'
D Σ
(w
Vi x. {s!? ) (Em, Xi, )S!! i1.
111
=E [x) my
Σ
(w
(t) Vim
x
(1())
B Bl x,, = 0.
$1,5
Based on previous results, by plugging [M) m.i = * ( (Vin x(,)) and utilizing the
definition of Pa b) in D.4 we can further derive that
D
D
(v X5, )010 ) v
4
K E
Σ
o
Σ (v (V (V
o' (E
(w
(0) Wim
!S!!
Σ
1=1
EG
11=1
D Σ 1=1
(Vin Xi, }p(()) )o!
(v x,,) - Kp(()
Σ
(v m. Xs, k1() +
p(t)K
o
D - K
EGG
EEGB
G
- Kp(t)
Σ (Vim X )1))+
Σ
(vi xn) xn)
D K
GG
Pa Ky(0?, (1- Kp(())) Kp(t)
D
(1-Kp(t))
DF
Kp(t) vn
D-K
The second equality is derived by fact that o(az) = ao (x) and d'(r)) = d')) a 2 0, and the definition of p((). The penultimate equality holds En,ea (Vin, xx))((() No0 (p())).
Enger (Vi x,, ) 1-Kp(t) ~ (o, (1-Kp(e) and they are independent. Then we can conclude the final result by the definition of F3 6) in D.4). Similar to the process of handling 41, we have the following for Iz:
D-K
D 1
ΣΣ
E
Σ
X
s(! 61.1
o
Σ
(w (0)
Kmm x,, (1 Im x ((O)
m
1 1- Kp(t)
(t)
Σ
E
(v Xz, :p(() + Σ
(v Xi,)
D - K
1=1
1- Kp(t)
Σ
(vin, X1,
Xi, :p(t)
(v x,,)
D - K
1- Kp(()
]
X;, pp(p) + Σ (vi x,,)
D K
EG
(t)(t) Kp(t)2+ + (1 (1- (1-K(())3
V.. (t) V.
D-K
where the last equality holds by Lemma F.I Plugging the calculation results for 41 and la which, as we stated previously, directly con- clude D.I In addition, we can further calculate that
into D.14). we can immediately derive (D.13
(ci(0)+ + Dm
ESD Kp(()
(100)()) - vin
(++1) w Vm
Ci(++1) vin
22

---

Published as a conference paper at ICLR 2026
which finishes the proof of (D.12) Next, prove that Ci(t) is always non-negative by induction.
Obviously Ci(t) 2 0, and we prove that Ct + 1) 2 by assuming that Ci(t) 2 0. Firstly, we define that
F
(t)
Fp(t)
Then based the definition of (L),, the iterative rule for C1 be re-written
Ci - = Ci(t) nDF3 (1- Ci(t)
Kp(t) (1- ci(t) 11-C6)
From the iterative rule above, it is clear that if Ci(() < ())), then C) Ci(),, and (1 +1 < Ci(() (i(1) > (T)). Notice that Lemma D.3 immediately implies that <
(t()) R. We can conclude that once (t) surpasses
R. then it starts to decrease until it
becomes lower than V
D K.
Therefore, we have
D Kp(t)² + + K (= Kp(t))
D +
F
Ci(t) S V K
Dn Kp(t)
SVK K
+ Dn
Ky(t)
D V K + nD (p(t) + VV K(D - K) Ko(t)) S
1
D
2nD
D-1 K
K
K
where the second inequality holds as F S Kp(t)² V K p(t Kp(t)) demonstrated in
Lemma F.5 and the last inequality holds by the condition of that n S O(D-5/5 in Theorem 3.1 Now prove that Ci + 1) 0 holds for both cases: (t()) (t) and Ci()) Ct()..
(t) s (t), then it is straightforward that C1 Ct() 2 0. Ci (t), then we
have
Ct C1(1) (()(()))
2 (1) (DCK()((Dpp()) - 2Kp(t) + 1)
K
D +
(DKp(1)2 2()((() + +1)
2 - DnV
K
D- K
D 10
21 - "V
K K 2
Here, the second inequality holds as ((() inequality holds by Ct())
DKplt) -aob))22Kppt+++ implied by Lemma F.1 The second D-K
D+1
and C1(t) Ci(t) 2 1. The third inequality holds as
DKp(t)² - 21 D-K when 1 S p(t) S k. The last inequality holds by the condition of
y that S (5/)) in Theorem 3.1 This finishes the proof that (t(t) is always non-negative. L.
F
In the proof above, we introduce the definition of proxy ()) =
Kp(t) and utilize this proxy
to provide an upper bound for (t). In fact, ()) can be regarded as a "stationary point" of the iterative rule for Ci (t) in (D.12). Inspired by the proof techniques proposed in Wang et al. (2024). we introduce the following lemma, which offers more refined upper bound for (t).. We demonstrate this lemma prior to Lemma D.8 as its conclusion will be utilized in the proof of Lemma
Lemma D.5. Suppose all conditions of Theorem 3.1 hold, and Ci (1) are defined in
Lemma D.4 In addition, define that
Kp(t)2
if (.)) is identity map; if a () is ReLU activation function; if (o()) (-) is Leaky ReLU activation function, (D.15)
Kp(t)
Kp(t)² 2x
arctan VK(D-K T-Kp(t)
A(t)
(1++))) Kp(t)
(Ix1) arctan 2x
(K(D-K)(()) I-Kp(()
23

---

Published as a conference paper at ICLR 2026
and
0
if (r)) is identity map;
B(t) 2x V D D
((t)((1)pp(t)
if (r()) is ReLU activation function:
(D.16)
(1xx² V D Kp(t))
if (-) is Leaky ReLU activation function.
Then it always holds that
4A(1)
1- Kp(()
C 1+
5(A(1) + B()) -1)
(D.17)
Specifically, when ptt) S 2VTDK' this upper bound can be tighter as Ci & (t).
Remark D.6. In fact, by checking the definition of F.! in Lemma D.2 and its calculated value in Lemma F.5 we can conclude that r A(L) - B()).
In addition, we also have the following lemma, which provides further calculation results when the conclusion of Lemma holds. This result will be utilized in the proof of Lemma D.8
Lemma D.7. Suppose C1 (t).. (t) as defined in Lemma D.4 and satisfying that
A(t) A(t) + B(t) Kp(t)(Dp(t) - (i())
Kp(t)
(1 = 1
for some scalar a < then it holds that
F(1)
(1- Kp(t))
E( (() (121 + +(()))())
(DKp())22Kpp())------ 2kp(())11) (1 a)A(t):
p(()
F(
(D- K)
Ci(1) p9(0)
(D (D-K)F22
1- Kp(()
a)A(1).
Kp(t) (1(t)(1-Pp(tt) 1 Kp(t))
1 Kp(()
Proof of Lemma D.7 We prove this lemma when ar-) is the identity map. ReLU activation func- tion, and Leaky ReLU activation function, respectively. When a(-) is the identity map, utilizing
Lemma Lemma
F2 Lemma F.5 Lemma [.6,, and Lemma F7 we can obtain that
(1 (
(1 KP(t))²
F-10 = Kp(t)² +
Fd9 p(t)²
D D-K K
(D --K)²²
rg = Kp(1); p/0 x1)²;;
p =0
(D.18)
Then combined with the definition of A(t), B(t) in Lemma D.5 we can derive that
F(()
((t)
KCT(t)
(1) F21 + =(t)
p(t)
E(0
FF())
(1+a 1- Kp((t)
Kp(t)(Dp(t) -1) (CT(t) F(O + +(t))=())
p(t)
1 Kp(t)
A(t) -(1+
1- Kp(t) A(t) + B(t) (p(()(D((() )
aA((t)
(1 (1-K)(t)) (Dp(t) -1)
Kp(())
(okp(t)--2K3)) +1 A(t)
Kp(t)
(1- Kp(t)) 2
(1 a)A(t).
where the first inequality holds by applying C1(t) = (1 + A(0)+(()) ko(t) (t() and
A00 1 1-Kp(t
(()) 0, the second inequality holds by applying the definition of (i) and the calculation results illustrated in (D.18). Similarly, we can also derive that
F
(D-K)F
(i))) ((()) (D-K)
Kp(t) Kp(t)(1- -
1-Kp(t)
24

---

Published as a conference paper at ICLR 2026
Kp(1) - 1 a 1- 1- Kp(t) 1-Kp(t)
A(t)
Dp(t) -1
(p(() -1) DKp(1)2 2Kp(()++)+ +1
A(t)
1- Ky(()
(1 At)
(DK)(() -2kp(() +
This finishes the proof when (.) is identity map. When o(-) is the ReLU activation function, utilizing Lemma E. Lemma
F.2
Lemmal
F.3 Lemma
F.6, and Lemma F. we can obtain that
Kp(t)²
(1- Ky(t)) 2
Fu p(t)2
F2 (1- Kp(t))² 2
F
2(D K)
2
2(D K)²
= Kp(1)²
Kp(1)²
(VK(D) k)p(t)
K
Fg
+
arctan
+ 2T V D - k k-D(t(11 Kp(t)) = A(t) + B(t):
4
2n
1 1 -t((t))
p(1)2
p(1)2
VK(D- K)p(t)
K(D - (1-(())(--ppi)))
F
+
arctan
4
2T
1 Kpt))
2T (mDKD(t)--223333 +1)
(D pp()2
= A(I) + DKp(t)² p()(()) + B());
F
p(t) -
K
(D.19)
2T ()(D((())2 2Kp(() +1 D-K K
Then combined with the definition of A(t), B()) in Lemma D.5 can derive that
(())
F
(CT(()( 11) + p(t))
p(t)
p(O)
aA(t) 1- Kp(())
(())
1+ A(t) B(1) (p(t)(D((t)
p(t)
K)
F(
P(U)
aA(t) 1- Kp((t)
(1- - (Dp()) -1)
1+ A(t) - B(t) (p(()(Dp(t) -1)
p(t)
DKp(t)2 2Kp(t) +1
1- Kp(()
A(1) (1 1+ A(t) + B(t) Kp(t) (Dp(t) - ),) (1- (Dp()) -1) Kp(t)
aA(I)
1- kp(((
- 2Kp(() DKp(t))-2Kp() +1 A(1)
(1- (1-Kp(t))(pp(t) B(t) (1+ aA(t) 1- Kp(t) (1 (1-K
A(())+((t)) (D(()(Dp(t)) )
(1-Kptt)
+ (20)())2KKp((+]]]
DKp())=-2kp((+1 B(t)
(I-K((1)²
(1 A(t)+B(t) ) A(1) a(())(())
aA(l)
A(t)+B(t)
(1-))(t))
(22((())-((((()))) (1-a)A(1).
Similarly, we also have
F (D (D-K)F K)F - -1(()) (00 (D-K)122
Kp(() Kp(t)(1 Kp(tt)
1-Kp(1)
1 KKp(t)
(D)((DK()222Xp())) B(t)
Kp(t))
(1+
aA(I) 1-Kp(t)
1Tt (Dp(t) -1
A(t)+B(t) Kp(t) (Dp()))))
1
11+ A(t) + B(t) Kp(t)(D((t) - Dp(t)-1 -
AA(t)
1 1- -Kpp(t)
A(t)
Kp(t)
+1
1-tp(t)
+
Kp(() K)())(DK(()322XX))))
25

---

Published as a conference paper at ICLR 2026
1+ (1+
Dp(t) - A(t)+B(t) (plt) (Kplt)2-2Kp(t)>
aA(tt)
1- Kp(t)
B(1)
1 Kp((()
(1- A(t)+
aA(l)
A(t) al(()B(()
A(t)+B(t)
1 Kptt(t)
(1 At)(t). (1-))((t)
This completes the proof of the scenario that (() is ReLU activation function. For the case that (.) is the Leaky ReLU activation function, utilizing Lemma F.I Lemma F.2 Lemma F55
Lemma F.6
and Lemma
F.7
we can obtain that
(1 + K)p((
(1+ (++x)(1= 11001110 Kp(t))²
2
2(D K)
(1 + x²) pp()p
(1 + k²2 - Kp(t))²
2
2(D KK)
(1 + kxp(2t
(1- --) x)pp((
VK(D-K
arctan
2x
1- Kp(((
(1 x)²
K
+
D K x2((((1 - -Kp(t) = A(t) + B(t):
2T
(1 + p(t)2
+ (1 (1-1))p((t2
VK(D k(p(tt)
(1()
arctan
4
2x
1- -x((())
(1 x)² > K(D - KK) k))(t))((( Kp(())
(D K))((t)²
+
=A(t) DKp(1)2 - 2Kp(1) +1 B(1);
2x (DKp(1)2- 2kp(1) +1)
K 2=(D - (DKKp(() 2K)(())(p))))))))))))))))))))2 + V D-K
(1 - (1-)((((((K(((()))
P (1)
(D.20)
Then the remaining proof entirely identical that of the ReLU activation function, when replacing the values of these terms demonstrated in (D.20).
Based on the conclusion of Lemma D.5 Dn5 and Lemma D.7 we are now prepared to prove Lemma D.8 We will address the proof of Lemma D.5 after completing the proof of Lemma D.8 Lemma D.8. Under the same conditions of Theorem 3.1 there exist time dependent non-negative, monotonically increasing scalars (2()) and Ct()), such that
D Σ
w KQ = C2(1) Σ Σ P.P.T - -Cs(1
Σ P.P. 1=1 GG
(D.21)
KQ demonstrated in (D21)) there exist a time dependent scalar p(t), such that S / = p(t) for all i [D] and in E G Otherwise, s(O) 1-Kp(t)
Due to the specific pattern of w(()
D-K K Additionally, 1 < p(t)= < K and p(t) is monotonically increasing. Based on the definition of p(t), C2(() and
Di
(t) have the following iterative rules
=C((t + C1
1 (())
F210) -C((() (C(((((FE++(()))) +D(()21)) (
(D.22)
VD
K
p(t)
(1130
1/00) (Kp(1) (p((((1 - Ky(()) -c((1) p(()) (D (D-K) K) F(l
(D- K)
vD(D- K)
K p(1)
(D.23)
In addition, based on all these definitions, the coefficients (t C2(t), and Cz(t) are essentially minimizing the following loss function by gradient descent
V 12 K(D K)
1 - -CP)) + +CP (1 (-kpp) - (FFCC.P) K
2
Z(C1, C2 C2) 2(D-
26

---

Published as a conference paper at ICLR 2026
where Co is an absolute constant such that Co 1 (o)-) is identity map) + 11 (o()) ReLU) + 1+s² 1 (o()) is Leaky ReLU3- In addition, Fo(Ci,p) is defined as
0:
af)is identity map
pC2
Kp
+ arctan
PV PVK(D-K) I-KP
1/
(I-KDVR TVD-K K
If a-))ss ReLU activation
F8
- x)-C
Kp
1 7
arctan
PV PVK(D-K) I-Kp
[1- KplvK
If (-)(s)ss Leaky ReL.U activation
FVD-K
Proof of Lemma D.8 Similarly, it can be easily verified that the initialization W isfies (D.2I). Assuming it holds at the i-th iteration, we aim to prove that it continues to hold at the t 1-th iteration. To do this, suffices show that
KQ = ODXD sat-
D w kQ) (1) Σ Σ P.P. + ez(() Σ Σ PuP.
D
C0. V Vxx C(W
(D.24)
GEG
where (2(t) and (t are two time-dependent non-positive scalars. By Lemma
D.I we have
VDV
L(WV)
KQ)
Σ
D
[r-(X)] m.t
Σ
D
(w (Wv. V.m X,, )5 )o
Σ
(w
V. xz, Ss
11.1
ΣΣ
V.m Xi, (s) S (p. (P:: Piz )p.
(t)
1=li==1
M
D
D
D
w V Xi
(1
(w V.m x,, ( Pi,
$1.3
17
E
[F(X]]
m,
o'
V.m
(t)
X
(5())
Σ
Σ
(w Wwm X1, /S! P P.P:
===
I
Σ
Σ
E
(w 11
V.m Xi, [s!!!
0'
Σ
(w
V.m X (1(00)
ΣΣ
(wl W. xi, ( s
t2. P..
11=1
11
Σ
M
D
E
0
Σ
D
x,, 15
or
(L
D Σ
(w
Wi x
(s(())
$112
Σ
D
D
(w mm x,,) ( s(O P.PP
11=1
11
16
(D.25)
In the next, we discuss the value of Is. Is, Is. and In respectively. For I3, can be calculated
1
Σ
Σ
R
Σ
(v
xx.))
Σ (0 Xi, (5!0
E(w )
P.P. Σ
s az
i=1
P=1
(t)
M
Σ
E
(Vmnxx))(()+
Σ
(Vire Xe, )p(t)
Kp(t)
m=1 t=1
PEC'
GEG
27

---

Published as a conference paper at ICLR 2026
(vm: x,,) 1 ( Kp(t)
o (v (Vvi.vx/// xx)p(() +
Σ
(vm x,,)x(1) +
Σ
(V Vyn:xx20(1 p
D - K
Gi(()
M D
Σ m=
Σ EE
(v X,, (((pp()))
Kp(t)
FE
o (Vin: x²)
1 - Kp(t)
Σ (V Xi, }p(()= +
Σ
Vin)in)
1 1 Kp(t)
(V (v)x)))((t) D.D.
D-K K
D K
GG
C((0M
F4 p(t)), (K - 1)p(t) (1-Kp(t))
D
Σ PPPD
p.P:
Ky(1)
D K
1=1 YE
C(O Kp(t)
Fs. Ky(t), (p)K)(t))² 1D-K-1)(C (1-pp((t)
D
ΣΣ P PgG
K)²
10-K)²
G(OME D ΣΣ P + G(OMF (1) D
EEPrl
Ky(()
Kyll) Pyg
PEE
Notice that (v xx) ~ N((0,p(1))) Enecs da (V x1, pl() ~ ()-KK)())))))) and
Engar ( x 1_Kn(t) ~ N(0,, ('-Ko(t)²) are three independent Gaussian random vari- D-K
D-K ables. Consequently, the first term in the penultimate equality derived by the definition of Fa(a, b) in (.5). Similarly, Eijecs m. x., ()(p(() NO Kp(()²), (v x2) - DーK (0, (I-p((t)² (D-K)= ), and Enggr (Vin X., D-X .(0, (D- K 1D-K-100-KP0)) kp(()² are three in- -)-
L-Kp(t)
dependent Gaussian random variables. Therefore, the second term in the penultimate equality is derived by the definition of (s b, c) in (D.6). Similarly, for 14, we can calculate it as
4
Σ
M
D
E
0
Σ E
(vm: xix))
)
(0(2
Σ
(w/ (1) Vim Xi, 1
(w
Σ
D
(w
Vin x x, 2917
Σ (
D
P..P
K
m==
i=1
11=1
12=1
(t)
M
D
1 Kp(t)
Σ
(v xn))(())))))))))))))))))))))))
0'
Σ
(vi X1y )p(p)
+ Σ
(vm:X,,)
Kp(t)
D - K
m=1 1=1
EG
Σ (v x()pp((
Σ X1,)
1- Kp(()
Σ
S
P.P
D K
22
GEG
G18G
C(())
KT(t); (1-K011K Kp(t))²
2
D
P3
ΣΣ Pr.P.
K
D-K
(1())((-KP()) Ps (KPTOR
(1-Kp(t)) 2
ΣΣ P.P.
(D-K)p((I)
D-K
G((
D PPPT + (1(1) - Kp(()
D ΣΣ P.P.
K
(D - k)p((t
1=1 iggg
The penultimate equality holds by the definition of Fa(a, b) in (D.4, as E,, Ea (Vin xi, p(()
(0 Kp(1)²). and En (Vin, - random variables. Additionally, 15 can be calculated as
1-Kp(() No, (1 (1-Kp(U))*
are independent Gaussian
D-K
M D
D
D
D PaP: Σ
Is = 2
ΣΣΣN
E
o
(v m x;, (S()
o'
Σ
(v Xi, (//0) (v x
s(D
m= P=
n=11
==
= -C,( 3
E
Σ
( X (1)
d
Σ
(v X
(1(())
( (rin
PvP.
m=li=lve
28

---

Published as a conference paper at ICLR 2026
+ ((1)
ΣΣΣ Σ
M
D
E
a
D
(v x s )
0'
Σ
D
(v X1, 1S!0 )
(v xx) 1 Kp(t)
P.P.
D - K
M=
=1=1
11
-MG(11)FF (11))
(1- (1-Kp(t)) Kp(t))
2
Σ Σ P+P.
D
i=1 PEE
(11- (1-Kp(t))
(D (D-K- 1D-K11)( - Kp(())
+MC((1)
(D-Ky² .K((t)++
Σ Σ PpP!
(D K))
i= rEGs
D 2
D
Σ
Σ P + MCG(t)
Σ P.P.T
1=1 PEG
1= PIRGC
where the penultimate equality utilize the definition of F2(a,)) in D.3
Similarly, for Is, we have
Is ( E
m- X,, 1s
0' Σ
! 11.
(v X1, S
Σ
2
Pip
D
1- (ptptt)
D
P(((()
P.P. + MC
Σ
Σ
PAP:
D K
1=1
12E
==1
12
Combining all these results of I3, I, Is. and Ig. and plugging them into
D.25 we obtain that
V (O: (O) ; KQ)
G() 1 F(())
D Σ pap.
-F!") - (1(() (F -)())()+))))))))
VD K p(t)
1 Feree NEG
C10M (100(1-KK(0) Ky(()
(
Kplt) (EDDE)
Kp(1)(1 - -P((1)) -C(())())) [F(00_ (D (D-K)F22
(D (D-K)F!
ΣΣ PrP.
( KVVD
1-Kp(t)
==(((t Σ E iEl
p.p. (t)
ΣΣ P.P.T. PEg
remains to show that (2 and ca(1) are always non-negative. Notice that Lemma D.5 guarantee
the assumption of Lemma D.7 By carefully compare the formulas and applying Lemma D.7 we can obtain that
K KvDea(1)
(1-Kp(()) MC1(t) (K)()DKK())-2K0(1)
A(t) 2 0.
Since we have proved that (t) is always non-negative in Lemma D.4 this result implies that c2 2 O. Similarly, for ca (t), we have
VD(D K))a(() (1-(0)(1-p((() 5K)0)(DKx()322200)))) A(t) 2 0.
Kp(()
This proves that ca(t) 0, and conclude that
C,((+1)=C2() + C1(t)M
1
(())
- -((0) (00/:+-p(())) +2(0)1)))
VD
K
p(()
-LUMLLL))) (1-K((t))
F(
(D- K)
(D (D-K)) (D-K)
Cs Cs(++1)=C3(0->) VD(D K)
Kp(t) Kp(t) (p(t)(1-pp(t)) Ci(() P((
1 Kp(t)
which completes the proof of D.21). (.22) and (D23). It remains to prove the conclusions regard- ing ((!), and p(t). By the orthogonality among the positional encodings p it is straightforward
29

---

Published as a conference paper at ICLR 2026
that for all i [D].
(2t(t
ih G
P w KQP:
Ct((t)
if $ G¹.
Then by the definition of ((0,, when in G
P. W
Ri
C2(1)
exp
/D
exp
vD
11
Σb= exp
Pl With
KoP.
K exp
C2(l)
+ + - K) exp
Ca(O)
Ca(t) (t(t) p()).
K + (D X)exp exp
VD
Since (t) and (t()) are non- non-negative and monotonically increasing scalars, we immediately con- clude that D S p(t) K. and p(t) is also monotonically increasing. Lastly, it remains to formulate excess loss into an expression of excess loss. By the parameter forms Lemma D.2
1
we have
E(CI, C2 C3) L(Wv, WKQ) Copt
R
M
D
Erear (Vm xx)
a
(Cip Cip
Σ
(v x,,) + Ci(1 - Kp)
Σ (vi m. xxi) x
K
D K
m=
REG
M
D
Eveeg (v Xv)
27
D
Ci Kp))
Kp)
27
E
E
Cip Σ (vm: xp)
Σ (vin x(x)
D G'g
PEG
Is
Ever (Vms Xir)
(v:n Gi(1-Kp
E2
σ CIP E
Σ (v xp:)
K
D-K K
NEG
PEG'
1.
For the term 1y and 1s. by the fact that B[o(x) a when x ~ N(0, a), we can directly calculate that
1, Co |l
K
C1 -p)11pp) 1 |
[ = Kp Cfl
D K
In addition, by utilizing the conclusions in Lemma F88
conclude that
co Cipll Dl 11 Fo(r,,))
2
Plugging all these results, completes the proof that 6, V 2(D K)
Z(C1, C2, C3)
K(D k) ( 1-
K -Cip)
+ +Cp (()1KKp) kD) - DIIV 2
Fo
Now, we successfully prove all the conclusions of Lemma D.8
Lastly, before we prove Lemma D.5 we first introduce and prove the following Lemma D.9
Lemma D.10 Lemma [.II and Lemma 1.12 which will be utilized for proof of Lemma D.5
Lemma D.9. Under the same conditions as Theorem 3.1 and p(t) as defined in Lemma D.2 it holds that
p(t) (1 -K((t) (AC2(t) + ACG(tt) < A v
D2p(0(1- - Kp((t) (ACz(t) + ACg(t)):
2VD
MD
D+1 1
AC2(1) ACa(1) S K2(D K)
K
DT2
where Ap(() p(1 ((1+1)-(())) AC2(1) C2(1 1) C2((), and ACs(() Cs(l Cs())
30

---

Published as a conference paper at ICLR 2026
Proof of Lemma D.9
By the definition of p(t) in Lemma D.2 it can be derived that
Ap(1) = p(t 1 p(()
K K K))xpp
Co(++))+
K + (D-K))e) Kxxxp
Ca(t)+
VD
VD
1
1
K+ (D-K))))t
Call)CC(t))
ACalt)+ACa(1)
K KD(D-K)ar exp
Calt)+Ca(1)
vD
VD
vD
(AC2(t) + (Ca(t)) - K) exp
Ca(0+Ca(l)
VD
VD + +D -K) exp
exp Ca(l)Ca()))
AAa())+Aaa()
K+ +D((- (D- (D-K)ens exp
Ca(t)+Ca(t)
VD
VD
vD
(D-K) - K)eN)) (- Ca(1)-Ca(1)
AC2(t)+ACss)
VD
VD - AC2(t) ACa(tt K (D) exp
C2(t)+Ca(11)
VD
AC2(t)
VD AC2(t) ACa(tt) p(t)(1 (1p(t))t(t)
Additionally, applying the update rules for (t) and (t( derived in Lemma D.8 along with
similar calculation to the one used the proof Lemma D.7 we obtain that
AC S (t)(t)
F(
(0)
1'3
KVD
p(I)
(t)(t)
Kp(() A(1)+
DKr(122 2k(()) +1 B())
(D- - K)p(t)
1 KVD
MC, (t)
(1ttt()
(l-K)()) (Dp(t)
A(t)
B(t)
KVD
Kp(t)
(p(t)(Dp(t) - + 11-++KKppp -Kp(t)
S1 (t) 1- - Kp(t), p((1) <1 S" KVD Kp(t)
(C) 1- Kp(t)
(KD(1² + 1
2 V 10-K - -K(()) ((((O)(-K((
K
KVD Kp(t)
M. D+1
S1 K2
V
DK
Here, the penultimate inequality holds Ci(() v also derive that
D+1 K
and 1/2 S p(() S R. Similarly, we can
(C&(t) S MC((0(1-kpl))
F
M
D+1
VD(D K)
Kp(t)
S
K(D - K) V DK
Combining these results, we have
MD
D+1
1
AC2(1) AC&(! S K2(D K)
< K
D
where the last inequality holds by the condition that 17 O(M -DD -5/2) in Theorem 3.1 Replacing these results, we finally prove that
( (AC2(t) (Cc(())
Ap(t) < VD(D2-1
On the other hand, since AC2(!) + ACa(() is sufficiently small, we can also have
Ap(t)= 2
K + (D K) exp
Ca(t) VD
+ 2/D
K + (D - K) exp
Ca(t)+Ca(t) VD
31

---

Published as a conference paper at ICLR 2026
(D - (D-K)ex) exp
Ca(t)+Ca(t)
+C2())+Caa))
VD
2VD
KD(D- K) K))xp) cxp
Ca(t)+Ca(t)
vD
p(t) Kp(t))
2VD
This completes the proof.
Lemma D.10. For (1) defined in Lemma D.2 it hols that (1) is monotonically increasing w.r.t t when p(() 2/ZDK'
Proof of Lemma D.10 As Lemma D.8 demonstrates that p(t) is always monotonically increasing. Consequently. it suffices to show that Ci monotonically increasing w.r.t p(t) when p(t) S 2V=DK' In the following, discuss the three scenarios where a(-) is the identity map, ReLU activation function, and Leaky ReLU activation function, respectively. When o(-) is the identity map,
(D K)((t) (t) DKp(t)2 2Kp(t))+ +1 DKp(1) + p()]
D- K
2K
is straightforward that (t) is monotonically increasing when p(t) S VDK as the denominator decreasing. When (-) is the ReLU activation function, have
<(D K)p(t) + 2(D K)p(t) arctan (VK(D- K) p(1) KP + + K) 1- 1-Kp(1)
VK(D-K K)
(t)
(TK (OKK()))-2Kp(()-1 +1)
By applying basic calculus, derive that
dti(t) 2 - 1) ()) - -2Kp(t) 1)
dp(1)
(DK(D2) 2Kr(t) + 1)²
2(D - K) (Tp(() + VK(D-K) Kj
4T<((D((t) -1
4.2 (DK(Dt) 2Kp(t) +1)
3T - K)(( 4t (DKp(tt)) 3T K)(1 2VTDKP(())
4 (DKp(1) - 2Kp(1) + +1) 4 (DK - pKp(() + +1)² 2
which is positive when p(t) S 2/EnK" Therefore, we can conclude that when p(t) S 2/=nk* (L) is monotonically increasing w.r.t. p(1). Similarly, when o(.) is the Leaky ReLU activation function, we also have,
+ -100----)))00000--0)0000-0000000 - 2(1+ ()) +1)
aretan (VK(D-K) p()) Kylt)
Ci(t)
(D=x)(D-K)(I-KP)))
/ 2Kp(t) +1)
and
(t)
2π(
2Kp(1) +
dp(()
((DD((()-22333)3 +1²
2(D - K) (P(()) + -K(D-K) K)
K K(Dp(1)- -1)
π (DKp(1)) 2K)(1) + 1)²
3r - -1
( 4mDp((tt)) 3r K)(1 2VnDKp(()
4x (DKp(1))-2222)) +1)² (DKP((1 - 2Kp(t) + +1)
which proves that (t) is monotonically increasing w.r.t p(t) when p(t) 1
32

---

Published as a conference paper at ICLR 2026
Lemma D.11. For (t) defined in Lemma D.2. it holds that
3DKp(t)Ap(() c DKp(t)² 2Kp(t) + CT()),
(D.26)
Proof of Lemma D.11 We prove (D266) for o(-) is identity map, ReLU activation function, and Leaky ReLU activation function, respectively. When o()) is identity map,
(D- - ())p(t++ + + (t+++)) DKp(t + 1)2 1) 1) 2 DKp(t)2 2Kp(t) - DK (2))(t)(2p(t)) + Ap(t))
(D - K)p(t)
DKAp(t)(2p()) + Ap((t)
3DKp(t)Ap(()
2 ((()) D(p( - p()(() +1 CI()) CI()) DKp(t)2 - 2K +1 (1)),
where the second inequality holds by Lemma F.9 and the last inequality holds by Ap(() S p(t)
implied by Lemma D.9
D.9 When (-) is ReLU activation function,
7(D - K)p(1 1) + 2(D K) + 1) arcian (VK(D K) 1)
p(++))
T- (ptttT)
2k 2n D p(l 1)2 12pl11)2---------- +1)+(+++)++)
(D (1 Kplt + 1))
v K)=D-K)x(DK2))) - + 11)
T(D K)p(()(() + +2( K) arctan (v KKD K)
p(t)
+2 K)
1-Kp(()
Kp(t)
VK(D-K)
(T(Dk(())-2Kkp)) + (K Ap(t)) 3DKp(()Ap()) (t() DKp(t)2 2Kp(t) +1 Ci(t),
where the first inequality holds as the numerator is monotonically increasing function w.r.t. p(t).
Furthermore, the second inequality holds by Lemma F9 and Ap(t) S p(t) implied by Lemma D.9 Similarly, when o-)) is the Leaky ReLU activation function,
(D=K)2((D-K) 1)++1) + 2D())))(((++++)) arctan ) (I-x)²
p((+1) (-p(ttt)))) + + [x(p(t+1)
(K(
(t+++
2x(1+x²)
(DKp(l + -)2-1Kp(11 1))
[I-K]²
(14x) (1-x)²
p(() + 2(D K K)p(1) arctan (v K(D K) p(t)
kpptt) 2(D K) kx(((t)
VK(D-K)
2x(1+e²) (DKp(t)2 2Kp(t) +1) 2x(1+x²) (1-x)²
DKAp(t) + Ap(t))
(1-x)²
3DKp(()Ap(1)
2 (t) DKp(t)2 2Kp(t) +1 Ci()).
This completes the proof
Lemma D.12. For A()), B()) defined in Lemma
D.5 it holds that
A(t + 1)
1 Kp(t + 1)
A(L+ + B( 1) Kp(l+ + 1)( 1) 1)
A(I)
1 1 Ky(()
A(L)
2DKp(()Ap()
(D.27)
A(t) + B(t) (p(())Dp(t) 1) A(t) + B(t) (2 (1)(())))))) -1)²
Proof of Lemma D.12
Notice that
B(t) is non-increasing function w.r.t p(t). Therefore, we can A(A)
derive that
A(t + 1)
Kp(t +1)
A(t + 1) + B(t + 1) Kp(t+ 11)(1(( +) -1
A(t + 1)
1
A(t + 1) + B(t + +)
Kplt + 1)( + +1)
Dp(t +
33

---

Published as a conference paper at ICLR 2026
A(t)
A(() + B() (x(()(Dp(() - 1) + (2(()(2D)(()) + DKAp(() - K) Dp(()
A(t) 1 Kpttt
A(t) Ap(t) + DKAp(() K)
A(t) + B(t) (ppt)(Dp(t) 11) A(t) + B(t) (²p(t)(Dp(t)) - 1)
A(t)
1 Kp(t)
A(t)
2DKp(t))p(()
A(() + B() (ppt) 11 A(I) + B() (²p(t))(Dp()) - 1)
where the last inequality holds 8p(1)= $ D
1 implied by Lemma D.9 This completes the proof.
Now, we are ready to prove Lemma D.5
Proof of Lemma D.3 As Lemma D.10
guarantees that (t) is monotonically increasing when Consequently, when p(t) I 2/DOK
p(t) Σ 2/DDK'
nDF 1 -((t)) (1- Kp(()C1(0) (CF0))00)00000
2 (1 (1-
nD K
(G(() -((0)) 2(1- 1D k (CT(0) ------C))))))
2 0.
F(1)
The second inequality holds by Kyft)
S k. and (1) 2 1. The last inequality holds by the assumption of in Theorem 3.1 and (0) = 0. In the next, we prove that (D.17) holds when p(t) 2 2/TDK by induction. We assume (D.17) holds at t-th iteration and examine the I +.+ 1-th iteration. Inspired by the separating strategy in Wang et al. (2024), we consider the fol- 24(4)
lowing two cases: (i). when Ci(() S (1 + 5(A())B0) Xp(C)Dp(0-T) (t) and (ii). when
1 1-Kp(t) (1 (((())+(0) Kp(e)(p(())r
2A(1) - (1)) < C(()) S (1+ B(O())B()) KP(O(P(()) (1). For
4A(1)
1- KPp(()
the first case, it suffices show that
1- + 1) 1+ 5(A(t+ B)+ Kp(t+))
4A(t + 1)
Ci(++1)
2A(1) 1 Kp(()
2 5(A(t)+ B(tt) (D(()(Dp((-1)) (i()).
(D.28)
This is because if (1)) < (t), we have
(1)) - 1t (1- kp(t)C7(t)
DDF
(CT0)) (G()) 20,
which implies that
4A(t 1)
1- Kp(t + C((+1).
< 1+ B(A(t+1)+B(t+1)) Kp(L+1)(D(((+))1)
The last inequality is guaranteed by (.28). On the other hand, if Ci()) Ci(t) s (1+
B(OJ) (P)))PP((-T))) Ci (4), then we can also obtain that
2A(1) (vp(()
2A(1)
Kp(t)
(((+1)sCC0)
1+ 5(A(()+ B(t)) Kp(t)(Dp(1) -1) (1(1)
s <(11 1+ 5(A(t+1)++(t+t)) Kp(++1)(Dp(t+))-1)
4A(t 11)
1 - Kp(t + 1)
Ci(++1). )
In the next, we show that (D.28) holds. By applying the lower bounds derived in Lemma D.I and Lemma D.12 we can derive that
4A(t 11 ((A(++)+B((++))) KD(t+1)DDp(+t))- ) )
1- 1-pptt+1 11)
34

---

Published as a conference paper at ICLR 2026
2A(t) 1- Kp(t)
2A(t)
1- 1- Kp(t)
5(A(1)+B(0) B()) (p(O((p(() -11 CI(1) + 5(A(1)+ + B()) (P())D Ci()
(1+ 5(A(++1)+ B((+11) (D((+1((p(t-))-ii DKp(t) (Kp(t) +1 (1
4A(t+1 +1
1 1 Kp((++))
3DKp(t)Ap(t)
A(t) 2DKp(t)Ap(t)
A(t) + B(t) (Dp(t)((p((t) - 1)2 Ci(t)
1+ ((A(t)+B(()) (D(()(D(()) -1) Ci(() + 2(1 2(1-K)(0)
2A(1)
1- kp((()
15KD())D0(0-1) (T())
3(4T + 1))(p(t)Ap(t)
10mDKp(t)Ap(t (t
Ci(t) (p(t)(D((t) -1)
Kp(t((p(()) -1) +1 +1kKp(t)
2A(1) 1 Kpp(())
(T(()+ + 1- 1-Kp(t)
_22T + +t)Kp(t)
1+ 5(A(t)+B(tt) (D(())
cT(()
(p(t))Dp(t)-1) 15 D
2A(1)
1 Kp(()
B(O)) (D(()(p(()-I) Ci()),
which finishes the proof of (D.28) In the derivation above, the second inequality holds as A(t) 2 and kp 11) I 5π when p(t) 2 2VEDK The penultimate inequality
A(e)
derived by Lemma D.9 As demonstrated previously, (D.28) implies that (D.17) holds at the + 1-th iteration for the first case. In the following, we consider the second case, where 2A(1)
(1+ s(A())+B()) (p(t)Dp((t-T] (1) S C((1) S (1 B(A Kp(tXDe(1)-1 (1). For
1-Kp(1)
4A(0) 1-Kp(()
this case, it suffices show that
4A(() 1- Kp(()) 1+ 5(A(t)+ +())) Kplt)(Dp(t) (i() 1D(1-Kp(t))
10K (Dp(t) -1)
1 1 Ky(t 1 ((A++ (((++) ko(t+1](pp(+t))-]]
4A(t+1
(1+1+)
(D.29)
This is because
Ci(1+1) + 7DA
(1- Ci(t)
CT()) (1-C(0)
Kp(()
4A(t) 1-Kp(t)
)))
CT(0)
DDF 2A(1)
1- Kp(t)
1+ ((A())+B(t)) (P(()(D((1))
Kp(t) 5(A(t) + B(t)) Kp(t)(Dp(t) - 1) Kp(t))
4A(t) 1 1- Kptt)
CT(1) 10K(Dp(t) 11)
5(A(1)+ B(U) (p(t)(Dp(t)
4A(t + + (p(t+1)(Dp(t+))
1 1- Kp(tt 11)
(1++1),
where the penultimate inequality is derived by F.! = A(t) Bt) and A(t) > Ky(t)² and the last inequality is guaranteed by (D.29). To show (D.29) holds, by applying Lemma D.7 refined upper bound for (C2(() and (C&(t) as follows:
4
we derive an
AC (t)
E((t)
F.( -KC((t)
10 +0++))))))
KVD
p(1)
3n (1-Kp(1))
A(t) < 3(4m + ()1 Kx(())²
3K/DD(K)())22Kp(0) CT()).
5KVD
KD()(DK)()=-2X000))
and
(MC(()(1-KK())
(D- (D-K)F(1)
(1(0) (D-K)
1 -- ) 1-Kp(()
ACa(1)
- Ci(1)
VD(D-K)
Kp(()
Ky(1)(1 - Kp(1))
35

---

Published as a conference paper at ICLR 2026
(1MC(((1--K(O))
1- Kp(t)
5VD(D-K) (DK(( A())
3(4T + (1) (1 Kp(t)) 2
(t).
5(D K)
Based on these refined upper bounds for (C2(t), ACs(t), and the lower bounds obtained previously, we can derive that
1+ 5(A(t)+ B(t)) (p(O)(D(()) -)()(D(()-1))C C7(1) )T0
4A(1)
1- kp(()
4A(1+1)
1 Kpll + 1)
1+ ((A((+))+ Kp(t+1/(D((+1))-1] (i(++)))
1+ +(A(t))+(()) (p(((Dp(()1)) Ci(t) ) (500(000+((())+B00)) ((()))+((tt)
4A(t) 1- Kp((t)
4A(t)
(P(()(P(()-1))) CI(1) C
1- Kp(t)
(22 + DAp((t (1(1)
Dp(t)= -1
(22m + 3)D D2 D2p(1t(( Kp(t)) (AC2(t)+ + (ACs(t)
AC3(t))
Dp(t)= -1 VD(D2 1)
3(22T 3) + D)nDD) Ky(()
K ((D)(t)1) - 1)D(D ()KK0))-)1000)))0/0000000- CT()2 C1))
nD(1 - Kp(t)) 60 nD(1-Kp(t)
K ((Dp(t)-11 DK S 10K (Dp(t))))
Here, the first incquality is derived by applying the upper bound of (1 + 1- K-K((t+1)) Ci + 1) obtained previously. The second inequality holds by applying Lemma D.9 The third inequality is derived by replacing the refined upper bound of (t) and (C2(t). The penultimate inequality holds as (1) I Kp(tj. and the last inequality
4A(++11)
5(A(++1)+ B(+++) 1p(t+1][
guaranteed by D 2 ((M) in the condition of Theorem 3.1 This demonstrates that (D.29) holds
in the second case, which completes the proof of (D.17.
D.2 THREE PHASES TRAINING
In the previous section, Lemma D.2 accurately characterizes the training dynamics of w and
w(())
KQ- Specifically, demonstrates that w(( C, (L)V,, where (I) is always upper bounded by (1+ + (() 1-Kp(t)
4A(1) CC (1). Next, will show that the update pattern of C1
differs across three distinct phases. In the first phase, (t) monotonically increases, approaching (t() while p(t) remains close to b- In the second phase, (t) remains in a neighborhood of
(1), while p(() monotonically increases. This increase exhibits modes characteristic of a tensor power progression, continuing until p(l) reaches Z In the third phase, the p(t) will eventually converges to K. and (t) converges to 1, leading the loss also to converge. The formal proof is provided as follows.
Lemma D.13. Under the same conditions with Theorem 3.1 there cxist t = Onn 1), such that (t,) 0.95 (t1), and p(t) < 1 for all t t..
Proof of Lemma D.13 Notice that when (1()) < (t), Ct is monotonically increasing, Let & be the first time such that (t( Z 0.95 (t). For the conclusion regarding p(t) with < f1. we first assume it holds and utilize it to demonstrate other conclusions, and lastly prove it by induction. Since p(1) almost remain unchanged for all / S L,, we can obtain that 0.975 -975--C (t') > 0.95 C (1") for all ,', 1" < l (This conclusion is proved in following Lemma D.14). Therefore for all / < 4
nDF Kp(t) (it (1-
Ci())
DDF
(1
(11 - 1)
DDF
Kt(t)
0.95
(11 -
40Kp(t)'
0.975
36

---

Published as a conference paper at ICLR 2026
where the last inequality holds by
(1(1-1) (1 I 0.95. On the other hand, it is straightforward that
C(t+ C(t)
2Dr Kp(t)
Additionally, when s p((,,((ti) s 1+D 1/4
we can obtain that
* -I) is the identity map, then
F1 Kp(t) = p(t) = = (1):
Eltt)
6(1)
(41) = Kp(())(())
@(1)
(f(()) is the ReLU activation function, then
P(
p(t)
p(t)
(VKID- (LLD-Kjp())
2T / K(D K) (1 (Kp(tt)) = -0
arctan
+
(VDK
Kp(t)
2x
1 -Kp(t))
psli Kp(t))F(')) = Kp(th)(D((t)) - 1) + - = =0(17) D
0 K)
(1)))
(fo()) is the Leaky ReLU activation function, then
F
+ p(t))
(1 x))p(t)
VK(D- VK(D-Kjo() -Kp(t)
-Kjp(t)
+
arclan
(1 Ky(())
Kp(t)
2x
2T / K(D K)
1
θ
VDK
e( R) Kp(t))F(1) = Kp(tt)(D((tt) - 1) + Kp(tz)
p(())
= =0( K
(41)
Therefore, conclude that
O() Φ(6
if af)(s) identity map;
@(v2)
4 0.95 (1) EE (C
O1)
if of ReLU activation function;
O 1)
( (-)is Leaky ReLU activation function.
θ
ny
Next we prove that p(4)) <
1 1/4
by induction.
Assume holds at 1-th iteration, then by
D
Lemma D.9 we can derive that
MD AC ACg()) S K2(D - K) V
D+1 DK
and consequently
Ap(1) < D²p(t)(1-pp(t) (AC2(1) + ACa()) < 3M
VD(D2-1)
D3
Therefore, we eventually conclude that
4-1
1 Ap(t) S 1 +
M
1 1/1
p(h)) S p00) +
Σ
D
DK
D
t=
where the last inequality is derived by our condition that = 2(poly(M)) in Theorem 3.1 This completes the proof.
37

---

Published as a conference paper at ICLR 2026
Lemma D.14. For all t',t" I t1, where t is defined in Lemma D.13 it holds that 0.975 (t') 2 0.95 . (t").
Proof of Lemma D.14 Notice that by the definition of (t),,, is entirely determined by p(t). And for all t', ""t" < t1, we all have 1 S p(t)), S 1+D-1/4 With Lemma F.1 and Lemma F,5 we can further derive that
If (-) is the identity map, then
(D- k)p(()
D-K K
Ci(t'= DKp(t)² 2Kp(t') +1
1+D-174 174 - K(1 - D-1/4) $1+D
(D -0-K K))))
D K
Ci(('= D 2K)(t))
2
=1.
D K
It immediately concludes that
1
0.975 (tt))) -
Ci(t) 2
D 2 0,
40
as D 2 Q(1..
(-) is the ReLU activation function, then
2(D K)
D-K
(D- Kin(t')
ci(()
5 (DKp(t - 2KK((t) +1)
x
K
1 --Kp(t')
(D- K)p(() Kp(t)(²
D-K + (D-R)p(t"))
1
D
K
+
K
2
(1 (-Kp(1'))²
1
D- K
K
2 K
D D-K
+
K
D
K
where the penultimate and last inequalities hold by utilizing Lemma F.10 and 1 < p(t)= <
14D D 2 (ppoly(K)) in the conditions of Theorem 3.1
Similarly, we can also obtain that
D
2(D- K) ci((") Kp(t") (DK 2KP(1") +1)
D-K
K
1 DKp(t")
1 K K
D - DKD('"
1
K
1 K -
1,
K
K
where the penultimate and last inequalities hold by utilizing Lemma
F.9
and 1 S p(t') <
D 1+--1/4 D 2 ((poly(K)) in the conditions of Theorem 3.1 is straightforward that
Based on these two results,
1
D
K K
0.975 (1") 0095 Ci(()
3 320.
40xt
D 2 ((poly(K)).
If (-) is the ReLU activation function, then
2(D - K)) (t)
(1=x)²
D-K
+ (1+*")(D-Koot Kio(t')
(t')
K
2
(1 + Kp)()(tt) (DKp(t)2-222) 2kp(t)+p
(1+ (1+x2)(1= (1-
(1 x))
D K
<
V
K +2;
+ k²)
2(D K)F(() (t)
[1xx²² D_KK
Ci(t")
(1 + K2)) (DKp(t))))()))) - 2Kp(t)+1] 2 (1+ + DKp(t")3)
38

---

Published as a conference paper at ICLR 2026
(1 x)² (1 + 2²)
VD= (D-K- 1.
K
Combining these results directly leads to
(1 [1-x)²
D K K
0.975 (es("T((")) - 0.95 cT())
3 320 0,
40 イ² π
D ((poly(K)).
This completes the proof.
Lemma D.13 successfully demonstrate that at the initial phase of training, C₁ (t) will mono- tonically increases until 0.95 (t), while p(t) remains smaller that L+D-1/4 Furthermore, D
once Ci(t) reaches 0.95 (t), it never falls below this threshold again. Combined with
the conclusion demonstrated in Lemma D.5 that C (t) is always upper bounded by (1 + 4A(e) 5(A(t)) B(t)) (p(t) (t),, we can claim that C1 will always remain inner neighbor-
1-Kp(t)
hood around Cf (t). The following lemma provides a formal illustration.
Lemma D.15. Under the same conditions Theorem 3.1 and with defined in Lemma D.13
for all 2 4, the following holds:
(1 2 0.95 v (1
4A(t) 5(A(() + B()) (p(t) 1T) )
1- Kp(t) CI())
(D.30)
where A(t) and (t) are defined same as in Lemma D.5
Before we prove Lemma D.15 we first introduce the following lemma, which will be utilized in the proof of Lemma D.15
Lemma D.16. For (t) defined in Lemma D.2 it always holds that
3(D K + (C)
1) +
+1 +1-11-1 (Ap(t) Ap(t).
(D.31)
2
In addition, (t) is monotonically decreasing when p(t) 2 VDK
Proof of Lemma 16 We prove this lemma by considering (-) as the identity map, ReL.U activa- tion function, and Leaky ReL.U activation function, respectively.
If (-) is the identity map, then
(D-K)pp(t +1)
- K)p(t
(D - k)p(t) (D- DKp(t)² 21
Ci(++1) DKp(t 1)2 2Kp(t 1) +)+1=
D- K + (C)(1)
sci(1) + D + +1 KAp(t) Ap(!).
In addition,
D- K Ci(t) DKp(1) + 2K'
which is obviously decreasing when p(t) 2 VDK
2
If (-)) is the ReLUU activation function, then
arctan
1-Kp(t+1)
+
1-Kp(e+1) VK(D-K)
(t
2T(DKp(t + 1)2 -2- 2Kp(t + +)
39

---

Published as a conference paper at ICLR 2026
x(D -))())+((-)))((())e))))))))))))) arctan (V((D-K)pC)
1-Kp(t)
2T D 2K + 2 / pp(()
2(D-K) 1- 1-Kp(tt -D
(K(D-K)
(DKp(()) -2Kp(t) + -
-1)-27KA Ap(1))
3(D 3 (KT((t))
S (1( +
-x + - -KA((t) Ap(1).
In addition.
(D K) arctan (K(D-K)p(t) 1-Kp(t)
D-K
K (1 (-Kpp(()
D-K (t) 2(DKp(t)+ plo) 22K)
π <DKp(tt) p(t) -2K)
(DKp(t)) p(t) 2K)
where all these three terms are monotonically decreasing w.r.t p(t), when p(t) 2 VDK demonstrates that (t) is monotonically decreasing when p(t) 2 VDK
This
2
If (-) is the Leaky ReLU activation function, then by similar calculation process,
1)
[11h-0/000-0000----22-------001- arctan (K(D-K)p(ttt)
arctan 211=))=((K6((+))3-2222))) + 11
I-Kp((+1)
(1=0)(D-=(((1-=p(tt))
(1--))/00(((+-322224444++++111/1 /K(--K)
<11 (1==))00-0)))2222--0000000000000 arctan VK(D-Kip(t)
1-Kp(t)
KK(D-K) (KD- 111--"D00)))/00000++>11111111 ((000--)))0000---(((00000000--00
(DK+)((XX())) ()) (
((D-K+K((())
(0 2(DKp(t)2- - +1- (KA((t)) Ap().
addition.
(1 (D= arctan (K(D-K)p()
(1+k)² K) (t)) 2(1 2(1+x2)(DKp()) + P(J 2K)
1-Kp(1)
(DKp())(DKP() p(1) 22)
(1-m)² D-K (1 - (1KKp(t)) K
#(1+x²) (DKp(t)+ p(T) 2K)
where all these three terms are monotonically decreasing w.r.t. p(t), when p(t) 2 demonstrates that Cf (t) is monotonically decreasing when p(t) 2 VDK'
2 This
VDK*
This completes the proof.
Now, we are ready prove Lemma D.15
Proof of Lemma D.15 We first prove the first part of (D30). i.c. C1(t) > 0.95 (t) for all
t > ti. by induction. To establish the conclusion, we consider two cases at the t-th iteration: (i). when (t) 2 0.975 Ci (t). (ii). when 0.95 (1) < (1) < 0.975 (t). For the first case,
40

---

Published as a conference paper at ICLR 2026
when p(t) S 1 Lemma D.5 shows that Ci(t) V (t), implying that Ci(t + 1) 2 Ci(t).
2V7DK' Then we can derive that
3 (t) - + 11 ktp(tt) Ap(t)
Ci ((0+1)20000 2 0.975 (t)) 2 0.975 CT(++
K) D 2 0.975 Cj(t - 2 0.95 Ci + 1).
0..97 - 11
where the third inequality holds by applying the lower bound of (1) demonstrated in Lemma D.
The forth inequality holds (t) 1 v and Ap(t) DJ guaranteed by Lemma D.9 The
penultimate inequality holds D 2 ((poly(K)), and the last inequality holds as (t) 2 1. When p(t) 2 2V the upper bound of (t) established in Lemma D.5 can help to derive that
(1 Ctt+ (t(( 5k²p(t)((Dp(t) - 1)
3(D - K)
An))((t)
20.975 C[(+++)
5K xx22p(t) (tt))ttt)((CC (Dp(t) 1)
20.97 1] 3(D - K)
D22 MV K 2 0.975 CG 0.025 2 0.95 CT + 1).
Here, the second inequality applies the previously obtained lower bound for (t). The third in- equality holds as A(I) s Kp(()², and p(l) 2 2VEDK. The penultimate inequality is derived by D 2 ((poly(K)) and S OOMD -5/2) in the condition of Theorem 3.1 These results demon- strate that under the first case, Ct + 1) 2 0.95 (t + 1). Let's consider the second case, where 0.95 (t) S (t) < 0.975 Ci (t). Under this case, is obvious that C1(t+ would be larger
than (t), and by the updating rule, we have
TDF
2 0.95 Ci(+) + 1DF33
3(D - K + (C)(t)) 40Kp(t) 2(DKp(t)) 2K +1 -KKp(tt) Ap(t)
(t (t) +
4t)p(t)
D
3M Ky(())
D+1
20.95 (t) np(()
V
80
VD(D2 - 1)(D - K)K2
K
D
4M
20.5 (t) np(t)
2 0.95 1.95-CT1 (tt).
80
K
Here, the second inequality holds by D.3I). the third inequality holds since F 2 Kp(t)² by D-K+ (t)
Lemma ES5
KApe) 2D, and applying the conclusion of upper bound of Ap(t) demonstrated in Lemma D.9 Besides, the last two inequalities is guaranteed by D >
(ppoly (M, K)). This finishes the proof of (l) 2 0.95 (L) for all / 2 4. In the next, 4A(t)
1-Kp(t) ) Ct)... In fact, we prove the second part of (D.30), i.e. (t) 2 (1 5(A(0)B(() Kp())D((t)-))
we only need to consider the scenario where p(t) Z VDK* This is because when p(t) S VDK'
1- ky(() 5(A(t) + B(t) (D(()(D)(t)-i)))) - 1) 5 (P(()) + 2 V D-K (1 Kp(t)))(D((t) - 1)
4A(1)
4(1 Ky(())
K
10v L0VDKp(U 2 0.05.
Therefore, Ci 2 0.95 CI(() guarantee that Ci > (1 5(A()) (p(-)( -)(()))))))))(((())) (T((
4A()) Kp((t)
holds when p() < 2
VDK When p(() 2 2
VDK we also consider two cases: (i). when Ci(1) > (1
4A(1) 5(A(1) B(J KptO( (t). (ii). when (1 B(A())+B()) ((t) < Ci(t)
24(1) kp(((t)
C
k-K((t)
5(A()) B() KP(0[DP(1-]] C (t). Then, for the first case, at the + 1-th iteration, we have
2A(() 1- Kpp(t)
(t) 4n))(tt(( -Kp(t) (tt2p(t?(DD(()) - 1)
41

---

Published as a conference paper at ICLR 2026
5(A(1)+ + B()) Ko(O((D(() (1) CI(1) 41DAM)(-KKPC kp(t)
2A(t) 1- kp((t
(Dx)(())((D)(()))))))))))))))))) -1)
4A(++1
1- Kp(tttp+)) + 1)
Ci(L+1)
((A(t+ 1) B(t+ 1)) Kp(t + 1)(p+ (1) -1
((A(t)+ + B(t) Kp(t)(Dp(t) (tt) 4nDA(t)(1- - Kp(t))
2A(t) 1- Kp(t)
A(t) D K)
(i(1)
((1)+B()) (2D(0((DD()-1))
4A(++1)
1- Kp(l + 1)
ci(t+1) 1)
5(A(t+ + 1)+ B(++BB++)))) 1)) Kplt+ + +)(pptt + 1) - 1)
1- Kp(t) (CT(1)
4D
((T(1))p(()
K (pp()) -1) 5
-15
5
Dp(t) -1
4A(t+) 11)
1- Kppltt + 1)
Kylt + 1)(Dp(t + 1) 11 (i(++1) 1)
((((++)) + +B(+++BD)) 1))
1- Kp(t)
4
2
K Kp(t) 11 5
5VD3
VD3
1- Kp(t+1) 5(A(t + 1) + 1)) 1)(p(( + -) ci((+1).
4A(t+ 1)
In particular. the third inequality is obtained by replacing the the lower bound of A(t) 1-Kp(t) A(0+B()) Kp(D(Dp()--) in Lemma D.12 and utilizing (t) 2 Cy + 1) when p(t) 2 VDK* which is demonstrated in Lemma D.16 The forth inequality is derived by the facts A(t) B(T)
2
2 2
when p(t) 2 VDK A(t) < Ky(t)", and utilizing the upper bound of Ap(t) in Lemma D.9 Lastly,
the penultimate inequality is derived as 1 S (t) S V R. Ap(t) S D and 1 S O(D 5/2).
This demonstrates that the second part of (D.30) holds at t ++ 1-th iteration for the first case. On the other hand, for the second case, Ci + 1) would be strictly larger than Ci (i),, and it can be demonstrated that
(('P(t)((D((t) 1) C1 + 1) (t)++ + 2nDA(0)(1-Kpp))
5(A(1) + B()) Kp(t)(Dp(t) - 11) Ci()) + 21 Kp(())
AA(1) 1- Ky(()
5K30(0R(DD(0-I)
4A(++1)
1- Kp(t +1)
5(A(t+ 1) + B(t+1)) 1)) Kplt+ + 1)(pp(tt + 1) -1) ci(++1)
-(D)((((-K((() Kp(())
-
A(t)
K)
()()
p (DD(1)) 11 A(1)+((11 B(1) (²p(1)((D((1) - 1)²
11(++1)
1- Kp(t +1)
5(A(t+ + + B(++ 1)) Kylt + 1)(pp(t + 1) -1) Ci(++1)
11 - Kp(t))
D
2DM(1 - 2DM(1- Kp(t))
(p(()(D(()) -
5
(Dp(t) - 1)
11
1- Kplt + +1)
+ 1)) Kp(t+ + 11)Dp(t + +1) -1) cG(++)).
5(A(L+1)+ + +B(++++(+++))))
where the last inequality holds as under the second case, we still have
2DM1 (I-K(()
vK (p(t)-1) < 00
VDM K3
S <O(D). This demonstrates that
1) (1- 5(A(t + 1) + B(t + 1)) Kp(t + 1)Dp(tt + 1) 1) Ci(t+ 1).
4A(t 1)
1- Kplt + +1)
42

---

Published as a conference paper at ICLR 2026
which finishes the proof of (D.30).
Lemmas D.15 and D.5 together establish matching lower and upper bounds for C1 (t) on these bounds, we can derive precise training time which p(t) achieves 2 This result is formally presented in the following lemma.
after &g. Based
Lemma D.17. Under the same conditions as Theorem 3.1 there exists T O such that p(T*) 2
KD
m- li
Proof of Lemma D.17 Notice that Lemma D.13 and Lemma D.3 guarantee that
0.95 (1)) S (1) S (4T + 1) (1)
for all 2 t1. The left hand side inequality is straightforward, and the right hand side holds because:
when p(t) S 2vnnk Ci(() < (t) (4T + 1). Ci (t); when p(t) 2 2VEDK*
ci(1) < 11 (A(()) + B(t)) Kp(t) (Dp(t)
4A(t)
1- Kpt)t
Ci() & (4m + 1)-
On the other hand, Lemma D.15 and Lemma D.5 also guarantee that
4A(L)
1 Kp(()
(t) - 1+ 5(A(1) + B()) (p(()(D((()- -1)
ci(1):
4A(t))
1- Kp(t)
2)
(i
(D.32)
5(A(t) B()) (p(()(D((t)- -)
These two lower and upper bounds of (1) allow us to apply Lemma bounds for AC2(1) + (I) as
D.7 to derive lower and upper
9(4 + (1 10K(D- -
I Ci()):
(C2 (C3(t) S
AC2(1) + ACs(1) 2 200K - 19Dp(t)(1 - kp(t)) EM 1v CT(),
(D.33)
where replacing M with E llv 12 to match the presentation in our Theorem 3.1
With
these bounds in hand, denote T the first time such that p(t) 2 Then for all S T by applying Lemma D.9 and the upper and lower bounds of (C + ACz(t) obtained in (D.33
it can be derived that
Ap(t) < DPp(1)(1 - (P)(0(1)p(()) (AC2()) ACa(t)) Sn
(ST + 2) MM m= v p(t)
VD(D2-1)
KVDK
p(1)11 -K((K) (AC2(t)+ AC3(t) 21
m=1 V 113
Ap(t) 2
p(t)².
2vD
50K SOKVDK
Notice that the iterative rules for p(() satisfying the assumptions Lemma Lemma F.1l with the initialization that D p(t1) B. we can obtained that
E11
By applying
5
D2K
T 4 <
100(8T T0 + 2) D - log K) s 5
"1
Σ m==
11
m=
||
D2K
D2
T - 2
(log D - llgg 2 0
35mH Em
||
>> m=1 |
This results demonstrates that T t + 0 proof.
DPK Su
0
D²k
This finishes the
live
113
live
L
In the next, we provide the analysis for the last stage that p(t) eventually converges to +. This result is formally presented in the following lemma.
43

---

Published as a conference paper at ICLR 2026
Lemma D.18. Under the same conditions Theorem 3.1 for any T T", where T
O
D²K
as defined in Lemma D.17 holds that
12mm liv
2DD K)
D(D K)
S p(() S K V 1 ΣM m= 1 T
(D.34)
K
2e n 2m=
1l (1((T P
In addition, it holds that
2
p(T));(T)
- (1 - Kp(T)) + (1 (Dp(p))))²
K
D
K
40K K)
400D (D K)-
(D.35)
1K > 2m II -T+
17 E m= v ( T*)
Proof of Lemma D.18
With the bounds established in Lemma D.17 and the fact that 1- I-Kp(t)
p(()
exp
Ca()+Ca(t)
the upper and lower bounds of (C2(t) + (Ca(t) obtained (D.33) can be
VD
rewritten as
8T ΣM
m== V
J0 (c21)+0310)
AC2(t)+ACA) Sn
K3VD3
+ 2 mm l|v
30 (c1000050))
200K3
The upper and lower bounds of AC2(t) + ACC(()) match the assumptions of Lemma F.12 By applying the lemma, we can obtain that for all T"
VD
17 E 1 (T') + teNB
Cs(T) 2 log
2
200K3
8
ΣM
VD
8Ti EM 1
C2(T) - C3(T) <
log
(T - T") + EEND
K3VD3
K3D2
Replacing this result into the formula of p(T), have
1
p(T)
K (D K)exp
C2(T)+Cs(T)
K (D- K)exp( ( log
?. = l 12
(T-T+))
VD
200 D
1
20D(D K)
K
1 Em v (2 T
On the other hand, we can also derive that
p(T)) <
K + KD- KXe - log
88n 18 A
- -T*) + tekv))
nia MM
live,
K
K
D(D - K)
< K
26 KK EM= | (T
This finishes the proof of (D.44. can obtain that
With this condition holds, by checking the definition of (T*)).
Kp(T) (D- - K)Kp(T}²
3 < 11+ 1 1- -pK(TT)
1 (1 - Kp(T))
Kp(T)
Kp(T)
Plugging this result into (D32). we derive that
2K D D SC(T) KD(())
2K
1 (1 Kp(T)) Kp(T)
Kp(T)
D
which immediately leads to the final conclusion of D.35
44

---

Published as a conference paper at ICLR 2026
Now, we ready prove Theorem 3.1
Proof of Theorem
3.1
We first prove the first conclusion.
D
D
S - -1 K -x(1)) + +DID - K) (1 (1- KT(T)) ())-K((())
((()) S(T---) 8*
Σ
S bi
K)²
11=1
D
DE
V (KID- K) (1-Kp(T)) =
E EM ( TT)
where the last inequality holds by applying the upper and lower bounds of p(T) derived in Lemma D.18 This finishes the first conclusion of Theorem 3.1 Notice that in Lemma D.18 we have derived that [C1(T) - 1| @ Kp(T)), which directly imply that
K
w ( -
- |r
r= D2
(T T-) TT-
||
n
Em
V
where the last inequality holds by applying the upper and lower bounds of p(T) derived in Lemma D.I8 This finishes the second conclusion of Theorem 3.1 For the third conclusion, notice that
M
)
L(W
W (T) KQ
Σ
E
Ymmi
V.me Xi )s (T)
2
i=1
M
D
( w (T) V.m x,, .n.² ) 27
Σ
Σ
E
(I5(X)) m. - o
Σ
(T)
+ E [[E
2
where the last term is essential Copt. and the last inequality holds by the independence between X and & and the fact that & is zero-mean. Since this equation holds, in the next, we directly deal with C( C(W(T) w k Lopt. We first prove the upper bound. By utilizing the fact that (o(x)) - ov)l s J y for all v. y E IR, we can derive that
L(W
KQ
Copt 2 m= 1=1
D 3Σ
M
m. - (Ir(X)
Σ
D
Ww V.m
X.,)
)
27
11=
(vi xi,) + C1(T) Kp(T))
ΣΣ E
Σ HEC
Σ
(vin X
D-K K
En GG
Notice that Ziam No ..m) where or. (0, o. ). where oz.m. bounds derived in Lemma D.18
K ( (T)) and z(T)
Iv 10 (1 and they are independent. Based on the upper D-K
we can finally derive that
L(W(T) W Copt
D
E
(Z
+ Z24.m
2
+
D
Σ
M
0 m
2
2
m= i=1
m=1
m=1
DK
M Σ
DC (T) (C- KP(T))
M
C((T)(())
V 122 +
Σ m=1
V
2
K
2(D - K)
m=l
KD1
<e (T - -T"
where the last inequality holds by applying the upper bounds for (1 - C1 and p(T) de- rived in Lemma Dhi8 completes the proof for upper bound. On the other hand, denote Z Za,m, Ei, (vi,, x:) (0, Kl[Vm. ll1))) and 74 74, -i, ear (v xx) NN(0, (D K) km) 1l].
45

---

Published as a conference paper at ICLR 2026
and (T)
1-Kp(T) m.. x y 1 1=20.020} can further derive that
5,mm P(T) m..
Then, by utilizing the fact that lo(x) o(y)l 2
D-K
L(W();
W kQQ
Lopt
2
1 2 i=1
Σ E
(x))
(2) V. X,, !S
M D
Za.mi K
(Cii(T) S.m.t) )
Σ
E
2
Wi t=l
( -i(TT)((7) Z3.....
C1(T) ())())) Za.ms
E
l|za.::00) 11
12 100m. 201
D-K K
i=1
(12 - Z3.m. (1())(1-p(())
E E
Zrmis
D-K
M D
((2 -C(()(T))
Zz.m.. Ci(T)((-Kp(T) Za.mss
Σ
Σ
[Zz (2amm.20))
2 M= 1=1
D-K K
1 M
D
K ((()))(T))
M D
Ci(T)((1 - Kp(T))²
E[Z 3,m,
E[Z2, mm.il
4
2 m=1 i=1
4(D 1))
m= 1=
G(TT) - -
E B[Z (.... 20)
E[ E m, (22.m.:200)
D-K
V/
(T) - Kp(T))
m=
8(D 8(D-K)
ΣM
||
(D) K -C((T))(T)( K2(T))/ -
2x(DD -K)
2m=1
M
|/ (T)(C(()(( KP(T))2 (CI(T)
2 BVR(D= K)
2(D
TD
MM
1/ kK((T)²
KD
2C c (T T
where the last inequality holds by applying the lower bound of Kp(T) demonstrated in Lemma D.18 This completes the proof.
E PROOF OF THEOREM 3.2 AND DISCUSSION OF THE WORST CASE EXAMPLE
In this section, we provide complete proof for Theorem 3.2 and a worst-case example can attain the upper bound in Theorem 3.2 We first prove Theorem 3.2 in the following.
Proof of Theorem 3.2
We first upper bound the OOD loss by the sum of three terms as
LOOD W
(7)
w
R2)
TF
w w
E
X +f
TF
TF
k
E - +
+ - -f w
-
5, E L - - + -10" W
46

---

Published as a conference paper at ICLR 2026
B -
TF W
where the last inequality holds by Cauchy-Schwarz inequality. Based on this decomposition, is critical to derive the upper bound for IE [ (x) and
( w KQ)
[[]]
For the first term E |||Y - f ("()) || we have
B (
2E IY + 2 2 (x)1
By the assumption that each column of Y satisfying that Elllym/ On the other hand, we have
< &, it is straightforward that |
E |IA - DE.
D
M
D
M
v
B[IL/
E
//1l///lll....1.1 xx,)
m.J.
K
m2
MEGi
D
M
E
D
v
1
K
MEG
For the second term (X]
( we can derive that
-
E
[[f X
TF
w
|
M
D
(1x
Ci(T( -)((K)
27
Σ
-((Pp(T)
Σ X
Σ
(Vin:
Xin)
K
D
EEG
in GG
D
12
-((TT(()
Σ 11
K
M
CIT)((1 - (1-Kp()))) K)(T))² 2
D
BIJ
kY²
m=1 i=1
4 G
KD'S
<0
(T -T+
Here the first inequality holds by (z(x) ov)] S |2 - y/- The second inequality is established by the fact (E-1 a;)² < DE of for all scalar a; 's and (vi, X < || 11 Lastly, the third inequality derived by replacing the conclusions in Lemma D.I8 Combining all these derived terms into the three terms derived as the upper bound for OOD loss, we have.
COOD (w R
-
-
so D
(D
o
K KEM l ITT - T*)
KD$g
nT - T*
Let the upper bound derived above smaller than e, we can derive that
= +0 KD Σ
m=1 |
o
6g2
Em=
M
|
ne
T7E2
This completes the proof.
In the next, discuss the construction of the worst case Y. such that LOOD W (T.)). w(?)) [ (X) [ > for some T. O MKPC) (assuming Vm.l12 and & O(1) for simplicity). In fact, this T. can be different with the T. defined in Theorem 3.2 but the same order w.r.t. hence matching result.
(T).
KQ
D2 By the conclusions in Lemma D.18 we know that π pT) O
( VTKMT K S 0( O ( Therefore, there exists an absolute constant such that ((T)C,(TT)
and |P(T)Ci(T)
1-Kp(T)
2
é' D. In addition, let An,i to denote the event such that En (Vin X;,))
47

---

Published as a conference paper at ICLR 2026
1.] We can assume the probability of Am,i is larger than abso- lute constant. In fact, such an assumption can be easily verified on many specific distributions like Gaussian distributions. With these notations in hand, we can design Y such that its (m, i)-th entry generates m,, sign (ig #f (Vi xx)) 1 (Am.:) + f (X)m.s Given this construction,
max Eneci EG (v:: (imx,,,)111
can deduce that
E
(*XX), f ( - TF Z w W2)
(Ym-- (X3)s)) ((x -
Σ
IF
Σ EG
(vin X
m=1 1=1
CI(T) KP(T))
Σ (Vin Xii))
D-K K
M
D
(1(T)(1-KA))) Kp((T))
Σ (v xx,))
((dmm))
D-K K
m=1
D3
MK
B[1 = =0 (03))
MK
2
n/T
TJT
Replacing the T with T. = 0 MKDD) ne²
) we can finally conclude that
LOOD W
W
(T,))
2 [[ fr [
(Z W
w H θ
(D'
D'
MK
11
O
D3
MK
This validates that the upper bound is indeed attained under our construction.
F TECHNICAL LEMMAS
In this section, we present and prove the technical lemmas we used in the proof of the previous sections.
F.1 CALCULATION DETAILS OF EXPECTATIONS
We introduce the details regarding
Let a N(O,a)) then holds that Lemma F.1 (Calculation of Fi (a) defined in D.2
If (-) is the identity map, then
E (Ir()) =
If (-)) is ReLU activation function, then
B [(xo(()) (x))
If (-) is Leaky ReLU activation function, then
([za(z)o' (x)) (1 + x²)a
2
Here, K is the coefficient of the Leaky ReLU activation function when the input is smaller than 0.
Proof of Lemma F.I. The first conclusion for the identity map is straightforward. When a(.) the
ReLU activation function, we rewrite that (z) Therefore, we have,
xl{2207 [|200 x² (z20)-
(x)o(()) = E|x² (a20)]
2
2
48

---

Published as a conference paper at ICLR 2026
Besides, when a(.) is the Leaky ReLU activation function, x² 1'1 (x20) + x² 2 (=<0). Therefore, we have,
rewrite that xo(2)o((x)
(1 E[("
(1 + K)
E [(za(z)) E| (c20)) + A/ L [m<o)]
2
2
which finishes the proof.
Lemma F22 (Calculation of (a, b) defined in D.3 ). Let ~ (O), ~ NV b) be two
independent Gaussian random variables, then holds that
If (+) is the identity map, then
E + 22)00 + x2)]
If (-) is ReLU activation function, then
E + 22)00 + +x2)] 2
II (+) is Leaky ReLU activation function, then
E[zjo(11 + x2)0 (x1 + x2)] (1 + k²2a
Here, K is the coefficient of the Leaky ReLU activation function when the input is smaller than 0.
Proof of Lemma F.2 The first conclusion for the identity map is straightforward. For the next two cases, we first introduce some definitions. Let 23 Iı + 22 (N + b). Then we have Cov(x1, [33) [[[x) x2)=a] a, and B[zi| a+6 ug. Consequently, when o(-) is the ReLU
activation function,
E[z101(11+ + 22) (a1 x2)] = B[rias (zx20) [[[E[[11 (x220) [=s]]
[zz20|] 2(a + b) Blz3]
In addition, when (-) is the Leaky ReLU activation function,
E[z10(<11 + 12)0 (m1 + x2)1 = E[ (m=20)) + x E[z [0a
E [[ (0) [2a]] + x [ [E[ (as [23]]
x²a
(1 + x3)a
a+b Blu |= +
El 1 (m)
+
2
This completes the proof.
Lemma F.3. Let ~ holds that
be two independent Gaussian random variables, then
If (+) is the identity map, then
B| + x22]
If (()) is ReLU activation function, then
(( (z) + x2)] = a
a
(v
a
Vab
arctan
(F.1)
4
2x
a+1
And there exist the following matching lower and upper bounds:
a
avab
v
a bvab
< (m
-
(F2)
4
2T + b)
2
2m + b)
49

---

Published as a conference paper at ICLR 2026
If (r))) is Leaky ReLU activation function, then
(1 + a)aa + 2tT
(1 --)
(v/) +
Vab
B[ (#1 + +x2)]
arccaa
(F.3)
+
And there exist the following matching lower and upper bounds:
(2) (x1 x2)1 (1 + x²)a
2
(lz1o(21) (x1 r2)] (1 k3 + (1 (I-n)avab)
(1 (1+12)a
(1 (1-x)) x1-x) (a + b)
2T + b)
2
(F.4)
Here, K is the coefficient of the Leaky ReLU activation function when the input is smaller than 0.
Proof of Lemma F.3 The first conclusion for the identity map is straightforward. When a()) is ReLU activation function, we can rewrite that (21)o' + x22) = zi (2120) (2)222 Let
21 Va
and 22 v
then we have,
B[zio(z1)0 x2)] a (1220) 4 +
(F5)
For I, by denoting
x
B we can obtain that
[1 / z/ (Azι) (221)d21
where (6()) and ((.) are the cumulative distribution function (c.d.f.) and probability density func- tion (p.d.f.) for the standard Gaussian distribution respectively, We can denote that [()) Jo 27 (Aea) (22)dd21. Then, by the Leibniz integral rule, we have
di(A)
dA = [
1
o ( = 2x + ^2)2
3
de
T(1 + X2)2
Additionally, since I(0)
1 1' we can derive that
arctan (Vo)
v ab
I = + 2m
arctan À + X2
=
(FF6
2m
+
Applying the result of (F.6) into (F.5), we finishes the proof of (.I).. In the next, we derive the upper and lower bound for 4t. By the property of c.d.f., know that 4(2) 1 for all E RR which implies that
1 E[2?] 2
I 21 =
Additionally, by Mills ratio, we further obtain 1 ((2) S o(>)/2 for all x > 0. Based on this result, we obtain that
12 270(z1) (1 (1- ((A21) dæi
1
=1 21 (()=)))()2)))
A
2 7
where the second term can be calculated by
224+22
1
210(21)6(Az)ddz)
zie
dz
zie
d
2T Jn
2 X2)
2 + X2)
Plugging this result into the preceding inequality, we can derive that
1
b
I
2
2x Vala +b)
50

---

Published as a conference paper at ICLR 2026
Combining all these results and (F.)), finally conclude that
(
Vab
b
1 2
<
(F.7)
2T (a + b)
2T 2
Applying the result of (F.7 into F.S). we finishes the proof of F.2 In addition, when o()) is the Leaky ReLU activation function, we can similarly derive that
B[11o(11)) (21 + x2)] [ (#1 (a120) 11 (Vāzi (bz. axl [2] (ar <o) (vās vbzz200
+ [x (s1 20) (aa + an B|e [ < (vas
(1 +1+ x 2m (Vaz + vo2220 + 2naB (0)< 'l (
where the last equality holds by the symmetry 21 and 22. By applying very similar calculation process, obtain that
L 2/ (Azι) ((2))(())))))
1
a (V
uub
B[i (s)<0) (Jan++be220)
arctan
2x
a
By replacing this result into the previous calculation, we can immediately prove be directly derived from (F2).
F3
And
(.4
can
Lemma F.4. Let N(0, ~ (), be two independent Gaussian random variables, then holds that
If (-) is the identity map, then
E (1220(11) (m1 + x2)] = 0.
If (-)- is ReLU activation function, then
bvab
(122 x2)) 2m (a
If (+) is Leaky ReLU activation function, then
(1-x) xbbbbbbbbbbb
(1)22 + +2)] 2T (a + b)
Here, K. is the coefficient of the Leaky ReLU activation function when the input is smaller than 0.
Proof of Lemma F.4 The first conclusion for the identity map is straightforward. When a(-) is
ReLU activation function,
rewrite that (2) (æ) 22) 2 21 (m120) (on t22-00))
Let 21 and 22
Ø1
da then we have,
vb.
E B (z) + x2)] = Vab E(21222221 (m120) I (váa + vb=220)
(F.8)
For T, by denoting
X VT i. it can be calculated by
[ 210(21) ( 120(22)d22 dz1
1
21 (212((1))((22ddl
Az1
d21 = / 210(21)
210(21)
2m
22E
d22
V2T
x2
c d22
=2(1-22)
21e
dz1
(F9)
2x
2n + X2)
2n((a + b)
Now applying the results of (..9 into function. In addition, when a(-) is the Leaky ReLU activation function, we can derive that
F8
finish the proof when (-)) is the ReLU activation
(x )a (an 2r2)
51

---

Published as a conference paper at ICLR 2026
= =VabEl21 (=) ( (va
+ Vabk ( +
+ avaabl 22 22 (=)20) - x² abE (s) <0) 1 (Vāz,
x2+ abb 21 22 [0) Vbzz + 2xv abb(21 22 (=Λ<0) L
where the last equality holds by the symmetry of 21 and 22- addition, by similar calculation process, obtain that
1
#a4x3
1 2T(1 + x3)
2| co) 1 +Vaan+vbsss 2
210
dæ
2 + b)
Consequently, finally obtain that
(1 2x + b²)
(1 - bb)
(( + +x2)] =
2T + b)
2x + b)
which finishes the proof.
Then, based the conclusions of Lemma H3 and Lemma F.4 we immediately obtain the following lemma as corollary.
Lemma F.5 (Calculation of Fa b) defined in (D4)). ). Let a), (N be two independent Gaussian random variables, then holds that
If (-) is the identity map, then
E[(2) + x2)o(1)) + x221 a
(-) is ReLU activation function, then
[[( + (2)o + 22)
a
+
arctan
a (VB)
+
ab
4
2m
2T
And there exist the following matching lower and upper bounds:
Vab
byab
a +
< + (2)0((1)0'(m1 + +x22 s
+
S
a
+
2
2m
2
2T b)
2
4x
If (-) is Leaky ReLU activation function, then
(1 + k)
(1 (1-k)aa
a1
(1
k=)² ab
E[(21 12)0 (æ1 + z2)]
arctan
2x
2x
And there exist the following matching lower and upper bounds:
(1+x)aa
(1+k)²a +
(1 (1-x)/aab
([r1 (a1 xx)) 2
2T
E[ (x) + x2)) < (1+12)a (1 (1-k)²bvab ab (1+x²)a
(1-x)2b
- 2 +
2
(a+bb)
4m
Here, & is the coefficient of the Leaky ReLU activation function when the input is smaller than 0.
Lemma E.6 (Calculation of Fa b, c) defined in (D.5). Let æi NN a). D2 NN(0,)) æ3 N/ c) be three independent Gaussian random variables, then holds that
If (-) is the identity map, then
E|z + x2)0 + x2 + xa)] = a.
If (-) ReLU activation function, then
(Vatb) + vla+bje)
[[ + x2)0 (æ) + 22 æa)]
arctan
2n
a+btc)
And there exist the following matching lower and upper bounds:
av (a + bjc
)
ac
+ 2m(a (((abb+c)
S E[aio(1 + x2) + x2 x3) s
2 + +(++bb+a))))))+
52

---

Published as a conference paper at ICLR 2026
If (r))) is Leaky ReLU activation function, then
B[r10(11 + x2) (21 + 22 + rs)] (1+ x) a)a
(1 h-n)aa
arctan (Vatt) + V b)c
4
2
a+b+c
And there exist the following matching lower and upper bounds:
Ba + +222) (x1 +2 + 13)
(1 k1+ [1 11-n)aa (a +bbc
( [1+ x1+)))
(1 - x)aac
V
2T (a b c)
2
2x Exva+blo b(a b
(1 (1+ (1+x2)a
E[z1a(1 + 22)0 (21 x2 + 0z)) v
Here, K is the coefficient of the Leaky ReLU activation function when the input is smaller than 0.
Proof Lemma F.6 The first conclusion for the identity map is straightforward. When (() is ReLU activation function, we can rewrite that (1 + o (r) D2 + Ta) = 2 (21 + 222 (21+0220) +222220)) Additionally, let 24 2ı 22 NO, b) and Va+b
(O,,)), Then we have Cov(z1, x4) = [x) + +2)) a, and E| [c:s] = a4 Therefore, we have
(a (2)) (21 + x2 + r3)]
=E (a + x2) (21-2220) (21t02t2220)]
= EE (2) x2) (21 203 [2]]
x (x) x2) (2) E
21 + 32
E E x x1J44 20, $
14
x4
VC
x (m)20
24
aB 2 1=20]
((x=)] = a / = ())))))())d
a +
where A =
o+b By the similar process in the proof of Lemma
F3 we can obtain that
1 +
aretan A +
4 +
arctan
(V
a aabb
+
V(a b)c
4
2T
1 X2
2x
a a b+b+
and
(1
v la +le
< 215 1
2T(a ar(a+b+t +b+c)
2T bla
Plugging these results into the previous equation of expectation, we finish the proof when a(.) the ReLU activation function. In addition, when of-) is the Leaky ReLU activation function, have
E[zoa(1 + +12) (x1 x2 + xa))
-Eir (a) + x2) (>) [=)) (m (z) nE| | (z) + x22 (m1 [ (11+0)
+ xE| pz (x1 + x22 (>) + = (2)+ " (#1 + x2) ( (x) (2))
=(1 + x² E (z) + xΦ) (m (2220) L tm220)))) + 2xE[x (21 + x2) L (2) (m)tat (m
By utilizing a similar calculation process, we have
E(a (mi + 22)) (m (01422--5200) =a
z 2 (
a
a
(Vatt l
Vla+bi
bJc
arctan
2T
b
+
Plugging this result into the previous calculations, we finish the proof. And the upper and lower bounds for Leaky ReL.U activation function be directly derived by comparing the formulas. n
53

---

Published as a conference paper at ICLR 2026
Lemma E.7 (Calculation of F5 (s b, c) defined in (D.6). Let #1 a(), x2 NNO, 23 () be three independent Gaussian random variables, then holds that
If (-) is the identity map, then
((220(21)0 +22 + r3)] = O.
(r)) is ReLU activation function, then
bva(b+
(z)2 + 12 x3)) = 2m + b
-If (-) ) is Leaky ReLU activation function, then
[[zza(11)0 (x1 + x2 + rs)] (1 - K) x) a(b c
2m (a +
Here, & is the coefficient of the Leaky ReLUU activation function when the input is smaller than 0.
Proof of Lemma F.7 The first conclusion for the identity map is straightforward. When
a1)) is ReLU activation function, we can rewrite that x20(x)) )o (ar + 22 x2)
2122 (2:20) t202000)) Then we have
E [[ (11 + 22 + 23)1 E[z1021 (m 20) 1 (an toa
E
2122
1
(m) 20)
(a toitott22200 21, 22
21 D2
E
T1
(2120) ①
Ve
31 + x2
T1I2
é(z)) ((2)))(z))dxdde))
Vc
00
21 + x2
1
11 V2Ta
T2
Ve
v2
duz
dri
We utilize the integral by parts derive that
11 +
22
=
b
T1 + J2
I1 + x2 vc
de
e
dΦ
2T
Vc
2n
V/
2T
00
-00
(122 o.))
b
2
Lo 1 dz
00
e
20
drg
215
2T
2T
-00
2x +c)
Now substitute this result of I back into the outer integral for the calculation for expectation, then we have
22 2 DIE
([I2a (x1 + 02 + xa)] = 2T /abb + c)
200
alb+c)
(athe))=?
(a b c)a 2a(b + c)
2a(b+e)
d
a a b
b/a(b+c)
(m +b(a)b+)))
This finish the proof when o--) ReLU activation function. addition, when a(-) is Leaky ReLU activation function, we can derive that
E[x20 )a (#1 + 22 + Ø3)] =B I (z1 203 (21+20 + KB(2 (ur<o) (ei (2:t2)
xB (0/20) 1 (m1zz==2<00) B(xna (m)< (e
54

---

Published as a conference paper at ICLR 2026
=(1 + x² ( (m:20) (an +22+2220) + 2kL[[1122 (ar <0j (mrtoat2200)
By applying similar caleulation process, we can derive that
E| <u) 1 t2a 2 + L I1e dæ
by ava(b
2m 2m(abb+c)
Applying this result, finish the proof.
Lemma E8. Let ~ (V(0,), ~ (0, be two independent Gaussian random variables, then holds that
If (-) is the identity map, then
(lo(s)) x2))
If (+) is ReLU activation function, then
(( + +2)2 = a + arctan
a
(1)0 a [bb)
v ab
+
4 2x
2T
If (-) is Leaky ReLU activation function, then
(( + x2)] (1 + x)²a
(1 (1-k)²a K)
aretan (V +
(1 x)2 Vab
+
(F.10)
4
2m
2m
Here, K is the coefficient of the Leaky ReLU activation function when the input is smaller than 0.
Proof of Lemma F.8
The first conclusion for the identity map is straightforward. When a() is ReLU activation function or leaky ReLU activation function, we can utilize the fact that zo' (z) (x) re-write that
(x (1( + x2)] [[( + (2))) + +2)]] E[[ + ( (a1 x2))
And this term has already been calculated in Lemma F55 Hence finish the proof.
F.2 ARITHMETIC INEQUALITIES
Lemma F.9. Let a, e be three positive scalars, it holds that
be
b
a
Proof of Lemma E.
c C
bc
bc
2
+ b
a (a + b)a
a²
This completes the proof.
Lemma F.10. Let a, b, be three positive scalars, it holds that
< + bc
b
(a 6)2
Proof of Lemma F.10
bc
bc
a
(a
a
b)
This completes the proof.
55

---

Published as a conference paper at ICLR 2026
F.3 SEQUENCE ITERATION BOUND
The following lemmas characterize the increase of a positive sequence with matching lower and upper bounds. Similar conclusions and proofs can be found Jelassi al. (2022); Cao et (2023);
Meng et al. (2024); Zhang et al. 2024a 2025c). We include the proof here for completeness.
Lemma F.11. Consider a positive sequence { satisfying the following iterative rules:
I++1 2 + 11-01 c1 21;; zt",
F++1 S mt +
where 02 2 C1 > 0 are positive constants. For any v > Tv. let T, denote the first index such that Ty v. Then, for any constant Ç > 0, the following bounds T, hold:
Te v 1+C
(1 + C) (g))
(F.11)
9-
and
log( (o (1+ c)
To 2
(F.12)
(1 + ()91ee220
Proof of Lemma F.11 To prove the bounds, let Ty be the first iteration such that 2; (1 + () Furthermore, define g* as the smallest integer satisfying (1 + c) zo 2 v. This implies
log(
log(
+ 1.
log( + 5)
log( 1 + ()
For 1 = Ti, the lower bound iteration:
Ti --1
TT.
2 zo +
n/12 2 3n
Tinera
t=0
from which we can deduce that
IT. 2o Ti
(F.13)
7/C1200
Utilizing the opper-bound iteration for æT, and the condition æT,1 S x0 + {), we get
IT, S 2T, n020y I xo + () + no2ag Ç)
(F.14)
Combining the results from F.13) and F.14) leads to
(1 ()+
T
The case for g > 1 is handled similarly. Using the lower bound iteration from Ty-1 to T, - 1:
IT. 2 ZT +
ncra / 2 IT,-; + (T + 1)00-11)
(F.15)
and the difference 2T, IT can be upper bounded using the upper bound iteration and TT,-1 S co(1 + < and IT, 2 xo + 9-1
JT, TT-1 < IT, + 1/0224 TT,- S ((1 c)9-120 n1222 + 9)
(F.16)
Combining F.15 and F.16), we derive that
(1 +1+51 8)
Tg S T
+
(F.17)
1 1)(0-110-11)
56

---

Published as a conference paper at ICLR 2026
Taking telescoping sum of the results of (F.17) from 9 1o g" and by the fact that T, I Tg' we finally get (F.IT). For the lower bound, we proceed similarly starting with = T1. We use the upper bound iteration:
2T. S +
Σ
nora & xo + TInoaa(1 C).
1=0
Substitute that DT Zo 2 fo we get
Tr
(F.18)
ne200 (1 + Ç)
similar derivation for g using the upper bound iteration gives:
Tg
IT, S TT +
nezai S TT + n02(TT - -
+ ()99.
(F.19)
The difference IT, TTg-1 also be lower bounded by utilizing the fact that TT S xo(1
c)a-1:
"T -T TT, TT,-1-1 1,225
5(1 + +)9--100 n)22z(11 1)00-119 (F.20)
Combining the results from F.19 and F.20p, we obtain that,
S
T T
(F21))
(1 + () 1)+1 (1 c)
Taking a telescoping sum of the results of F21 from = 1 to = g" - 1 and by the fact that T. 2 -1 we finally get (F.12).
Lemma F.12. Let Te be positive sequence for 2 0. Assume 2, satisfies the iterative formula
I++1 = It + eie -22
for given constants C1, C2 O. Then, for all t 0, the sequence ia is bounded as follows:
log(cie2t cx200) x S eie cza0
0250
log(c)c2t + ex2z0).
Proof of Lemma F.12 First, we establish the lower bound for Iı. We introduce a continuous-time
sequence Σa, 2 0 defined by the integral equation with the same initial value.
r. = C1 J"" cagrdr dr,
Eo = T0-
(F:22)
Observe that v, is clearly an increasing function of t. Hence, we obtain
I++1 I4 +
d7
ptll
$ x, + C1
F. + exp(-c220
for all E N. By comparing the preceding inequality with the iterative formula for {at}. the comparison theorem implics that v, z, for all t N. Equation (F22) possesses an exact solution given by
log(c10zt + e00))) 2
Thus, have
log(cre2t 2200)
24
C2
57

---

Published as a conference paper at ICLR 2026
for all E N. This concludes the derivation of the lower bound. Next, we derive the upper bound for 2g. We have
32 20
22
log(e1earte<200)
S 250 + 61
20 01
C10sT
C1
20
+
C1
C1027 + 60200
S 20
C1
C102T ecazo dr,
eC220
where the second inequality utilizes the lower bound for 24 derived in the first part of the lemma's result. Consequently, we obtain
It S 30
+ t0200)
log(e
C2
1
ere
log(c102l + 2200)
This completes the proof.
G PROOF OF THE CASE WHEN D K
In this section, we provide the theoretical results for the special case D = K. Under this setting, the ground-truth softmax scores reduce a trivial rank-one structure that S = 1nll Consequently,
the initialization w KQ ODx D already yields ( I achieving exact recovery of S at the start of training. As result, the gradient with respect W KQQ remains zero throughout the optimization, and the problem effectively reduces to optimizing the single parameter matrix Wv- Under this reduced setting, the loss becomes strongly convex in Wv, and gradient descent enjoys a linear convergence rate, which is much faster than the (11/T) rate established in Theorem 3.1 Since Theorem provides matching upper and lower bounds and is therefore tight and not be improved, this linear convergence phenomenon is exclusive to the degenerate case D = K. This explains why the proof strategy for Theorem 3.1 does not extend the D) K setting.
Now, we present the following Theorem G. to characterize the loss convergence when one-layer transformer is supervised by teacher model f* (+) with S* 1 b.
Theorem G.1. Suppose that S 2. then for any t 0, the excess loss is minimized as
M
[
L(W
w(O (1) KQ) Lopt
(t-1)
Before we provide the proof for Theorem Lemma G.2. Under the same conditions of Theorem 3.1 there exist a time dependent non- non-neeative
G.1 we first provide and prove the following lemma.
scalar C(I), such that
V. C(I) vin for all [M];
(
(G.1)
and C(t) has the following closed formulation:
C(t) = (1 11)
where the function Fi(.) is defined in D.2 In addition, w (1)
KQ remains zero throughout the training.
58

---

Published as a conference paper at ICLR 2026
Proof of Lemma G.2
W.L.O.G.. we assume that V is already normalized, and m Ivm Em.2.... Sm.d) E Raxd be an orthogonal matrix with Vm being its first column. We prove this lemma by induction. Since these two conclusions holds at initialization with C(0) 0. It is
sufficient to prove that Vwv.. L(WW; w KQ) = c(t) v and V Wwxo w KQ = 0, when
assuming w V.m (()) v and W KQ O. Notice that W kQ - 0 implies that S 1/ D for
(1)
all i', E [D] By the gradient calculations demonstrated in Lemma D.1 we have
(w V.m x )s )
V ww.. L(WV, W KQ)
E
Y m.
o Σ
E
WVV. Vim Xix
Xi,)S!!0
11:0t
Xi !!!!!
D
(w Vime Xin)
X11
EE
Y mi.io'
D
D
(w
Vim xi,)
D (w())
Vimi x
(1)
E
o
Xi
(G.2
D
i=
For I1 we have
D
D
(w
Vim X(n)
(1)
I Xii
Σ
E
Y
mi
D
(w!))
v.m Xxx))
(vi Xin)
E
[X)
v
m.
D
D
D
D
d
D
(w xix)
(Sm,,e xin)
E
(X)())
Em,k
D
D
E
k=
D
D
(t) Xxx)
(v Xi,)
Σ
Σ
(X)
D
D
1=1 =1=1
The first quality holds as E is mean-zero and independent with X, and the last equality holds the orthogonality between vi and Em,k implies that (Vin Xix) is independent with (Em,k xi,) for all
i i2 E [D]. Notice that [x*(X] m. to Eipzr (ym x,,)) and o' o'
(w (wrf (imiit))) x.))
D
C(t) o' D Eh-n(ve (v% x.))) = o' (V x,,)). Consequently, (Em,k: X(,) is a mean-zero (w(t)
Gaussian random variable, and independent with both [f*(X)) and o' EEp= multancously, implying that
xx,) si-
D
{ (t) V.m xi,)
E
(x))
mix
(Sm,, x.n)
D
1=
D
(
(t)
Vom Xin)
E
[r(X)]
Bl(Em,ke x;,)]=
mit
Based on previous results, by plugging [f(x]] m.
5o Eii=iiiivv xx,))
and utilizing the
definition of (a(a) in (.2) we can further derive that
h= ΣΣ
D
E
Σ (vin x
o'
(w(t)
(v Xin) V.m Xix!
v
D
=1
59

---

Published as a conference paper at ICLR 2026
E D
E
D Σ (vin x of
(E m. X x
D
(v
Σ 11=1
D
(v x1n)
vn = =i(1) Vin-
The second equality is derived by fact that ((ax) ao can conclude the final result by the definition of Fi (a) in (.2).. Similar to the process of handling 1,, we have the following for 12:
(x) and d'(xx) = d'() if 2 0. Then we
W (0) Vim XT,
1
Σ
E
o'
W
Vim
X
T m M
D
1=1
C(t) EE
(v X Jo
Σ 1,=1
(v
Xis)
D
v = C(t)Fi(1) Vn- X1.) (v' m
D
Plugging the calculation results for Iı and I2 into we can immediately derive that V Vw.... L(WV); V w k c(t) v which, we stated previously, directly conclude (G.1). In
addition, can further caleulate that
(++1) W V.m
= 1) V (c0)+ v
which implies C(t) possesses the updating rules as:
C(t + = C(t) + n) C(t))
Subtracting 1 on both sides of the equation above and rearranging the terms, can obtain that
= = - (1 This proves the closed formulation of (tt)). In the next, we prove that Vwxa ((W w = 0.
1 - C(t 1) (1- (1
= (1 ()))
By Lemma D.I
we have
DV ke
L(W()) ((() KQ)
M
D
(X)
D
(w
Vime
(t)
Xix
0'
(w Van Xix
D mh m==i=
D
D
11=1
D
ΣΣ
(w(t) (1)
V.m Xz,
-P
4=16==1
I (x))
Vimi
Xin
Vam X P..
D
m.
D
11
M
D
E [x())]
Σ
D
(w
Vim Xin)
(t)
Σ
D
Σ
Vin X
D
D
m=1 1=1
f==1==1
14
D
M ΣΣΣ 1=1
D
P
0
Σ
D
(w
Vom
(t)
D
Xi,
o'
(w Vim Xi,)
Σ
D
Σ
D
Vim
(t)
x4, P.,P
11=1
:1=1
1=1i==
E
0
w
V.m
(t)
Xin
o'
(w V. X.,)
(w
Wm X ).,p
(t)
D
D
+==1
(G.3)
In the next, discuss the value of Iz. La. 1s. and Ta respectively. For Ts, can be calculated
60

---

Published as a conference paper at ICLR 2026
M
C(1)
Σ
E
0
V me Xii o'
Σ (v (vinxxx) xix)
E(vmo
(v Xin Xin)
Pij Σ
Pi
i=1
=
C(t)
Σ E
M
(v
Xx,))
D 0'
Σ
(V Xii)
Σ
D
V Xz,)
Σ
D
Pii
D
P;
D
11=1
=1
11=1
11=1
i=1
D)
(1)
Σ
Ph
Σ
P
The second equation holds as E o (Ei=iia
x,,)),
(v..
Eb (v x takes
identical value for all it E [D] as they follows the same distribution. Similarly, for 14, we can calculate it
C(t)
M M
E
0
Σ
v
x
of
D (Vinn
E(vi
D X
(v
Σ 12
D
Pis Σ
D
P
11
D
(M
Σ
Pia
P
1==1
i=1
This implies that I3 = 44. Through similar calculation, we can also get
D
13 1== (M)(()))) Σ P
Σ p.
1=
Plugging the results that 1 4. and Is = Is into G.3
we immediately concludes that
C(W V w(() kQ) 0. This completes the proof.
With the conclusions of Lemma G.2 ready prove Theorem G.1
Proof of Theorem
G. Since we have demonstrated in Lemma G.2 that W(t) KQ = 0, implying We can decompose and simplify the loss
S(0) I.
M
(w (1)
Vin X ))
C(W()
KQ)
Y Ym,i
Σ
2 m=1
D
t=1
1 Σ
Σ
E
(IF(x)]
Σ
(w (4)
V.m Xi,)
(omxxin))
))
27
2
m,i
D
where the last term is essential Copt. and the last inequality holds by the independence between X and E and the fact that E is zero-mean. Since this equation holds, in the next, we directly deal with
((W w(()) KQ) Copt. By utilizing the fact that (a((x) - o(v)l s jæ yl for all x. y IR we can derive that
L(W w(()) xQ)
Lopt
M Σ
(lr(x) mit
(t) w Vim X ))²
27
$
E ((1 C(t))
Σ
(v me x(.)
27 }}
(1-((t)
3
ΣM
M
|m 12
2D
2
m=1
Notice that we have derived - C(t) = (1 - < e-1F100001) < Se e-(tt-1)/2 in
Lemma G.2 where the last inequality holds as (i 2 demonstrated by Lemma F.1 Plugging this result into the upper bound above, then we complete the proof.
61

---

Published as a conference paper at ICLR 2026
H ADDITIONAL EXPERIMENTS
In this section, we present additional experimental results on transformer learning of bilinear teacher models under more general training data distributions.
Each batch of training data (X,, Yn) is generated with X,, drawn from either (i). a Student-t distribution with df = 5; or (ii). mean-centered Gumbel distribution with loc = 0 and scale = 1. We then repeat the learning experiments for the six types of teacher models described in Section + Except for the change in the input data distribution, all other configurations remain identical to those in the Gaussian-date experiments.
The results are demonstrated in the following Figures 6 and Figure and report the results when training data are generated from Student- distribution, while Figure and report the results when training data generated from mean-centered Gumbel distribution. We could observe that all these results seems almost identical to those demonstrated in main body. Specifically, for both different distributed training data, approximately -1 on their tails, and the curves of O.O.D. loss have slopes approximates -0.5. These results empirically shows that the (1/TT) convergence rate for training loss and O(1/VT) convergence rate for O.O.D. loss still hold, even the model are trained Gaussian data. In addition, can also observe that the trained softmax attention scores ((T) perfectly replicate the patterns of S", almost identical to the results obtained on Gaussian data.
still observe that the curves of training loss have slopes
Excoss tautions
ttrations
mihrrty
can
FeLl DON
Loaky Phoci
Locty fall
Onk
12
2000
-0000
00000
5000
10000
(a) Excess training loss (log-log) (b) Excess OOD test loss (log-log)
(c) Cosine similarity
Figure 6: Excess training loss, excess OOD test loss (both in log-log scales), and cosine similar- ity between the value matrix Wy of one layer transformer (2.4). and ground truth value matrix V*. These results are presented for experiments where training data is generated from Student-T distribution.
62

---

Published as a conference paper at ICLR 2026
Abention
Atvetion
Aftent
122
01
-42
-02
-00
(a) ReLU CNN
(b) Leaky ReLU CNN
(c) ReLU GCN
Aflee tion soors 5
0.75
011
0.00
0.00
0.00
-01
-01
-02
(d) Leaky ReLU GCN
(e) Sparse token selection
(f) Group sparse linear predictor
Figure 7: Heatmap of attention score matrix ((T) when the training loss converges. These results are presented for experiments where training data is generated from Student-T distribution.
tacexx oas over fasabana 2og 9g
Excruee COD loea ovar darabarz Oag-lgg
Creat ralrrry selunen Wyard
Lodyy MNL COM
072
12
(a) Excess training loss (log-log) (b) Excess OOD test loss (log-log)
(c) Cosine similarity
Figure 8: Excess training loss, excess OOD test loss (both in log-log scales), and cosine similarity between the value matrix Wy of one layer transformer (2.4). and ground truth value matrix V These results are presented for experiments where training data is generated from mean-centered Gumbel distribution.
63

---

Published as a conference paper at ICLR 2026
Attention soore sin
Adetiion soom 5TD
Atten score
87
0.1
6.1
(a) ReL.U CNN
(b) Leaky ReLU CNN
(c) ReLU GCN
score
score
Altert soors ..
-02
0.0
0.5
000
0.1
-11
-07
-2
(d) Leaky ReLU GCN
(c) Sparse token selection
(f) Group sparse linear predictor
Figure 9: Heatmap of attention score matrix (T) when the training loss converges. These results presented for experiments where training data is generated from mean-centered Gumbel distribution.
64


Reviews:
[{"review_id": "xXxOfE9syH", "reviewer": "Reviewer_sWSz", "summary": "In this paper, the authors study how simplified one-layer transformers with position-only attention trained via gradient descent can learn a class of teacher models defined through a bilinear structure combined with a non-linearity. Included in this class of models are, among others, convolution layers with average pooling, graph convolution layers, sparse token selection, and group sparse linear predictors. Several experiments on synthetic and real-world data are carried out to confirm the theoretical results.", "questions": "I would like the authors to address my main concerns expounded above. In particular, I'd like the authors to discuss how their work is still relevant despite the strong theoretical assumptions (in particular, the input-independent attention), and to discuss more thoroughly how their assumptions compare with those of the previous works they compare to (Wang et al., 2024; Zhang et al., 2025). \n\nAs a minor side question, I see that S in (2.1) is assumed to have all non-zero entries to be equal to 1/K. Can your approach be generalized to non-uniform S. What are the main difficulties for this case?", "limitations": "", "rating": 2, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well written and well organized. The proofs are complex and they seem sensible to me, even though I could only superficially go over them.", "weaknesses": "My main concerns with the paper are the strong theoretical assumptions and what I perceive is a somewhat unfair comparison with previous works. \n\nRegarding the first point, the assumption that the attention matrix only depends on the position of the tokens is extremely strong. In fact, this simply means that the attention matrix does not depend on the input at all, since the positions are the same for all inputs. (In Eq. (2.4), P is fixed and independent of X). With this assumption, the paper neglects what is arguably the main feature of attention, that is, to attend to other tokens depending on their value. With this assumption, the architecture precisely mimics the teacher class considered (compare Eq. 2.1 with Eq. 2.4), therefore it is not surprising that the simplified architecture can correctly learn this class (albeit, I concede that proving this fact is still challenging). \n\nRegarding the second point, I find the way the authors compared to some previous work quite misleading. Specifically, the authors claim that their work generalizes (Wang et al. 2024), which also consider how a simplified transformer architecture learns the sparse token selection task. However, the two works differ on two important points. Firstly, in (Wang et al., 2024), the fact that the attention only focuses on the positional information of X is shown to be learned during gradient descent, and not assumed from the start. Secondly, and more importantly, in (Wang et al., 2024) the attention matrix depends on the input. In fact, in their work, the sparse subset of indices over which X is averaged is sampled randomly and given as an input to the transformer. The transformer then selects the correct indices of X through the attention layer based on the input. In this paper, instead, the subset of indices over which the average is taken is presumably fixed before training for all inputs. This is a much simpler task, which explains why the attention is chosen to be input independent. Comparing the two works, and in particular their convergence rates, seems unfair.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "oEOIS4WF4Q", "reviewer": "Reviewer_MrQP", "summary": "This work presents a theoretical analysis of the optimization dynamics for learning a class of teacher models via a single layer transformer. This class of models satisfies a bilinear structure, and is general enough to recover the results of previous theoretical analyses. In particular, when optimizing over the population loss of the least squares error, the authors show a tight convergence rate, improving upon previous analyses. An additional result on mild OOD generalization as well as experimental results are provided.", "questions": "- Is there any example of a teacher model in this class that does not do sparse selection as a subtask? If not, what are the challenges in characterizing such a class that does something other than sparse selection?\n- I know this does not change the takeaway much since experimental results are provided, but what are the challenges in getting sample complexity guarantees? I suspect it will be some online SGD type analysis; would the authors elaborate on this?\n\nI am mostly curious about the first question, and I am happy to raise my score if the authors address it sufficiently.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 2, "strengths": "- The results provide a tight convergence rate of $\\Theta(1/T)$ when training on the population loss, improving previous analyses while generalizing the setting.\n- Experimental results for different examples of teacher models are provided, which supports the theoretical analysis.", "weaknesses": "- The setting is restricted to the population loss, as well as Gaussian data.\n- The generalization of teacher models is nice, but ultimately it is still the idea of sparse selection from the input.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NS5L9lfVrb", "reviewer": "Reviewer_xeJq", "summary": "This paper provides a theoretical analysis of one-layer transformers trained via gradient descent to learn from a class of teacher models with bilinear structure. The teacher models covered include convolutional layers with average pooling, graph convolutional layers on regular graphs, sparse token selection models, and group-sparse linear predictors. The authors prove that transformers with simplified \"position-only\" attention can recover all parameter blocks of these teacher models, achieving optimal population loss with a convergence rate of Θ(1/T). They also establish out-of-distribution generalization bounds and validate their theory through experiments on both synthetic and real-world (MNIST) data.", "questions": "See weaknesses.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Unified theoretical framework: The paper identifies a fundamental bilinear structure shared across diverse learning tasks, enabling unified learning guarantees. This is a significant conceptual contribution that connects previously studied disparate settings.\n\n2. Tight convergence guarantees: The paper establishes matching upper and lower bounds for the convergence rate of Θ(1/T), improving upon prior work\n\n3. Empirical alignment with theory. Synthetic experiments match predicted slopes and show early directional alignment between learned value matrix and teacher weights.", "weaknesses": "1. **Limited model complexity** The analysis is restricted to one-layer transformers with simplified \"position-only\" attention. While the authors justify this simplification empirically, the gap between this architecture and practical multi-layer transformers with full attention is substantial.\n\n2. **Notation density**: The paper is extremely dense with notation making it difficult to follow. A notation table and more intuitive explanations would improve accessibility.\n\n3. **Scalability concerns**: The convergence time $T^*$ grows quadratically with sequence length D, which may be prohibitive for longer sequences. This scalability issue is not discussed.\n\n4. **Gap between theory and practice:** Theory–practice gap in positional encodings. Assuming a fixed  D×D orthogonal positional matrix departs from learned or sinusoidal encodings used in practice; further clarification is needed to justify this assumption.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "yBJ0ZRKTFV", "reviewer": "Reviewer_hPEb", "summary": "The work demonstrates both theoretically and empirically that a modified version of Transformers can mimic a collection of well-known models (GCN, CNN, etc.). Briefly, this modified version replaces the separation between query and key matrices in traditional Transformers with a single parameter and introduces position-only attention weights. Importantly, the authors provide proven polynomial-time convergence rates for the gradient descent-based minimization of the population loss for Gaussian data. On top of this, theoretical claims are well-positioned within the broader literature and clearly supported by a comprehensive empirical analysis. \n\nAll in all, the main document is nicely written and provides apparently significant contributions for the theoretical understanding of Transformers. Nonetheless, the supplementary material is very dense and it is difficult to assess the correctness of the theoretical results. Lemma C.2, for instance, states that the rows of the learned value matrix are always a positive multiple of the rows of the true value matrix, and the proof is long and convoluted. Also, it is unclear why the data’s dimensionality should be larger than a polynomial function of $M$ and $K$, and why many of the presented inequalities break when $D = K$ (See weaknesses below).", "questions": "Please refer to the questions above.\n\nOn top of that, calculations heavily rely on the Gaussianity of the data. When the data deviates slightly from a Gaussian - a t distribution or a Gumbel distribution - how do the authors believe that the results would behave? Would we keep the polynomial convergence rate, for instance?", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 4, "contribution": 4, "strengths": "1. The main document is very well written; the notation is clean and easy-to-follow.\n\n2. The proposed convergence bounds are, AFAIK, the tightest known for this specific type of problem (which, as remarked by the authors, has been extensively studied lately).\n\n3. The presented empirical results strongly agree with the introduced theoretical analysis.\n\n4. Assumptions for both the Theorems and experiments are clearly stated (e.g., Gaussian data for in-distribution learning, exponential data for out-of-distribution learning, hypothesis class for optimization, Gaussian noise for the target, etc.).", "weaknesses": "My main concern is regarding the lack of a clear explanation for the results obtained through extensive and dense calculation, as I will detail below. I will be more than happy to raise my score in case the following points are addressed. \n\n1. Convergence rates of $\\mathcal{O}(1 / T)$ (non-tight) are well-known for strongly convex problems (or, more broadly, for Polyak-Lojasiewicz-type problems). However, this is clearly not satisfied by the squared objective function in this work. Which other regularity conditions do the authors think have allowed these tight bounds to exist?\n\n2. Lemma C.2, as I referred earlier, seems to be a strong result, and its demonstration is broken up into several pages with many hand-waved computations. Could the authors provide a short and intuitive explanation for this result? What is (if there is) a meaning for $C_1$, $C_2$, and $C_3$?\n\n3. Also, demonstrations look invalid when $D = K$, that is, when $\\mathbf{S}$ is a simple pooling matrix - this might be the case, e.g., when $g = [D]$ in Example 2.3. What is the fundamental reason for this constraint (if there is such a reason)? \n\n4. On the same note, the bounds also seem to fail when $K = 0$, i.e., when all entries in $\\mathbf{S}$ are identically null. As I understand it, this should be the easiest-to-learn scenario. Could the authors briefly discuss this restriction? \n\n5. In line 1088, it is stated that an inequality follows due to Lemma E.1, which only derives a few equalities. Could the authors please elaborate on this regard?\n\n6. Theorem 3.1 bound for $T$ depends linearly on the inverse of the quadratic norm of $\\mathbf{V}^{\\star}$, while Theorem 3.2 depends linearly on the quadratic norm of $\\mathbf{V}^{\\star}$. That is, the dependence is inverted in both results. Why is there such a difference between in-distribution and out-of-distribution results?\n\n7. Could the authors provide the slope for the linear curves fitting the tails of Figures 2(a) and 2(b)? Also, do the authors plan to release the code?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)