Transformer Learns Optimal Variable Selection in Group-Sparse Classification
Authors: Chenyang Zhang
Xuran Meng
Yuan Cao
Venue: ICLR 2025 Poster (2025)
Keywords: transformer
self-attention
variable selection
group-sparse classification
Primary area: learning theory

Abstract: Transformers have demonstrated remarkable success across various applications. However, the success of transformers have not been understood in theory. In this work, we give a case study of how transformers can be trained to learn a classic statistical model with "group sparsity", where the input variables form multiple groups, and the label only depends on the variables from one of the groups. We theoretically demonstrate that, a one-layer transformer trained by gradient descent can correctly leverage the attention mechanism to select variables, disregarding irrelevant ones and focusing on those beneficial for classification. We also demonstrate that a well-pretrained one-layer transformer can be adapted to new downstream tasks to achieve good prediction accuracy with a limited number of samples. Our study sheds light on how transformers effectively learn structured data.

Published as a conference paper at ICLR 2025
TRANSFORMER LEARNS OPTIMAL VARIABLE SELEC- TION IN GROUP-SPARSE CLASSIFICATION
Chenyang Zhang'. Xuran Meng', Yuan Cao
f The University of Hong Kong i University of Michigan, Ann Arbor
wtaphhnngcooneenn w.h..k.hh hk, xuranm@umic ch. edu, yuancao@hku.
ABSTRACT
Transformers have demonstrated remarkable success across various applications. However, the success of transformers have not been understood in theory. In this work, we give a case study of how transformers can be trained to learn a clas- sic statistical model with "group sparsity", where the input variables form multi- ple groups, and the label only depends on the variables from one of the groups. We theoretically demonstrate that, one-layer transformer trained by gradient de- scent can correctly leverage the attention mechanism to select variables, disre- garding irrelevant ones and focusing on those beneficial for classification. We also demonstrate that well-pretrained one-layer transformer can be adapted to new downstream tasks to achieve good prediction accuracy with limited number of samples. Our study sheds light on how transformers effectively learn structured data.
INTRODUCTION
The Transformer architecture (Vaswani et al., (017) has emerged of the most popular models in the modern deep learning, demonstrating remarkable success across a wide range of real-world applications, including language processing and language modeling (Vaswani et al. 2011 Radford
et al. 2019 OpenAA, 2023), computer vision (Dosovitskiy et al.] 2020 Rao et al. (021). and rein- 2022222220aeee Janner et 2021). Despite its empirical achievements, the underlying mechanisms of transformers remain poorly understood due to their complex architecture, especially the mechanism of the self-attention layers.
forcement learning (Jumper et al., [202]] Chen et al.) [202]]
As the core of the Transformer architecture, the attention mechanism has consistently been the primary focus of research aimed at understanding how Transformers work. Some empirical re- search (Xu et al. 2019 Vig & Belinkov 2019 Rao et al. [022]] Yao al. | Chen et al.) 2022) demonstrated that the attention layer can effectively extract structure information by assigning dif- ferent weights to different imput tokens. To fully understand such property theoretically, researchers explore the expressive power transformers across various aspects (Edelman et al.) [022] Bai et al.) 2024). Edelman et al. (202)) study the capacity of single-head attention to approximate sparse func- tion by presenting the covering- number of attention function class. Bai et al. (2024) demonstrate that the transformers can implement a broad class of standard machine learning algorithms con- text with various data distributions.
In addition to examining the expressive power of Transformers with custom-designed parameters, some recent studies have focused on analyzing the training dynamics of Transformers to determine whether these favorable properties can be attained using popular optimization algorithms in the deep learning community. Zhang et al. (2024) provide a global convergence analysis of gradient flow on the linear regression in-context tasks, where they consider linear layer for attention calculation. Huang et al. (2024); Siyu et al. (2024) extend this result to one-layer transformers with softmax attention layer. Wang et al. (2024; Chen et al. (2024b) study the capacities transformers to learn sparse linear regression tasks. Specifically, by exploring the optimization trajectory, |aag et al. (204)) show that one-layer transformers can effectively extract the positional information, and pro- vide rigorous convergence analysis. At the same time, fully-connected neural networks fail in the worst case. [Chen et al. (2024b) study the mechanisms of multi-head attention, revealing that multi- head attention will exhibit a sparsity in multiple-layer transformers. All the preceding works focus

---

Published as a conference paper at ICLR 2025
on the regression tasks, and for the classification task, Jelassi al. (2022 theoretically demonstrate that the one-layer vision transformer trained by gradient descent can converge on spatially structured data and the attention map exhibit patch association. Tarzanagh et al. (2023a b) prove that one-layer attention can converge in direction to the hard margin solution of SV M. Li et al. (2023a) provide a generalization enror bound for vision transformers trained by stochastic gradient descent. Although these papers provide some insights into mechanisms of one-layer attention for classification tasks, their theoretical results rely on some specific initialization which is practically infeasible, or strong assumptions that input data follows some particular patterns or certain parameters of attention ar- chitecture remain fixed throughout the training.
To further explore how transformers extract structure information in classification tasks, we consider a classic statistical problem with "group sparsity". In this setting, input variables are generated from multiple groups, while the true label of this input is determined by variables from a single group. We investigate the properties of one-layer transformers trained by gradient descent on this data model. The major contributions of this paper can be summarized as follows:
1. We establish a tight global convergence analysis with matching lower and upper bound for the population cross-entropy loss of a one-layer transformer trained by gradient descent (Theo- remb 2.2). All parameters in the one-layer transformer are jointly trained with the same learning rate from zero initialization, without imposing any prior knowledge of the desirable patterns of the parameters. Specifically, by precisely characterizing the global optimization trajectories of all trainable parameters, we decrypt the working mechanism of each component of the one-layer transformer for learning the "group sparse" data model. Our results theoretically verify that transformers can learn the optimal variable selections.
2. We demonstrate that the well pre-trained one-layer transformers on "group-sparse" inputs can be efficiently transferred to a downstream task sharing the same group sparsity pattern. Specifically, denote the number of variable groups as D, the dimension of variables in each group as d, and the downstream sample size as n.. Then we show that the one-layer transformers fine-tuned by online-
SGD on the downstream task can achieve Ō ((+DD) generalization error bound (Theorem 3.2
3. We conduct numerical experiments, empirically show that training loss will converge, and verify our conclusions regarding the optimization trajectories of trainable parameters. Specifically, the sparsity of the attention score matrix empirically demonstrates that one-layer transformers can ef- fectively learn the optimal variable selection. Additionally, we transfer the pre-trained one-layer transformers to downstream tasks, and empirically show that it can achieve good generalization performance with a small sample size. All these empirical observations back up our theoretical findings.
Notation. Given two sequences [Ir, } and {yn). we denote Xn Oyn)) ) if there exist some absolute
constant C1 > 0 and N > such that [xn] S C1 for all n 2 N. Similarly, we denote In (yy,) if there exist C2 > 0 and N > 0 such that |x,|| > C2lyn1 for all > N. We say In = (yx) if Vn = O())) and In = M((,)) both holds. Besides, we denote Fn, = o(yn) if, for any 0, there exists some N(c) > 0 such that [æ,|| < clym2l for all n 2 N(c). and denote on (y,,.) if yn o(In). We use Ō(-). (k-), and θ(-) to hide logarithmic factors in these notations respectively. Moreover, we denote I, = poly().) if O(v.)) ) for some positive constant D, and Fn = polylog(y.) if In = poly plog()))),..). For two scalars a and b, we denote b ma (,, bj and a A min(a,b) b}. Finally, for any E N., we use [ri] to denote the set (1,2,.... m.].
2 TRANSFORMERS LEARN "GROUP SPARSE" DATA
In this section, present theoretical findings that one layer transformer can learn "group sparse" data by implementing an optimal variable selection. We first introduce the definition of "group-sparse" data distributions and the one-layer transformer study in this paper.
Group sparse learning problem. Assume the feature vector X E RP, the label y of feature vector & is determined by labeling function x(-) : R R as y o (*),, where 3 E R" is pre-defined ground truth. Furthermore, we assume there exists a predefined D disjoint partitions of (p). specifically, [p] G2 and n Gj 0 for any i * j. Then refer to this learning
2

---

Published as a conference paper at ICLR 2025
problem "group sparse" if the ground truth vector 3 satisfies that
supp( := (k / + 0; k E [p]} G
for some E [D]. In particular, if the labeling function (xx) = I. we denote this learning problem
as group sparse linear regression. If the labeling function ((x) sign(x), we denote this learning problem as group sparse linear classification.
The definition above regarding group learning problem motivated by Huang & Zhang (011
Li et al. (2023b). In this paper, we focus on binary group sparse linear classification problem. For simplicity, We consider the setting where all G5 s are of the same size d, implying p dD. Besides, we convert the feature vector & E R into matrix X E Rd where each column xj is the collection of the variables from Gj. Let be the index of the label-relevant group, i.e. s)pp(33) C Gy.. Then we can denote v² the d-dimensional vector obtained by restricting 3 on Gy- so that (B* X) (v",x,-) Lastly, assume the features are generated from Gaussian distribution. We provide a formal definition of the data model as follows:
Definition 2.1 (Group sparse inputs following Gaussian distribution). Let E Rd be a fixed vector representing the parameters of the labeling function, and j E [D] be the index of the label-relevant group. The binary classification data input pair (X,y) IR x {1,+1} generated from the following distribution D:
1. The features are generated from Gaussian distribution, i.e., X = [x1, X...... XDl while each col- umn J E [D] represents group of variables and is independently generated from (O, Le).
The label of this input X determined by the features from the label-relevant group indexed by j", defined y sign bign((xy..) v*)).
Notice that the label y is determined solely by the direction of v", while the norm of v does not affect the data distribution. For simplicity expression, we assume that ||v*||2 = 1 the following.
One-layer self-attention transformer. For each x, with j E [D) from data distribution D defined in Definition 2.1. we define the corresponding positional encoding P, as
2j D sin Dii T
(DIDTT)
Pj = sin (DDT)
(2.1)
The definition above motivated by the fact that Pj. E [D] form an orthogonal basis. We generate new training input by concatenating the Gaussian feature X with positional encoding matrix P. which is defined as P = P1. P2. .Po]- Specifically, we define the new training input as
= [2,, Zz. zp] E D)xD with z; [x] .p.] for all j E [D]. Now we introduce the one-layer self-attention architecture applied to the input Z as:
D
f(Z, W. v)
ZS (2,) v ZS1p.
(2.2)
Here S(-) DRD
and S R denote the softmax mapping S(h), ΣP_
[S(Z T Wz1), S(Z Wzz), sZ Wzp)l E RxD Besides, we denote the entry in the j'- th row and j-th column of S as S. Compared to the classical single-head, one-layer self- attention structure in Vaswani et al. (2017); Dosovitskiy et al. (202)), we make some mild re- parameterizations the architecture:
We combine the query and key matrices into one trainable matrix W E R(d+D)xid+D) We replace the value matrix with one trainable value vector Rd+D
The consolidation of query matrix and key matrix is commonly considered in recent theoretical analyses attention structure (Tian et al. 2023 Wang et al. 2024 Huang et al. 2024 Zhang et al. 2024). The simplification of the value matrix into vector because only one parameter vector v needs to be studied, and the value vector v could regarded as a combination of the value matrix and second layer with uniform weight. Similar simplification is also considered in some recent
3

---

Published as a conference paper at ICLR 2025
studies about understanding the attention mechanisms in classification tasks. Tarzanagh et al. 2023a).
Jelassi et al. 2022
Loss function and training algorithm. We consider minimizing the population cross-entropy loss to train the above one-layer attention transformer in (2.2). Specifically, the loss function
[(v. W) (x))) [1 f(Z. W. v)]l,
where ((a) = log(1 +exp(-a)) is the cross-entropy loss function for binary classification. Although the choice of population loss implicitly assuming infinite number training set is not feasible in practice. It can significantly simplify the training dynamics and enable us to focus on an analysis of the global optimization trajectories. This objective loss is considered in large number of recent
theoretical works (Jelassi et al., 2022 Zhang et al.) 2024; Huang et al. 2024; Wang et al. 2024] Nichani et al., 2024) Chen et al. (2024b) focusing on the optimization analysis of transformers. Besides, one of the most popular loss functions classification tasks, cross-entropy loss is widely considered recent theoretical works concerning transformers on classification tasks (Jelassi al,, 2022 Tian et al. 2023 2024). Compared to the hinge loss considered in Li et al. (2023a), it is more common and general in practice. We utilize the gradient descent algorithm to optimize the preceding loss function. We consider the joint training regime, where all trainable parameters W are updated simultaneously with the same learning rate 17, i.e.,
((++)) = (() - -VVCCWW ((1);; w = w( V (()).
(2.3)
(2.4)
Besides, we consider zero initialization (0), w = O.
Tight convergence bounds. Following the preliminaries, the next theorem presents our main results regarding convergence bounds.
Theorem 2.2. For any 0, suppose that D 2 (1/e)) d s Oppoly(D)) o 17 O(1) with a 1/3 and let 7 @(D* v Do.)). Under these conditions, holds that
1. The self-attention extracts the variables from the label-relevant group: for any new sample (X, ~ D with the corresponding softmax output matrix S(T") at the - iteration, with probability at least 1 - exp ( ((VD)) it holds that
s $ 2 1- exp ( (DD)))
for all j [D].
The first block of value vector aligns with the ground truth: v(T) [T -0%]T with v (T") Rdd where
ScDexp (-0(VD))
v
3. The loss is sufficiently minimized:
Ci ci(eN = w(r') s eΛ 1
D
where Ci is positive constant solely depending 1) and with Ci < 1.
Theorem 2.2 implies that lne-l attention model (D.4) can learn the group-sparse data with an assumption of an infinite number of training data. It can achieve arbitrary small loss by imposing mild assumption concerning the scale of sparsity in our data model D. Specifically, for any given loss tolerance e, need the number of variables group be larger than the logarithm squared of the reciprocal for loss tolerance, i.e., (0g(((/c) Based this assumption, Theorem 2.2 can establish tight upper and lower bound on the convergence rate of population loss.
In addition to providing tight convergence bound, Theorem 2.2 decrypts how each component
of one-layer self-attention (2.2) takes effect when learning sparse group data, by precisely identi- fying the scale of the parameters at - iteration. Specially, for any new sample (X,y) D.
4

---

Published as a conference paper at ICLR 2025
the attention score regarding the label-relevant group at - iteration, i.e. $ is approximately I with high probability. This observation indicates that the self-attention layer almost attends to the label-relevant group, effectively selecting the variables mostly beneficial for classification while disregarding all other irrelevant variables. Besides, the second block of v(T retains implying positional encoding is only involved in the calculation of attention weight. Following the proce- dure of optimal variable selections, positional embedding is never presented in the final output of one-layer attention. This observation benefits our learning tasks as positional embedding can not (1)
provide additional beneficial information for classification. Finally, V1 is almost aligned with the ground truth direction v", verifying that the value vector learns the optimal direction for the binary classification problem. These observations collectively guarantee that the one-layer transformer 2.2 can correctly implement the classification task on the "group sparse" data 2.
Notablely, recent work (Wang et al., 2024) studies the sparse token selection of transformers on a similar data model to ours. By characterizing the entire training dynamics, they demonstrate a global convergence of transformers on the sparse token selection task without further assumptions on initialization. Besides, they do not impose fixed positions on target tokens, which is more general than our settings. However, they assume the positional encoding of the next token an average of target tokens, which alleviates the technical challenges, is somewhat impractical. What matters most is that they consider the l loss function in their settings, making their learning problem more akin to regression task, which significantly distinguishes optimization analysis from theirs.
3 DOWNSTREAM TASKS
The variable selection mechanism of one-layer self-attention enables be efficiently transferred to downstream task sharing a similar structure. In this section, we provide a generalization er- ror bound for the downstream task to theoretically verify this conclusion. We first introduce the definition of downstream task data distribution we study as follows:
Definition 3.1. Consider a downstream task with training data (x()) E RdxD x (1,+1), i [n] generated from a new distribution D satisfying:
1. Linear separability with a margin: maxv:l|vllas () (V. (v)) 2 almost surely for E [n]. Besides, the lahel-relevant group index j is the same as the distribution D in Definition 2.1
2. Each entry of (()) is independent sub-Gaussian, with | Mell
S Sxx for all i E [n], [d] and
j E [D].
Compared to the data distribution for pre-training in Definition 2.1 we make downstream task distri- bution more general by relaxing the Gaussian data to sub-Gaussian data. Besides, the assumption of linear separability with margin is commonly considered in the analysis concerning generalization
risk bounds of classification tasks (Wu et al. 2024 Ji & Telgarsky 2020 Cao & Gu 2019).
Fine-tuning with online stochastic gradient descent. Similarly, we concatenate each x (8) with
positional encoding Py as defined in (2.1). and generate the new training raining data (Z(), ()), for i E [n]. Then, we consider fine-tuming one-layer attention model with the same structure defined in (..2). and denote the new parameters as v and W. We fine-tune the new parameters by online SGD, i.c.,
g((+1) -
Z (W(), (())
(3.1)
w(++1) w( 7V C. (1), v(0)
(3.2)
for all i E [n]. where L, (WO, ()) ((g(t)) (Z W W(,) The initialization are set as (0))- = 0 and ()) = WT)) which is obtained from the pre-trained model in Section 2 In the following theorem, present generalization error bound on this downstream task, which concerns an average of all iterates ((W(), (0))) of online SGD.
5

---

Published as a conference paper at ICLR 2025
Theorem 3.2. Suppose that D 2 w(log" d O(poly(D)). oz O(1), and f
O (dvD))=)
). Under these conditions and for any > holds that
Σ P PX (v f(Z, W v00) S 50) s (d log' O log (0g(1/0)
n2n
n
with probability at least 8 nexp( (Q(())) over the randomness of {(X(), ())
Theorem establish generalization error bound composed of two terms. The first term concerns the ratio between the practical dimension d D and sample size n. Here consider d D as the practical dimension since the positional encoding is also involved in the learning process.
The second term standard large-deviation error term. If we take 7 ( as most literature assumes, we can obtain a sample complexity on the order of ? d+D + log())) In comparison, the existing lower bound for the sample complexity of linear logistic regression vectorized X is O d1 + log( )). according to the classical PAC learning (Long 1995). This superiority in sample complexity theoretically demonstrates that one-layer transtormers tramnd on "group data" can be effectively transferred to similar downstream task.
4 PROOF SKETCH
In this section, we provide a comprehensive explanation of how a one-layer transformer implements variable selection group sparse data. Furthermore, we share insights into the reasoning that led to the conclusions presented in Theorem 2.2 along with a proof sketch for the convergence bound.
For clear illustration of each component in the one-layer self-attention, we first separate v(*) and (t)) into the following blocks.
W
v(0) = (vv W
1. (8)
W (t)
W 2,2
2,1
2,1 E Pxd and ((t) PxD 2,2 In the following, we will show that a precise characterization of these blocks at 7 -th iteration plays key role in our theoretical convergence analysis.
where v
RR.V ER", w
1.1 Raxd ( 1.2 E Rd w
A low-rank structure of (T') beneficial for attending to label-relevant group j The con- clusion that the attention score S 1 with high probability for any sample (X, y) DD serves as a foundation for transformers to correctly implement classification on group-sparse data distribution D. Without such guarantee, the loss cannot converge regardless of the direction in which the value vector converges, since the product of the label and variables from irrelevant groups still follow a Gaussian distribution.
The following lemma accurately characterizes each block of W(T"), providing insights into the rational of S/"") Rs
Lemma 4.1. Under the same condition with Theorem 2.2 ((") satisfies that
( W (T*)
1,2 2,1 0;
TT + ((T)) W(("")
W.
D +
w B2
Σ
(D P.))
Σ
p
(("")
2.2
j
where |311 ((VD), 32 = 0 (z)). The error terms W(T")
1.1.emo and W (T")
2.2 are small such that
W V II: W. I exp (vD)v))).
Lemma 4.1 indicates that W(T) exhibits particular pattern: (i). the off-diagonal blocks w
and W!l 2.1 remain 0; (ii). the leading component of left-top block w(T") aligns directionally

---

Published as a conference paper at ICLR 2025
with the projection matrix v"v";; (ii). the leading component of the right bottom block aligns directionally with (E Eez (D,, - p3)) (E)) In the following, we explain why this pattern implies that S for all j with high probability, which benefits the label-relevant variable selection.
First, Wi.2 W = suggests that the features xj x,'s, do not engage with the positional em- bedding Pj s when calculating the attention scores. Besides, due to the orthogonality among the " positional embedding Pj's, the position-position interaction term Py 2.2 Py takes a large value only when j' = jr. Furthermore, by standard concentration inequalities, it can be shown that (1")),
x.W. 1,1 xj < p.. 2,2 Py for all j,j E [P] with high probability. All these observations col- lectively conclude that p... w((") Py will eventually dominate the softmax calculation, which finally implies that S3 x 1.
Accurate characterization of the alignment between v (T") and ground truth v'. Based on Lemma(4.7]an preceding discussions, we have shown that one-layer attention can effectively extract the features from the label-relevant group. The remaining challenge is to accurately characterize the alignment between v(T") and ground truth v² in terms of both direction and scale. The following lemma presents this result.
Lemma 4.2. Under the same condition with Theorem 2.2 holds that,
(T") a (T") + Vienori (T") O,
(T*)
where the error term Vi,eror satisfies that (viton) v*) and ||vi,otoo l2 S exp (- ((V)).
Besides, the coefficient of the projection of (T")) onto the direction v", which is denoted by (*) follows that
(2T + DD S (T) S ( + D),
where C2, C3 are both positive constants solely depending on o₂ and r7.
Lemma 4.2 reveals two important facts. Firstly, the fact that V2 (T') O implies that the positional embedding part only contributes attention weight calculation, and is not presented in the output. Furthermore, v (?"") is well aligned with the direction of v*, and the scale of its projection onto this direction is accurately characterized by a matching lower and upper bound. Besides, the second
2 (T ll2
conclusion Theorem 2.2 is a direction corollary of this lemma as ||
1|2 v.
Upper and lower bounds of y f(z, W(T") v(T")) A tight upper and lower bound of [Z, w("), v(T") is undoubtedly of utmost significance for the analysis of global loss conver- gence. The accurate characterization of (1 in Lemma 4.2 enables to get this matching bound. Before we present the result, we introduce a notation ET- to denote the event that the conclusion in Theorem 2.2 holds, i.e., S)T" 2 exp(O(D). Then, we present our results regarding the matching bound of y f(Z, w(T"), v(²")) in the following lemma.
Lemma 4.3. There exist two i.i.d. random variables B1, B2 with B2/02, B2/o X2. such that
(Z (T"), (T))) (D B -D D || [B
Moreover, under the event Ers, it holds that
(T"
2 v(v" x,.) 1 (Z, W(T") v(T')) < (T(TT) +1.
Based on these lemmas, ready to prove the conclusion of convergence bound in Theorem 2.2
Proof of Theorem 2.2 We separate L(v(T") W(T)) into parts as
(() w(T") E E[[lyf(Z, w(T)) ((T)) (1p-)))))))))))))) (E.- 1 E [Llyf(Z, W(T), v(T)))
I1

---

Published as a conference paper at ICLR 2025
For [1. by the fact that exp(-r)/2 Σ {(x) exp(-2) for all >> 0, and utilizing the results of Lemma 4.3 for Er- occurs, it holds that
(T(
2
2
11 - E
eexp
viV (VVV,x).)) 18T-
S V n o Da
2 V ,, Da(
1 2 E
exp
Da yyv'" , E
2e
The last steps of both formulas are derived by basic integral since y(v", x,-) follows folded nor- mal distribution with mean 0 and variance o² (See more details in Lemma E.5 and Lemma E.6 Furthermore, for 12, we provide an upper-bound as
E log? (1+ exp (Da(T))B + D M.v.f.oocl122
12 S
((Eg.))
2/20, DE (a(tT))
s 2 2D2 (aD((a())) + 2D | 1, [22] P(E
(T
eO(DD)
The first inequality is from Cauchy-Schward inequality and Lemmal 4.3 by the fact that log(1 + a) S 2 2og a for large o, and (a + b)² < 2 a
The second incquality holds
12). The third inequality is
derived by applying the definition of B1, By in Lemma 4.3 and the fact PEE
I ee (-((D)) from the
first conclusion in Theorem [.2] (We provide a rigorous detail Lemma C.3
By the upper-bound
of (T") in Lemma 4.2) the definition T in Theorem 2.2 and our condition 2 (1/e),
we know that ( <<a e e indicating that 12 < J.. Let T = (v Caas ) V DAz), by
applying the lower and upper bounds of (T") in Lemma
4.2 we can finally obtain that
[(v(1") w((")) < 21 S
2 4e
1
S
C2a. D(T+)1/3 1D2
C2 π 4eC3o, D(T.)1/3 + D 32c²Cs
2
((v(r", (2)))222
2
c.
This completes the proof.
NUMERICAL EXPERIMENTS
In this section, we conduct numerical experiments on synthetic data to verify our theoretical conclu- sions. Two types of synthetic data are generated following the group sparse data distribution defined in Definition] 2.1 and downstream task data distribution defined in Definition 3.1 respectively. For each type of synthetic data, generate two samples with different sample sizes n, variable group numbers D, and variable dimensions d. Besides, the variable group numbers D, and variable di- mensions d are matched across these two different types of synthetic data to enable the parameters of one-layer self-attention to be transferred to downstream tasks.
In the experiments where transformers learn variable selection, we generate data according Def- inition 2.1 setting 0g 0.25 and j" = 2, with two different pairs of (n, d, D): (500, 4, 6) and (200, 2, 4J. The vector v randomly generated and then fixed. We set the learning rate 7) 0.5 and train the models for 400 iterations. During the training process, we plot the training loss, the
||- After the training loss converges (at final iteration), we calculate the attention score matrix for each sample and display the heatmap of the average attention score matrix across all samples.
cosine similarity [ |||v=I|
(? >v)
and the norm ratio ||v l/
As shown in Figures aligns with the direction of v* early in the training and continues in that direction for the rest of the iterations. When the training loss converges, we can also observe in Figure that ||v!"'| is relatively small compared to | When we examine the attention matrix, as shown in Figure 2
the training loss converges to 0 after sufficient iterations, and v rapidly
8

---

Published as a conference paper at ICLR 2025
Traiting Loss n Cradiem: Descort
Cc Angs and
Vorm Rate [1e5 (0871)
Egcth
Oporn
(a) Training Loss
(b) Angle vi and
(c) Norm Ratio || (1) |1 (1)
That-ing Crudient Descort
Coe toge of and
Norm Ranis ( P'D''
0200
04
0275
kp4ch
4pocn
(d) Training Loss
(c) Angle in vi and V
(f) Norm Ratio [[v (t) |/ [|
Figure Figures on training loss, cosine similarity and norm ratio. The first line presents the training results with sample size of 400, 6 variable groups, and variable dimension of 4. The second line shows the training results for a sample size of 200, with variable groups and a variable dimension 2.
Heatmap of Attention
Healmap Attention Matrix
3.861
0.971
0.922
3.912
0.213
0.011
5.033
0.030
0.002
2.035
0.902
0.903
0.5
5.000
0.002
0.003
2.005
3 3.205
0.005
0.02
0.005
0.4
0.205
(a) Hcatmap of Attention Matrix
(b) Heatmap of Attention Matrix
Figure 2: Heatmap of the average attention matrix. Figure 2(a) shows the heatmap of the attention
matrix corresponding the variable groups, and Figure 2(b) shows the heatmap of the attention matrix corresponding to the 4 variable groups.
we observe that the second row significantly larger than the others, indicating that the attention predominantly focused this row. This phenomenon well matches the prediction from theory.
We conducted additional experiments on the downstream tasks by generating two new Gaussian samples, each with different v* from the previous setup. These samples follow Definition 3.1 with parameters set to J. = and 7 = 1, while using the previous configurations for d and D. To ensure linear separability with margin, The test accuracy is calculated after each iteration using 1000 test samples. Both experiments use a sample size of 400, and the learning rate is set to 10-3.
applied projection to adjust the data accordingly.

---

Published as a conference paper at ICLR 2025
Test Accuracy
Test Accuracy
Test Accuracy
Test Accuracy
0.8
0.7
100
150
200
250
300
350
50
100
150
200 Epoch
250
300
350
400
Epoch
(a) (d,6)=(4.6)
(b) (d,D)=(2,4)
Figure 3: Test accuracy in the downstream task performance with different variable group numbers and variable dimensions.
As shown in Figure 3 the test accuracy initially oscillates during the early stages of training. How- ever, after longer training period, the test accuracy stabilizes and converges to 1 in both settings. This oscillation primarily due to the transition the direction of the vector (+) during training.
According to analysis, once vi (t) completes its direction transition, the test error decreases, lead- ing to stable learning in the downstream task. These trends are consistent with the theory, confirming that the pre-trained transformers perform effectively downstream tasks.
CONCLUSIONS AND LIMITATIONS
In this paper, we study how one-layer transformers trained by gradient descent learn binary clas- sification group sparse data, without imposing any prior knowledge on the initialization. We pro- vide accurate characterization of optimization trajectories. Based on this result, theoretically demonstrate that the one-layer transformers can almost attend to the variables from the label-relevant group, and disregard other ones by leveraging the self-attention mechanism. Besides, establish tight convergence rate with a matching lower and upper bound for the population loss. Furthermore, also propose that pre-trained one layer transformers transferred to downstream task with the same sparsity pattern. We theoretically demonstrate this conclusion by providing an improved generalization error bound for one-layer transformers, which surpasses that of linear logistic regression applied to vectorized features. Our numerical experiment observations support the theoretical findings. Although our theoretical findings provide insights into the mechanism of self-attention in variable selection, we focus exclusively on one-layer self- attention model. Future research could be more intriguing if explored deeper transformer architec- tures. Additionally, investigating the integration of self-attention with other modules, such as MIPs, ResNets, and normalization layers, presents a promising direction for further work.
group sparse data be effectively
ACKNOWLEDGMENTS
We thank the anonymous reviewers for their helpful comments. Yuan Cao is supported by NSFC 12301657 and HK RGC-ECS 27308624.
REFERENCES
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers statisticians: Provable in-context learning with in-context algorithm selection. Advances in neural information processing systems, 36, 2024.
10

---

Published as a conference paper at ICLR 2025
Alina Beygelzimer, John Langford, Lihong Li, Lev Reyzin, and Robert Schapire. Contextual bandit algorithms with supervised learning guarantees. In Proceedings of the Fourteenth International Conference Artificial Intelligence and Statisties, 2011.
Satwik Bhattamishra, Kabir Ahuja, and Navin Goyal. On the ability and limitations of transformers to recognize formal languages. arXiv preprint arXiv:2009. 2020.
Satwik Bhattamishra, Arkil Patel, Varun Kanade, and Phil Blunsom. Simplicity bias in transformers and their ability to learn sparse boolean functions. arXiv preprint 1xXiv:2211112316 2022.
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. In Proceedings of the 33rd International Conference Neural Information Processing Systems, pp. 10836-10846, 2019.
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 Annual Conference Learning Theory, pp. (6699-7753. PMLR, 2023.
Dexiong Chen, Leslie O' O'Bray, and Karsten Borgwardt. Structure-aware transformer for graph rep- resentation learning. International Conference on Machine Learning, pp. 3469-3489, PMLR, 2022.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Picter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:1504: 5084-150977 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.
Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2019.
Qingxiu Dong, Lei Li, Damai Dai, Ce Zheng, Zhiyong Wu, Baobao Chang, Xu Sun, Jingjing Xu, and Zhifang Sui. survey on in-context learning. arXiv preprint arXiv:2301.0022 2022.
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 at scale. arXiv preprint arXiv:2010.I 1929, 2020.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation self-attention mechanisms. International Conference on Machine Learning, pp- 5793-5831. PMLR, 2022.
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 2024.
Processing Systems,
Shivam Garg, Dimitris Tsipras. Percy S Liang, and Gregory Valiant. What transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems, 35:30583-30598 2022.
Junzhou Huang and Tong Zhang. The benefit of group sparsity. The Annals of Statistics, pp. 1978- 2004, 2010.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. In Forty-first International Conference Machine Learning, 2024.

---

Published as a conference paper at ICLR 2025
Muhammed Emrullah Iidiz, 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 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, 34: 14:1273-1283 2021.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-37836, 2022.
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve ar- bitrarily small test error with shallow relu networks. In International Conference on Learning Representations, 2020.
John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger,
Kathryn Tunyasuvunakool, Russ Bates, Augustin Zidek, Anna Potapenko, et al. Highly accurate protein structure prediction with alphafold. nature, 596(7873):583-5588) 2021.
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is 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.
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.
Jiangyuan Li, Thanh V Nguyen, Chinmay Hegde, and Raymond KW Wong. Implicit regularization for group sparsity. In The Eleventh International Conference on Learning Representations, 2023b.
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, 2023c.
Zihao Li, Yuan Cao, Cheng Gao, Yihan He, Han Liu, Jason Klusowski, Jianqing Fan, and Mengdi Wang. One-layer transformer provably learns one-nearest neighbor in context, Advances in Neu- ral Information Processing Systems, 77:2266-822044 2024.
Valerii Likhosherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices. arXiv preprint urXiv:2106.0377 2021.
Bingbin Liu, Jordan T Ash, Surbhi Goel, Akshay Krishnamurthy, and Cyril Zhang. Transformers learn shortcuts to automata. In The Eleventh International Conference Learning Representa- fions, 2023.
Philip M Long. On the sample complexity of pac learning half-spaces against the uniform distribu- tion. IEEE Transactions on Neural Networks, 66)6,1556 1995.
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 (21/0720024-277772222 Vienna), 2024.
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. In Forty-first International Conference Machine Learning, 2024.
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, et al. In-context learning and induction heads. arXiv preprint arXiv:2209. 1895, 2022.
OpenAL. Gpt-4 technical report. arxiv 2303.08774. View in Article, 2(5), 2023.
Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than sgd for transform- ers. arXiv preprint wrXiv:2366.00204 2023.
12

---

Published as a conference paper at ICLR 2025
Jorge Pérez, Pablo Barceló, and Javier Marinkovic. Attention is complee Journal of Ma- chine Learning Research, 22(75):1-35, 2021.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, llya 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 Hsieh. Dynamicvit: Efficient vision transformers with dynamic token sparsification. Advances in neural information processing systems, 34: 13937-13949, 2021.
Laura Ruis, Jacob Andreas, Marco Baroni, Diane Bouchacourt, and Brenden M Lake. A benchmark for systematic generalization in grounded language understanding. Advances in neural informa- lion processing systems, 33:19861-19872, 2020.
Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Un- raveling attention via convex duality: Analysis and interpretations of vision transformers. In International Conference on Machine Learning, pp. 19050-19088 PMLR, 2022.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. Advances in Neural Information Processing Systems, 36, 2024.
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, pp. 4573-4573, PMLR, 2024.
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 lection attention mechanism. Advances in Neural Information Processing Systems, 36:48314- 48362, 2023b.
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 in Neural Information Processing Systems, 06:719111-71441 2023.
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Shaolei Du. Joma: Demys- tifying multilayer transformers via joint dynamics of mip and attention. In The Twelfth Interna- tional Conference on Learning Representations, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention 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 ACI. Workshop BlackboxNIPP: Analyzing and Interpreting Neural Networks for NLP, pp. 63-76, 2019.
Martin Wainwright. High-dimensional statistics: A non-asymptotic viewpoint, volume 48. Cam- bridge university press, 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.
Colin Wei, Yining Chen, and Tengyu Ma. Statistically meaningful approximation: a case study on approximating turing machines with transformers. Advances in Neural Information Processing Systems, 15:11271-1208 2022.
Jingfeng Wu, Peter L Bartlett, Matus Telgarsky, and Bin Yu. Large stepsize gradient descent for logistic loss: Non- monomoonotni of the loss improves optimization efficiency. arXiv preprint arXiv:2402. 15926, 2024.
13

---

Published as a conference paper at ICLR 2025
Da Xu, Chuanwei Ruan, Evren Korpeoglu, Sushant Kumar, and Kannan Achan. Self-attention with functional time representation learning. Advances in neural information processing systems, 32, 2019.
Shunyu Yao, Binghui Peng, Christos Papadimitriou, and Karthik Narasimhan. Self-attention net- works can process bounded hierarchical languages. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 3770-3785, 2021.
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are functions? In International Con-
transformers universal approximators of aequence-to-sequee ference Learning Representations, 2020.
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 Hlghhdimmensioaal Learn- ing Dynamics 2024: The Emergence of Structure and Reasoning.
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:333311533, 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, 2024.
14

---

Published as a conference paper at ICLR 2025
A ADDITIONAL RELATED WORKS
Expressiveness of transformers. There has been a line of works studying the expressiveness of transformers, Yun et al. (2020); Dehghani al. (2019); Pérez et al. [222]: Wei ct al. (2022) inves- tigate the universal approximation results for transformers on a general class of functions, Likhosh- erstov al. (202T) ) Bhattamishra et al. (2022) evaluate the expressive power of transformers
boolean concept c classes by sample complexity. Bhattamishra et al. (2020); ; Liu et al, (023) explore the the capacity of transformers recognize formal languages, (Sahiner et al. (2022) study the equiv- alent finite-dimensional convex problems of transformers by the lens of convex duality. Olsson et al.)
2002 Dong et al. 2202)): Garg al. (2022); Ruis et al. (2020) investigate the capacity of trans- forms to learn in context and interpret the attention layers as gradient descent iterations, [Sanford ct al. (224) discuss the strengths and limitations of the representation power of attention layers, by
demonstrating the necessary intrinsic complexity of parameters on two different tasks.
Optimization of transformers. Several recent works study the optimizations of transformers. Zhang et al.] (2020) compare the performance of transformers trained using SGD and adaptive meth- ods, attributing the subpar performance of SGD in training language models to the heavy-tail dis- tributed noise it introduces. In contrast, by adjusting the batch size in experiments, Kunstner et al. (2023) empirically demonstrates that the varying performance of Adam and SGD can be ascribed to the differences in their geometric trajectories rather than the noise introduced by stochasticity. Pan & Li (2023) compare the convergence rate of Adam and SGD on transformers, and argue that Adam has better directional sharpness of the update steps than SGD. Through two-stage training regime, Li et al. (2023c) investigate the optimal parameters of transformers applied to a masked topic structure model. Ildiz al. (2024) explain the mechanism of attention from the perspective of Markov chains. The output of transformers is generated by context ceneeett-ondttioned Markov chain, with weights determined by the attention structure. Nichani et al. (2024) demonstrate that two-layer transformers trained by gradient descent can encoder the latent causal graph in the first attention layer when solving in-context learning tasks with latent causal structure. TTia et al. (2023) analyze the training dynamic of transformers with one self-attention layer and one decoder layer trained by SGD, revealing how transformers combine the tokens and attend more to distinct tokens. Tian et al. 2024) further investigate the training dynamics jointly with MLP layer, predicting particular pattern of attention map and theoretically decrypting the hierarchies of tokens. Li et al. vestigates the training dynamic of transformers when performing one-nearest neighbor selection.
2024) in-
Gao et al. (2024) studies the global convergence of transformers under particular conditions. CCen al. (2024a) investigates the in-context learning capacity of transformers by studying the training dynamics of a two-layer transformer on 7-gram Markov chain data.
B DETAILED COMPARISON WITH JELASSI ET AL. (0022)
In Jelassi et al. (2022). the authors provide a theoretical guarantee that one-layer vision transformers can learn the inner structure of data, which is defined as "patch association" by the authors. This result offers novel insights into the training dynamics of one-layer transformers.
However, both the conclusions and proof techniques of Jelassi et al. (2022) can not be directly extended or applied to this work due to the distinctive settings among these two studies. Specifically, Jelassi et al. 2022) considers a one-layer vision transformer defined as
f(X) ov XS(A))LD.
where X is the sequence of tokens/feature groups with each column x, denoting token/feature group, a--) represents the activation function, indicates the value vector, S(-) represents the softmax function, and A is input matrix of the softmax function. Unlike the common design of transformers, they directly treat the entries of the matrix A as trainable parameters and consider training A with gradient descent. In contrast, we consider softmax attention with the formulation S(Z T WZ). and treat the coefficient matrix W as the trainable parameters, which aligns with the general design of transformers. Besides, the initialization of the value vector inJJelssi et al. (2022) is assumed to strictly align with the direction of ground truth v*, which is a strong and impractical assumption. In comparison, we consider general zero initializations.
In addition to the distinctions among the settings of problems, our theoretical results are more pre- cise and refined. While Jelassi et al. (202)) provides upper bound on the number of iterations
15

---

Published as a conference paper at ICLR 2025
needed to achieve population loss of 1/poly(d), offer both matching upper and lower bounds for iterations to reach arbitrarily small population loss. Furthermore, we present a sample complex- ity analysis for transfer learning, which surpasses the conclusion of linear logistics regression vectorized inputs from the PAC learning theory. In contrast, Jelassi et al. (2022) does not include such sample complexity analyses.
C PROOF IN SECTION
In this section, we provide detailed proof for lemmas in Section Before demonstrate the proof details, we first introduce some basic gradient calculations of ((v, W) w.r.t v, W as,
0f(Z; W.
V.L(W,v) v) (lyf(z:
δv
D
D
=E(e(yf(2 w.), ZS1p = flff(z Σ j=1
y Z2
Σ
S
y==1
(C.1)
VwC(W,v) = E (luf(Z: v)) 0f(z: w.
OW
explz ) exp(2). (2;;)
E elyf(Z, W. v))
Zz'
(zj- -zn 2zn1zz
= exp(xx 23.00 W2,)]
j=ly'=1
3"+1'
D
E eluf(z: (. v))
ΣΣΣ J
Sy Sy' 22 2j" )2}
(C.2)
For the following derivation, we first introduce a shorthand notation that (())
(y, W v(9)). Besides, we also use to denote an orthogonal matrix A with v being its first column. We denote .2,..., Ea the rest columns in A, i.e.
A [v", E2,
Ea] Rdddd
(C.3)
where (v", O, (E:,E)) = 0 and ||E:122 for all i 1..
C.I RESTATEMENT OF LEMMA 4.2 AND LEMMA
For the sake of conciseness and coherence in the presentation, we rearrange some content from Lemma 4.2 and Lemma 4.1 Specifically, we consolidate the conclusion that v ((() w(()
1.2
into Proposition C.I The remaining conclusions of Lemma 4.2 are presented in Lemma C.2
We
also include the remaining conclusions regarding W(!) from Lemma 4.1 in Lemma C.3 and the
2,2 from Lemma 4.1 in Lemma C.4 We present these new lemmas in this subsection and proof them respectively in the following subsections of this part.
conclusions about W!!)
We first introduce our new Proposition C.I
Proposition C.1. For iterates v()) and W(O of gradient descent in 2.3 and 2.4 it holds that
V2 (4) O. (( and (!!)) for 2
The next Lemma C.2 a restatement of conclusion of v in Lemma 4.2 Lemma C.2. Under the same condition Theorem 2.2 holds that
(4) Vi,erorr
for all 3 < < T", where the error term satisfies that (v", ) 0 and
| (c)
V. 1l S -G-D
(C.4)
16

---

Published as a conference paper at ICLR 2025
for some positive constant C₁ solely depending 0z. 17. Besides, the coefficient a satisfy that
(V 2 3m 20 D (1
(1-3)
2 oz 2 Do)
à
S
a
S V
π +
(V 2 6n or
(1 3) + 2
02
- D3
2
125 π
D3
(C.5)
In the following Lemma C.3 and Lemma C.4 we present the conclusion concerning W and W from Lemma 4.1 respectively.
Lemma C.3. Under the same condition as Theorem 2.2 holds that
W 1,1 = 81 1iv*v* + W 1,1,emrr
(C.6)
for all 3 < / < T*, where |31| S ei D for some positive constant solely depending on 1), a.
Besides, the error term satisfies that
1, ll Se-D
(C.7)
for some positive constant C7 solely depending 0,., 77-
Lemma C.4. Under the same condition as Theorem 2.2 it holds that
/ 2,2
B2
Σ (p (p,)) $$1
[p]] ())((_))))))))
( omor
(C.8)
for all 3 3tts T", where 09 S 132
Da
2
eg for some positive constants C22 os solely depending on D
7i Ons and the error term satisfies that
2,2 ll S CtVD
(C.9)
for some positive constant C7 solely depending o.. 1/.
Lemma C.5. Under the same condition with Theorem 2.2 with probability at least - exp ( CsVD), it holds that
$ 2 - D
$ S exp(-CoD)
for all 3 s T v [D] and j' j. The coefficients Cs. Cg are all constants solely depending on o₂ and 77.
C.2 PROOF PROPOSITION
We will prove Proposition C. by induction. And we first introduce several lemmas, which will be used for further proof of Proposition C.I
Lemma C.6. If ((()) (
1,2 2.1 0, then S (1)
Junt can be expressed function of { {y-x1,'. (..
and is independent with y for all j1.32 E [D].
Proof of Lemma
C.6
When (
(t)
O, we can re-write (t) ji.d2 as
1,2
2,1
exp (2)] 23))
exp
X + P ((())
2.2
J2
= exp [I], W W(1) (112),)))))))
Σhe exp [x] W i, + P.J. w wT
exp (y J // X,2 + P. (!!) 222
Σh exp (y
1,1 Xj2 P.J. w 2.2Pin)
If we omit all non-random components, we have (/0 (1)
jn.j2 (y "1," , xp). Moreover, by
j1.ja Lemma R.2 we obtain that S S(!)) is independent with y.
Li
Lemma C.7. If W
W 0 and v 0. then ((()) can be expressed as function of {v XI,"., . xp), and independent with y-
1.2
17

---

Published as a conference paper at ICLR 2025
Proof of Lemma C.7
When V2 (1)
0, express (1()) as.
1
1 1 + exp EE yxz) Eソ= S
(1(t) (yf w (0)) + exp (us(Z, Wio, v(()))
3J is a function of {v x1, . XD) for all j.J E [D]. Therefore, (((t) can also be expressed as a function of {yr ",,".,, y xn) when omitting all non-random components. we have ((()) is independent with y.
By Lemma C.6
((t)
By applying Lemma
E.2
Lemma C.8. If W
W 2,1 v = O. wi. av'v' and vs = bv for some scalar a, b, then !!0 and ge(t) can be expressed as functions of {u(v", xI), ; ... ,(v",xD)3 for all j1. E [D],
1.2
and they are independent with x().,x,) for all i E (2,..", ,d and j E [D].
0 and w(. = av"v", similar to the proof of Proof of Lemma When w w 1,2
Lemma C.6
we can re-write S 5a as
exp (z), 23.)
exp (ax), v'v' x22 + +,, 2.2p).) D Σ exp (ax). v"v X/2 Pj, W 2.2
(!))
ji.j
Σ 5= exp [ W()2)))
exp (ay
X D/ 2)
-== exp (ay(v*
x J/ 2.2P
This is a function of {y(v", x1): ... u/v" xD)} when omitting all non-random components. Be- sides, when (t) 0 and V1 bv", we express ((()) as,
(((() = Lof( w v()) = 1+ exp (u/(Z, W(O, v(oi)) 1 exe (bE ulv", x5) E3'=1 /
Since have demonstrated that (11) is function of (u(v", x1), ,0(v", ,xv) for all j,j
J1.j2 [D]. Therefore, g(0) can also be expressed as function of (y(v", x1). ,(v",xn)) when omit- ting all non-random components. By Lemma E. we obtain that y(v" xx²) is independent with y for all j E [D]. Furthermore, by the orthogonality among v and E...... Ed. we can also obtain that y(v", X,) independent with (E,.x) for all (2....,dd .d.. While for j + j'. y(v", x1} is independent with (E:.x)) since x, is independent with Xy. Combining all the preceding results, we conclude that y(v", xy) is independent with v(E,, xy.) for all i E f2.... d] and j.j [D).
Since (10 's, ge(t) are functions of {v(v*, (x))"""x)"""">"" x1)s"">, "bys""" ,x)), we finally prove that they are in- dependent with all y(E., xy)'s.
Now, we are ready to prove Proposition C.I
Proof of Proposition C.I Since the iterates of (1) and (()) start with (0) 0 and ( 0, it
suffices to show that V ((v(), w(t) = 0, V <(v))(v()) w()) = 0 and Vw,,, ((v(t), w(t))
given v O, (())
1.2 = 0 and wi(t)
2,1 = 0 by inductions. We first prove that Cvv(()( w(())
0. By we have
D Vv w()) = Σ M
D P ye DJ = - 3=15²=1
EBy
g()(()
j' D =0
The second equality holds because ((())(()) (t)
j.j is function of {y X1, . ...xx. xn} for all j.j E (
by Lemma C.6 and Lemma C.7 when v2 0. w(O 0 and w(i 2,1 0. Then by Lemma E.2
have y is independent with (())
y And the last equality holds since Fly) 0. Next, prove
18

---

Published as a conference paper at ICLR 2025
that W1 ((vl), W By C.2 we have
O, and skip the proof for Vwa, C(v(") w since it's similar.
D
D
Vw, [(v w(0))
= Σ
E
ye
(00 of S ,x)),, -x
D
D
3'J S
S (XX) wx D]
ΣEΣ =ij-"
Σ EyjE
(Us) 93 S ( - exz") |p]
S
The second equality holds because (v(1) 23-) = (v1 Xyy when v2 = 0, and y² = 1. The third
equality holds since ((() is function of {y X1, y XD} for all j. " E [D] by Lemma C.6
and Lemma C.7 when v22 0, w 1.2 and w O. Similarly by Lemma E.2 we have y
independent with ((0) The last equality holds since Ely] = 0. This finishes the proof.
C.3 CALCULATIONS FOR INITIAL ITERATIONS
Since we have demonstrated that V2 W w
throughout the training process in Proposi- tion we will focus our analysis solely on the iterates 1 w and w
1,2 2.1
2.2 Before we further
present the global characterization of v((), ( of these parameters in the following lemmas.
1,1 and w we first calculate several initial iterates
Lemma C.9. Under the same condition with Theorem 2.2 the iterates vI W of gradient de- (1) scent defined in 2.3 and 2.4 satisfy that W(1) 0 and vi (1) with (1)
n2
V2
Proof of Lemma C.9
At = 0, we have (v 2 for all j' E [D] since v = O, and correspond-
ingly Vw (()(()
W(0)) = 0. Therefore, we still have W(1) = 0. On the other hand, we calculate
the gradient of ((v(), W()) w.r.t at t = as
Vv,L(v() w()) = -
Σ
B
(0) [
ΣE
E [yx,] - E[xx--]
Or 2nT
j=1y²=1
The second equality holds since (T()) by the fact f(z, v()), w = 0, and S / yo for all
j.3 [D] by the fact W)() 0. The third equality is by the independence between y and x, for j + j The last equality is from E.4 By this result, we obtain that
(1) = (v, w()) = 1H0 v = vv.
2
which completes the proof.
Next, we provide the calculation for the second iteration.
Lemma C.10. Under the same condition with Theorem 2.2 the iterates v of gradient descent
defined in 2.. satisfies that ((2) a (2) with a (2) (D).
Proof of Lemma C.10 By the fact (1)) 0. still have 5/ 1 for all j.J (D]. Therefore can calculate the gradient w.r.t vi
(4)
as
D
D
Vv, C(v(1, w()) = Σ
ΣΣ
Σ
E
y
= y[eW)xx(1] ] + Σ
E E yx,
1=
19

---

Published as a conference paper at ICLR 2025
We analyze the value of E
(((()) yx,- and E
( yxy| respectively. For ((e()) we have
E. yx AA E ( yx,+] AIE p(()) [ x3-) (62, ...) (Ed, x->]]
[ (())
AK
[(v" (E2, (Ea, x,)]
= =B(eW)yvv x,-) )
y xy. E
B B y) E [(s. y(e
where A is the orthogonal matrix defined in The penultimate equality holds since (((1)
1texp (a(1) Ey= /v" wxz)) By replacing y with sign ((v* (xy))), we can notice that yl((1) only
contains the projection of x,, the direction of v", ie., (v*,x)-). Hence by the orthogonality among v and E2. ... Ea and properties of Gaussian distribution, we have (5.,x.)) is independent
with yL(()) for all E {2....,d]) ,d]. The last equality is simply by E (((,),)))) = O for all i e
f2. .2...",, Through a similar process, we can also derive that
B[eWyx, = E[eOy(v" xz)]v'
for all j # j". Moreover, we could notice that y' have the same distribution for all j + j", and correspondingly E ((0""("" (o(vv,x))) xy)] take the same value for all * j" By carefully checking the distribution of ((()) yvv,, xy.) and g()) y(v (xx)")))))))))))) xy) with j + j", we notice that Lemma E. and Lemma E. apply to the calculation can derive that
S[eW)yyv x5-)] and IE E[eaay(v" x,)] with j 7 j We
Vv,((y w()) = (E(e x + Σ BleWDlv ((v))(("
j411
And for the coefficients of v², we have
1
2D
2t
E[ xy-1]
E y( x,)] Jx
2 O(VD),
π
3V
eT
exp
v 2n
11²
$
and
B[eOy(v" x
Σ
E Ble y(v x,)] <
O2
4(1 + no
D
n' 2x
< <Q(VD).
exp
2
110
Applying these results into the gradient descent iteration of v we have
v v() (E(e x
Σ E ye (x)lv, xx)l)r* ====0 =al
where -== This completes the proof.
Lemma C.11. Under the same condition with Theorem 2.2 the iterates w((.) of gradient descent
defined in 2.4 satisfies that (22) fi V The coefficient &n satisfy that |31| S vDD for some negatege constant e₁ solely depending on 17 and d...
Proof of Lemma C.Il We first demonstrate the calculation details of Vw,, w()) Since (1) (1))r" and v. O, we always have (y(1), 2,) = (()))v* in the following calcula-
tions. For Vw,, ((v(), W(I)), we can obtain that
Vw C(v() w(1) B (()) y. ΣΣ
Σ
(1) J
S
( (y (2)) xx)
20

---

Published as a conference paper at ICLR 2025
D D a (D - D2
(1)
D
D
ΣΣ
B B(e y. (v' xx)xxx
ΣΣ
Σ
B y (v"
B
D
We analyze the value of I and 1 respectively in the following. For 11, we can obtain that
(D
D D ΣΣ AE
(2,x,y/
4
y(v xx²'
x(v",,)),
(Ed, A
D2
: (Ea, X,:)
5=1
BAd' ERdxa
We denote the entry in the i1-th row and i2-th column of the expectation matrix B B ii By utilizing Lemma C.8 examine the value of B for two cases: j' and j # jr.
Case I: j j'.
1. B B[ew u(v" xx,"),
2. B E( y(v" " x,)) B[eW) xy)] Xy)] xy) = ( y(v" xx))],
bi.d
for all i # 1.
1,12 = E [ 1 "yll" x,) (Sig x,)] E E| y(v x,) E (S: x,) = 0, for all i2 + 1.
4. 11.1 ( x,) (E Xز}] = E E( x,) B x,)] = 0, for all i * 1.
5. B
11.d E ( y(v xzy (En x,) xy) (E1,, x,)] EE (1)) y(v (EE)) xB) (S x,>| 0.
for all i2 # and i * i2.
Case II: jr.
1. BI = E E y(v, (v" x))].
2. Biy ylv" xy. (E x,)] B[eWD(('" xy) | x,)] 0, for all is # 1.
3. B,J"" ( y(v", (En, xy)(""""'''"v.... xx)] = E[ x (v' (V) xyy] = 0,
for all i + 1.
B
E u(v Xy),,, x = (, =
for all i1, 12 #
By previous discussion of Bi 11 we derive that
== (D 1)
Ble yly" x,)) ",) a - 1)0
ΣΣ
Σ 2
E Bye" x))l,,T
D2 =l
D
1=2
(C.10)
Similarly, for 1g, we have
(v",Xy" xy" (g2))
11 D
1 DE
ΣΣ L
AE
a elyylv
(E2,x)), (Exx)))] A
(Ee, xy")
21

---

Published as a conference paper at ICLR 2025
We denote the entry in the 11-th row and i2-th column the expectation matrix CAd' as C.d J"
And we examine the value of Cy.5 fd" for two cases: j = and j # j"
11,02
Case I: j j-j=) j"
1. C =E(eW)ylv (v", xy)]]
2. Co = E BleWyyvv"
xz)
x,) =
=
(B(LeWyl' x for all i 1.
3. Chr's = E ([( ylv (,, x,y]] = B[eWD) .xy/vv">>)Vv>>v>> xxy)xv,, xx)x [{Er,, = 0, for all l2 + 1.
4. C E ('a) xx) (SI, = E y( (x))(())))v"'v"""" x,)) E x,ر] = 0, for
all 1 1.
5. C = EE ( (En x x,>| (("') | (Ea, xz)) B =
0, for all in, i2 # I and in i2..
Case II: jT:
1. Ch = EE
x,)]
EE ((( u(v"
x/)] E[e(il(v' (v*, (()))
0, for all + 1.
3. Chs'd" E (((2) u(v' xy
xz)] e(e x (v*, xy) (E xyr)
1
0, for all + 1.
Ch E ( ylv' xyy (E., Xy" (Ex)
E5 (((()) x x (E:, (E,,,
0, for all i,, + 1.
By previous discussion of Gi 5" 11 we derive that
(1) D D
+ - 1)0 d D
(D E "
1/ D2 -
ΣΣ
Σ
E yle
Σ
D
2=
(C.11)
Notice that the coefficients of E are all equal in both 1 and Ig. Besides, we define two sets: J1_ {G. j") E [D]; j.J + ij", j" # jr} and Jn ((5.3', J""'""""""""""""" e [D].j.j' $J
j' # j") Then by using (C.10) minus [C111. we get
V (L) W(1) = 1 - Tz
(1)
(D- 1)
Σ
D
Σ
E ",
ΣΣ
((e "yy x vevar
D
J=1\'=1
j=
(!
(D-1) Σ
BLEO ,,,'", ,,)]
E
Σ
( x())) (x" xx)]
I5.1
13.2
22

---

Published as a conference paper at ICLR 2025
xxy xj)] v-veT.
+ (D - 1) E ((e
Σ
B[eWa(vv
(Gd')4
(hd' (1'' (
Ja.a
By carefully checking the terms inner expectation of 13,1, we can utilize Lemma E.12 to obtain that
2 4 + 1mo
an2ad
D5/2
13,1 S
ven
2x g0
D5/2
48TT 5
π
π
Similarly, we obtain that
16
22
Ox
n0
Pot
5/ = 1322 S
16
0x +
n2ot no
47
3
D
4m
4n
n
2t
2T
170
2T
/2T
by Lemma E.13
and
2
13,3 S 6V
0
Applying all the preceding results the gradient descent iteration of W!()
by Lemma E.14 finally obtain that
1,1
w wffl //w,,fC(y's w = =BIv'VT.
And the coefficient 31 satisfy that |&1 S ScvD for some non-negative constant e₁ solely depending on and σz. This completes the proof.
Lemma C.12. Under the same condition with Theorem 2.2 the iterates w 2,2 of gradient descent
defined in 2.4 satisfies that
D ΣpT (p,- - p,))
(2)((_))))
w 2,2 = B2 Σ
The coefficient B2 satisfy that Ax pending on and dx.
$ B2 1 D
for some non-negative constants c2 Cj solely de-
Proof of Lemma C.I2
For Vw Vwa:C(v(1), ((v(1), W())), we can derive that
D D ΣΣΣ ド"
C(v(C( w(1) = E (()) y
S by)) - pp
aD
D
D
ΣΣ
B y. (v x
ELC y. x
D2
D
We discuss the value of I and In respectively in the following. For J, we can obtain that
1 (1)( 1) E y(v" pp)
(L)(D)) 1)
Pj.
p
+
Σ
ELE
E
y(v
(x)v D
ΣpJ Pp.
D2
D2
j=l
While for Ig, we can also obtain that
(1)
D
D x
12
D
E "iv' x
Σ
PM
P
D2
y("
ΣE1 p
j'
1=1
12.1
12.2
23

---

Published as a conference paper at ICLR 2025
( +
D Σ
Σ Σ
yle, xy>lpp
P
D
1"
j=1
12.3
As we discussed earlier, y(v" takes the same value for all j F jr Therefore, we can obtain that
V C(v()) w()) 1 - -1== -h1,1- (1,1)) + (11.22 - 12,3) -12.2
(11) B (B y(v" x--) - B[e ((", Σ (p.- -P.- p.)) ドデ
1
Σ P/
D2
Furthermore, we can utilize Lemma E.10 and Lemma
E.11 to obtain that,
2
4(1+ no²)
1/001 1 2TT
B[e x--)]
B ((( x;)l 0r
exp
no DT
and
yle xxx", - B x1)l
S
Ox
2V
Applying all the preceding results to the gradient descent iteration of W
2,2 finally obtain that
w 2,2
w/2 1/ W.. V ((v(2), w())
(1
(-B[P() y(" + B[eW :xy)1) ,(r",,
EE<:; (p,, - -pi)) ())){pp) 345"
D
= D
= 32 Σ (Dy- - -pi)) (2
ΣΤ,] p.)
341
where
0²00" ?on S
1200
4(1 + n10)
211a0²
32 S
0g +
2
<
D2 4t /
D2
no.
n2D
This completes the proof.
C.4 PROOF OF LEMMA C.5 LEMMA
C.2 LEMMA C.3 AND LEMMA C.4
In this subsection, provide complete proof for Lemma C.5 Lemma +2 Lemma C
and
Lemma C.4 We first prove Lemma C.5 given that the result concerning W in Proposition C. Lemma and Lemma C.4 holds. Then we use Lemma C.5 to prove Lemma
Lemma C.3
and Lemma C.4 by induction. We would like to clarify that this is not circular reasoning, since are utilizing induction. It's reasonable to assume that all conclusions hold for each iteration and verify all conclusions still hold for the next iteration, long conclusions hold at the beginning.
rigorously demonstrate these
Proof of Lemma C.5
By Lemma C.3 and Lemma C.4 there exists constants C1, C2 solely depending
on 17, a. such that
|1 < civ
D and 83
2 C2 Then further combining with Lemma 1.
Lemma E.16 and Lemma E.18 with probability at least
20c102 62
De 28-1o7
VD.
De 2
can obtain that
1,1 + p.J. (t) Z W (t) 2, = X w
2,
D + 1)²
(3((vixy)( xx)
82 + X ((())
em + D. ((
2.2. P
<VBi xx)))) + [
L. 1 ||||,||2||x||/ 1|2 +| 2.2.2oo wap.llllpllle lallplll2lpp |2
24

---

Published as a conference paper at ICLR 2025
2 Do
ecsVD (1
savD
oz
+
130,0
C2 D <
12
holds for j'l * j and E [D], and
2) w/ 2, x W 1.1 x3 + P..
T. W (2)
2,2 Py
31(v" x,- (v x,) +
(D
L)(D + 1)²
82 + x W
1,1, P.. (
2,2,
4
ogD
Bi(v" x (v' الحو
W (7)
1.1
||||||||>>*|1| w
2,2,emor @elpy|lll|ppp1
2
22D
-VVD (V
1330,02 o Do
vDJ² eCavD
D D+1
eCDD
2 D
6
holds for all [D]. Therefore, we obtain that
21 De
De iD
e + (D D)e
$ S 2 e
S
+ (D 1) e
holds for all # j and j [D] with probability at least
26c D3/1 200
VD De
c2T
1 1-e 1-e-Cave -CvD for some non-negative constant Cg solely depending on d.-
by induction. When we prove Lemma will assume that Lemma C3 and Lemma C44 hold current iteration. The same situation still holds for proof of Lemma C.3 and Lemma C.4 As we discussed earlier, this is not circular reasoning by the essence of induction. Rigorously, all the conclusions from these three lemmas could be composed into a big induction. However, for simplicity and consistency, we present them respectively. Besides, we denote E, the event that ((v", xxy>| S D14 and ||xx S az (v for some constant C4 solely depending 1) ax and all j E [D], and denote FG the complement
Next, we prove Lemma Lemma C.3 and Lemma
C.4
event. By Lemma C.5 the occurrence of E, can imply that // 2 D exp(-Cg and s
3's S exp(-Ca for all j * j and j E [D]. And the probability of E, follows that PP Eb()) 2 1 - cxp( DC.
In the following, we present the proof for Lemma
Proof of Lemma C.2 By Lemma
C.11 and Lemma C.12 we have w B V with |31 s
C.D. and w!?
B2 Σ jデン (py- - ps) EP p.
with 02 D S 32 S ca Da aligning
1
2,2
with the formulas of (itt) (!(t) Lemma C.3 and Lemma C.4
We assume the conclusions of
Lemma C.3 and Lemma C.4 still hold for any T", then by Lemma C.5 we have P(E! >
1 exp( ( 11-CCvV CD) Based this result, we define proxy gradient g. which is calculated by assuming (()) = 1, i.e..
c DE(E((Da (Da (x,)) + (t) L. x;-)) xxx])
Besides, since v. is perpendicular v* (which is 0 at 2). it is inner the linear subspace spanned by the . and denote its decomposition v! (1)
V. Eiz a; (1)
E.. By the
25

---

Published as a conference paper at ICLR 2025
orthogonality among $ . we have Σi= (a )
| 1. 1| We also denote Pe the
projection matrix of the linear subspace spanned by the (E2.- Ea). Then, we can decompose
g( by a similar process in Lemma Lemma C.
g. - DE (Da 10) yDa""yvv y(v,, + Dy (c) 1,emor x5.)) ))
DAATB l' (Da (0 y(v x5-) Dy
Σ ai
(1) (Ei, (Ei x3- () [033]]
d DAR (yD .xp)+ + Dy
(0)
(Ei x ))1 x (E2, (xr)") (Ea. x3-1]
ai
1=2
= (Da xy)) + D Σ
x,, ( ")r(", 28-))
1=2
Σ DE l' (Da xx-) D
u(E v() xx.)
And we can upper bound the difference term as
D D
Vv, S(v(), w())- - -200 = =Ir xx,) + Dylvi. I, Xj. x,,)) 1)25
Σ Σ
e y
]
ド'=1ø=
D
<E
DE DE (Da x5-) + Dylviienor x,.1)
(t)
Σ
Σ
(0
yx,, 580 [(En)
j'= j=1
12
E
DE (Da y(v" xx,)) + Dylvi. lemmors x,.)) 1)x
(1)
Σ
yxy. ///// ) [(7)}
3'=1j=1
< DE
8
x5-) + Dylvilo eeron x,-)) p(())
||x, 11 11 (E,)
+E E
D -
Σ
((()
||x 11 2
j=1
D Σ 1=1
D
Σ ド"
1 j",s
[E))
+ DE |x 1 [E:)
+
E
|x
l1e (6:))
Ta
where the inequalities hold by triangle inequality and the fact that 1e') S 1 and [Syryl S 1. Next, demonstrate our analysis on two cases: t = 2 and 2 3.
When t = 2, v( (2) v implies that v (2)
V.ener = O. Therefore, we have
g = DE l' (Da (2) u(v", x,, ()))))(( xxv,))(v"
And we provide the upper bounds for each term of Vv W(2) - g. respectively. Spe- cially, for I1, we can derive that
h
DR
(Da y(v", x5,) + Dyllii.oor x).)) (12)
(1)
|| | (E{)
D
D
E
Da (2) y(v x (2)
Σ
Σ
v(v x/)
| 2 1 (Ei)
|ー11==
D
221 E (D- -S
(
l
D 1a1221
Σ Σ
1| ((E))
j=1
8' 1=1
26

---

Published as a conference paper at ICLR 2025
CD ( (Vd VD) D0
5n
The first inequality is because ll is Lipschitz continuous with 4. The penultimate inequality holds
since 1(v",x;) S czD4//4 ||x||22 s oz(vd + vD). S 2 1 Dexp(-CoD) and و S exp(-Cg D) when ((5)) = 1. And the last inequality holds since a = --(VD). Similarly, we also have
2 ك 0 D(d+ + VD)
and IS DvD* VD)
1
eCoD
CIVD
ECD
For 1a and 4s, by applying Cauchy-Schwand inequality, we have
DVd
1 S D
B x
12
[ET)
S
Cs vD
6
D2
1
1 < D
E |
E
1
VD
Combining all preceding results, we have Vv. ((v(2),
- -.-h S eCaV for some constant
C4 only depending xx and 17. Furthermore, we can derive that
(3)
1 - 1/ - (2) w - 17 Vv. .w(2) - -000
(2)
a( + DmB l' (Da y(v x "() x3 + +vv', (r'ng V, C(v(2), wrel))
+ ripe (((2) Vv,C(v(2) w(s))
(3) v²
(3) iieoorr
For (3), we can utilize Lemma E.10 derive that
(3)
1
D + a - //| V., ((v(2) w(2) - 0 > 10,
2 D; 5
4
a (3) < n0
2
D + a (2) + 1 Vv, (2) w 00001 S n m
2 D.
And we also have [ 2 < n| Vv, (2) w(2)- - () || S e C.VD which completes the proof f = 3.
Since we have derived that (C.5 and hold at / = 3. We will use induction to prove the case when 3 < < T Instead of directly proving C.4), we prove the following inequality
ne
| ecvD'
(C.12)
where e is some positive constant solely depending on and 17. Then we assume C.5 and (C.12)
T - 1. Similarly, we can bound I in [()(()()), w(l)- hold at any 3 < t tts T
c(t)
as
4 <DE
(Da 99 ylv(v" + Dylvi.omor x,-)) - ((())
||x 1121
D E
Da x,,) + f.mor Σ
D D
Σg
(00 + v. x,, S ||x,,
4
j²=1==1
()
D D
+ (a)l))
Σ
/ x|=) ((,)
E
ll |/ا 1 [
4
j=1
j' j=1
D E
D-
Σ
(1) jrز)
|
| 11111111111 (Ee)
+
D
Σ
E
||
12
1l (E.)
4
1=1
8'
1=1
27

---

Published as a conference paper at ICLR 2025
cgo, D t0 (Vd+ + VD) + 403 1.cmorl ll max D}
5eeVD
The first inequality is because E' is Lipschitz continuous with 1 4' The penultimate inequal- ity holds since (v",)))) S D, |xx|la & o=(vd + VD). 5J 2 1 - Dexp(-CgD) (())
and fJJ I exp(-C.D) when 1 [ = 1. And the last inequality holds since (t S
(1)) + DD < eXp(CyD) by (C.5) and our definition of T in Theorem 2.2 For other terms
in D(v(), w(t)- g: we can obtain the same upper bound with f = 2. Therefore, by
combining these results, we obtain that IV,, w - 9. 12 S -VV for some constant only depending on og and 17. Now, we are ready to derive the gradient descent update for ((++1) as (((+)) =v(") (1) V [(v(), w()) = v'v +
(9
Lemorr -ng ' ((v(), w()) - -o!)
= -(a1)) + DnB - El (Da (1) y(v' ,x"',,,+++++ + Dylviiemmn emorr x ()))))((r'. x3 + (r', :n!! - n Nvnl))
a (0) DE E' (DDaalyl x + D
(1) a vl., 1 E: + riPe (P( -vv,C)
(+++)ly +
(++1) Lemor
For For (t+1), we can utilize Lemma E. to to derive the following iterative formulas
2
(1-1) 2 a0) + +nd
(2D4(1))))) m/| /v. w -
(al()))
π
2 "
2 (1) + V
20 D (al1))
2'
and
(1+1)
< <a
nD
2
1
02
n/ V L(v(), w
D2 (a(0)²
2 2n
Sa
π 0 D) (a(t))2
hold for all t 2 3. Then combined with the fact about initialization that 1g V s a( 5
no V
2 D. Lemma E.19 and comparison theorem, we finally get that
(V
2 3n
2nj 2 D3s
2
$ S <a(tt))
2
(V 2 6n (1 2)
2
2 D3
$12
(1 (1- 2) 125 V
2 D3 +
π D
π
π
(C.5). Next, we demonstrate the upper bound for to provide this result, first analyze the coefficient of each &,, and W.L.O.G, we assume a (() 2 0. Then by Lemma E.IT we have
which completes the proof for
|| (tttt 1|2 In order
E - CDD "lV", + Dy momer Xj. xx-)))(v", x,.)
}] s 0,
and
E -
-Da ,xy") + (1 (f
2 oa Lenrrr x,- xx-)))((v x 2 -V e(
-20² w1 l|
60 0
(t)
n
where the last inequality holds by (10) 2 ag
D. Therefore the coefficient of & follows that
d Σ
(t) + nDE - C'Da x + D
1 x
s $1a.9/
28

---

Published as a conference paper at ICLR 2025
by ox 1 3/8 from conditions of Theorem 2.2 Based on these results, finally provide the upper-bound for [|v]).mmrr ll
as
d
| |.
ll <
+ DD
el (Da ylv x5. D
Σ qi a yli, ( E:
1=2
v,, [(v(t), w
17
d
2 u v((x:xx-))))(;xxxx))
Σ
+
nD
(Daaa)((' x5 + D
Σ
+
eoVD
1=2
7
| ll +
n
(++ 1)m
D
e
ecvD
which completes the proof for C.122 Further by our definition of T in Theorem 2.2
obtain
T
1 emoor 12 -
e S e CCVD
This completes the proof for
C.4
Next, we present the proof for Lemma
C.3
Proof of Lemma C.3 By calculations in Lemma
C.11
we have obtained that 1,1 follows (C.6)
with 1,1,
(2)
0. Instead of directly proving (C77
prove the following inequality,
[w()) (t)
L.
n
(C.13)
envD"
for some constant e solely depending on and 0n. Therefore by induction, we assume it holds for 1, and prove it still holds for 1. Actually, it suffices to show that / [(v()) w eevD By C.2 we have
V W1. ((v(t), w(1)
D
E (()))
ΣΣ
ES
Σ
SiS xx) xxx)x]
j=
D
((())
ΣΣ
(0)
(1 S 12) u(vl ,
E (((()
ΣΣ
s
(1)
D
D
D
E
(())
ΣΣ
((vl x 1(m))
+ ((0)
ΣV
sit
(1
( 1(m))
j=1f=1
j=1
D
D
E
ΣΣ
Σ
S
uiv
x 1(B))
E
Σ
(0) (0)
ulv? x 1(EF)
For 11,1, we have
D <
(,''
x + +|Iv l 112)) |x ||al| -
/=1/'=1
29

---

Published as a conference paper at ICLR 2025
8D2 max(d, D} (cgal() D1/4 | 2| (v a vD)
1
4eoVD"
The first inequality is by triangle inequality and (u]. 1e1 < 1. The second inequality holds because
I(v", xz)| S e3 D1/4, ||xy|l2 S VD), S 2 1 - Dexp(-cD) and j' S exp(-cD)
when [ 1. The second inequality holds since (1) < O((T*)1/3 + D) < exp(CoD) and
lyy 2 < e
C.VD by Lemma C.2 and definition of 7 in Theorem 2.2 Similarly, for 12,1,
can obtain that
D
D
D E
- ΣEΣ
E
Sp (vl)((r + Vifemerllll x 1xx 112
5 3-13-113 デリ
9D2 maxfd, (cga D1/4 livfomell (Vdd vD))
ecD
4
In the following, we also denote v by & in summation calculation for simplicity of expression. Then by applying Cauchy-Schwand inequality, we can obtain an upper bound as
x6.xy)
D
D
(Fa.xy)
14111222 =
AE
Σ
S
Jd (1
uyy xx))
[(E,, (E2,xx))) .. (Ea. x,)x 118))
A
j=1'=1
2
(Ed,xy))
(((()
Σ
Σ
/ (1 (1-1)1/)yurt s/y))(() (0) (E, x3) (EE)
E:EY
j=1
d
D Σ
E
(((e)
E
(1 - / Xy)
D
D Σ
27 (3)1)) VP(E.)
E
Σ
(
j's
(1 - /// (v XX). )
P=1
5==1
d
D Σ (Er, x))
E
Σ
(v (t)
x,. (2, )
i= P=1
d
2 )
x )2 ) (5:x)))) )
E
KSrr xلحز
Σ
P=1
j'=1
< (a + (V2orooll22) /
(2)
Σ
E
(3,, v (P((F))
To (a + |
1
CCOV
4eeVD
The first and second inequality is derived by Cauchy-Schward inequality, and the facts that
Σp= (S < 1 and (1 - s .2? e² I 1. The last inequality holds since (t) <
(((T))1/ 1/3 + D) exp(CavD) and I I c -vVD by Lemma C.2 and definition of
T in Theorem 2.2 Similarly for 12,2, we also have
X61)) (Ez, Xpr)
D
1/221/2 =
AE (T(E)
ΣΣ
Σ S xxy)
(Ei,))) (E2, (Ea, x;) 1(27)
AT
2
E.
ΣΣ
(G. (EE) E:EN
30

---

Published as a conference paper at ICLR 2025
D
D
E
Σ
Σ
D
[ΣΣ yー
E
VP((F)
d
D
D
EE
EE X1)
Σ (v x2,))
Σ
(E, x,x))
/P(EF)
="=1
Σ
d
E
KErr xلرو| (
Σ (v xxx))
(Ei, x1))) )
(
VP(EF)
P=
5"=11
D D D D Σ1ΣMMM
d
= =|
(2, xp))2 (P(F))
- 1"=
(a + +1I1 112) eCOVD
4eo/D
Combining all the preceding results, we have
(t-1) 1.
|/
((t)
+ +1 Vw, [(v(1, w(
e.
(
+ +n| 1 112 + n/| 2121122122 + 1/122,112 + 77) 1222222
1,
nt
4n
n(t+1)
e
4e
ecVD
which completes the proof for (C.13). Further by our definition of in Theorem 2.2 we can obtain
( emoor
T'n 2 e
-
S
-CrD
This completes the proof for C.7
Next, we present the proof for Lemma C.4
Proof of Lemma C.4 By calculations in Lemma
C.12
we have obtained that w
(2) 2,2 follows (C.))
with W (2)
0. Instead of directly proving C..
prove the following inequality,
2,2,emor
((())
nt
(C.14)
coVD
for some constant e solely depending on n and ag Therefore by induction, we assume holds for t, and prove still holds for 1+ Actually, it suffices to show that V/V [(v(1), w & e evvD By
C.2 we have
V W (t)( (()))
D
E
(()))
ΣΣ
ES xy)(p pp
Σ
D
(!))
ΣΣ
(
(1 -S- - [
B e(0
s L
j=y
31

---

Published as a conference paper at ICLR 2025
D
D
D D
Σ
(1 - / ( xxy)
Pypp
Σ
Σ
E
e (1 - .x)) [
pyp]
3
h1
112
D
D
E
s (1 - -1-
((t) y S S .x)ll
(x)) p.
2..
For 11,1, we have
D
|
f p (//5) ((1' xx,)) + || ||2| (22 1 (Fa)
E
(0
||py
j=
D2 + (t D2/s D2/4 + llvll.mmlll 21 (Vd vD)
4e
The first inequality is by triangle inequality and |ul. 18 < 1. The second inequality holds because
l(v", x,) S eg D2/4, |x, la a. (vd+v s 1 D Dexp( ex and $ S <xxp(-CgD)
when [(F)) = 1, and |p3112 V(D + 1)/2 by Lemma E.18
The second inequality holds since
a < <0( 1/3 + D) < exp(Cg and [vi.nnee l|g Se Ce-CDD by Lemma C.2 and definition of
T in Theorem 2.2 Similarly, for 12,1, we can obtain that
D
|221112 s
Σ
E
(
((t) j" (v', xy) + 1| 1emmellllxy/ / |1 ||, |P22
/=1ド==
2D + (c3a DJ/4 + || L (VV vV)+VD)) 1
<
DCo
4eaVD"
By applying Cauchy-Schward inequality, we can obtain an upper bound for I as
D
D
1121122 =
Σ
E
e
S po (1 - (11-))))utt Jyvvi xy) (Epps pypp
F=1
D
Σ
(1 Divv xy
2
| 1
2'
(D (+1))P(EE))
D Σ
D
(s)))) (v (v x,)33
2
j=1ド'=1
(D+ (1+1)(a(0_ + +||v (12) v P(F))
D E((v() /ll11'
Σ
2
3d + 1)D (a + 1/v10 112)
DeD
4eoVD
The first inequality is derived by Cauchy-Schwarz inequality, ||p:1|,
V(D+ + 1)/2
by Lemma E.18. The second inequality is derived by Cauchy-Schwanl inequality, and
Σラ= (S
s 1. The last inequality holds since (t)) < <O((T))1/3 + D) < exp(CsvD)
and [[ |
I E CaVD by Lemma C.2 and definition of T in Theorem 2.2
Similarly for 12.2
we also have
D ΣΣ
T2,2112
Σ
B S
(EE P
32

---

Published as a conference paper at ICLR 2025
D
D
erOS Pysp x 1 ||y 11 1
E
DD 11+1)
Σ
D
D
E
(v)
xz.)3
Σ (5 })
2
j=1 y==1
8"
(D 1)(1)( + [ (12) v P(E))
D
Σ
[( (1) //Iv2"122.xx)) )4]
2
3t0z (D- 1)D$ (a || l12)
2eD
40orD
Combining all the preceding results, we have
(t-1)
2.2.
2
(t) W
2norr
2
+
+1|
V w fl12
(
2,oor etrrr
(t)
+1 11.1 11 + 11221,2122 + n/|22all2 + +7] 22,2122
+
2
ntt
4n
nit 11)
eovD
4eeVD
eovD
which completes the proof for
C.1
Further by our definition of T* in Theorem 2.2 we can obtain
wi!h)
2 ecVD S -CND
T'n
2m2
This completes the proof for C.9
Finally, we provide the proof for Lemma 4.3
Proof of Lemma 4.3 By Lemma C.2
and Proposition C.1
we can re-write vf W(1), v(t))
(Z, w()), v()) (1)
Σ u xy)
3 ( L Xy xy (5
By the fact that Σp= (s(t)
2
< 1 and Schwarz inequality, we can get a worst-case lower
3'
bound for y (Z, (()) v(t))a
D
D
()
yl w(0), (())
2- (1)
Σ
Σ
(v
x²)
Σ
j=1
/'=1
D
D
- ||v 11/
L.erree
Σ
(s
(1))))
||
v.eerelll2 2
(t)
1=1
y'=1
7/4
D
(t)
))
2 1
2 - (t)
x
- DD| -l fomorll
Viemr
|| Ilv.fmoolls ll2 xx.
D Since (v*, (1) v = O, we can derive that ΣB=1 xx)) and
v
2 xy
are i.i.d.
[ Te
random variables and J/ Ey-i x
57 -1 liv
1
Vi
X3y)
x2 by the properties of
Gaussian distribution. Besides, when E, holds, we also have
D
(Z, w v(0) 2 "00''v' xx,) E
(0)
|x
112
Σ (1)
||
Σ
X
1|
j'j
/=1
8'
j=1
j'=1
33

---

Published as a conference paper at ICLR 2025
D
D2
ozD(Vd+ VD) eoNVD
2 ()
1
,((",xx))
cCoD
eCoD
D v((
2 Da
1 ccD
The second inequality holds because Eg implies that S
1 De De C.D s
C.D and
[xxy lla S on (v + vD) by Lemma C.5 and Lemma E.16 Besides, Lemma 4.2 guarantees that
[ ll S eonV. The last inequality holds because (0) I 0((77)1/3 + DD by Lemma C.2
which is much smaller than eCuD by definition of T in Theorem 2.2 Therefore, the last two terms are much smaller than 1. Similarly, we can also obtain that
lZ, W(, v(9) S a u(v", x5-) ES +a (1)
Σ
1
(0) 3' Σ
Σ j'=1
xxy 112
Es
2
8'4j"
5-1
jm
D0 y(v" x5-) 22a D2
(vd+vD)
DzD VD)
DCDD
cCoVD
s Da x
This completes the proof.
D PROOF OF THEOREM 3.2
Before we demonstrate the proof for Theorem 3.2 we first introduce and prove several lemmas which will be utilized for further proof.
This section provides complete proof for Theorem
3.2
The following two lemmas are very similar to Lemma C.3
Lemma C.4 and Lemma C.5
Lemma D.1. Under the same conditions of Theorem over the randomness of (()) i-1 it holds that
3.2 with probability at least - -CyDD
TT
W
B2 EjEs (Dz- p,)) (E), p.)
+ W
(D.1)
0
where 13:1 < CI Da B2
2
C2 D2
for some constant cı, 02 solely depending on y/ oz, and the error
term satisfies that
error 112 S exp (- C1oVD).
(D.2)
The coefficients C3, C10 are both positive constants solely depending on σ2, 0 and 77.
Lemma D.2. Under the same conditions of Theorem 3.2 then with probability at least 1 ne CavD over the randomness of (x()) ())) it holds that
jjj > 1 - -Dexp(-CC1))
s $/ S exp((CCiDD
for all i E [m] .j E [D] and j' * jr The coefficients C3,, C11 are both positive constants solely
depending az On and
Proof of Lemma D.2 By Lemma D.I there exists constants C1, C2 solely depending on v), ar such
that |81| < ci D and B2 2 D Then further combining with Lemma E. and Lemma E.
with probability at least ne
CIvD we can obtain that
1D+
1)2
(0)))
(4) (8) 2
(v
B2 + +Z)
wii error
2
4
< (v', x
+
(1)
12 11221112
12
34

---

Published as a conference paper at ICLR 2025
+ (v + vD)²
S savD
13 o.
D2
13
e
D
12
holds for all all $ j j E [D] and i E [D]. and
(2)2) T w() = (i(v" vv' ()) +
(D - (D-1)(D + 1)²
32 + (z.
(3)
T
W
4
error
2 c2D Bi(v" x (1)) (v" ,x,
W
12 2 112111 1 ll2
error
2
2
vVV /
02
20 Vd + VD2
13e00 o
CIoVD
C D 6
holds for all j E [D] and i E [D]. Therefore, we can obtain that
s j'f e
2
+D 1) 21 - De
- 3D
D
/ <
Se
et + (D
2
+(D
1)e
This completes the proof.
Proof of Lemma D. By Lemma C.3 and Lemma C.4 we have obtained that W(1) = W(T")
follows D.I with [ emor ll S max
[||
(T*)
(T")
2,2 [la] S exp - OrVD). Instead
L.rrr 2
of directly proving (D2]., we prove the following inequality,
[ venorlly 112 S exp ( CrvD) + exp(cD)
7p
(D.3)
where e a constant solely depending on az Tw, 17. Based on our previous discussion, we know holds for 1. Therefore by induction, we assume it holds for i S n 1, and prove it still holds
for + Actually, it suffices show that (((0), w S cc And we firstly provide
an upper-bound for | 112 as
11
e 1000022000 .Z S0
1'=1
< (v2ō,= + 1) VD) < (V2ō, + +1)) VD)
n
The first inequality is because |, (V2o,- +1) (vd+ vD) by Lemma E.17 and Lemma E.18 The last inequality holds by the scale of îj from the conditions of Theorem 3.2 Based this upper- bound, we can further demonstrate the upper-bound of VwC,(v() w Then by CC. have
we
(V(O) w())
D
Σ
Σ
D D
D D
{1)
Σ
2 (1- (1- 2606))))))))
p())
ΣΣΣ
ΣΣ Σ
J=1 11 f"
35

---

Published as a conference paper at ICLR 2025
For I1, we have
| EE syy) 1 ( |2l| 11l (1) 11 - 2(V20, + Dm)
D
Dm (on VD) 2 1
CCD
2eeD
j=1y'=1
The first inequality is by triangle inequality and |u|. |e s 1. The second inequality holds because || < (v2o, + (vdd + VD), و 2 1 - D exp(-C11D) and / I exp(-C11 by
Lemma D.2 and
and upper-bound for [| 2 The second inequality holds by the condition concerning
n from Theorem 3.2 Similarly, for 12, we can obtain that
D
2(V25 + 1)² Dm Dn( + vD)²
1 2 S EE
Σ
S
1blk 12l 12
S D
eCiD
j=
Combining all the preceding results, we have
W error
l12
W s 2 + /|| (w(( w
W emor 2 + 77 11 1| + 77 12 2
27
+ 7(
S exp (- (-CVV + eDD
+
exp ( - CVVD)
2eeD
ecD
which proves that D.3p. Then by conditions concerning n from Theorem 3.2
we further have
[ Worlll S exp ( (CCVDD + ecD S <epp(-CovD).
Tp
This completes the proof.
For further proof in this section, we introduce the following several notations. We denote V argmaxv.llvllas1 minggjn) y() (v, > and v [( [c VT) OD 7 E Bd+DD And we define W an idealized matrix such that s/n = and s /6 = for all j' j and j E [D] (Such a matrix exists by Lemma E.18 Furthermore, we define a proxy loss by -l' as RV, W) = -( f(z(), (1, W, V)) and R(v, W) = E {xx,)>= (y f(Z, W, v))]- Based on these notations, we can finish the remaining proof. Firstly, we prove that (V, W) is nearly convex.
Lemma D.3. With probability at least ne Cav/D it holds that
L.(V, (V(( w()) (V.L(VV) V- 7e
1
for all E [n], where C12 is a constant solely depending on oa Tz 7)
Proof of Lemma D.3
By convexity of e, we can obtain that
LV, W) - (V(),, W())
-(2 f(Z (), W, V)) - ( f(Z() Wo, v(0))
(0 (Z wi ()))))(()))))))))))))))))) W. v) 1(2(), W(1) v(1))
D
-0
Σ
S
(v 50, -00)))))) ,()(()))
D Σ
(v. )
(())
Σ3
S (
35 (v,
1=1 f==1
y'=l
1/1"
= -VV.L,(V() W (7,.-((1) + ((0)(0))
(D- Σ
(i)
(i. 2 + g (
Σ Σ
S 5/0 (,, (8)
j'
y'=1
And we can bound
[T] as
7 < D
10-5s/ ='=1
Σ
||
+ Σ
S 9.
11/ 12, 11.
341
36

---

Published as a conference paper at ICLR 2025
2(V2ō= +
1
yeCu.D
yeCizD:
where the penultimate inequality is by Lemma D.2 This completes the proof.
Lemma
E.17
Lemma
E.18
and definition of V.
Based on the preceding Lemmas, we are now ready to prove Theorem 3.2
Proof of Theorem 3.2 By updating rule
3.1 we can obtain that
[((+1]]
112 = [(0) 112 + 2 ((+)) v(0 7(0 - v) + [|
1| - 2
[(0 - 41/2 27 V.L.
v(0) - B) V L. 1
< [4 - 2/l + 27 LV, W) L,( W()) + 7e
1
+ 7 [V((
[ 3/6 + 27 LV, W) LV(()) W()
1
+ - W()
ye
112 (.(( w())
27
2π
(D.4)
7e CraD
The first inequality holds by Lemma
D.3 The second inequality holds because
D / W 1 - 77 (160 f(z vi, wo)) )
S ( $0/
S < F(Z W()) )) (V26, + 1)² (Vd+ vD) < (T, W(O),
where the first inequality holds by the facts that (x) S ((x). -((x) S 1 and |2] 12 $ (V25
1) (Vd++ + vD), and the last inequality holds by our condition that 7j S (V26. DDV And the third inequality is because
CV, W) ((y S W. v) = =6( |log n
y (V' (0))
< Ellogn) S -
7
By rearranging and take a telescoping sum on both side of D.4). we obtain that
CVV w( s || - v12
2
2 <
2( 2(v26, + 1) log" n((d + DD)
2 2
CIZD +
+
T
nn
n
72n
D0DD
n
Further by applying the fact that 7 S minseln] can get
|xxx S 22 ( (Vd+ + vD) and -f(x) S ((x)) we
W S 1 Σ (v()) w()) < 4(V26, + 1)2 log² n(d + D)
2
+ CiD
n2n
1=1
(D.5)
Notice that the quantity Σigs (R((V()) W()) - R(V(), W forms a martingale w.r.t. the fil- (111, 1)) (k), The martingale difference RVO,, W(0)- -
... ... R(V(0, W() is bounded by since -f' is bounded. And we also have
tration of ((Z(), yy(")
E (V0), W(O)- - R(VO) [000
=E (V00, w( [00 - -2R((() W() w0)oo(01 + R(V W(y)
-E
-E(R, (V!0) woiy - -R(V() w() < B ( (.00 w - R((V)) w(0 S R(V(), ((((),,(.))
37

---

Published as a conference paper at ICLR 2025
Therefore by Theorem 1 in Beygelzimer et al. (201) (a version of martingale difference concentra- tion inequality), with probability at least 1 for any positive 6, it holds that,
(R(V W - RVO, W())
e 2
E (R.(v()) W - R(V W())) 2 [000 +
log(1/0)
e 2
R(v(), w log(1/6)
Rearranging on both sides and applying the results of D.5 we get
+ (log(1/0)
Σ R(V w()) v RV()
==1
16( (0(v22_ 1)2 log* n(d D) 4 log(1/0)
8
7
CraD
Since ((x 112, W V 0) 22 W, v) holds for any v and W, we finally derive that
Σ P PX (v. f(z W x(0) 20 1
R(V W o
log "(dd + D)
O
(1g((10))
log(1
2
which completes the proof.
E TECHNICAL LEMMAS
E.1 INDEPENDENCE AMONG GAUSSIAN RANDOM VECTORS
Lemma E.1. Let 21, 22 be two Gaussian random variables with zero mean, and y = sign(z1). Then is independent with y and 22. Moreover, y also follows the normal distribution, which has zero mean and the same variance with 22.
Proof of Lemmal E.1 W.L.O.G, we assume that 2 22 are standard Gaussian random variables. We first prove that y is independent with y 21. is clear that 21 = |21]- Then for any x 2 0.
PPy 21 x) P(==1 S x) P(==1 < xly):
which indicate that 21 1 has the same distribution with y- ly.. Therefore, y is independent with y- Next, we prove that y is independent with y 22 and y Z2 follows the standard normal distribution. We denote (--) the cumulative density function (c.d.f.) of the standard normal distribution. Then for any x R.
PP(y 22 xx) = =P(22 S <xy=1) P(22 2 -; ---y-1 1)
= P( S x) Ply = 1) + P(22 2 -x) Ply = 1))
1 (x)) + 2 -1 --=)(())) ((x).
2
where the second equality holds by the independence between and 22, and the last equality holds
by the symmetry of standard normal distribution that ((-x) ((r). This proves that y 22 ~ N (O, 1). Moreover, it is obvious that
IP 22 S xly = 1) F(22 < æ = 1) = 4(x),
and
Ply 22 S x/y = 1) = P(z2 2 - 1)
=1 =1
(<-x) = ((x),
which indicate that y 2g has the same distribution with y 22ly. Therefore, we complete the proof of Lemma E.
Lemma E.2. For y and X1 X2, " defined in Definition 2.1 it holds that y is independent with y Xj for all j [D]. Moreover, y X; N(O,, L2 for all j 7 j".
38

---

Published as a conference paper at ICLR 2025
Proof of Lemma E.2 For simplicity of expression, we assume 02 1 here, and the proof for ox # is the same. We consider the following two cases: j + and = =j"
When j * j" y is independent with x,. Then for any l [d]. we have y X [k] (N, 1)
and y/ xy xy[k] is independent with y by Lemma E.1 Therefore, y xj is independent with y since each coordinate of y xj is independent with y- Moreover, E[x,[A [k] x, [A'] = for any k + k', which finally proves that y X; ~ (2)
have E'y x, [k] y x3 [M]]
When = j" let A be an orthogonal matrix with v being its first column. We denote E ... Ed
the rest columns in A, i.e., A = (v"E,, .,, ... ,.a]. Then we have
y X AA Xj+ = y. [v", E.. .Eal [(v" xx)), " (Saxx;;))1'
= (v' v' + Σ (E) &
By orthogonality among v", -,...., ,dd. we obtain that (v", Xy.) and (E,, x,.) for E (2.... d} are all independent random variables with standard normal distribution, which further indicates that y is independent with (L,, xy-) for all E {2, (2, .d. Then by Lemma E.I we can obtain that is independent with > 2i=22 V (E,, v.. Besides, note that y is the sign of the standard normal random variable (v", Xz+). From Lemma E.I we have that y is also independent with y (v", xy.). Hence y is independent with y XXy-. This completes the proof of Lemma E.2
xx")
. Lemma E.3, For and X1, X2, , defined in Definition 2.1 it holds that y-x1, V- VyXD are mutually independent.
Proof of Lemma E.3 W.L.O.G., we assume j2 * Jr Then by Definition [2.1] Xj, is independent both with and Xj hence it's independent with the product y Xji- On the other hand, by applying Lemma H.2 y X is independent with Therefore, we finally prove that X is independent with y Xj2
For any j1.32 [D] with j1 * f2. then least one of ji and ja is not j"
E.2 CALCULATION DETAILS OF EXPECTATIONS
Lemma E.4. Letx (0, a Ia), ) for any fixed V E R&,, holds that E [xsign ((x, v))] oz /
[
-..... .E similarly with Lemma E. Then by similar method, obtain that,
Proof of Lemma E.4 Let A =
//
El be an orthogonal matrix, which defined
sign ((x, V)) sign ((x, v))=
sign ((x, v)) (5% x) E.
| 2
Note that , x) and (E., x) for i E (2.... d) are i.i.d normal random variables with mean
and variance o2 therefore we have
E sign ((x, } E sign ((x, V))
V M
x)
2 oz
v 2
π
This completes the proof.
Lemma E.5. Let 21 ~ N(O, of) and a fixed scalar a, then holds that
1. If < O, then
2
[[exp(c])11) < min {2,
2 1 v
max {V
_ oia
-o
0ja
2. If a 2 0, then
-703 2
<
[[eep([||=11] < 2e
E
39

---

Published as a conference paper at ICLR 2025
Proof of Lemma E.S The probability
density
function for 121 is f(al(())
2-=x" (220f)))
(220}.
By this density function, we can calculate
E [[exp(a|z11]] as.
1
!
tax dx
1 01
1 /
2 of
(= dr 2e 2 P(z -o1a),
of
E [xxp(a|211]]
-
01 V
2
where z is a standard Gaussian random variable. If 2 holds that S P(2 2 -01a) s 1. Then we can obtain that
< [e(ep((alaa)) < 2e
When a < O, by applying the tail-bound of standard Gaussian random variables and Mills ratio simultaneously, we have
SP(z 2 -01) se Lia"
1
min
2m
-1a
-0/a3/
2tT
0ia
Therefore, we derive that,
2
2 1
S E [xxp(a|aa1]]
< min 2,
01a
-osa3
-011
Besides, we also have
(-01041)2
P( 2 -01a)
e dr
22T -01a
22T
which further implies that B[ [[exp(a|z]]] the proof.
2
V
me
Combining all preceding results, finish
Lemma E.6. Let 21 ~ N(O, of) and two fixed positive sealars a, b, then it holds that
2
1
6² ab
)
E [[exp(-a/aa)
π
σia
ora
σia + 6/01 exp
207
Proof of Lemmaa
E.6 The probability density function for |21 is fx²
2 ((2o))
1220)-
By this density function, we can calculate EE[ exx((=ala 11
as,
1
2 of 20 dr T V a
of22 1- (
E
dr
01
Jo
ofa -
2e
(P(2 0la) P(= 2 aia + b/o1)).
where: = is a standard Gaussian random variable. By applying the Mills ratio, we have
P( 2 -1)a)= 2
V2T
e
01a
o7a3)
and
P22 2 aia t/0)
/ b/oo))
Therefore, finally derive that,
E (expp(alaal))aamp
2 π
exp ( ab
20 )
12
σia
aia²
a1a + b/o
σ1a
which completes the proof.
Lemma E.7. Let 21 NOO, af) and fixed scalar a, then it holds that
40

---

Published as a conference paper at ICLR 2025
1. If 2 0, then
1b (01 S
<E[[z1/ [xpp///xppaa/ai1)]
(E.I)
π T
where 1b = max @i 20ja,
via
2v
01a
0.a
2. If a < 0, then
sia²
2
ofu²
o ofac
E
21
exp(-alz11) S 01
20,
(E.2)
Proof of Lemmal
E.7
Similar to Lemma
E.5
calculate E exp( - 1)
1
2
E |||=|| exp( =1
20
dr
01
2
a
20
-22 +o1a) dr
e
DE
2
=(otoaa) dr
01
n
2
010 (o c c-2(0to10)) dr / (o (x+ - dr
)(0+010)2
π
2 π
oru² P(2 )10)
20
(E.3)
where z is standard Gaussian random variable. When 2 0, by applying tail-bound of standard Gaussian random variables and Mills ratio simultaneously, we have
1
off P(= 2 mi)) se 2 min
1
1
3
01a
12TT
01a
oga²
ofa5
(E.4)
Applying the result of
E.4
into
E.3
finally get
}}
2 2 20aa,
3
< E|a exp((=pp(]])))) [[[xxeesp(-00 s
$
2
max 0iV
V
of
(E.5)
Besides, we also have
2 o7a²
E E|aa|eep((aa/a]]
1V
-6
re ((+ dr
π
2
o7a²
xe- (((+oia) dx
σ1
01
(E.6)
1/2
2
Combining the results of (E.5) and (.6)) we finishes the proof of (E.1) When < it's obvious
that 1/2 < P(2 2 σ10)= < I. Replacing this results in (E.3). we finish the proof for (E.2)
Lemma E.8. Let N(0, of) and a non-negative fixed scalar then it holds that
m 2 (a)21)))))] < 20 =
Proof of Lemma
E.8 Based on the density function of |z1/, calculate that
E 2 exp((a=21)]
2 00
ze 207
+az
dr
41

---

Published as a conference paper at ICLR 2025
2
x²e -t dr
01
π
(x (0) (0e - dx
_
-aa
ofa2
B[(2 + ao) = 20 a² 1)e =fa²
$ 20
where: 2 is standard Gaussian random variable. This finishes the proof.
Lemma E.9. For any Gaussian random variable with mean 0, it holds that El-[((z)] 2 1/4.
Proof of Lemma E.9
It can be derived that
1 2 B [
(=51 -
- - which finishes the proof.
B[ - (=)) 2 B[(
2
4
Lemma E.10. Let 21 N (0, of) and 22 (N, o2) be two independent Gaussian random vari- ables, and a, b be two scalars. Then it holds that,
1. If a 2 0, then
1b - + b22))2a1] S ubi((1, (0,, a, b);
(E.7)
where 1b1 (01, = max { /o 20fa.
or c
3 ofa"
as defined
2/2ex
010
a2 in Lemma E.7, and ubi (01, a, min 2
5,01 2
2. If < 0, then
01 2 - S (a][(all))++)aa))) S saiv= 2
(E.8)
4
Proof of Lemma E.
By the law of total expectation, have -E[E b=22))))))))))
E E
((aazzl bzz)) 22] Since 21, 22 independent, we still have 2 22 NNO, a1). We first
prove the case for a 2 0. The lower-bound can be calculated as
|2g|
E E - (a(2zl + b22 (21 [2] =
22
+ exp
(a|z1
+ bz2)
121
E
exp (a|21))
+ cxp (a)23 + bz2)
22
> ( exp(px/aaa) [xxxp(])]]
lbi (01, a),
(E.9)
where the last inequality holds by applying Lemma E.7 and Lemma Lemma E.7 On the other hand, we have
E.9 and lb, (a), a) is defined in
|2
- E[ b22)) = E
<
E
+ exp (a/2a) + b22)
exp (a 22
bag)
2
| exp((a//21
exp ba2)
V
(E.10)
where the last inequality holds by applying Lemma E.7 and moment- -ew.ennttngeeeetingggg function of Gaussian random variables. Besides, we also have
/
2
-E [al(all bz2)(=1[]
< BE [
(E.1I)
42

---

Published as a conference paper at ICLR 2025
Combining (E.9), (E.10) and (E.1Ip. finish the proof for (E.7). Next, we prove the case for O,
and (E.11 still holds for a < 0. Besides by Lemma E.9 we can obatin that
-B(e (a|22 b2g b21)
[[[((ba>)
01 4
(E.12)
Combining E..1 and 12 we finish the proof for
Lemma E.11. Let 21 NNO,, of), x2 V(0,02) and 23 N(0, σā) be three independent Gaus- sian random variables, and a, b, e be three non-negative scalars. Then it holds that
1by(01, 03 b, E (ale) bez + cza)zal S ubz((11, 02, a,b,,
(E.13)
where
co3
03
2 ozz afuP-a3e
2 oze
-20 2
02,(01,02,03,0..... max
21
4
v 12
π o2b
π 010
and
3rozb (V -
max
{ 2 (+t aia
{
(
1 1
oru
20362 2032
Proof of Lemma
E.17 Similarly, by the law of total expectation, we have -E EE + baz + c2)) = E (a|z1| + bag + cz3)23| 21, =2] Notice that - is a decreasing function and
2g has zero-ceree ceroecentered symmetric density function. Hence, we obtain that E (a(1|21 + bag + - 0. Specifically, we have 21, 22
cz3)
23
E
(a 21| bzz + cz3) (2g)23 23 21, 22]
[a(a| 1 + bzz + cz3) 23 1 (zs20) 21, 22 + + (a|221 + b22 + c23) 23 (x) 22]
E (a|=1 bb2 cza) + ll (( |2|| + bz2 + eza)) 23 (x,20) |1 #2
(E.14)
Since L' is Lipschitz continuous with 1 and the value of - l" is always in (0, 1). we have
E - l'(a/21l + bag c2g) (a(222) + b2g c22) 23 (=020)
21,
22
2 min 2mmn [E|c2]1 (zs20) 21 22
23 (z>20) 21 22 max {
ca3
22 j (E.15)
03
4
On the other hand, we also have
(a/2221 + bzz czs) + (a| + bzz + cz3)
calzittbs (6 -053)
- 1+ colaal+b22
+ calai|+begten
e1 =0l11/b22220 + calz1l+b2
aa 22 e e22))
(c01es) colzs|+ = = -//1
C - -2z
(c0121 c0/21|+b22 0z3
GCt e- + c c baz
c00 c2g
26- -a/al
(2/21/c) 1 (12a|51/6)
6o3g + -0-0s + bes + bez
6 e- e-
-aaa
1 (x>21/c3 (12)))1/1)
[ -
(2211/) (1aals1/6)
(E.16)
2(e + e-1)
By applying result of
E. 16 in E.14 and Lemma
E.5 we obtain
E
El (a|211
bag
c23)23
--
Ele
(221) 1/6)
3mozb
max {V
{/2
2
01a
c
-1a
2
0ia
1
07a3)
)
20
203
(E.17)
43

---

Published as a conference paper at ICLR 2025
x² 203 dx
where the last inequality holds because El 23 (2521122//e)
S1 de
aa
26t03
= 03/2n
and P(|z2) - 1/b) o2ver S-1/0 e 20 dr 2 b02v22
2
20102 Similarly, we also have
(a/211 bb2 c2s) + (a + bog c2s)
1
calzx|tbez (e c
1 tba
+ eat 2i +bzz-aza
(1 + calzs|+hs2
(1 + en 462gt22
(c c c=s) (x3)
c c
=ca/21l
(1 (12-02)) (1 + cb22+023)
ecoa + e e
bzz + e b2g
Sealai1 c c
د e bbb S e
b/22| ( e cRs))
(E.18)
By applying result
E. 18. Lemma E.5 and Lemma
E.7
14
we have
B (a)221 baz c2a)zs]
-E(ea==1]] T blza B
(e
2 oze o2a240302
22 o2b
(E.19)
Moreover, we can also obtain that
e + b22 cz2) + (a 21| + bog c23)
1
(e c00ss)
+ ealz1|+bz2--a
1 + ea 21 +bz2-ora
(1 + ea|z1| czs
+ ea ea tbast
(e
e c2a)
alale bzg c2a
(E.20)
bzz -33g
By applying result of E.20 Lemma R.5 and Lemma E.7 in
we have
E[ f' [a|21] + ban + c23)23]
2 Eeee-
-a/=1| E
E
(e
[ss20)]
oze
-3.0 -2
(H.21)
Combining E5.15
E.17
E.19
and E.21
we finish the proof for E..3).
The proofs for the following Lemmas are very similar to that of Lemma E. while we still include them here for completeness.
Lemma E.12. Let N(0. ~ N(O. oz). N(0, og), and 24 independent Gaussian random variables, and a1 02 0z a₁ be four non-negative scalars. Then holds that
N(0, o be four
lby (os(O1)00 02 03, 01 01, a2 az a2, a1) -E[['((1122 02 + 4323 2 s ubs abb( 02, 03, 04, a1 a2, az, a1): (F.22)
03,04,
where
010302
2
o0 1) of
lby(01,02, 02, 04, 02, 01) = max
4
e
4
V2m
π
0202
and
udg (a), 02, 03, 04, a1, 022 03, a4) 1927332a33020
0
max -01a1-10 1/2
"8020] "eggun 202-3.
0101
ofa
Proof of Lemma E.12 By the law of total expectation, we have (a|21 + bzg + C23 + 234] E B (a|zıl + bas + ezs + dz4) d1,) 22] Notice that -6' is decreas- ing function and 24 has a zero-centered symmetric density function.
Hence, we obtain that
E El-a'( + bag + c2s + d2))2a 21, 22, 23] < 0. Specifically, we have
E ( | + 22 + a3=s + a424 24 21 22, 2a]
44

---

Published as a conference paper at ICLR 2025
=B{ (1 02 + a34s + (=120) 21 22, 23
+B[ (aii 02 + 0323 + (4 24 ( 21, =2, 23)
E (11| 0222 0323 a1z4) (a1 21 + 02 + 0323 + aa22)) 24 (2120) 21 22, 23
(E.23)
Since -6' is Lipschitz continuous with & and the value of - ff is always in (0, 1). we have
- B ( 0222 0323 a1=1) l' (((11 |21 + 02 + az23 + a1=1)) 21 (z:20) 22,
2 - min B[aazq (2200//2122 E 24 (2, ž1 #2, 23] = max { { a40 2
04
V
Then we have
- ([ + 0222 + 0323 + (424)2z24) 2 max
0407
04
R[=2.
max {
a40307
oz00
4
/2T
4
/2T
(E.24)
On the other hand, we have
[' 21 + 0222 0323 a424)= + (a11/11 + 22 + az23 + a124)
1
1 +0
+ + _0l1
Ga1 +02 (c - c44444
+ (1+c00ll-+a22to0s--=aa 240
2 (1laa -424)
(enilal +2// +03333333)))))3333--33 (en)
2 -4424
0424
-n/al
ga424 +
20 +
2222
Ze -q1/ai
-
(222)) (1aal 2
g0424 + 4 20232
-222
0222
0453
-a/al
2(e + e-1)
(2
(21 (1=215 20g
1
(
(E.25)
By applying result of R.25 in E.23 and Lemma E.5 obtain
E - (ar//1 22 + 2z2s + 2424)21
S -
E a |2
E
(2 2,
[
-100 20g
P((k
P
2az
04 max {e 1-----1112
sogus
2020d
<
(E.26)
19273a2a30203
01
22
where the last inequality holds because R. (2211/0c)
J//a, xe
20
dr
24401
2n
/2
E
Zug
20 dr
80gu3 and P(
1
aa
Zug
320
202
2
20g
dr S 20
3o303 Similarly, we also have
22
-((a1/211 + az23 a4z1) + e' (a1/21) + 0222 + 0z23 + 242))
1+ 14le 1 + cu..
@ 4444 --44
7 == (1 |
45

---

Published as a conference paper at ICLR 2025
+012: (c C-0424)
c -0444
c01/211
(1
00252tta33 -424 @02s2-tassstaaz)
4424 + c 02 -322
6 8444
S <eailzile l
(c ( < e ( a
2 + -2-
(E.27)
By applying result
E.27
Lemma
E.5
Lemmal E.7
and Lemma E.8 (E.33). J, we have
El-((an1 + 0222 + 0323 + 6424
2 -B [[ [ [23 aall |eal (e
(
- a + 1) to
0202
(E.28)
Combining (E.24), E.26 and E.28 we finish the proof for E.22
Lemma E.13. Let ø1 NO0, af). 22 NO oï), 2q oz), 24 N(0, and 25
(0, og) be five independent Gaussian random variables, and (1, (22 d3 04, 45 be five non-negative scalars. Then it holds that
([1'(a1 - 0222 + 0323 + 0424 0525)
2
of
ai
+
ris to
O202
2m
1=3
2
B[(((a1)a1 02 2323 2124 a5 232125]
S
+
a,e
T
0202
2t
(E.29)
Proof of Lemma
E.13
By the law of total expectation, we have
B(
02 + 23 24 + 2525)
E
(a
1|||=+++0000 + + 0323 + 04 +
24. 26]
E
B (l-2((ai121 0222 + 0323 44 o5 23
|,3 (
24, 26
+ ( - a222 + 0323 + a4d4 + a523) (0) 24, Ø6
Since -L' is a decreasing function and 25 has a zero-centered symmetric density function. We can conclude that 1 < 0 and 2 0, and
1 < -EL (a1/<1l + 02 + as²s + a4 + 05 S 2
Next, we provide a lower bound for ① and an upper bound for 2 Similar to the proof of Lemma E. have
1 2 - aul B
22|22| E 23 laal
+auleal [2022
E 250
[s0
- 1 |ell
E
E Zze R3
|>>20)
E 24
1z1201 E 230 28
12-207
2 o
5
σ: + of aie
2 -8
π 0202
2
t=3
The last inequality is derived by applying Lemma hand, we can obtain that
Lemma E.7 and Lemma
E.8 On the other
2 - - a
e
02|22|
E 23
asl +al
(sse)<0) B 2
[e>00)]
2E e 1
E
E 220
23
(z>20)
E 24 (2220) E
25c 88
25 1 (20200)
2
2
a202
<8 π T2a2
a
22
This finishes the proof.
46

---

Published as a conference paper at ICLR 2025
Lemma E.14. Let 21, 22, be three standard Gaussian random variables, and 24, 2j are either independent 24 2j for i + j. Then it holds that =| 2
Proof of Lemma E.13 There are three possible cases:
1. 21, 22, 2g are all independent. Then E [| 222a|] = E
-B[[21]]
222
E
23
||
2
2. Two of 21, are equal while another one is independent with these two random vari- ables. W.L.O.G. we assume 21 22 and 23 is independent with both 21. 22.
Then
E 22
E |s1| =V²
3. 21, 22, 23 are all equal. Then 21 22
= 1 2
This finishes the proof.
E.3 CONCENTRATION RESULTS
Lemma E.15. Let 22,
"2, FD Lid. o(0,02).
Then with probability at least
/2 Da/a c/VD/2 it holds that |2 D1/4 for all i e [D] and any constant c1.
Proof of Lemma E.15 By Mills ratio, obtain that
P P( D10
2 π
cDVD
ei D1/4 exp
2
By applying a union bound over all E [D], we complete the proof.
Lemma E.16. For X1,
Xp defined in Definition 2.1 holds that
||x1122 S a Vd + DD
with probability at least De + for all j E [D].
Proof of Lemma E. 16 Since || |la is I-Lipschitz continuous, then by Theorem 2.26 in Wainwright (2019) we can obtain that
P P(Ix)l12 - E[[xx 2 ooVD) s -
Besides, by Jensen's 's inequality, we also have
E | [la]
S V E = =
Applying a union-bound over all j E [D] completes the proof.
Lemma E.17. For x defined in Definition 3.1 holds that
I(v xfy) oō 11/1
1x1 ||, I V2o, Vd VD)
with probability at least 1 - ne </VD for all i E [n] and j E [D], where c is any positive absolute constant, and é' is positive constant solely depending on
47

---

Published as a conference paper at ICLR 2025
Proof of Lemma E.17 By definition of || 11,,,, we have |[xx S 22o and [112 2 for
all i [m]. [D] and & E [d]. By Hoeffding inequality, we have
p
(o 1 [K
2 DV,
2exp (- @VD
Besides, by Bernstein's ss inequality, we also have
[ E
|| 2 202
S 2exp(_
(E.30)
where e" is an absolute constant. the proof.
Applying union-bound over all i E [n] and j E [D] completes
E.4 ORTHOGONALITY AND NORM DISCRETE SINE TRANSFORM
Lemma E.18. For positional encodings PL, P2, Pp defined in 2.1 it holds that ||pyll
P for all j [D] and Py) 0 for all j # j' e [D].
Proof of Lemma E.18
This lemma is equal to prove that
π Tkj
Tk'j D+1
D 2
Σ
sin
sin
OEK'
kk,k's E [D].
D
1
1=1
where 6FA when k k' and OR 0 when f K. Applying the form sin(xr) (exp(ix) exp(-irr)/a we can rewrite the left side above by
rkj
Tk'j D D+11
in +K exp
iT(k: k')
Σ j=1
sin
D + 1
sin
exp
D+1
D +
1=0
i (- exp
iπ (k k'),
exp
D ++11
D
When l: + sk' # O, the geometric partial sum shows that
exp + it(k + 2) D+11
expl + sk')] 1
Σ 1=
exp(+in(k + sk:) s//)/(( )/DD)//(/))/ + 1)}
where E (+1,-1} In this case, it is required that 1: + sk' + 0. We have that when ho + sk' + O, if + sk' even, the terms will vanish and then the equation equals When tk is odd, must hold that one of l and k is odd and the other is even, under this case exp(Lim (k + sk')) 1, the term compensates as
1
exp
(z))
exp
ir
This indicates that the equation above equals O. We conclude that when & + sk' # 0.
D
rkj
Tk'j D +
Σ
sin
sin
O.
D
+1
We consider the case k sk' O. By the definition, it only holds when -1 and k'. In such case, simple algebra shows that
Tkj D + 1
Tk'j D +
D+1 + 2
Σ
sin
sin
f=0
which completes the proof of Lemma
H.18
48

---

Published as a conference paper at ICLR 2025
E.5 SEQUENCE ITERATION BOUND
The following lemma is inspired by Cao et al. 2023
Meng et al. (0024) Zhang et al.)
Lemma E.19. Suppose that positive sequence Ir, 0 follows the iterative formula
T++1 3t + cæ2
for some positive constant > O. Then it holds that
((a + 1)ct + x =gtl)at air
S as S sro" + +(aa+ 1)et + x = atr
for all / 0.
Proof of Lemma E.19 We first show the lower bound Hτ. Consider a continuous time sequence Et 2 defined by the integral equation with the same initialization.
[
Eo
dr
Z0-
(E.31)
Note that n, is obviously an increasing function of Therefore we have
L++1 = Et +
E dr
S 2, + c
E4 dr
2, + C2
for all & N. Comparing the above incquality with the iterative formula of {x²), conclude by the comparison theorem that xa 2 I, for all E N. Note that (E.31 has an exact solution
F ((a + 1)ct + x (22) xgt2)att पोr
Therefore we have
2 ((a + 1)ct + mo =tt
for all E which completes the first part of the proof. Now for the upper bound of Ft, we have
T=0
(la 1)e7 + - atr
$ xo
7=0
-ंr -afr Σ ((a (a+1))es 1)er + = =g+1)=
20
+ [ ((a+ 1)er + x x oी dr,
< 00
where the second inequality follows by the lower bound of ae as the first part of the result of this lemma. Therefore we have
2e < 20 + cr + ((a + 1)ct + x
at:
T0
= cr a + ((a 1)ct + x = xtr
which completes the proof of Lemma H.19
49

---

Published as a conference paper at ICLR 2025
Gradient Descort
Coe Inye i 3nd
Opoen
Takesisda
(a) Training Loss
(b) Angle in vi and V
(c) Attention Matrix
Trwining Low Gradient Descert
Con lngg and
27
(d) Training Loss
(e) Angle in v(( and v
(l) Atlention Matrix
Figure 4: Figures on training loss, cosine similarity and attention matrix. The first line presents the training results with j* 30. The second line shows the training results for j = 70.
F ADDITIONAL SIMULATION RESULTS.
In this section, we present additional simulation results with larger numbers of variable groups D and higher variable dimensions Specifically, we generate data according to Definition 2.1 setting o₂ 0.25 and (n, d, D) (10000, 100, 100). We consider two scenarios with j 30 and j = 70, respectively, and investigate whether the attention matrix effectively concentrates on the specified jr, even under nigh-dimensiona settings designed to mimic the scale of image data. The vector v* randomly generated and kept fixed throughout the simulations.
Due to the large amount of sample size n, consider the SGD training with batch size 64. We set the learning rate = 0.01 and train the model for 100 epochs. During the training process,
we plot training loss and the cosine similarity [v(*)_v") After the training loss converges at the final epoch, we calculate the attention score matrix for each sample and display the heatmap of the
average attention score matrix across all samples.
As shown in Figure the training loss steadily decreases, eventually approaching zero after suffi- cient training iterations. and v(") rapidly aligns with the direction of v even when d and D have higher dimension. More interestingly, when the value of j" varies, the attention mechanism consis- tently adapts to the target index. For instance, when jL 30, the attention matrix sharply focuses on the 30th row and effectively isolates the label-relevant group. The focus shifts with the same precision j* 70. This behavior highlights the model's ability to redirect its attention to the specified index, even in high-dimensional settings designed to mimic the complexity of image data.
G REAL DATA EXPERIMENTS
In this section, we conduct experiments using the CIFAR-10 dataset, where each image has shape o x 32 x 32, representing three color channels (RGB). For this experiment, we select two labels, "Frog" and "Airplane," and use 500 images from each label. To prepare the input for our framework, each CIFAR-I0 10 image is embedded as either the first patch (positioned at (1.1)) or the 25th patch (positioned at (4,4)) in grid, while the remaining 48 patches filled with noise. These noise
50

---

Published as a conference paper at ICLR 2025
(a) "Frog" Images at position (1.1)
(b) "Airplane" Images at position (1,1)
(c) "Frog" Images at position (4,4)
(d) "Airplaae" Images at position (4,4)
Figure 5: Examples of embedded images. Figure 5(a) and Figure 5(b) show images labeled as "Frog" and "Airplane," respecivelll embedded at position (1,1) with a token index of 1. Figure 5(c) and Figure ((d) show images labeled as "Frog" and "Airplane," respectively, embedded at position (4,4) with token index of 25.
patches, each of size 3 x 32 x 32, are generated using random values sampled from a Gaussian distribution with mean of 0 and standard deviation of 1/3. This arrangement forms 7 grid (49 patches in total), resulting in final input with dimensions of 3 x 224 x 224. Examples of the processed images are shown in Figure
We apply one-layer transformer to learn from the processed images. Each patch, including the original CIFAR-I0 image and the 48 noise patches, flattened into a vector of size 32 x 32 3072, resulting in a total of 49 vectors corresponding to the 7 x 7 grid of patches. These vectors, each with a dimension 3072, form the input sequence for the transformer model. As shown in Figure and Figure 5(b) images positioned at (1.1) correspond to j 1, while Figure 5(c) and Figure 5(d) show images positioned at (4,4), corresponding to j = 25. In this setup, d 3072 represents the dimensionality of the data, and D 49 denotes the number of variable groups The transformer model initialized to O, and we train it using a batch size of 64 and a learning rate of 10-3, For comparison, also present results from directly applying logistic regression to the clean
51

---

Published as a conference paper at ICLR 2025
57
Tranfomer
0.7
Trrefsma
Logest Regrowon
Ligats flg tout
26-
06:
Traiiing LOGG
Trahing Loss
0.5
Esoch
(a) Training Loss position (1,1)
(b) Training Loss at position (4,4)
Toum trdee
5.6
Tokan Index
Tokan Index
(c) Attention Matrix position (1,1)
(d) Attention Matrix at position (4,4)
0.55
Teetiro fecurayy
Testing fccurecy
0.75
Logat: Segraiion
Logudi
80
107
1.00
(c) Testing Accuracy at position (1,1)
(f) Testing Accuracy at position (4,4)
Figure 6: Experiment results on training loss, attention matrix and testing aceuracy. The first column shows the results when images have position at (1,1). The second column shows the results when images have position (4,4).
data points, where each CIFAR-10 image is flattened into single vector of size 3072, and logistic regression is performed on these single vectors.
The results of our experiment are presented Figures Figures ((a) and ((b))boooon the training loss over 100 epochs using gradient descent. The plot demonstrates steady decrease in the loss during training, which closely aligns with the decreasing trend observed in the clean logistic regression model. This indicates the effectiveness of the optimization process and the transformer model's ability to focus on the true images. Figures 6(c) and ((d) display the attention matrices for the trained images, further confirming the model's 's ability to focus relevant features. For images positioned at (1,1) with j" = 1, the attention matrix shows that the model focuses predominantly on the first row, corresponding to the original CIFAR-10 image. Similarly, for images positioned at (4,4) with j* = 25, the attention matrix highlights that the model concentrates on the 25th row, again corresponding to the true image. Figures 6(e) and 6(f) demonstrate the generalization performance
52

---

Published as a conference paper at ICLR 2025
of the transformer unseen test dataset. The results show that the transformer maintains strong performance, further verifying its ability to effectively generalize to new data.
53


Reviews:
[{"review_id": "Ai5NKU76Au", "reviewer": "Reviewer_CzbC", "summary": "This paper investigates one-layer transformers trained on specific dataset, where the input variables are generated from multiple groups, while the true label of this input is determined by variables from a single group. Based on these simplifications and assumptions, it theoretically demonstrates that the one-layer transformers can almost attend to the variables from the label-relevant group. Moreover, it provides a tight lower and upper bound for the population cross-entropy loss of a one-layer transformer trained by gradient descent. It also shows that the well pre-trained one-layer transformers can be efficiently transferred to a downstream task sharing a similar “group-sparse” structure and further provides an improved generalization error bound for one-layer transformers fine-tuned by SGD, which surpasses that of linear logistic regression applied to vectorized features. The numerical experiment observations support the theoretical findings.", "questions": "the input x_j is concatenated with position encoding, does Theorem 3.2 still hold if using addition instead of concatenation?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The paper is well written and easy to follow. \n\nThe theoretical analysis is thorough and the results align with the group sparsity assumption, although I did not check all the math and proof in detail.\n\nIn a sense, it provides new insights into sparsity analysis on the workings of attention based models.", "weaknesses": "(1) too much data assumption and model simplification. According to Definition 3.1, each patch x_j is i.i.d from Gaussian, and its label is determined from a given v (which can be learned from samples later). These data assumptions make group attention trivial.\n\n(2) And the model is one-layer transformers, which may cannot capture attention-based models to a sufficiently satisfactory extent. With these simplifications, the contribution of this paper is limited, especially considering the pre-work from Jelassi et al. (2022).", "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": "wIqhoeeBJY", "reviewer": "Reviewer_BJ8K", "summary": "This paper analyzes the variable selection of the Transformer theoretically.", "questions": "See weakness", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The author presented a novel theoretical analysis on how the transformer can learn to select variables for the group sparse classification, which is interesting and novel.", "weaknesses": "1. Given your previous definition of the notations, pls define $\\Theta$ in advance （otherwise people would get confused about this common notation.What is the 'variable' in the paper? Are they some features or some attributes of the data that define the label? What is the $v_1$ in the Thereom 3.2-2？ The weight corresponding to the label 1 or the value vector (output from value matrix in Transformer)？\\\\\n2. The proof stretch should be placed in the Supplementary instead of the main paper. \\\\\n3. The experiments are conducted on synthetic data, however, in the higher dimensional images (e.g. 3*224*224), things would change. I have concerns on these too simple experiments. \\\\\n4. The Lemma 5.1, I do not see why the $W^{(T^*)}_{1,2}$ holds and $W$ is diagonal. In NLP, it seems that the feature is correlated to its position and would correspond to the output, is that a too strong assumption?\n5. Can you show a more detailed demonstration on why the inequality of $\\alpha^(T^*)$ holds (or is just a definition?) and why Lemma5.2 can be incorporated into Lemma 5.3? Also, I do not think Lemma 5.2. holds as the underlying assumption is the Transformer is naturally with low error. Lemma 5.2 is more like a strong assumption and definition to me. \\\\\n6. How you define the group sparse? What is this different from the standard classification?\n\n\nI would like to raise my scores if more details are provided.", "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": "cJU0mOAbJE", "reviewer": "Reviewer_tEUD", "summary": "This paper investigates the theoretical understanding of transformers, which have shown success in various applications but lack a clear theoretical foundation. The study focuses on how transformers can learn a statistical model with group sparsity, where only a subset of input variables (groups) influences the label. The research shows that a one-layer transformer can use attention mechanisms to select relevant variables and ignore irrelevant ones for classification. Additionally, it demonstrates that a well-pretrained one-layer transformer can adapt to new tasks with limited samples and achieve good accuracy. This study provides insights into how transformers learn structured data effectively.", "questions": "Please refer to the weaknesses section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. This paper investigates the theoretical understanding of transformers, an important and interesting problem that may greatly benefit the deep learning community.\n2. This paper is well-written and has a good motivation.\n3. The conclusion that a one-layer transformer can use attention mechanisms to select relevant variables and ignore irrelevant ones for classification makes sense to me.", "weaknesses": "1. Though the problem is interesting, the conclusion, a one-layer transformer can use attention mechanisms to select relevant variables and ignore irrelevant ones for classification, seems to be well-known in the deep learning community and has already been verified by the original paper (Attention Is All You Need).\n2. The theoretical understanding of transformers given by this paper mainly focuses on the well-pretrained one-layer transformer. However, it will be much more interesting to study the multi-layer transformer since the scaling law shows that the loss of transformer scales as a power-law with model size, dataset size, and the amount of compute used for training, with some trends spanning more than seven orders of magnitude. \n3. The study focuses on how transformers can learn a statistical model with group sparsity, where only a subset of input variables (groups) influences the label. However, the concept of group sparse is not properly defined, making me confused to fully understand it. Further, what is a similar “group-sparse” structure as mentioned in the Introduction section? What kind of “group-sparse” inputs are similar?", "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)