Implicit Bias and Fast Convergence Rates for Self-attention
Authors: Bhavya Vasudeva
Puneesh Deora
Christos Thrampoulidis
Venue: Accepted by TMLR ()

Abstract: We study the fundamental optimization principles of self-attention, the defining mechanism of transformers, by analyzing the implicit bias of gradient-based optimizers in training a self-attention layer with a linear decoder in binary classification. Building on prior studies in linear logistic regression, recent findings demonstrate that the key-query matrix $W_t$ from gradient-descent (GD) converges in direction towards $W_{mm}$, which maximizes the margin between optimal and non-optimal tokens across sequences. However, this convergence is local, dependent on initial conditions, only holds asymptotically as the number of iterations increases, and leaves questions about the potential benefits of adaptive step-size rules unaddressed. To bridge this gap, we first establish scenarios for which convergence is provably global. We then analyze two adaptive step-size strategies: normalized GD and Polyak step-size, demonstrating finite-time convergence rates for $W_t$ to $W_{mm}$, and quantifying the sparsification rate of the attention map. These findings not only show that these strategies can accelerate parameter convergence over standard GD in a non-convex setting but also deepen the understanding of the implicit bias in self-attention, linking it more closely to the phenomena observed in linear logistic regression despite its intricate non-convex nature.

Published in Transactions on Machine Learning Research (03/2025)
Implicit Bias and Fast Convergence Rates for Self-attention
Bhavya Vasudeva' University of Southern California
beasudev@@usc.ett
Puneesh Deora*
punceshdeormdsecceuuu ubc..
University of British Columbia
Christos Thrampoulidis
clhrummo@cce..ke Bew..ub................oo
University of British Columbia
Reviewed on OpenReview: https: //tp:///ppeeeeeeeeeeeeem net/ forum? id= id=pKslnj@ebb
Abstract
We study the fundamental optimization principles of self self-attention, the defining mechanism of transformers, by analyzing the implicit bias of gradient- based optimizers training a self-attention layer with a linear decoder binary classification. Building on prior studies in linear logistic regression, recent. findings demonstrate that the key-query matrix W. from gradient-descent (GD) converges in direction towards Wmm. which maximizes the margin between optimal and non-optimal tokens across sequences. However, this convergence is local, dependent on initial conditions, only holds asymptotically as the number of iterations increases, and leaves questions about the potential benefits of adaptíve step-size rules unaddressed. bridge this gap, we first establish scenarios for which convergence is provably global. We then analyze two adaptive step-size strategies: normalized GD and Polyak step-size, demonstrating finite-time convergence rates for W, to mimy and quantifying the sparsification rate of the attention map. These findings not only show that these strategies can accelerate parameter convergence over standard GD in a non-convex setting but also deepen the understanding of the implicit bias in self-attention, linking it more closely to the phenomena observed in linear logistic regression despite its intricate non-convex nature.
1 Introduction
Self-attention serves as the fundamental building block of transformers, distinguishing them from traditional neural networks (Vaswani et al., 2017) and driving their outstanding performance across various applications, including natural language processing and generation (Devlin et al., 2019; Brown et al., 2020; Raffel et al., 2020), as well as computer vision (Dosovilskiy al., 2021; Radford et al., 2021; Touvron et al., 2021). With transformers establishing themselves as the de-facto deep-learning architecture, driving advancements in applications seamlessly integrated into society's daily life at an unprecedented pace (OpenAI, 2022), there has been surge of recent interest in the mathematical study of the fundamental optimization and statistical principles of the self-attention mechanism; see Section 6 on related work for an overview.
In pursuit of this objective, Tarzanagh et al. (2023b;a) have initiated an investigation into the implicit bias of gradient descent (GD) in training a self-attention layer with fixed linear decoder a binary classification task. Concretely, the study paradigm implicit bias soeks to characterize structural properties of the weights learned by GD when the training objective has multiple solutions. The prototypical instance of this paradigm is GD training of linear logistic regression separable data: among infinitely many possible solutions to logistic-loss minimization (each linear separator defines such solution), GD learns weights that converge in direction to the (unique) max-margin class separator (Soudry et al., 2018; Ji & Telgarsky, 2018). Notably, convergence is global, holding irrespective of the initial weights' direction, and comes with explicit rates that
"Equal Contribution.

---

Published in Transactions on Machine Learning Research (03/2025)
Train accuracy
Test loss
Test accuracy
1.5
0.5
0.8
SGD
0.5
1.5
0.6
SNGD
0.5
0.4
1
0.4
SPS
Adam
50
100
150
50
100
150
50
100
150
50
100
150
Figure 1: Comparison of train and test dynamics of various optimizers SGD, stochastic normalized GD (SNGD), stochastic Polyak step (SPS), and dataset; see App. C for details. SNGD and SPS, employing adaptive step-size rules, demonstrate significantly faster training, closely resembling the performance of Adam. Motivated in part by this observation, our work establishes fast convergence rates for NGD and PS for single-layer self-attention.
fine-tuning a pre-trained BERT model on the MNLI
characterize its speed with respect to the number of iterations. Drawing an analogy to this prototypical instance, when training self-attention with linear decoder in a binary classification task, Tarzanagh et al. (2023a) defines a hard-margin SVM problem (W-SVM) that separates, with maximal margin, optimal input tokens from non-optimal ones based on their respective softmax logits. For this, they show that the key-query weights W, found by GD converge locally to the solution W Wmm of (W-SVM) as the number of iterations grows to infinity.
Despite the intriguing analogy between the two settings, the findings of Tarzanagh et al. (2023a) are highly non-trivial not only because the nature of the max-margin solution differs, but also because of the intricate non-convex optimization landscape introduced by the presence of self-attention. The non-convexity, induced by the softmax operator in the self-attention layer, complicates the analysis and is the reason why the convergence result of Tarzanagh et al. (2023a) is: (i) local, holding only for an appropriate initialization direction, and (ii) asymptolic, applicable only as iterations t approach infinity.
Identifying these limitations, this work is motivated by the following questions:
Q1: Are there settings under which GD iterates converge globally to the solution Wmm of (W-SVM)?
Q2: Is it possible to oblain finite-time rales of convergence to W Wmm?
Additionally, motivated by the practical benefit of using adaptive learning rates in transformer optimization (see for example Fig. 1), we pose an additional open question:
Q3: Can using adaptive learning rates in self-attention optimization accelerate the convergence to W Wmm? ?
Contributions. Our work addresses the above open questions, thus contributing fundamental insights on the optimization properties of the self-attention mechanism.
Concretely, we study a single-layer self-attention model 4(X,0) = u X (TTT((W)))) where, sp)) is the softmax nonlinearity, X [.",""",,T]]] is the sequence of input tokens, W the key-query matrix, and u is a lincar docoder. Following the sctup of Tarzanagh et al. (2023a), for each sequence X, we associate cach token #, with a score 7+ := yu x, and let opt [T] be the index of the token with the largest score. Given training set of sequences X,, this defines a hard-margin SVM problem with solution Wmm that separates the optimal tokens with maximum margin; see Section 2 for details.
Motivated by question Q3, we study here the optimization properties of training ();0) with exponential loss using normalized GD, which sets the learning rate π at iteration adaptively as 7) = [v.[(..] for some constant 17 > 0. Our results also extend to Polyak step-size, another adaptive step-size rule.
Our first set of results (Section 3), for fixed decoder (similar to Tarzanagh et al. (2023a)), answer questions Q1-Q3 as follows. In response to Q1, we begin by identifying sufficient conditions on the input data under which GD converges globally, in direction, to Wmm- Then, simultaneously addressing Q2 and Q3, we establish fast Snite-time rates of convergence for normalized GD by proving that the iterates W,, at any time t, satisfy W. Wmm. ((-112) We identify two key ingredients towards establishing these results. First, we [W.T m
2

---

Published in Transactions on Machine Learning Research (03/2025)
show that the Euclidean norm |Wall of the iterates grows at rate (t)). Second, and more intricate, we demonstrate that even if the iterate at any time i violates the constraints of (W-SVM) corresponding to any training sample, the softmax score of the optimal token for that sample must be non-decreasing. In turn, this establishes that the itcrates, initialized in any direction, remain inside a cone around Wmm- Our convergence results also imply an explicit. rale al which max-tttttnnni gets sparsified as softmax scores of optimal tokens converge to at an exponential rate O(cxp(-ntt).
Our second set of results (Section 4), raises the limitation of fixed linear decoder of prior work and applies to joint training of and of the attention weights W. In response to Q1, we construct a representative data model with Gaussian-distributee tokens and prove that GD converges globally, in direction, to Wmm- To address Q2, we show that for normalized GD with an aggressive step-size, the iterates W,, at any time 4, satisfy W. Wmm
[WI [Wam s (1/llog))) Further, prove that the línear decoder converges, in direction, to Vmm: the solution of the hard-margin SVM problem (u-SVM) that separates the examples with maximal margin, using only the optimal input tokens at a (t rate. Finally, to completely characterize the training dynamics, we show fast train loss convergence at a O(exp(-(')))
rate. We highlight three key technical
contributions towards proving these results: (i) growing token score gap Topt.t - Vr. > O, between optimal opt and non-optimal tokens with time; and (iii) constant loss ratio across all training sequences valid for any time. These properties pave way for PI.llike inequality which is crucial to show the loss convergence rales.
opt tokens, for any time t > 0; (ii) non-decreasing softmax scores of the optimal
Throughout, we validate our findings on both synthetic and real datasots (sec Section 5). We end with Section 7 discussing the implications of our results and the open questions they raisc.
2 Preliminaries
Lowercase and uppercase bold letters represent vectors and matrices, respectively. ||a|| and |A]| denote a and A [ respectively. [a1, = {a : Ba, + (1 (1- B [0, 1]} denotes the line segment between an and a2, and concal(.,-) denotes the concatenation operation. a A b denotes the minimum of numbers a and b. a v denotes their maximum. We use standard notation O, Z to hide absolute constants, and Ō, Z to hide poly-logarithmic factors. All logarithms are natural logarithms.
Fuclidean norms. à and À denote |a]
The output of a single-head self-attention layer, parameterized by key, query and value matrices WΩ. WK Rdxdd, Wy Rddxd2, is given by (XW W X')xx where X ==[ Rxxd is the input sequence, and the softmax map sp(.) RT + RT is applicd row-wisc. We can compose the output of the attention layer with a linear projection head to obtain the final prediction as
$ (X:0) u X (XWWx,),
(1)
where 0 = cu, W) denotes the set of trainable parameters. Here, similar to Tarzanagh et al. (2023a); Deora et al. (2023); Tian al. (2023a); Oymak et al. (2023) we reparameterize the key-query matrix as W == Wo E Rexd, use the first token for prediction and subsume the value weights Wy within the prediction head u e Rd. Let a(W): = XWai denote the vector of softmax logits, then ('(a( e RTxT denotes the Jacobian at a( (W).
Given training data ((X,,))) 91) with labels y 6 [+1] and X, == [xx,1 Di2,.,,,,... and decreasing loss &: + R,, the empirical risk is I(0) = + Eiemi ((x; We focus on GD optimization with adaptive time-dependent step-size T/- Concretely, we study two variants: normalized gradient descent. (NGD) (Hazan et al., 2015; Nacson et al., 2019), and Polyak step-size (PS) (Loizou et al., 2021), which set
n xc (1)\oI((e)II. and n=
with İ = I((), 0
respectively. Our results are applicable to both update rules. For specificity, we present them for NGD and include remarks for PS.
"Note the rate here is slower compared to the fixed-decoder case. As detailed in Section this is due to the additional (t+ 1) factor in the step-size, which results in slow down of the rate of growth of [W Our proof requires this additional factor to account for the non-smooth objective.
TThis is without loss of generality as our results hold for any taken 7 [T].
3

---

Published in Transactions on Machine Learning Research (03/2025)
We follow Tarzanagh et al. (2023b;a) in defining token scores follows. Definition 1 (Token scores and Optimality). Given a fixed prediction head u, R", the token score vector for a sample (y. X) is given by Y= yXu,, The optimal token index! is opt = rery Yr, where 7 = ya, denotes the token score for the token at.
Similarly, opt; denotes the optimal token index for sample (y,, is [[]]. Intuitively, these are the tokens that minimize the training loss upon selection (Lemma 2 in Tarzanagh et al. (2023a)). Given a set of optimal token indices OPT:= (opt,) define the following hard-margin SVM problem, which separates, with maximal margin, optimal tokens from the other tokens for every input sequence:
W mm argmin W
W s.t. (Ii,ops, (t.,)'Wa,,, 2 1 Ví [m], e [T] > {opt,}-
(W-SVM)
Throughout, we assuine that (W-SVM) is feasible, i.e. softmax logits Topp. Wæ,,1 of optimal tokens can be separated from logits TI,- W of the other tokens + opt,.
3 Training Dynamics of W
Here, let fixed prediction head u = u. This allows focusing first on the dynamics of token selection induced by training key-query weight matrix W, which is the only trainable parameter:
V I(W)) t>
W = W. - VV T(W.)l
(2)
We start by setting up (mild) assumptions on the data and initialization, and then present our main results.
3.1
Setup
For convenience, first define A W m B max 7= = min Yi.r: K4 == max exp("Hopp - -%;,, K_ = reopt,
1,7
log(n=) min exp(Ti.pps - -)), and Y K4 log(n_)) Note 1/A is the margin achieved by Wmm separating optimal tokens from the rest, and B is a uniform bound on the tokens' magnitudes. The parameters N₂ represent the largest and smallest degradation factors across sequences for each sequence's 's individual loss term when suboptinal tokens are selected. Respectively, Y can be interpreted as conditioning parameter for the problem, measuring the variability in token gaps across different sequences.
Our first technical assumption towards ensuring global convergence requires that tokens are nearly orthogonal. This is often the case in high digeesiiensional settings; see Example and the references below.
Assumption 1 (Nearly-orthogonal Tokens). For any i,jj, k [m], and for any T, r' e {TT),
|..- 11 4n (T
j+ i + opt,
(pr op,, 4n (TT 2k.))) )
y: = y # yk.
We will use Ass. 1 to prove that softmax scores of optimal tokens are lower bounded by a constant throughout the optimization trajectory. Note that Ass. itself makes further guarantees on softmax scores of optimal tokens approaching or even increasing during training, which is essential for global convergence and we prove separately. Furthermore, we impose no assumptions on the direction of initialization Wn. In particular, our main results hold even when Wo is aligned with "bad" stationary directions W that could asymptotically saturate non-optimal tokens, i.c. V Vw Tw[((W) + and Li,ept, (a + 0 as a + co. Thus, the only mild requirement on the initialization Wo regards its scale, rather than its direction, and is formalized in Lem. 2 in the App. See Fig. in the App. for numerical validation of global convergence despite initializing in a "bad" stationary direction. The following example illustrates the above assumption.
Similar to Tarzanagh et al. (2023a), we assume unique optimal token which holds for almost all datasets. Thhis is analogous to linear separability assumptions iu linear models, which are fundamental to understand the implicit bias in GD-based methods. For linear models, fessibility of hard-margin SVM is guaranteed when the ambient dimension exceeds the number of samples. Similarly, for our self-attention model, (W-SV)) feasibility is guaranteed when the ambient dimension exceeds both the number of samples and the context window size (as shown in Theorem of Tarzanagh et al. (2023a). Oorr results also extend cases where tokens are antipodal and the remaining takens are nearly orthogonal.
4

---

Published in Transactions on Machine Learning Research (03/2025)
Example 1. Let μr e RR, pe pr,, ||,,|| -1e U, and data generated follows:
y Unif( (+1)))
opt ~ Unif([T]), v ,~N(00,0'))
Mr - 1 Topt u+
ar -.-N(O, VT opt. | . yy=-1²
-U-Pp - It is easy to show that when d= (1(2T22) the tokens are nearly orthogonal with high probability. Further let high enough signal-to-noise ratio U22 2 Si(( p)RTTvd) and appropriate u. such that the loken scores salisfy T= 0(1,, then, the dala model satisfies Ass. 1 (see Lem. 6 in the App. for details). We remark that similar data models have been considered in prior works on the analysis of linear overparameterized models (Muthukumar et al., 2019; Chatterji & Long, 2021; Wang & Thrompoulidis, 2022; Cao & Gu, 2019; Wang el al., 2021), NNs with Leaky ReLU aclivation (Frei el al,, 2022a; 2023), CNNs (Cao et al., 2022; Kou et al., 2023a), and self-attention models (Deora et al,, 2023; Li et al., 2023a).
where E:= Ia I-- -
Our second technical assumption is similar to Tarzanagh et al. (2023a) with two key distinctions: Firstly, it applies to self-tttention rather than their simplified attention model. Secondly, and most importantly, under this same assumption, we will prove stronger results for global convergence and finite-time rates.
Assumption For any E [m], and any T opt, [T], the token scores satisfy "t, = 7/--
Note that in Ex 1, Ass. is satisfied by choosing an appropriate u. (see Lem. 6 in the App. for details). It. has been recently shown that the optimization landscape of attentttntion can lead GD to converge to local directions that are different from Wmm (Tarzanagh et al., 2023a). Thus, global convergence necessitates additional conditions. Assumptions and 2, as outlined above, serve this purpose. While being only sufficient condit ions, they lead to the first-known global convergence result with finite-time guarantees, contrasting with the local and asymptotic convergence in Tarzanagh et al. (2023a).
3.2 Main Results
We now present our main results. For case of exposition, we use exponential loss, but our results directly exlend to continuously differentiable decreasing loss functions. First, we establish the rate of NGD's directional convergence to W um- A proof sketch is given in Sec. 3.4.
Theorem (IB rate). Under small initialization scale (Lem. in the App.) and Ass. 1-2, using the NGD updates in Eq. (2) with Q=)(B-)) it holds for any t2 (1, B,A, A,,,),,
W W mm) ) 2 C (logt)²
where C:= C(1 B,A, poly (17, to). In particular, assuming B= ((1),T=0(1) and T=0(1), we have A=O(1). Thus, for any t to = =(n),, it holds [W W mml| 0
Thm. establishes that normalized iterates W. of NGD converge globally to W mm: starting from any initialization direction. It also provides a lower bound on the rate of W. approaching W mm- In Soc. 3.3, we establish an upper bound for the rale of convergence of GD, demonstrating that NGD's adaptive step-size provably accelerates convergence. This is the first global convergence result and finite-time convergence rate for self-attention using either GD or NGD. To the best of our knowledge, this is also the first demonstration of directional convergence of NGD and its superiority over GD in a non-convex setting (Sec. 3.3 for comparisons with existing results in convex linear settings).
We now discuss an implication of Thm. 1 concerning the evolution of the attention map during NGD training. Specifically, we demonstrate that as training progresses, the attention map increasingly becomes sparse, selecting optimal tokens at an exponential rate.
Lemma (Softmax score rate). Under the setting of Thm. 1, it holds for any e [m] and any t 2 poly (n, B, (o)) that pi,op,, (1 (T- -1)e n88B3
1)8= 1 13)²1)
5

---

Published in Transactions on Machine Learning Research (03/2025)
1.8
6x 10:1
48
(W. Weell
0.0
( stop
Sx 1
0.5
0.0
MCD-mon
18-4
MCD-theory
2000
2000 4006
2000
4000
Tterations
Trerations
Iterations
Figure 2: Training dynamics of a single-head self-attention model (Eq- (1)) when optimizing only W on synthetic data with nearly orthogonal tokens (Example 1 with o 0). The observed softmax score saturation, norm growth of W and directional alignment with Wmm closely match with our theoretical results (Lemma 1, Eq. (9) and Theorem 1, respectively).
To understand why the attention map becomes sparser during NGD training, recall from Thm. that for sulficiently large iterations, W, W Wmm: and that W separates optimal tokens from ot In proving Thm. 1, we also show a norm-growth rate [W/ @t).=((t. Combining these yiclds that softmax scores of optimal tokens Phopt, approach 1, while the rest approach 0. Concretely, given W W mmm the opt softmax score Pi.opt, 2 i+ E. exp(- exp(-Iwaml
[W&]
IWil
approaches as Wil grows. Lem. formalizes this intuition and rrops; [Wam!
Wel
pairs it with an explicit sparsification rate.
Finally, we validate the predictions of the above results using synthetic data generated as per Ex. with o (see App. C for details). In Fig. 2, plot the train loss, the norm growth of W., the softmax score (ii,opt, for the opt token (averaged over i [n]), and the alignment of W with W Note that since is fixed, the train loss does not converge to 0. Observe that the directional alignment for NGD is closely predicted by Thm. 1. Similarly, the softmax-score saturation rate in Lem. 1 closely aligns with the empirical rate in Fig. 2. Moreover, note that NGD significantly faster than standard GD, which we formally prove below.
3.3 Remarks
NGD is faster than GD. We demonstrate that there exists dataset, satisfying the assumptions of Thm. 1, for which the correlation of vanilla GD parameters to Wmms (We. t. W mm). is upper bounded by I (((o())))))) In contrast, according to Thm. 1, the rate for NGD is (l)g))))))) thereby highlighting significant gap between the two.
To construct the datasel, let n= 1, == = y= 01 = [1,0]", 22 [0 and u. [1,0]*. Clearly, Ass. 2 holds since there is only two tokens; specifically, opt = 1, Topt - 1. Also, Ass. 1 holds trivially since n For this datasot, the max-margin classifier can be casily expressed in closed form Wmm (x) xz JT]
Using this, we observe that V WwI,W) -[Wi (1-)t Wmm. which implies
W411 Wo + G1G W Wmm, where we set G := Σ VV #=0
Leveraging this, we now show that [.11 s O(logt): Note that [ 0o as l + oo which leads softmax saturation P'pe and hence, loss minimization. This implies that G, grows with t, and there exists some to > 0, such that for any t 2 to, holds that [Wall a nG2 + W where W is the (1,1) entry of the initialization matrix. Using this, we can show for & 2 to that. sl1 Irexp((<))=e)) Wonn S exp(-(x1 222"T (W+11)) exp(-1GG. W exp W I) Thus,
Wenl & W f
Σ T(W)) exp (- W O
Σ P20
exp exp(-IWW.l)
where we use T(W,) 0(1,, since the decoder is fixed. This proves that W, Ologt), which use to upper bound the correlation as follows:
G W
Woll (W!l)²
Wol 2 (WW!))
W W mm) V(nG.+ W1y
(3)
W. (W!1)²
2(1 + W!1)2
2 W/

---

Published in Transactions on Machine Learning Research (03/2025)
Thus, (Wa, 1, W mm) m 1 (11((((ogg))))) equivalently [W W Wmmll ((llot))=)). This shows that using adaptive step-size accelerates convergence to W compared to standard GD.
Extension to Polyak step-size. In Thm. 7 in the App., we establish an analogue of Thm. 1 specifically for the Polyak step-size rule. This extension required us to further establish a lower bound on the gradient norm. This is because PS has an additional [vo[(0,) factor compared to NGD. To accomplish this, we used the variational form of the Euclidean norm for a "good" direction Wmm, yielding [vo[(0.) 2 poly ( I(O.).
As seen in Fig. 2, PS leads to even faster convergence than NGD across all metrics, including the train loss. In our implementation, we hyperparameter Mmax threshold the step-size value, which has been shown to stabilize training (Loizou et al., 2021). Further, since the selection of opt lokens minimizes the loss, I is calculated by setting Pi.opt, 1 for all i [n]] Note that after some time, when 2 Maax, the PS updates used in practice are identical to standard GD with stop-size mmax- Hence, the rates are much faster in the beginning, but slow down as training progresses.
3.4 Proof Sketch of Theorem
The first key intermediate result we need to show is that after sufficient iterations, the negative gradient of the loss at W; is more correlated with Wmm compared to W, itself. Formally, we show that for all e > 0, there exists R. 2=c log(Ce")) with C=poly (T, nT), such that for every t for which [Wll 2 Re, it holds that.
(-W.T(WW.) W mm) (1- (-W. W.).
(4)
A proof sketch of this follows; see Lem. 7 in the App. for details. First, under Ass. 2, for any W:
n le, - pipp,, (1- (1-Phhope,)hig
((W.)(WW)) W) Σ
(5)
Exery, ph,por where he opt Errope Ph.- respectively. Then, to prove Eq. (4), it suffices that
- -x;-XX1 KWWaa 2.11 Analogously, define h, h", using a-AXW 11 d'=XW,,, -X
has s (1+ where V= v=e(1-e)"
(6)
Suppose that W W mm x since otherwise the claim follows easily. We categorize the training samples into three subsets based on the satisfaction of the constraints in (W-SVM). Concretely, define the "minimum SVM-gap per sample" it[n] as ouiin: miin ās,ept, a.,-- 1 and based on this consider three sets of samples as follows:
I (ie [n] i <0.3v), for samples where the constraints are violated or weakly satisfied,
-===(ie [m] 0.3v <o s 0.8v), for samples where constraints are satisfied with small gap,
Ia {ie [n] 8miin 2 0.8v), for samples where the constraints well-satisfied.
Since | W mml is large, there is at least one point e [n:] that violates the constraints, i.e. In is non-empty. We first show the following for examples in I₁ and Iz
J s (1 0.5 i Zı, and hh, S (1 (1+ E
(7)
Note that we prove tighter inequality for I₁ than required by (6) in order to use this residual to get the desired inequality for examples in Iz. Specifically, when H. is large enough, we wish to show (following Eq. (5)) that
0.5v
Σ Eli((riopp,- -Wiphoppe (1- (1-pp.ope)) ) SET9
Σ Tez.
Ce,((T - -niphopp (1 phope))
(8)
ABB2
7

---

Published in Transactions on Machine Learning Research (03/2025)
Intuitively, for the LHS, any sample i<Iz satisfies (W-SVM) constraints with large gap giving Ps,opt, close to and growing with [We]. Similarly, for the RHS, we expect Pi,,pt, to be smaller. Thus, for a sulficiently large WWlll and controlling Pi,ot,, for any iεZi and t>0 would allow us to show the above inequality. This is achicved by establishing a non-trivial lower bound on the softmax scores for i I that holds throughout training (Lem. in the App.). Finally, since ht* 2 and h we combine Eqs. (7) and (8) using the formulation (5) and complete the proof of Eq. (4).
The next key ingredient in the proof of Thm. 1 is showing that the iterate norm grows at rate ( Specifically, we prove in Lem. 5 in the App. that
2np 2 2
The upper bound follows by applying Sauchy-Schwand For the lower bound, we write the iterate W using its gradient updates, the dual norm characterization of the Ez-norm, and select the W direction get.
+- Σ 5
Will 2 (Wo, W )
Fw I(W) I(W.] [Waml
Wimm
where we use the key property that for any W the gradient correlates positively with Wmm (formalized in Lem. 4 in the App.), i.e. (- V L(W)
Vw T(W) Wmm (2B²A) O.
Next, using Eq. (2) to substitute the gradient update Eq. (4), we can show that W
W 1-- W 2
W. W W 6e W W.,
WWi1- W to Wmml
2 (1- cc) WW.. Wt, [W.1]
W.
2 Wel
[WiT]
1/2
2 W 1 W
Well -C7)
We then select e, = A (nt) and use Eq. (9) and Lem. to obtain a O(llott)| rate. Using Eq. (9) then finishes the proof.
4 Training Dynamics for Joint Optimization
We now study training dynamics when jointly optimizing prediction head u and attention weights W. Compared to Sec. 3, the key challenge is that the token scores of evolve with time, driven by the changing nature of ur Additionally, the objective function becomes non-smooth in this context, given the dynamic changes in u, Addressing these challenges necessitates additional technical considerations compared to Sec. 3, which we also discuss in this section.
4.1 Setup
We consider the following updates for us and W..
U4+1 u,, V.I(O). n²= 7((+ 1)-23 2/3
n((+ 1)-1
Iv.I(0.)] W vw th IV.
(10)
Here, we focus on exponential loss. Note that in nh, the factor (t -2/3 can be replaced with (t+ 1) where e [2/3, 1). In our results, we present. findings using p 2/3, as it yields the fastest rate. We also remark that common adaptive learning rates like Adam (Kingma & Ba, 2014), AdaGrad (Duchi et al., 2011), etc. similarly set per-parameter learning rates that vary dynamically.
Similar to Def. 1, given a trainable prediction head u, at any iteration t> 0, the token score vector for a sample (y, X) is given by -
Data Model. We study the following data model. Within each example X, single opt ~ Uni[([T) token and an additional T tokens are sampled i..d. Gaussian as follows:
Topt a² ær N d²La), VTES, v=sign(u,
(DM)
8

---

Published in Transactions on Machine Learning Research (03/2025)
where S:=[T] (opt), u. Rd is fixed vector that generates the label ye (=1). Here, p controls the token norms, and a>1 separates the opt token from the other tokens. We consider the following mild conditions on the initialization, token norms, overparameterization and step-size.
Condition 1. Let B = apv/1.5d 2 1, and Co 1, C2,C2 be some absolute constants. We consider the following conditions for any 8 e (0, i) Zero inátialization: | 0. W 0; ü) opt token has larger norm: Co: tii) Sufficient overparameterization d2 Cia'n" log 10n tv) Small step-size: n s (C2/TBB(BV d))=
The condition on the initialization is for ease of exposition; our analysis and results extend to random initialization with small scale (similar to the conditions of Lem. 2 in Sec. 3). We also remark that under Cond. I and data model DM, the tokens are nearly orthogonal, similar to Ass. in Sec. 3. We formalize this in Lem. 9 in the Appendix.
4.2 Main Results
We now present our main results for joint training. For ease exposition, we consider T=2, but our results directly extend to T>2 assuming token scores 7 satisfy Ass. every iteration is also possible extend results to the case where Ass. 2 is not exactly satisfied, but the non-opt token scores show small deviation.
Our first result shows that train loss goes to 0 at ((exp(-t!/)) rate; see Sec. 4.3 for proof sketch. Theorem 2 (Train loss convergence). Under Cond. 1, data model DM, using the updotes in Eq. (10), it holds for any t>0 that I(Ow.) exp (- "B"BaC (t+ (++1)1/3)) vn
The next theorem illustrates that W, converges to W Wmm in direction at an O(1/ logt) rate. Theorem 3 (IB rate of W). Under Cond. 1 and the data model DM, using the updates in Eq. (10), for any 12 te exp(poly (n, B, n, d, log gog ()),,
W., W mm) 2 121---pap poly (1), e, 1
logt
In particular, assuming p= (2 (nA)=") for any t2t, 12tt=f[(cxp((( (xxx=((
W.-W mm (0 (logt)-1).
Intuitively, since the updates for W, have an additional (t 1)+1 factor due the smaller step-size Th the iterate norm ||Wil grows as ((logt) in contrast the @(t) rate in Sec. 3, where we were only optimizing W. Consequently, the convergence to implicit bias is slower (11/logt) instead of Ö(t-)). The additional 11+ factor in the updates for W, helps to account for the non-smoothness of the objective, which will become apparent in Sec. 4.3.
We now discuss the implicit bias and convergence of u. From prior work (Soudry et al., 2018; Ji & Telgarsky, 2019), one could guess that u, converges to the max-margin predictor separating the set of samples [(vl.ops, Ti,opt, (1- -po.oot,)) ) ).).r::))] But as shown above, when f t- co, Popt 1. This motivates the following hard-margin SVM problem,
Uum marmaargminlu s.t. you otp 21 Vi [i]
(u-SVM)
The following result confirms that the learned model attains margin that is asymptotically constant factor (1/4) of the maximum margin / min; you Di,opt. Theorem 4 (IB rate of u). Under Cond. 1 and the data model DM, using the updates in Eq. (10), for any t2t, v exp(C(n, B, ce v (8B²))*)), it holds that
7 4 1 (x=xp(((333)))))111ggg
min you oit 2
4.3 Proof Sketch of Theorem
The main challenge in establishing the loss convergence rate in Thm. the con-smoothness of the objective function. To overcome this, show three key properties that hold throughout training: i) constant loss ratio
9

---

Published in Transactions on Machine Learning Research (03/2025)
for any two samples, ti) growing token score gap between opt and non-opt tokens for every sample, and fif) non-decreasing opt-token softmax score. We formalize these in Lem. 11 in the App. Property (ii) is crucial in proving the implicit bias rate in Thm. 3. Property (iii) is crucial in obtaining PL-like inequality which is the key challenging step in proving Thm. Specifically, using propertics (i) and (ii)) we demonstrate that the training loss satisfies the following (see Lem. 13 and Rem. the App. for details):
1. PL-inequality-like result: [V.I())) 2 Q(Bn (1)==,)))(
2. controlled loss between current and next iterate: max0r(00.0.... I ST((O,)
3. second-order self-boundednes::
max 0r.10..0..11
= ((((())vv((O)))
where
To prove the first point; the dual norm characterization of the fz-norm,
| a w 1-1
[vI[()]
sup
x.). v I(o,)
sup min Bi (V oppp + (1 - Pi,ope, )
=///-11
To proceed, we lower bound the softmax scores for the opt tokens for every i [n] and t 0 as
ph.opt, 2 2
(11)
Using this, we show that the set of samples (0.5((z). optt (. where T # opt,, are separable using û Eie[n] opt, (formalized in Lem. 12 in the App.). Then, selecting == à proves the first point.
We now brielly discuss the process to prove Eq. (11). We can show this by induction. Since Wo = 0, then for any e [n]. vh.opt. 1/2 at initialization. To show the inductive step, note that for any i [n].
(Di.opt,
>0,
(12)
8 141 /
exp ((zj,opt,- (.:+))((..... W.)
pj,opt /499
We consider two cases. If (fyff.o Ploop, 3/4, Eq. (12) follows vi is small. Specifically, for any j E [n], opt;,
phopp, exp((x1,r Bjoopt, " W, 3/4
(0 1 D.joptt) Wi-13,, " -
After the gradient step, since n (2B²) (3), we have
(I,,r 1h.--ajm 35,opt,) W1+13,,1 (I,, jopopt ) Wi Th-1 (2, -x,,) (-W. 2nB² 2 log(3)
If plopp, $3/4, we proceed as follows. Consider the expansion of the LHS of Eq. (12),
(Ii - xIiT)) V I(..))), lest (11 - -7))s%. (1 -li )|m,, 121120 xi, 12
le. (es - -25)%% (1 shj )a(,1 1z (m.jopt, - -x,T)) (Ii a.,,))) Σ j=1
To use this to prove Eq. (12), we first show that the token scores o satisfy Defn. (1ope, > >nt for any ie [m]],
and that the loss ratio max.,ee[]]
hi is bounded by a constant throughout training, f.e., for any i [m]. ley
lett SC,
Mop oo 2(t2/3) and max [ E,j
(13)
where C> universal constant. This is formalized in Lem. 11 in the App. We then prove Eqs. (11) and (13) jointly by induction.
10

---

Published in Transactions on Machine Learning Research (03/2025)
For the second point, first show that for any 0,0',
() , |u x x'XX'"XX""XXXX"""
$ X 2, u-
X 2,00 1/ W - W'
Since we are using exponential loss, we use this to get,
(VD(O, + X(0++1 - x)), s max exp [2 ا| ulll W Wil + X| 2,00 u4+1 uel)
max Ad
((v0(o,
xc
Further, we show that the iterate norm 14l grows as O(11/8) (formalized in Lem. 10 in the App.). Then, using the updates in Eq. (10), we get
( (0, A(0 0,), x:)) I exp (6n) BB)B)))
max max jefn] Ae[0,1]
(O,,
Then, using Cond. 1 for 7 finishes the proof for the second point. For the third point, first upper bound I | as @ xI 6dB ull -VdB3. Using this and the upper bound on /obll, we show that for any 0' [0, 0..1.
VVT((') max((1o (,)1 (1, ())
ie[a]
Then, we leverage the second point to finish proving the third point.
With these ingredients, we prove Thm. by working with the sccond-order Taylor expansion of L(O.11). Specifically, using the updates in Eq. (10) and the three key properties of the training loss established above, we show that
BB
8n² I(O.) + (1+1)48 (w. v ((((:)))((()
TT(.))
10/7(L 1)2/3
Further, since ||uz|| s O(11/3), we show that w(6,) ((133)) Using small enough 7, and telescoping, gives the advertiscd result.
5 Experimental Results
To complement our theory, we present experiments on synthetic/ real-world data demonstrating that (S)NGD- based training leads to faster convergence for various metrics compared to vanilla GS)
Synthetic Data. We first consider the self-attention model in Eq. (1) and data generated as per data model DM. Fig. 3 shows the train loss, iterate norms |/u||| and |Well, average softmax score for the opt tokens, and alignment of iterates u, and W. to the respective max-margin solutions Wmm and Umm- We compare standard GD, NGD updates in Eq. (10) (without the additional 2/3 PS updates, and joint NGD updates. We observe that all other algorithms are faster than GD, with PS being the fastest. We also see that the train loss converges at similar rate for NGD and joint-NGD, while for the other metrics, NGD is slightly faster than joint-NGD.
2-2/3,("* _ factors in the step-size),
Language and Vision Datasets. Fig. compares the train and test dynamics of various optimizers while fine-tuning a pre-trained BERT (Devlin et al., 2019) model on the MNLI (Williams et al., 2018) dataset. As evident, both gradient normalized step-sizes SNGD and SPS train significantly faster than SGD, Fig. the App. shows similar results for the CivilComments dataset (Borkan al., 2019). In App. C we present additional results on vision datasets. Interestingly, in this setting, while SPS still outperforms SGD, SNGD trains slower, which could be explored further in future work. It is also worth noting that in this setting see significant gap between Adam and SGD; this is consistent with the observations reported in Xie et al. (2023); Kunstner et al. (2023).
11

---

Published in Transactions on Machine Learning Research (03/2025)
MX2
Polynt xtep
Mab-jaint
15
1.0
1.0
19
12
PVL
090
0.4
Lterations
Lterations
L1ur aiirt
LLeration
Ite Iteransers
Figure 3: Training dynamics of a self-attention model (Eq. (1)) with data generated using model DM.
6 Related Work
Implicit bias of NNs. There has been extensive work on the implicit bias of GD for both linear predictors (see Sec. 1) and NNs. Despite non-convexity, our results are more closely related to those for linear predictors since give an explicit formulation of the solution in the convergence limit (compare to implicit formulations in termss of KKT points for nonlinear NNs; see App. D for review).
Thm. 1 provides a non-trivial extension to self-attention of corresponding results for linear logistic regression. Specifically, Ji & Telgarsky (2021) demonstrate that NGD on logistic loss under separable data converges globally, in direction, at rate ((1tt) to the max-margin separator of data points belonging to distinct classes, While the convergence rate of Thm. 1 for self-attention is slower, holds under a more intricate non-convex landscape induced by the softmax nonlinearity; see App. A.4 for a discussion on the lightness of our results. As we have seen, in addition to NGD, our convergence results also apply to Polyak step-size, which empirically outperforms NGD (see Fig. 2). We also present results for NGD with momentum in Fig. 2, which results in performance boost. In the simpler setting of linear logistic regression, this method is proven to have a faster implicit bias rate of Ø(1/1²) (Ji et al., 2021). Extending our theory to incorporate momentum in the self-attention setting is an interesting direction for future work. The loss convergence rate in Thm. 2 is analogous to the (ex)(-()) loss rale for linear predictors on separable data with т/- = n/ v/IvI(0.) (Ji & Telgarsky, 2019). The difference O(exp(-t1/3) vs. ((ex(-tt) shows up as we use smaller step-size. Nacson et al. (2019) similarly show slower rate of ((cxp(-+*/2) for lincar predictors with 7)- 1(t++)))////l//// Contrary to all these works, to the best of our knowledge, are the first to prove parameter convergence of NGD and PS in non-convex settings. See App. D for discussion.
Transformers theory. To understand the optimization and generalization dynamics, (Jelassi et al., 2022) shows that Vision Transformers (ViTs) learn spatially localized patterns in a binary classification task using gradient-based methods, while Li et al. (2023b) shows that attention maps sparsify as SGD training progresses. Oymak et al. (2023) studies the initial trajectory of GD for the closely related prompt-attention. Additionally, Tian et al. (2023a) studies SGD-dynamics for the next-token prediction task in a one-layer transformer with a linear decoder. More recently, Tian et al. (2023b) extends this analysis to the joint training of multi-layer transformer with an MLP. As detailed in Sec. 1, our work is most closely related and improves upon Tarzanagh et al. (2023b;a), which adopt implicit- piii view to attention training. After completion of this work, we became aware of contemporaneous work (Sheen et al., 2024) that also studies implicit bias in the same problem setting (single-layer self-attention, fixed decoder, Ass. 2). Unlike ours, their results hold only for gradient flow (i.e. GD with infinitesimal step-size), and are asymplotic. On the other hand, they study separate optimization of key-query matrices, albeit with appropriate initialization assumptions. See App. D for full review.
7 Conclusion and Future Work
We studied implicit optimization bias of GD for self-attention in binary classification. For both fixed and trainable decoder settings, we identified data conditions under which global convergence holds and characterized, for the first time, the convergence rates. Our convergence rates hold for NGD and PS, which are shown both thcoretically and empirically to outperform vanilla GD. In future work, we aim to further relax Ass. 2 and identify and investigate other sufficient conditions for global convergence, such as overparaneterization. Additionally, extending our theory to incorporate momentum and explain its fast convergence also an interesting future direction. Finally, motivated by our experiments, we want to further investigate why NGD trains faster on language datasets, yet appears to be slower on vision datascts. This may involve extending our findings to the next-token prediction setting, which is also of interest.
12

---

Published in Transactions on Machine Learning Research (03/2025)
Acknowledgements
PD and CT were supported by NSERC Discovery Grant No. 2021-03677, the Alliance Grant ALLRP 581098-22 and a CIFAR AI Catalyst grant. BV was supported by NSF CAREER Award CCF-2239265. The authors acknowledge the use of the Discovery cluster by USC CARC.
References
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithmi is in-context learning? Investigations with linear models. In Int. Conference on Learning Representations, 2023. URL
46
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. In Advances in Newral Information Processing Systems, volume 32, 2019. 46
Pierre Baldi and Roman Vershynin. The quarks of attention. arviv preprint or arXim:2202.0837 08371, 2022. 46
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer:
A memory viewpoint. arXiv preprint ar Xi::2306.00802, 2023. 46
Daniel Borkan, Lucas Dixon, Jeffrey Sorensen, Nithum Thain, and Lucy Vasserman. Nuanced metrics for measuring unintended bias with real data for text classification. CoRR, abs/1903.04561, 2019. URL
http://arxiv org/ abt//9////11//.////// 04561. 11, 42, 45
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert- Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, 1lya Sutskever, and Dario Amodei. Language models are few-shot learners, 2020.
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. Advances neural information processing systems, 32, 2019. 5
Yuan Cao, Zixiang Chen, Mikhail Belkin, and Quanquan Gu. Benign overfitting in two-layer convolutional neural networks. ar preprint orXiw:2202.00:: 06526, 2022. 5
Niladri S. Chatterji and Philip M. Long. Finite-sample analysis of interpolating linear classifiers in the overparameterizzed regime. Journal of Machine Learning Research, 22(129):1 30, 2021. URL http:
//jmlr.orre www.ooogg//p//..... pv2prs/v22/20000 974..97...0 html.
Lenaic Chizat and Francis Bach. Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference Learning Theory, pp. 1305-1338. PMLR, 2020. 46
Puncesh Deora, Rouzbeh Ghaderi, Hossein Taheri, and Christos Thrampoulidis. On the optimization and generalization of multi-head attention, 2023. 3, 5, 37, 39, 46
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova, BERT: Pre-training of deep bidirectional transformers for language understanding, In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short. Papers), pp. 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 1..18653/v1/N19 1/213/111..33. URL
wwwpatthhooogoooo///eeeeeee org/N19-14222 1, 11, 42, 45
Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pp. 2793 2803.
PMLR, 2021. 46
Alexcy Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale, 2021.
13

---

Published in Transactions on Machine Learning Research (03/2025)
John Duchi, Elad Hazan, and Yoram Singer. Adaptive subgradient methods for online learning and
stochastic optimization. Journal of Machine Learning Research, (2(61):211122136) 2011. URL http:
wttppapprs/vv//duueeeee http://///cchhhaaaaaaaaaaaaaaaaa htm.. 8
Tolga Ergen, Behnam Neyshabur, and Harsh Mehta. Convexifying transformers: Improving optimization and understanding of transformer networks. arXiv preprint arXiv:2211.110 11052, 2022. 46
Spencer Frei and Quanquan Gu. Proxy convexity: A unified framework for the analysis of neural networks trained by gradient descent. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.),
Advances in Neural Information Processing Systems, 2021. URL https: http:p/:///peeeeeeeeeeee/e//e net/forum?ide
WVpGLJuuPx5. 37
Spencer Frei, Niladri S Chatterji, and Peter Bartlett. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. In Conference on Learning Theory, pp. 2668-2703. PMLR, 2022a. 5
Spencer Frei, Gal Vardi, Peter L Bartlett, Nathan Srebro, and Wei Hu. Implicit bias in leaky relu networks trained on high-dimensional data. arXiv preprint ar 1:X::2210.07082, 2022b. 46
Spencer Frei, Gal Vardi, Peter Bartlett, and Nathan Srebro. Benign overfitting linear classífiers and leaky relu networks from kkt conditions for margin maximization. In Gergely Neu and Lorenzo Rosasco (eds.), Proceedings of Thirty Sirth Conference on Learning Theory, volume 195 of Proceedings of Machine Learwing Reseorch, pp. html. 5
PMLR, 12-15 Jul 2023. URL httpss:/proceeeding mlr, www..press/vve press/v195/frei255
Elad Hazan, Kfir Y. Levy, and Shai Shalev-Shwartz. Beyond convexity: Stochastic quasi-convex optimization. In Proceedings of the 28th International Conference on Newral Information Processing Systems Volume 1, NIPS'15, pp. 1594 1602, Cambridge, MA, USA, 2015. MIT Press.
M. Emrullah Idiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawal, and Samet Oymak. From self-allention to markov models: Unveiling the dynamics of generative transformers, 2024. 46
Samy Jelassi, Michael Eli Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. In
Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022. URL https: : /oppreeeeeeeeeeeeeeeeee
12, 46
Ziwei Ji and Matus Telgarsky. Risk and paramcter convergence of logistic regression. arXiv preprint
arXiv:1803.07300 2018. 1
Ziwei Ji and Matus Telgarsky. Characterizing the implicit bias via a primal-dual analysis. In International Conference on Algorithmic Learning Theory, 2019. URL https: //api. sem.ppicccaanicllllrrar wwwooorrg//rorrrre 226245936. 9, 12, 46
Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. Advances in Neural Information Processing Systems, 33:17176-17186, 2020. 46
Ziwei Ji and Matus Telgarsky. Charactorizing the implicit bias via a primal-dual analysis. In Algorithmic
Learning Theory, pp. 772 804. PMLR, 2021. 12, 46
Ziwei Ji, Nathan Srebro, and Matus Telgarsky. Fast margin maximization via dual acceleration. In International Conference on Machine Learning, pp. 1860-4868 PMLR, 2021. 12, 46
Diederik P Kingma and Jimmy Ba. Adam: A method for stochastic optimization. arXiv preprinl brXiv:1412.6988 2014. 8
Pang Wei Koh, Shiori Sagawa, Henrik Marklund, Sang Michael Xie, Marvin Zhang, Akshay Balsubramani, Weihma Hu, Michihiro Yasunaga, Richard Lanas Phillips, Irena Gao, et al. Wilds: benchmark of in-the-wild distribution shifts. In International Conference on Machine Learning, pp. 5637-5664. PMLR, 2021. 42
14

---

Published in Transactions on Machine Learning Research (03/2025)
Yiwen Kou, Zi-Yuan Chen, Yuanzhou Chen, and Quanquan Gu. Benign overfitting in two-layer relu convolutional neural networks. In International Conference on Machine Learning, 2023a. URL https: 5
Yiwen Kou, Zixiang Chen, and Quanquan Gu. Implicit bias of gradient descent for two-layer relu and leaky relu networks on nearly-orthogonal data, 2023b. 46
Alex Krizhevsky. Learning multiple layers of features from tiny images. pp. 32-33, 2009. URL https: //www. ca. toronto................ wearonto.edu/eee
pdf. 44
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam transformers, but sign descent might be. In The Fleventh International Conference Learning Representations, 2023. URL https: https://openreevir wwwpeeeewwwwieeeeeeeeeeeeeeeeeee net/forummitt a65YKOcqH8g. 11, 44, 46
Yann LeCun and Corinua Cortes. The mnist database of handwritten digits. 2005. 43
Scung Hoon Lee, Seunghyun Lee, and Byung Cheol Song. Vision transformer for small-size datasets, 2021. 43
Hongkang Li, M. Wang, Sijia Liu, and Pin- Yu Chen. A theoretical understanding of shallow vision transformers:
Learning, generalization, and sample complexity. ArXiv, abs/2302 abs/2302.06015, 2023a.
Hongkang Li, Meng Weng, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In International Conference on Learning Representations, 2023b. 12, 46
Yingcong Li, M. Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithmss: Generalization and stability in in-context learning, 2023c. 46
Yingcong Li, Yixiao Huang, Muhammed E Ildiz, Ankit Singh Rawat, and Samet Oymak. Mechanics of next token prediction with aelf-attention. In Sanjoy Dasgupta, Stephan Mandt, and Yingzhen Li (eds.), Proceedings of The 27th International Conference on Artificial Intelligence and Statistics, volume 238 of Proceedings of Machine Learning Research, pp. 685-693. PMLR, 02-04 May 2024. 46
Zhiyuan Li, Yuping Luo, and Kaifeng Lyu. Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning, 2021. 46
Chaoyue Liu, Libin Zhu, and Mikhail Belkin. Loss landscapes and optimization in over-parameleriieed non-linear systems and neural networks. Applied and Computational Harmonic Analysis, 59:85 11185---6666--1-- 2022.
Nicolas Loizou, Sharan Vaswani, Issam Hadj Laradji, and Simon Lacoste-Julien. Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence. In Arindam Banerjee and Kenji Fukumizu (eds.), Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 Proceedings of Machine Learning Research, pp. 1306 1314. PMLR, 13 15 Apr 2021. URL https: //proceedings wttpsssmmrppppreeee..eeem wwwpppeeweeeosieeseeeeee htmml. 3, 7, 46
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In
International Conference on Learning Representations, 2020. 46
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multi-head attention transformers. arXiv preprint Xrxv:22306.0 02010, 2023. 46
Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot. Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar. Attention with markov: A framework for principled analysis of transformers via markov chains, 2024. 46
Vidya Muthukumar, Kailas Vodrahalli, and Anant Sahai. Harmless interpolation of noisy data in regression. Xiv preprint Xiv: 909909099999 2019.
15

---

Published in Transactions on Machine Learning Research (03/2025)
Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent separable data. In The 22nd International Conference on Artificial Intelligence and Statistics, pp. 3420 3428. PMLR, 2019. 3, 12, 30, 37, 46
Quynh Nguyen and Marco Mondelli. Global convergence of deep networks with one wide layer followed by pyramidal topology. Advances in Neural Information Processing Systems, 33:11961-11972, 2020. 37
OpenAl. Openai: Introducing chatgpt, 2022. URL https://openai. caapaaai.cmmmbaaia hhtp///ptgghhatp......
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023. 3, 12, 46
Adam Passke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Köpf, Edward Yang, Zach DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang, Junjie Bai, and Soumith Chintala. Pytorch: An imperative style, high- high-performanc deep learning library, 2019. 42
Mary Phuong and Christoph H Lampert. The inductive bias of reflu} networks on orthogonally separable data. In International Conference on Learning Representations, 2021. URL htttp:://opp http::////peeeeee.ee/eem net//
forum?id=krr2/TxxU22 46
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and llya Sutskever. Learning transferable visual models from natural language supervision. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8748 8763. PMLR, 18 24 Jul 2021. URL https:/ http:://rroceedngee mlr.
press/v139/ httv////3//raar/20/ html. html..
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with a unified text-to-text transformer. (1),,, 2020. ISSN 1532-4435.
Arda Sahiner, Tolga Ergen, Batu Oxturkler, John Panly, Morteza Mardani, and Mert Pilanci. Unraveling attention via convex duality: Analysis and interpretations of vision transformers. International Conference on Machine Learning, 2022. 46
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. arXiv preprint Xrx:222306. 02896, 2023. 46
Fabian Schaipp, Ruben Ohana, Michael Eickenberg, Aaron Defazio, and Robert M. Gower, Momo: Momentum
models for adaptive learning rates, 2024. URL https: : //ttx//xxvvv///////////e//aaaa org/ abs//3//5/255/.. 07583, 46
Heejune Sheen, Siyu Chen, Tianhao Wang, and Harrison H. Zhou. Implicit regularization of gradient flow one-layer softmax attention, 2024. 12
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasckar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1):2822 2878, 2018. 1, 9, 46
Hossein Taheri and Christos Thrampoulidis. Fast convergence in learning two-layer neural networks with separable data. Proceedings of the AAAI Conference on Artificial Intelligence, 37(8):9944 9952, Jun. 2023a. doi: 10.1609/aaaii 2aa////aavv37/ 23718.2686. URL https://ojs http:///os.aaaaa...aaaaaaaaaaaa. aaai awai.arr/aaaaaaaaaaa.a.aaaaaaaaa index. phpp wwt.///aaaeeeeeeee/////aaaaaaa//
Hosscin Taheri and Christos Thrampoulidis. Generalization and stability of interpolating neural networks with minimal width. arXiv preprint arXiv:23022.09235, 2023b. 37
Davoud Atace Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines, 2023a. 1, 2, 3, 4, 5, 12, 24, 27, 29, 41
16

---

Published in Transactions on Machine Learning Research (03/2025)
Davoud Atace Tarzanagh, Yingeong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism, 2023b. 1, 4, 12, 24, 27
Christos Thrampoulidis, Implicit bias of next-token prediction, 2024. 46
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in l-aayer transformer, 2023a. 3, 12, 46
Yuandong Tian, Yiping Wang, Zhenyu Zhang, Beidi Chen, and Simon Du. Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention, 2023b. 12, 46
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Herve Jegou. Training data-efficient image transformers distillation through attention. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 10347-10357. PMLR, 18 24 Jul 2021. URL
https: / https::/pprcete /pprce////occdngs... mlr. pre.sppes///eeee wasss/vv3////v wtt////oovroonaaabeeaom html. . 1
Gal Vardi. On the implicit bias in deep-learning algorithms, 2022. 46
Gal Vardi and Ohad Shamir. Implicit regularization in relu networks with the square loss. In Conference on Learning Theory, pp. 4224 4258. PMLR, 2021. 46
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Johannes von Oswald, Eyvind Niklasson, Ettore Randazzo, João Sacramento, Alexander Mordvintser, Andrey Zhmoginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. arXiv preprint ar arXin:2212.07677 2022. 46
Ke Wang and Christos Thrampoulidis. Binary classification of gaussian mixtures: Abundance of support vectors, benign overfitting, and regularization. SIAM Journal on Mathematics of Data Science, 4(1): 260-284, 2022. doi: 10.1137/21M1415122 5
Ke Wang, Vidya Muthukumar, and Christos Thrampoulidis. Benign overfitting in multiclass classification: All roads lead interpolation. Advances in Neural Information Processing Systems, 34:24164 241799, 2021.
5
Adina Williams, Nikita Nangia, and Samuel Bowman. A broad-coverage challenge corpus for sentence understanding through inference. In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguisties: Human Language Technologies, Volume (Long Papers), pp. 1112 1122. Association for Compulational Linguistics, 2018. URL http: org/ org/anthology/ 110-1101. 11, 42, 45
Xingyu Xie, Pan Zhou, Huan Li, Zhouchen Lin, and Shuicheng YAN. Adan: Adaptive nesterov momentum al- gorithm for faster optimizing deep models, 2023. URL http:///ppppppooeeee/ee//e/eeeaa 11, 44
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank J. Reddi, and Sanjiv Kumar. Are transformers universal approximators of sepuence-to-sequenee functions?, 2020a. 46
Chulhee Yun, Yin-Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawai, Sashank J. Reddi, and Sanjiv Kumar. onn) connections are expressive enough: Universal approximability of sparse transformers, 2020b. 46
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context, 2023.
46
17

---

Published in Transactions on Machine Learning Research (03/2025)
Appendix
A Training only W
18
A.1 Preliminaries
A.2 Key Lemmas
A.3. Proof of Theorem
26
A.4 Optimal Rate for NGD
28
A.5 Proofs for Additional Results in Section 3.2
28
B Joint Optimization
31
B.1 Preliminaries
31
B.2 Key Lemmas
32
B.3 Proof of Theorem 2
39
B.4 Proofs of Theorems and
40
C Additional Experiments and Settings
42
C.1 Initializing a Bad Stationary Direction
C.2 One-layer Self-attention Model Trained on Binary MNIST
45
C.3 Validation of Assumption 1
45
D Related Work
46
A Training only W
A.1 Preliminaries
| denotes the spectral norm and (A) denotes the minimum cigenvalue. Also, | max|a:l, where a, is the fth row of A. We list the useful notations used in the Appendix in Table for convenience. Further, let (0(v):: == Vyp(v= diag({<v))) (p)) denote the Jacobian of the softmax vector yp (pvv) at
th token in the th sample
Figr a( a(W)
W
XW
yXu.
(.)
softmax vector
20) r [T] for sample at time t Wumll
A
B
max; |
(nju(7ri.spt -X)
min exp(7so))))))
r
ks max (7x
max max (re.. -%) maxexp(1x)
Ç
Table 1: Useful Notation.
18

---

Published in Transactions on Machine Learning Research (03/2025)
E R We note that under Assumption 2, for any W, we have
(-VwI(WW,, w)
4, (y) WWW X.). W) Σ
n 1=1
Ersopt, Pia Ersopt, Pir
(Ti,oop - -.ppiopp (1 phope,
ahopp,
(14)
where a; X,Wx,,, x1,1.
Intuition for the Constraints in Equation (W-SVM). The intuition behind the constraints in Equation (W-SVM) is follows. Consider the loss for the 2t sample:
4 = (y,u ((X,,,,,,.)) (()) ((,,) = WWW, Proope, + +7 (1,,opt)))
Since l is decreasing function, it is minimized when Pi,opt, 1. Next, consider the softmax weight for the optimal token:
exp( W,ar)
exp( War1) 142 1oor (xp W.ns)
For loss minimization, this score should tend to 1 the iterate norm WW, tends to infinity. This only happen when the exponent is negative, that is, W, satisfies (x,, Topt) translating the constraints to the SVM for training linear models on linearly separable data, we can obtain Equation (W-SVM) for self-attention.
" Waa 0. Then, analogous to
A.2 Key Lemmas
We first prove some key Lemmas that are useful in the proof of Theorem 1. Lemma 2. Let wi.v ~ ((0,oo) for all u, u [d], with the initialization scale
To s 2d 1/2 (C" log T-1 nd 1/
KK-1)
for some 1 > 1, and absolute constant C' 0. Further, if 2 log(2/5) for any 5 (0, then with probability at least 8, we have vt. olo, ET for every i [n].
Proof. Let wh. N(0,o)), for any u, E [d]] Then for any i e [n] and e [T], for any E (0, 1), with probability at least - d we have
(Ri, Fipopt,) ) Wo s (2 Fipopt ) Waal
$ |iir - oi, Wo Wo 11112
2B ( ((oog2/))) (20g()
GCB3ooVd B² 'C
Using this, we have
Phopt,
1+ Σ exp((xi,r otpopt.) ) Wo reopt,
1 (T (T=1)xxxp
AT
(using do C0-1 112 log
KT
T-1
Also, W. 200d (ABB2A))-
19

---

Published in Transactions on Machine Learning Research (03/2025)
Lemma 3. Under the conditions of Lemma 2, Assumptions 1 and 2, using the updates in Eq. (2) with loggT- at any time i 2 O, for any sumple J [n], the softmax weight for the opt, loken satisfies
S3 2 Pjoopp, T
(15)
Proof. We prove this by induction. Clearly, Eq. (15) is true for t= O. Next, we assumne that it is true at time 1-1. 1. Combining this with the conditions of Lemma 2, we have 10-11
Pjoopt, 2 2T
Further, for any j [n]] we can have one of the following two scenarios:
Scenario 1: At time t L, Phopt, 1 2T This implies
> 1
t-1 jiopt
1 Eroopt, exp( (Zi,r Fjoopt, Soor,)(W,,.. (-111
2T
1
Σ exp( (z), )joopt W1 2T
T+opt;
(2)r - )opp, ) W1,1 S
Vr + opty.
After the gradient step at time t - 1, for any T opt,, we have,
(D Dsoopt,) )
"' Wa (æj,r jopt, W 1-1 7h- 1 (-F lo - 1) + 7/ u5 opt, æ,,11
- log(2T - BnB²
0,
where the last step follows since n s - 1). This gives
Σ
exp((xjrr - )jooptt) TT W s T
reapt,
1 joopp, 2 T
Scenario 2: At time 1, 1/jopp, 5 2T This gives
2T 1
Foop,, (1 - (1-4//ppe)) )
(16)
4T
4T"
We will use this to show that
(2)jopt, (-Vw(((-.)) 2,,1 20
(17)
Let a' = a((x)oop,, 2z,))1)11) X X( (I)oopt, Fz (20,) x1 Using similar calculations Eq. (14), we get
(I IJ,T) I (-VwI((--))))
1 a
E Ercopt, 1it d/ Eroope, s
Σ 241, - -wpttope (1 (1-ptoppe)) aliope
(18)
n 1
There are two cases:
20

---

Published in Transactions on Machine Learning Research (03/2025)
Case 1: opt; 1. In this case, by splitting the sum over [n] in Eq. (18) into i and i J. and using non-negativity of the softmax weights, get
(Bs,opt, (-Vw
4-1 (1j.opt, 7/5 73))93.0pp,, (1 - 110ooppp |1 r'sopt,
min (Tjjopt, Fjjr) (Fojo " - xz,+))
n
max 4-1,1 (H.opp V6 p (1 Pioope )12 و3,1
max (0z),ops, - -z,,r) =,,,))((Z,,op,,, - -I,r)).
(19)
r'oopt; Next, we use Assumption and that nTT2 2 to get
izj
max max (2tj,pt,, - -.j,)) (Fs.oop, Bi,r)) ) max jopt,, Droop iyy r'sopt,
+33 max la [ [.ttej.e([]] ['ppt]
tzj
12
$ 4
1zs,opt,ll 11?
(20)
4
Similarly,
min (Ij,opt, (osr)((as.ppp))) xirr) joppp, s'eopt,
12 - max j rc(T]. r'e[T]> (opt,)
22-3 j.ppppt 12
j.opp,,
x j.opp, 11
(21)
4
Combining Eqs. (16), (20) and (21), we have
( === n maxi4j la (7H,op - 75) (1 -iippp, =axp, (2.joot,, xzj,)) ) (FFi,opt, - -xi,r') min,'=opt, (o),, - -z,77 (TI)))()) (2), -Zj,7()
l (eslop -75) -35)4 (1 phopp,
max; l- 1. (7H, 7;) 1/4 pap 1?
n
min, 1,115 (.Koppp Khopp, - 74) 1/(4T) le0 12/4
< 4n TK max; 4-115
min; f max; exp(-(xff,pp. Nropp, (1 (1- )v:) sp/ips min; min,exp( exp(-(xff)) (iopt Ni.opt p/ )r1)
4n'To TK
max; cxp(-7)) = TT..
4n'TK
(22)
min; exp( -h,opt,))
Using Assumption and Eq. (22), we have
|, [ 12
which implies that Eq. (19) is non-negative.
Case 2: opt, = 1. In this case, let I = (i [n]] : us A opt, 1}. Using Eq. (18), we
get
(Ij,opt, - (-W
l- n
(Thoops, - -75) (1 -1--yj..pe || 12 min (I - -xjr) (x - -0jrr' roopt,
min l 4e-1,(7i.ppt
-:) 1-1 (1 pl.ope, )ar a Fga min (Ijoopt, 2z,r)" (Ii ( Fi,5t)
stj,ieI
proope,
max (7h,opt Y/ (r:)yp.- (1 sp. )|ا) m max (T)j,pp, - =zjr)) ( o(,)()) (23)
stj,if
roopt
21

---

Published in Transactions on Machine Learning Research (03/2025)
Using Assumption the second term is non-negative. For the remaining two terms, proceed in similar way as the previous case to show that using Assumption 1, these terms are also non-negative.
Using Eq. (17), we get
(zsopp, -3,,) Wi-1)a5, = (Ij,opt, (-VW))(-VWWW--))) [WwI(W
2 20.
Then, for any o have
P)oopt, /1
= exp ((zj,opt, x,,r)'(W, Wi-1)a5,1 21 1.
j.opt, /50
By telescoping, we get Len 2 which implies that phopp, f.oop,,
Pioy,
Lemma 4. Under Assumption 2, for W mm defined in (W-SVM) and any W,
VwI(W)
W mm
2 ~2B²A)~~ 0.
[ww)(l)
W m
Proof. First, using Eq. (14) for Wmm. we have
1 n
(-VwT.(W), W Wmm) 2 min(dhoop
max (4,7) rtopt;
Σ
le (H,oop, - - Pt.oo, (1 pi.
1=1
Σ
l (Ti.opp - Psooop (1 Pi.opt.),
(24)
where a; = X2 W 1,1, and the second inequality follows by the definition of Wmm.
Let W [(W)l
VwI(W) and à. = X, KWW.Y 1i,1, then by similar calculations as Eq. (14), we have
VV (wwT(W)) P (-VTT(W), W)
Extope, PPrr di.r
(7i,opp - Pi)pp, (1- Pi.ope,) Σ
di.opt, Ertope, Pir
n
n
s d
max
min âi,r) reopt,
Σ
L CiiVi, - mi)spropp, (1 - :Pi.opt;)
[=]
1
2B
(Nhopp, - Proop, (1 Ps.opt;);
(25)
n
where for the last step
max (ax.api.
min (i.r) s 2 max |āi.r| = 2 max |.. Wx,,1
reopt,
mmax X211,,0 |x;, W 2B²
Using Eqs. (24) and (25), we get
VwI(W) WwI(W)]
W mm
(2BRA)-1.
W
22

---

Published in Transactions on Machine Learning Research (03/2025)
Lemma 5 (Iterate Norm). Using the updates in Eq. (2), under the conditions of Lemma 2 and Assumption 2, at any time O, we have
An)(4B²A)~ 1)t2 2 (1B2
(26)
Proof. Using Eq. (2) and Lemma 4 with W Wt we have
W = Wo Σ
te
V
w
T T(WW))
P=0
1-1 Σ
I(W,)
W Ymn
(o,
W mm
I(W))l
W
f'-0
(2B² Wo W mm))
Under the conditions of Lemma 2,
( W,, W mm)) | Wollssn(4 (ABBBA)
which gives the required result for 1>
Similarly, for the upper bound have
11 VeV V T((W))
W
Wo
Σ
(-1
[Wol Σn "
1'=
$ nn4B² A) nt.
Lemma (Data Example 1). Using u. #, v_) and d 2 M(n log (2x)) data generated as per Example 1 satisfies Ass. 1 and 2 with probability at least 1- for any d (0, 1).
Proof. Using u. H (u u_), for all ie [n]] and 5 opt, we have
Thop yupp Hh = yulw 0.
Using this, we have
max max(Th,ppt, (Hoopt, ()) min ein(7x777ippp.
T min (Hi,pp, 76)) max exp(1:) e.
Next, by using Bernstein inequality, we know with probability at least we have
In for every i [n] and 7 # opt,
4n
o²d d
|10 12 s0²2
co²
dlog
2
/dog || 10 p²d| cp² =
(n
p²d
2
where e> O is absolute constant and we use d 2 4ca log
in #x)
Further, for any ú, [n], j and any T, e [T]. By applying Bernstein's inequality, each of the following is true with probability at least 1 2n
4n ( s
4n d
(xt,, 2jooot,)) (My, M 202 V dlog g
8
4n 10 Dirid 20 dlog
23

---

Published in Transactions on Machine Learning Research (03/2025)
where r' opt,, opt,.
In order to satisfy Assumption 1, ignoring the log factors, we require
oU²+o'd) p'd & nI(p opvd) +²+ nT pdv V
The above inequalities hold if we have d = T²) and U² á(( v p) The proof finishes by the application of & union bound.
We now state the key Lemma used in the proof of Thm. below, with the expression for R.. Lemma 7. Under the conditions of Lemma 2 and Assumptions 1 and 2, for any e (0, 1), there exists
R. = 2AE 2.- (log(An( -TT 1 500g(2
such that for every t where Wil 2 Re;
W mm
W.
V ((W,,
2 (1-
- T(W.)
(27)
W mmml
W.l
Remark 1. Lem. improves Lem. 10 of (Tarzanagh et al., 2023a), which only holds for n= i.e., for a single training sample (same restriction as (Tarzanagh, et al., 2023b)). Our key idea to extend the result to n21 is to divide the samples into three sels, based on whether the construints in (W-SVM) are violated (I, ) satisfied with small margin (Iz) or well satisfied (Is). We prove a tighter version for somples in In and the residual allows us to prove the resull for the sum over all samples.
Proof. Let ā; := X, W aj X;W Ii,1 We consider two scenarios based on whether IW W Wmml S 2Ba (Scenario 1) not (Scenario 2), where
W W mmmll v:= (1-e)
W.
(28)
[
Scenario of the proof is the same as (Tarzanagh et al., 2023a). Below, we consider Scenario 2.
Using the definition of v and similar calculations to Eq. (14), Eq. (27) translates to
12 Σ le (Thhop, popp (1 Ploopr) S (1 - Σ (oooo, - -n)phopp (1- ppopop, )i,
(29)
n
Exarge air and hi, is defined similarly. Ererpe phr
where he, ai.pp,
For any sample i E [n], let
max (c,w):= == max 0i,opt; - āj.r 1, r((opt,)
guai (e,w) min āi,opt, 0i,r 1. rf(opt,)
Define the following sets
I = (ie [n]
8min s $..vv),
Σ = (ie [m]
0.3v 5mmin 0.8v),
I3 = (i e [n] gmii 2 0.8v).
Intuitively, I₁ is the of samples for which some (W-SVM) constraints are either violated or all constraints are barely satisfied. Similarly, Ia contains samples where all the (W-SVM) constraints satisfied. Finally,
24

---

Published in Transactions on Machine Learning Research (03/2025)
Ig makes up the rest of the samples.
First, we will show that for any i It,
Ersopt, Pir 0i.r (1 + 0.5v) opoo
Etsope, Pip Ersope, Pi,
@r.opt Ersopt, P/ttr
(30)
Next, will show that for any ie I2, have his s (1 (1+2) ie.
Erropt, Pi,
Ersopt, Pi
dit. Ereopt, Prr,r
s (1 (1+v)) "hope, Erropt, Phr
(31)
Finally, notice that in order to complete the proof for Eq. (29), using Eqs. (30) and (31) it suffices to show that
Σ (iioo, -x)+p% (1 pho (hre - (1 ()
vell
S 0.55v
l,,(Tpopp; - p)pp pl.optt)) hi, 1EI1
Σ
(32)
Using N7 2 1, Eq. (32) follows when
Σ M
(T,(To ph.op. (1 (/u, 1- 0.5v E
Ces(Hopp -x) (1-0,,,,1-,,,,,, (33)
Case 1: Sample i I, i.e. 6rmin z 0.3v.
If 8max < 0.5v, (30) follows directly. For the rest of the samples i where gmna 2 0.5v, let
N {T [T] ai,opt, - Ai.r S $1+0.4v)
W. Let R Wmml
then by definition of N, 8iii, we have
Σ Pi,
seopt, +N
Tea sir
T Texp(-R(1+0.pp) exp(-R'(1 +8min)
Teopt,
$ Texp(-O.1Rlv).
Combining these, we have
Σ (ūi,opt, āi,))p!.
Σ ot
Σ
T4opt,
(1+0.4v)
TEN
8max)
Teopt, . TFN
Σ P/r Teopt,
Σ
Σ 10t reopt,
reopt
Σ
ett Teopt,eN
Σ
Σ
fot
$ (1+0.4)
reopt,
(1 5mmax
TEE
Σ vo
r+opt,
Toopt
$ $1+0.42 0.40 (8maax 0.4p))exp(-0.111
satisfy (30), we need (8max 0.4v) (xtp)))11p-- S 0.1v, which is true when
R' (og
This is true since R' > 10e-1 (0gg10mmmxxff gmax/ where we use the definition of v from Eq. (28).
25

---

Published in Transactions on Machine Learning Research (03/2025)
Case 2: Sample i é Iz i.c. 0.3v gmin $ 0.8V.
If 0R v, (31) follows directly. For the rest of the samples where 6 2 we can proceed in a similar way as Case 1. We use threshold of 0.9v to split the tokens, and obtain
Σ (di,opt, ai,r Per,
Teopt,
$ 11+,
Σ Pe, roopt;
when R' 2 10e- Loge= max //nax/e).
Case 3: Sample i Iz, ie. p. 0.8v.
As W. mmml 28 at least one sample i [n]] violates the (W-SVM) constraints, i.e. [ 21. 1. We have
Σ H.lhHopp, wlplopp, (1 pioop,
min &, (hoppt, 7)) minn(pi )² Σ
T=opt;
min ieZ,
4,, (Moppt 7)) exp( 0.3v) (R.32)RR)) R'),
(34)
where the last inequality follows by using Lemmas 2 and 3, and that 6:uin z 0.3v for any i In. Also,
Σ L.i(Toppp; - 78) 1i)sph.pp,, (1 ppoope,) s (n max l,, - -%) max Σ IeIs rtopt,
ieIs
$ (n - 1) max 4, ("h,oo (T exp( (1 R'), (35)
ieIs
where the last inequality follows because 0.8z minnez, 8miin Combining Eqs. (34) and (35) gives us (33)
when
maxieTa (,, (,,(7i, 71)
exp(0.5vR') 2 <0001>/))3..5228110000--00--1-00
(36)
minser, lass (H.opp, - -71)
which is true when
exp(0.5vR*) -
which is true when R' 2 2c log((n(( (og(A((""))--TTT>>
Combining all conditions on R', and using 8maa s 2B²A for all i [n], we get the desired result since
-='=2A4 R. A- = 2Ac 20 (log(4((B"A)) 220g(20BB'A
7
A.3 Proof of Theorem 1
We first restate Theorem 1, this time with the exact constants.
Theorem 5 (IB Rate). Under the conditions of Lemma 2 and Assumptions 1 and 2, using the updates in Eq.
(2) with s 1ogg2- 1), for any 2 to == 10BA vn
3
v log TTTTT(B 2T)
20A
W.
W
C(n, B,A, to (logt)
2
W.
W mm
where C(n. B, A, 1o) 4nj A W
(W to
W mm)) 2B 2B2A(40A +1))
26

---

Published in Transactions on Machine Learning Research (03/2025)
Remark 2 (Comparison (Tarzanagh et al., 2023a).). Thm. 1 estoblishes finite-time convergence of NGD to Wo starting from any inilialization direction. This marks clear improvement over prior work ('Trraanahh et al., 2023a), which demonstrates that GD converges asymptotically (t + 00) ) to Wmm only under on appropriate local initialization direction. The only previously known global convergence result in Tarzanagh et al. (2023b;a) requires Ass. 2 and n 1, limiling ils applicabilily to training with a single sample. Thm. addresses this limitation and additionally provides convergence rates. For completeness, we demonstrate in Thm. 6 in the Appendix that GD (with constant step size) also converges globally to Wmm under the conditions specified in Thm. without requiring 1.
Proof. First, use Lemma 5 to show that Lemma 7 is true for any t 2 to; by selecting (80B²A² A log) Since logt < 1.25t for t> 0, 6 $ 1 for t (10BAn-1/3)) 1/333. We also have 2A (log( A "T 1) v 55 Ae )) s 10Ac logg( (4m ABB"A) A) 'T v A0T
10A
Ty
log (4nT(B2A)-TT> 20TB²A) S0B²A² log
80B²A² logt
((4BAA) log/ log logt
(T( 5)
20A
$ (4 -
Thus, |W 2 R for 2 to and we can use Lemma 7 with c= ces i.e.,
(- (WW))
W mm
W VV T(W.),
2 (1)
W mmm
Using the update Eq. (2), we have
W
W
W+1 W.
2 (1 -c) c2)
W - W.
W mm
W
1
W. Will
[[W+1ll
2 WWoll | Woll3) Et
W411
W.
2|
1
(IW
1|
[W,1)
+ + Vw W.)
(using Eq. (2))
2|
2 W:
n/2
a²
2 W 11-1l-1 - W..l 2 W,
CJ
(sinc
2 a- ba-bya aabbva,>>0 0)
2b
80B² A logt
WW [W||
A2 (1+1)-1
(substituting et, using Lemma 5).
Summing over t 2 to. we get
WW. W mm) 2 Wil - W 1 (We., W Wum) S0 A2 Σ (logt')(t)) BnB²A Σ (t' + 1) (37)
t'oo
t'=o
We bound the last two terms as follows
Σ (tggt
s logt Σ (t')- s (logt)?,
Σ (1' 1) s log
Using these in Eq. (37) and dividing by Wil throughout, we get
W. W mm
2 |Will (lW.,ll - (W for W mm) + +2B 222AA(40A+ +23²A(40A+))
[|W]|l [ mm
21 C(n. B.A, to) (logtt)
t
where the last step follows by Lemma. 5.
27

---

Published in Transactions on Machine Learning Research (03/2025)
A.4 Optimal Rate for NGD
In Th. 1, we showed an upper bound on the correlation W t: W mm/ s s1- for NGD. In this section, we consider the example used in Sec. 3.3 to derive a lower bound on the correlation.
We restate the example here for convenience. Let n 1, 2, y = 1, æ1 [1. 0]", 22 [o, and u. [1,0]*. As we saw in Sec. 3.3, opt L, Topt -7 7-1, 1, Wmm (x) 1 -x2) X and A A= Using this, we can write = ([()))-I(W))))
(1- oppe ) W Wmm: mm: and W Wo + n,G mms
where G mm Σi=0 IwI(W (t + 1) mm..
IV IWW,))
mm-
Clearly, [Will nt + [Wol s Ott.. Using Eq. (3), we have
[W W mm) 1- Woll (W W.
$ 11))())).
Comparing this with the upper bound, that the rate in Th. may not be the optimal rate for NGD. However, it is the first finite-time convergence rate for self-attention. Improving this rate is an interesting direction for future work.
A.5 Proofs for Additional Results in Section 3.2
Lemma 8 (Softmax score rate). Under the conditions of Lemma 2 and Assumptions 1 and 2, using the updates in Eg. (2), for any i [n] and any 2 2 (BBAA)) B,A, v to, where C(n, B, == BB (1I (Wio. W mm) 2B2 222 1B2AA + m)), to = 10BA
v log (K
20A
Piopp, 2
1+ 1T((=-)exp(
Proof. Using Theorem and Lemma 5, we have
(x,,r T(o,op.)) T (W, W Well W [W1))a,11 2BB W. |ll| W
W
mnll
Wil
W [mmm]
22 C(1, B,A,
(2nt) - -
/1/2
AnB² /2C(n, B,A,t)) (.AB
(1BB
S s-n(8B2A2)
(38)
since 2 2" 2BB B,A,
We use this to find the softmax rate as follows
1
hhopt, Eroops, exp( Ti.opt,) ) W
1 Ertopt, exp(( Dipt (W, W mm W. W mm
W.
2
1 (T- 11 11)=1)xppppppp (x)3
28

---

Published in Transactions on Machine Learning Research (03/2025)
Theorem 6 (Asymptotic Convergence of GD). using the standard GD updates with 1j s BCC
Under the conditions of Lemma and Assumptions 1 and 2, for any 2 te such that [[:]| 2 R. 1/2,
log(2T-1)
W
W
2 c- C(n, B,A,)
W
W mm
[W
where C,,, 254 (1 e) Wt. (1-(()))) ( t. W mm) 2n) W., 1 I(O..))
Proof. First, show that Lemma 3 holds for standard GD updates, i.e,, for any 12 0 and any j [n],
Sioopt, 2 Pjoop,, A T
Following similar steps the proof of Lemma 3, consider two scenarios, Scenario 1: At time L, Pjopp, 1 2T This implies
Pjopp, S - 1ogo - 1), Vr V
After the gradient step at time L 1, for any T opt,, we have,
(0 ks,opt,) W = (æj,r - - Foopp, ) W - 77 (.ooo - (-W,
s log((1 1) 7/ 5, x jot, 20, V
$ log(2T 11) (2B²)) ( (B*T(W,-) wax 74 phopho (1 sp(opp.)))
$ logg(TT BB max exp( (-7)) max max - 76)
$ 0,
where the last step follows since 17 s
Log(2T-1) B4G
This gives phoopt, T-
Scenario 2: At time t 1, phops, $ 1
1 2T
The proof for this part is the same Lemma 3.
Using this, we can easily show that Lemma. 7 is true for GD updates, ie.. for any e (0, 1), there exists
R. 2Ac (log(4((B"A)) v 5 llog (oog(20TB*<>*)),
such that for every t where |Wl 2 Re,
W mm mmml
W. (1- (Vw-e)-WWW(W.))
IwW,))
[W
Once we have this, the remaining proof the same as Theorem in (Tarzanagh et al., 2023a).
Theorem 7 (IB rate for Polyak-step). Under the conditions of Lemma 2 and Assumptions 1 and 2, using GD updates with Polyak-step, with S (2/32)-1un (og(2T-1), where w1 = 3(1 - (8AA min;(7H.opt, - -7;,, for any To 2 22 to w2y LOBA
(T(TT(BAA)" T)5) )) where W2 0.5 max; (Tr,opt, - -7), such that for 20A
log
/
any ie [n], vf.,op, S 538,
W.
W mm
21 Cl1, B,A, (n),),,)))lo0)1
W.l
W mmml
where C(n. B,A, A, 1o) Anj B Awzz (1
-(W.,
W mm)) 2B² Aw (40A meue (
to:
29

---

Published in Transactions on Machine Learning Research (03/2025)
Proof. To extend our analysis and results for NGD updates to the Polyak-step, we mainly use the following lower bound on the gradient norm,
W mm
V ((W)) L(W)) min VV
[
Exeope, Pi. Exsope, Plan
= - - min( 7:) (1- piops) ) oiops.
2 2A min(m..pp 76) min Phopp,, (1 vioope)) i,
2 2 9)A min(W..ps, - -70 (W.)) WW.).
Note that similar lower bounds have appeared when deriving loss convergence rates of lincar predictors trained with NGD on separable data (Nacson al., 2019).
We also have the following upper bound,
VVII(W)) s ((W.)(1/4) max(<,,ps -
= -
First, show that Lemma 3 holds for Polyak-step updates, i.e, for any 12 0 and any j e [n),
pyap,, 2 pjoop, ^ T
Following similar steps as the proof of Lemma 3, we consider two scenarios,
Scenario 1: At time - 1, Phopt, 1 - 1. This implies 2T
Poope I -- log((T Vr opt,-
After the gradient step at time t - 1, for any T opt;, we have,
(æ - Bj,opt,) ) WW = (x)r )g,op; ) W1-1 1/--1 ( x3,5) (-v=))(-Vvwll) (-Vw-((:--)));;; T(W,11) D. V VV [(W 1:03.7- - Fj.opt, -
$ -loog(2T-1) - 1)+ 1
I(W(-1) (2B²)
looog wi I(W,-))
2B
= logg(2T
W1
<0, 0,
where the last step follows since 1 1og(2T-1)
This gives &hoopt, +
2134
Scenario 2: At time - 1, Pjope, s
The proof for this part the same as Lemma 3.
2T
Using this, we can easily show that Lemma 7 is true for Polyak-step updates, i.e., for any (0, 1), there exists
R, 2A ( v v (
such that for every t where Wil 2 Re,
W
W. (1 (1-6) v.
V.w..W))
(39)
W mm
W,
30

---

Published in Transactions on Machine Learning Research (03/2025)
Next, following similar steps the proof of Lemma 5, the iterate norm growth be characterized follows. Using Lemma 4 with W Wt, for any i <T,
W.
Wo
Σ M=0
Vw I(W.)
11 I(W)-
VwI(W, [ww((W))
W W
Wo, W mm
Σ Vww((W))
mm
2 (2B² Aw2u)
+ (W, W Wmm)
n Awzu))
since W2 1. Using this, and following similar steps as the proof of Theorem we can show that Eq. (39) is true for any 22 ty. Using & (80B² Algg (nt) we have,
W Wmm
W.
W+1 W b,
W - We 2 (1 62) (2)
W mm
[W.1
172
(I(W.) )z + 62
I(W.)- -T- WWII(W))
([W -1,l)
(w.,),,). W.
w|
W [(W:)l
W W..
172
47
2w7 We
W1
2 W WWI 80B² A w (logt)t 2T Aut (or2/r(t 1)-1
Then, telescoping over t 2 to: we get
W
W
2
Well (I (Wi.. W Wmm) + 2B2 A2, (40A + '40A
+neww (2) (log t)
W1
W ml
21 C(n, B, to
to (logt)²
B Joint Optimization
B.1 Preliminaries
We first state the complete version of Condition below.
Condition (Complete version of Condition 1). Let = apy 1.5d. We consider the following conditions for any ò 1),
Zero initialization: |uoll = 0, Wol 0,
opt loken has larger norm: o 2 6,
Sufficient overpanameteriaation. d 2 Coa'n' log
10n
where Co is an absolute constant,
Small slep-size: S 18a4 A ap log(2)
128 1.5mm where 13(B ()
B
160m
3B
Lemma 9. Under the data model DM and Condilion the following evends are brue with probabilily al leust 8,
Ej := { papp'd | 1oo.op,,22 12 pdd Vie [n]].
Fg = (e'd
||
3d'd - [ opt),
31

---

Published in Transactions on Machine Learning Research (03/2025)
E3 == op op, 2a²p" dlog 10n²
Vi+ j [m]],
= ( 3y ) 2p² dlog 10n2) Vi+ j [n],-' T' opt, opt,),
Es == of 33 2ap2 dlog (10m²) Vi j 7
Proof. First, by applying Bernstein's 's inequality, with probability at least 1 on for every ie [n]]
10n a pdpdd
opp a2= a²p²d ca²p² p² dlog
5
2
where e> an absolute constant and we use d (2c)2 log (10m) in the second inequality.
Similarly, since d 2 (2c)² ( (L0m), ), we also have that with probability at least
on for every i é [n].
6
T + opt;,
10m
p²d 2
- - cp² d
i + and any T, e [T] since (0), 2z,rt) is zero-mean. By applying Bernstein's Further, for any i,j [m], inequality, each of the following true with probability at least 6m2
1002 )
10n2
".ooptt 5 20 p² dlog
(1 s 2p"
dlog
8
10n2
(a F. s 2ap²
dlog
where 7' + opt,, opty.
Applying a union bound over all these finishes the proof.
B.2 Key Lemmas
We first state some intermediate Lemmas that are useful in the proofs of Theorem 2 and 3. Lemma 10 (Iterate Norm). Using the updales in Eg. (10), al any lime L >> 0,
Ju s 3nt 2/3
Proof. Using Eq. (10) and triangle inequality, get
v..((.)))
[-1 in (1+ & 33 v.I(O))
uu S (u,11 +
22/3
F=0
Lemma 11. Under Condition and the dota model DM, for any j e [re], any i> 0,
Sopp, - 3 nu2 Σ (t+ 2/3 pi. 2 1 and max s C, i.je[n] Eej
les
1=0
where wa w2(a,, p, oe V 10n²
and G 24 is universal constant.
32

---

Published in Transactions on Machine Learning Research (03/2025)
Proof. We will prove this by induction. First, we prove the three parts at t == 1 as follows. Since ug O, for any j E [n]. fo.; L. Using this, We have
(Thops,,
- - -===(((..... (xF,)))) - -.j.7)
7u
( (*f,))t)) Fzrr)
"
77
(f/),op, 2soopt phops, )zi ( (z)jopt,, - -xr,))
7
170
Σ (i,Wi ( (Ph,opt, Vi.opt, (1 - -i,op, )abirr) ( - -(
(#)jopt,,
12
70
1 a² p2p
3p2d ap² V dlog
10n²
2(n 1)Mpp² a2 Vdlog
10n?
2 2
2
8
8
70
a
pdd
π 16
since a 2 d (16n)² log
10n" 8
Also, (voo(()) $ mmxx Xi112,00
s a.v/1.d.. |
Using these, we get
16n a² p² Vd16n log
10n²
(1)ops, -7,) >
apvi..5 1.5d
nap
10n²
log
2 1J42:
V1.5
8
when n Next, we have
si.opt phopp,11 /49,
exp (1t)oopt, - -. (W, - Wo))os.)
= exp (arjoopt, - -Dj,r) (-V
since up O. Therefore, shopp, shopp,, 2 1
log(2) Further, since 7 3apvI.5d we have
max E1 i.j[n]] li
s max (xxx s (xp(2)u1
$ (ssp(naap/1.5d)) s exp(2log(2)) 45C.
Next, we assume that these are at itcration t. Let M, max; b,, and m, min, las- We will first show that for any [n],
175-1 := (7),ops, -23+1) - (75,ope, - -2ÿ)
(+1
= (-V T.(,)) (I.),op,, - -,,)) 2 n7(1 + 1)///
(40)
33

---

Published in Transactions on Machine Learning Research (03/2025)
We have
<< n(-vuI(0.)) (#j.opt, -Z,,)= (1, (((aa)))))(( (W.)))( (ojoott =x
(1((9s. Rhopp, + +1 +11-11).opp Pl.ope, (z,,r)
(I - -xz,))
+ E liiViii (e..iissPPh.ppp. Poooop,
Di,opr, Pt. (1 piopet ( - -zj.))
a 3pdd ap dlog
10n²
10n²
2 m4
2(n Map
dlog
2
8
10n2
10n2
a p²p²d
3p"d
ap? v aloo
(10
=2p² Mig dlog
2np² a mi
dlog
2
m
2mp²a² a dlog
10n²
0
10n2
10n²
M
2 220))
dlog
2mp² a m&
dlog
8
m
16m
10n² log
5
10n²
2 22'' 2(())
dlog 8
(41)
where the definitions of M,, me and the first inequality follows by using Lemmaa 9 and the second part of the IH, the second inequality follows since M, I(O.), a 2 and 2 4 4lg 10n
_or ), and the final inequality
follows by using the third part of the TH and that d (16Cn) log 10n"
Next, have [V,]((,)] max;{[]] X 1, I(O,) s apv I.VVI. Combining this with Eq. (41), we get
( - -1
2pa
10n
log
2 W1,
((.))l v1.5n
5
which implies Eq. (40). Using Eq- (40) and the first part of the IH, it follows that
1+1 Hopt, -33 2 Trop,, -7s + nw/wr 2 1pl2 Σ (t' + 1) 1
P'=0
Phoppt, 2 2 We consider two cases: log(3) Case 1: v),ope, 2. In this case, we can directly use the steps in Scenario of the proof of Lemma 3, and s+ll get that jopt, since n Pdd Case 2: v)o
Next, we will show that for any j [n], t-1
3. Using Lemma 10, we have
Mope, -7/ 2/ max| X2 s 6n
Σ
(1 (') 2'1- 2/3
opv 1.5d. 115d.
(42)
#=
34

---

Published in Transactions on Machine Learning Research (03/2025)
Next, using similar calculations Eqs. (18) and (19), for any j [n]], we have
(ajpopt, (-v
Ru (1),opt, - -))) (1- Phope, | (e),a11123y,ept - =,,T1
(n m--1 1)
max 2,,1 ( (hoopp - - (1 ll,op Ma0, a 13,11 -
#4
me (nuz Σ
411
(t' 1) 2/13)
3 p²d
p'ppd
p²d
20p2
dlog
10n²
16 2
2
2
5
(n- 1)
1-1
10n²
10n²
M1 (7(7.5nppV
Σ
(t' + +1)
+11-2/31
2/3 ) p
dlog
(a² + 2a + +)p²
dllog
5
6
--=
a²p d312
1-1
3
10n²
10n?
(17
Σ
(t' 2/3)
dlogg
mr
32a (n Me log
n
P00
20n
320 3 2//3/2
11 (t' (1)+)) 213) log
10n²
d
Mr
n
(n
P-0
[((6.)
nmm
210
log
10n
my
32a peasp&aao
L-
1102
(11
Σ (1'++1)
(1'+ 1)/ log
T.(O:):
(43)
where we use the definitions of M, and my, and the second inequality follows by using Lemia 9, the first two parts of the IH, and the final inequality uses the third part the IH and follows when d z C2²"Ca""²) log 10n") This implies that 8
Phoopt, 09o0p,11997
exp (<xj,pt, =xjr) Wi)a,,)) 2 1.
Ph/ops, /49/
Therefore, Poope, 2 "pppp 2 1/2
Finally, we will show that A maxx,eelm Ben
fass C. Let sp (W.)((W)).. For any i [n], we have
L = (x.==exp(-uup) y-u x
= -exp(-y;,, X. (vp - -pt) eee(ppyppp (y; X Xo X.
= exp( - yyauh x? (vp/+1 .pi)) exp
n Σ [((0))))))))))3(1
(44)
XT s0,) u Xx of can be expanded as where : == Σ Lej (y3 X
1-1
S₁, l 44.illpioo aill Tipopt, + (1 (1-v%, pl.opt,) )a:;r
+ Σ Sesiiills (1.iiii( (4s Rjoopt, + (1 Pj )airrr ( (0h Dioopt, + (1 - (ii.,,pp.))) )xir)
(45)
Using Lemma 9 and the second part of the 1H1 we have
3xx 3dpp2d
I
+ liphopt,mr..pt, Ti,opt, - -pt. )ar,-1 |=
2 pd
12 pd
10n² 2
a² pdd
lipi,op,, Ti,opt, + (1 - (1-19% )xr,r 122
+
2appp
2
dlog
2
22
2
22
2
4
8
10n²
140%oop, Fioopt, (1 s0/opp,)) )xr,)))) (vpl.op,, Si.opt, + (1 pi.opt, Dii )l s 2²2²
dlog
35

---

Published in Transactions on Machine Learning Research (03/2025)
Using these in Eq. (45), get
3a p'd S l,, na? p² V ddog
10n² Me. St, a²p²d l,s na? p² V
10n²
p dlog 8 Mr
(46)
8
Next, for the second term, using similar calculations as Eq. (53) and Lemma 10, we have
R == y.u X ( co & | x (vp!t1 xp!)
2 max X 1,,00
W 2-1 W.l
6n(apv (1 n[ 1) = (app () +1-2//3
(47)
Consider the loss ratio for any two samples i, E [m], A == 4.. b.
Using Eqs. (44) and (47), we have
A+t1 Ar exp(Ret1.)) exp(- 2h St.)) St.))
2h St.i) S A exx(112 (apv). 1)/3/// exp(
"
(48)
exp( 1 exp( S1,)
n. St,)) exp(- "
We consider two cases:
Case 1: A, < C/2. In this case, we have
The Se, s (l( + 1)-2/3 max 2, & napv/1.5d(( + 1)-2/3.
2/3
n
Using this in Eq. (48), and that n
3apv11.3d 1.3d 17 S 18a2 we get
(2)
A $ A, exxp(12r)) (app: exp(()ap 1.5d)
log(2)
1
$ A, exp
(180 pddd eex(( exx((2log())) (2)33))
1apv/1.5d 18 p2d
-A,
Case 2: A, 2 C/2.
In this case, using Eq. (46), we have
exp(- 2h' St.1)
pd
302
3a2p2
10n2
S exp
4,
fes
2na2 p²
p2 dlog
Me
exp(- n Stj)
71
8
2
mhaa
ppppp
dl
3d
10n²
A pmax ))
$ exp
A.
2m
dlog
n2
2
Further,
a p² p2
A: 3d
10n
17t 0 22222mm p²me
d C-
3d 2m
10n²
2n
dlog
Amax
d
dlog
C
11
1/t a pm
OVd
10n²
Vd 32m
log
16
8
Th a² 22 ///dM,
na² 2 p² Vd
T.(0.)
(t+ (1+1)-2/3
32n
32n
apv 1.5dL (0,)
7ap (t 1)
12///3 2/3
40n
36

---

Published in Transactions on Machine Learning Research (03/2025)
10n2 (2 and [vI((.)] s apv (..d. Further, since 17 s ap as C² 24, 2 (32m)² log
160n we have
exp 1) exp (- yap
(t+ + 1) 2/3 s 1,
40n
which gives A+1 & A s C.
Lemma 12. Let u == Σ Yo Under Condition 1 and the data model DM, for any j [n], T opt,,
se(n]
ap d
(1j(0o,pp,, #zj))
4
Proof. Using Lemma 9, we have
0j (o(po + xz,))
Σ
1zy,ott, 12+
Σ
yop jjoppt + Uj Σ
y. Zj.r
ie[m]
im]9[n]
[c[m]
10n²
10n2
.0.aa2p' 1 2(n p)
dlog
2napp V /dlog
8
10n²
0.5aa²ppd 4a²
dlog
8
a² p²p²d
since d 2 (12n)² log
10n ?
Similarly, we also have
m Di,opt,
|| n² max vej
(Ii Fjoopt,
10n²
n n(L.5a2p'd) 2n²a²p²
dlog
I app/1.75nd
since 2 (8n) log
10n² 8
Combining these and using 3 finishes the proof.
Lemma 13. Under Condition and the data model DM, using the updates in Eq. (10), exponential loss has the following properties:
|vo[((.)) 2 [V.I(()) wi)())),
I((0) &T.))),
max 0..1
max 0'((0..0.111 [2I[(0)1 s 8(w(0,) v (((O))))((())),
where W1 = =. B (0 13(B 1VV 111), B ap/11.dd
8vI.5n"
Remark 3. We note that the first property is the most challenging to show in the analysis and it yields PL-inequalityy-tin like form, e.g. for neural networks (Frei & Gu, 2021; Nguyen & Mondelli, 2020; Liu et al., 2022). Analogous properties but for simpler settings, e.g. to obtain loss convergence results for linear predictors using NGD on separable data have been shown in (Nacson el al., 2019). The second point is an analog of the Log-Lipschitzeess property shown in (Taheri & Thrampoulidis, 2023a) for the analysis of two-layer newral networks trained with NGD and controls the loss on a line between current iterate O, and the next one 0.-1- The thind property of second-oe order self- boundedness has been seen previously in the convergence analysis of muulti-head head self-atlention. (Deora el al., 2023) and MLPS (Taheri & Thrampoulidis, 2029b) trained with GD.
37

---

Published in Transactions on Machine Learning Research (03/2025)
Proof. We first obtain the lower bound the gradient norm follows.
1
a Σ 1-1
1 VoluiVo (.), x.), 00
voT((.))| = sup 1//111
1
sup 2 vl [[,,1l
(O..
1|111
2 I(())) sup miny ( X.)
1//|=1
2 I(,)) min vv. X.),
(49)
where ū= Σ Vo ie[n]
Using Lemma 11 and 12,
á min y V00, X.), |
= min vi (' opt (1- (1-P ) u
0.5 yi Ti,r) u ap d
w1.
(50)
n
Using Eq. (50) in Eq. (49), we get the first bullet point.
We show the second bullet point as follows. For any 0, 0,0', we have
x1) x1) n X pp XW" 21)1 1y, X) (006'' |u X ((X''x))
(u u') X (XXX' Iu" X (p(XW(),) -x((X"')))))
(51)
These two terins can be bounded as follows.
luu u') x (XW) (aTaaa)) I u' Iu- |1 11
X 20 ||
(52)
|
)))
s | 1.2
(XW s0 (X W 1
2 2 2,00
XW X x1
$ 2 X
I(W W æ1
$ X 1
[
W W'I
2 X 2,00
W WWWI..
(53)
Using Eqs. (52) and (53) in Eq. (51) for 0, and 0, + A(0 - 0:), we get
(O, + A(O..1 0.), X))
max xe[o.1]
(()0(0,, X))
max xc(0.1]
exp(1y A(O - -.), X) -0, X))
< max exp (2X X|] 1e[0.1]
3 00 1n
W
Will + X|| 2,00 [uet - well)
exp (2mh
W | V T(
13
(using Eq. (2))
1yt 1/3
S sexp
61)
X -/200
17 X 2,00
(using Lemma 10)
(1 + t)
$ exxp
(6n²1
X 2,00
X 2,00
38

---

Published in Transactions on Machine Learning Research (03/2025)
Since this is true for every X,, and 1 30pv/I.5d
log( 2)
log(3)
we get
3o+p"d
I(0')=
maxexp (6n² | | 2,00
+ 22,00 I(0.)
max 0+.1)
re[m]
6 log(2) log(3) log(2)
exp Bd + BB) L)))))(O))) s exp
I(0:)
+
2
3
s (T.)
Next, obtain the upper bound on the Hessian norn follows. Using Prop. 3 from (Deora al., 2023), have
X), X
2,00 +
+ 2000 X ues oo.
VV.00,, x)] 6d x |
Xu, 00
+2/ |1,00
Using these, we get
Vo, X:)l | |2,0 2 2, uell
1v, x.)] = 6d x 13
[u
+ 2 X 1
2,00
2,00
(o)( 1 X,)I max 4 X 1 12, 2 6d X 12,000 1n +2 X X 12,00 + +|
:c(n)
=4
2 | + 2V 13²
s 13
=
Using this, we get
VQI.('')
max w(0')T((')) 0'.10,0..11 $
max 0'e(0r.0..1]
B.3 Proof of Theorem 2
We finst restate Theorem 2, this time with the exact constants. Theorem 8 (Train loss convergence). Under Condition and the data model DM, using the updates Eq. (10), for any >> 0,
I(0(.)) so
00(ep)) exp
Iul (t+ 1)
where Wi
Proof. First, using Lemma 10, we have
((t)) v w((..)) ( 115
$ (13(
1 s 1w)(t+1)
(54)
where wo 13(B ())BB1))(BB (Bv)) d), since n 1/3.
39

---

Published in Transactions on Machine Learning Research (03/2025)
Next, for some 0' E [,, 0++1] using the second order Taylor expansion of I(Or-1), we have
I(0)-1) - T(,)) (voT(O.), l++1 - 0.)
2 (0 - -.)'v (( -- 0.)
sT(0.) + (VoT 0.) + 0 - 0il max
VzT(('))l
0.1(
sI(0)) - ((00)=='(V W V |V L
* 1. T...) (b.)))V
> (1, \ V T()
(oTT(0')
max
2
sT(O)) (++ 1)2/3
"
V v.t(())):
1
'7
max
VoT(''>
(using Eq. (10))
2
(t+1)443
(++ 1)²
01.10..0.11
8r7²
s si(0.)
I(Oi)+ (1+ 1)4/ (w((,) vw(Ot-1)EE(O))
(using Lemma 13)
(L+ 1)2/3
wn (1 + (1) 1)2/3)
M/w1
T.(O.)
(using Eq. (54))
(1+ 1)2/3
(1+ 1)43)
1
IHWI
[(0.) S sexp
TW1 (++ 1)-2/3 I(0.)
(using Condition 1)
2(t + 1)2/3
A mw Σ
$ exp
(1+ 1) 1)-2/3 I(00)
(by telescoping)
2
M=0
?wi
(1+ 1)2/3) T()O))
$ sexp
2
m
B.4 Proofs of Theorems 3 and 4
We first restate Theorem 3, this time with the exact constants.
Theorem 9 (IB Rate under Joint Training). Under Condition 1 and the data model DM, using the updotes
1/3 - v exp (BaA) ()
in Eq. (10), for any 2 te == exp (10BA)²
V
V10B2A)
vexp
250 log 10m2 &
We
W mam
C,A,,,
21-c-
W
mm
logt
where B apv/.5d, C2 24 an absolute constant, and
C(r, = 2B² A (1 e) (1--
W Wmm) 2n/ W., 1
1
Proof. In this case, Lemma. 4 follows directly, using Lemma 11. Since |Wol 0, using Eq. 10 and similar caleulations as the proof of Lemma 5, for any t> 0,
Ww 2 (2112A) logt.
(55)
Next, using Lemma 11, we will show that for any C (0, 1), there exists
4Cn2 2d
R 2A
log
v vllog (40B²Ac
(56)
log
10n?
such that for every t where W 2 Res
W
W,
(1 (1-6) - -V.T(()))
(57)
W mm
W.
40

---

Published in Transactions on Machine Learning Research (03/2025)
Using similar calculations the proof of Lemma 7, in the first two cases, since R' R mm 10 Ae Egs. (30) and (31) follow, respectively. Using Eq. (36), we can show that when
2
exp(0.5e(1-e) Rt) 2 0n 1)(2)² max l,d (H.opp, - -1) minke, l, (Hoop,, - 76)
(58)
Eq. (29) is satisfied. Using Lemma 11, Eq. (58) is true when
4Cn² 2d
exp(0.5e (8)
log 10n")
which is satisfied by Eq. (56).
Next, using Eqs. (55) and (58), we can show that for any 12 te, [W 2 R. v 1/2. We have
V2d 22
2Ac
log
-1
log
10n²
1/3
Cn2v200 2d
$ 25A
V 10BBA
-4/3
250 log 10n
$ 17( ()4B2A) logt [Wll.
In addition, |W.| 1/2 for 22 exp in (Tarzanagh et al., 2023a), we get.
B'A "
Combining these and using similar steps the proof of Theorem
W
W mm
2 2 ((2B2A) C,,,BBB
W
W
mmll
Welll
21 C(n, B,e)
logt
where the last step follows by using Eq. (55).
Next, we restate Theorem for convenience.
Theorem (IB Rate of u). Let 7 maxululs! min; y,u y,unaar oippp, Under Condition 1 and the data model DM, using the updates in Eq. (10), for any 12 be exp(C(m,
B,A,,,, (, (8B*A))),
π mjn y. oi, 4 [
Proof. First, by definition of 7, we have
Ivo((.)) 2 I(()..
3
Then, following similar steps as the proof of Thcorem 2, we have
Σ(O.) exp (-T (t+ 1) I((o).
Using this, we have
Di.opt + (1 Pi, (xipp,,).ir)) ){ir)) u.) Tu.) S sexp (-14 17
(1+ (1+1)118)
min (i Dioope, (1 Piooop, (t.oot,)) T u4 2 (t+ 1)3
17
4
41

---

Published in Transactions on Machine Learning Research (03/2025)
Then, using the proof of Lemma 10, get
7
min ya (Pr,opt, oipo + (1 Pi,oop,, )biir)' uux 14.2 2
4
?
min Vilfi,opt, oopt,, Ui 2
-axx
Pi.ope, x
4
7
PV 1.5d
2
4
1 expepp - ( logt)
where for the last inequality, we use the following lower bound on the softmax scores,
poop 2 11 1 cxp(7((8''))) logt)
which is obtained by using Theorem 3 and following similar steps the proof of Lemma 1, as (xirr - opt ) (W,
W W. + W
aa, 2B Will W.
W mml
W,
W mm
C(n, C.A,c
$ $4B²
E
(2nlogt) - (24 logt
logt
$ - -
since 22 exp(C(n, (, ()))
Using Pi.opt. $ then finishes the proof.
C Additional Experiments and Settings
This section includes some additional experiments and details about the settings for the results included in this work. We use the PyTorch (Paszke et al., 2019) library for our code, which is licensed under the Modified BSD license.
Fig. 1 In this case, use the MultiNLI (Williams et al., 2018) dataset, which contains sentence pairs belonging to one of three classes: entailment, neutral, contradiction. The task is to predict whether the second sentence entails, is neutral with, or contradicts the first sentence. It relcased under the ODC-By license. We use the Hugging Face pytorch-transforme model, with pretrained weights (Devlin et al., 2019), released under Apache License 2.0. We use batch-size train all models. Learning rates are Adam: 2e 5, SGD: 3, SNGD: 0.01. The Tmma for SPS and SNGD is set to 0.1.
implementation the BERT bert-base-uncased
Fig. In this case, the samples are generated following Example 1, with U 1 and p= 0.05. We set n 20, d= 100 and ==6. 6. We cap the adaptive step-size, Mmax 100 for NGD and NGD-mom, and Mmax 10 for Polyak-step. The learning rate 1) is set as 0.025 for NGD and NGD-mom, whereas for GD, it is set 0.25. For NGD-mom, the momentum parameter is set to 0.9.
Fig. In this case, the samples are generated based on the data model DM with p= 0.1 and a = 3. We set n 10, d= 100 and T= 2. The learning rate 7 is set as 0.002 for NGD and NGD-joint, whereas for GD, it is set as 0.02.
Fig. We use the CivilComments dalaset (Borkan et al., 2019), which consists of online comments and the task to classify if the comment is tozic or non-toric. It is released under the CC BY-NC 4.0 license. We use the BERT bert-base-uncase WILDS package (Koh et al., 2021). The learning rate for SGD is 10-3, 10-2 for NGD and 10-5 for Adam. The parameter Maaz is set to 10 for both SNGD and SPS. All models trained with batch-size 32.
model, with pretrained weights (Devlin et al., 2019) using the
42

---

Published in Transactions on Machine Learning Research (03/2025)
Train loss
Train accuracy
Test loss
Test accuracy
5GD
0.5
0.6
0.5
2.6
SNGD
0.4
0.2
Adam
0
0
100
200
300
400
100
200
300
400
100
200
300
400
100
200
300
400
Figure 4: Comparison of train and test dynamics of various optimizers- SSG,, SNGD, SPS, and Adam fine-tuning a pre-trained BERT model on the CivilComments dataset.
7x10---
Softnax score opt
1.0
1.8
1oss
6x-00-
48
8.3
TD
[W
5x10
20
8.4
(Wr.
Polyek step MG0-thcory
40
0.2
1890
2800
3900
4000
1830
4008
iterations
Iterations
iterations
iterations
Figure 5: Training dynamics when optimizing only W synthetic data with antipodal opt tokens.
Train loss
Train accuracy
Test loss
Test accuracy
2.5
2.5
SGD
1.5
1.5
0.6
SNGD
3
0.4
0.4
0.5
Adam
0
0
10
20
40
20
30
40
10
20
30
40
10
20
30
Figure 6: Comparison of train and test dynamics of various optimizers SGD, SNGD, SPS, and Adam for a ViT model on the MNIST dataset.
accuracy
Test
accuracy
2.5
0.8
2.5
0.8
2
1.5
1.5
0.4
0.2
0.5
0.2
SGD
SNGD
Norm Growth
SPS
Adam
70
20
20
15
40
1.5
15
10
90
10
Figure Comparison of train and test dynamics (lop row) and parameter norm growth (bottom row) of various optimizers SGD, SNGD, SPS, and Adam for ViT model on the CIFAR-10 dataset.
Fig. Training dynamics of single-head self-attention model (Eq. (1)) when optimizing only W on synthetic data (Example 1 with 0). In this case, the opt tokens antipodal instead of orthogonal. The remaining settings are the same as those for Fig. 2, except we set Tmax 5 for Polyak-step.
Fig. We use the MNIST dataset (LeCun & Cortes, 2005), which contains gray-scale images of handwritten digit 0 0-9. 9. It is relcased under the CC BY-SA 3.0 license. We use the ViT for small-scale datascts proposed in (Lee et al., 2021), referred to as ViT-small hereon. The implementation is available at https: //github. www./oo//ccaaeeee com/lucidrains tataa/vvttptytoo
under the MIT license. We use patch-size 4, and set depth as 2, number of heads as 8 and MLP width as 128. All models are trained with a batch-size of 100. Learning rates are set as follows. SGD: 0.1, SNGD: 0.001, Adam: 0.001. Mmar for SPS and SNGD is to 0.01. We observe that unlike language datasets SNGD trains slower than SGD, and Adam achieves significant, gain
43

---

Published in Transactions on Machine Learning Research (03/2025)
in terms of training speed over SGD (also reported in (Xie et al., 2023)), showcasing behaviour similar to CNNs (Kunstner et al., 2023).
Fig. 7 We consider the CIFAR-10 dataset (Krizhevsky, 2009) which a benchmark dataset for object recongnition tasks and contains colored images from 10 classes. It is released under the MTT license. We use the ViT-small model with patch-size 4, and set depth as number of heads as 32 and MLP width as 512. All models are trained with batch-size 100. Learning rates are set follows. SGD: 0.1, SNGD: 0.001, Adam: 0.001. The parameler Mmaa is set to 10 for SPS and 0.1 for SNGD. As evident, SNGD portrays slower training speed similar to the observations for MNIST. Additionally, we plot the norm growth for various layers of the ViT and see growth in parameters similar to our observations on synthetic datasets.
Compute and Runtimes. The experiments using synthetic data were run on Google Colab. The experiments on vision and language datasets were run an internal cluster with two NVIDIA V100 GPUs with 32 GB memory each. We train for 40 epochs MNIST and 80 epochs CIFAR-10. The runtime for the latter is about 2 hours for each setting. We fine-tune the models on language datasets for 10 epochs, which takes about 32 hours for each run. In Figures 1 and 4, we plot the dynamics for every 200 iterations.
C.1 Initializing in Bad Stationary Direction
Fig. 8 shows a synthetic setting where we initialize in a bad stationary direction W one that gives a higher softmax score to the non-opt token. To be precise,
1
0.2
-25
0.5
X1
X²
-0.2
2.5
-0.5
with labels V1 -1,92 = 1, respectively. We set u, [0, 1] Using token optimality (Definition 1), this gives opt, opt, 1; see Fig. 8 (left) for illustration. The chosen initialization Winit violates the (W-SVM)
constraints for i = 2, that is
(I2,opt, roop,, W 22,1 < O.
This can be seen in Fig. 8 (left), where the - line correlates more with the non-opt token than the opt one. This consequently forces the softmax score
22pt
<
exp( o2g ) W i2 2
which can be seen in Fig. 8 (plot 2). We call Wanit a "bad" stationary direction as starting with a Winit
and a + oo would result in op, (iii)) + O, and Tw-WWammss) O. On the other hand, for X1, Winit behaves the opposite way (see plots 1-2 in Fig. 8).
1.0
1.0
0.8
Wmnk;
Aiignnent
Train loss
WmmXi opt
0.7
Attn-
0.5
0.8
non-opt
Sample
0.6
-2
-2
0
101
10a
10a
10a
10a
10a
Iterations
Iterations
Iterations
Figure 8: The training dynamics for a synthetic setting with 2, = Attn-score denotes the opt-token softmax score pi.opt, and alignment shows (W W mm). Plot shows the initialization Wini,, and the (W-SVM) solution W Wmm projected onto the first token, respectively. For sample 2, the W Win is more correlated to the non-opt token which consequently gives a small softmax score (plot 2). Plots 2-4 show that despite this "bad' initialization direction GD converges to W.m. and achieves the loss minima. L.
Despite this "bad" starting direction, plots 3-4 in Fig. show that GD still globally convergences to Wmm: and minimizes the train loss. This behaviour aligns with our thcoretical results (Thm. 1) showing global parametric convergence starting in any direction Woa, with a small enough norm avoid the bad stationary directions in the limit as Winiel
44

---

Published in Transactions on Machine Learning Research (03/2025)
Train Loss
Avg. Softmaa Weight on opt Token
Cosino Similarity with W
0.05
0.04
0.8
0.8
0.03
0.6
0.4
NGD
0.2
0
1k
2k
36
4kc
5k
1k
2k
3k
4k
5k
1k
2k
3k
4k
5k
Figure 9: The training dynamics of self-attention parameters W. for a binary classification on simplified
MNIST. For details on the task see Appendix C.2. Results show parametric convergence W mmm+ and fast convergence for NGD aligning with our theoretical findings.
Distribution of cosine similarity
Civil Comments
MuitiNLI
50.21%
36.97%
Percantage
3.9
774
-0.6
-0.5
-0.4
-0.2
0.0
0.2
0.1
0.6
0.6
-0.8
-2.5
-04
-02
0.0
0.2
0.4
0.6
0.6
Cosine similarity
Cosine similarity
Figure 10: Validation of part 1 of Assumption 1. We compare cosine similarity of token embeddings for tokens within a sentence, and plot the distribution over multiple sentences. See details in Appendix C.3. We observe that a significant portion of the distribution lies around 0, indicating near orthogonality.
C.2 One-layer Self-attention Model Trained on Binary MNIST
In this section, we conduct experiments using a simplified MNIST setting as follows: we created a binary classification task where each input consists of multiple patches, with one patch containing either a 0 or digit (the 'signal' patch) and the remaining patches containing Gaussian noise. We trained self-attention module that processes these patches before feeding them through fixed linear classifier (pre-trained on 0/1 classification). As shown in Figure 9, the results in this setting empirically validate two key theoretical predictions: (1) the self-attention parameters W converge in direction to W mm,, and (2) normalized gradient descent achieves faster convergence than standard gradient descent. In this setup, the signal patch naturally corresponds to the token with the highest score, aligning with Ass. 2.
C.3 Validation of Assumption
We conduct an experiment verify the first part of Assumption on two language-based datasets. To do this, we obtain token emboddings from pre-trained BERT tokenizer (Devlin al., 2019), and, for every sentence, we compute the cosine similarity between every pair of token embeddings. We consider 50k text examples from the Civil Comments dataset (Borkan et al., 2019) and plot the distribution of these similarity values. We repeat the experiment using 100k samples from the MultiNLI dataset (Williams et al., 2018). The results are shown in Figure 10. Values close to 0 indicate that the tokens are orthogonal.
We observe that a significant fraction of all the cosine similarities fall within a small range around 0, suggesting nearly orthogonal behavior. Note that while Assumption 1 is only on the similarity of non-optimal tokens with optimal or other non-optimal tokens, whereas the experiment considers the similarity between every pair of token embeddings. Nevertheless,
that the similarities concentrated in small region.
45

---

Published in Transactions on Machine Learning Research (03/2025)
D
Related Work
Implicit bias of NNs. Since the first few works characterizing the implicit bias of linear predictors on separable data (see Introduction), there has been an abundance of work studying the implicit bias of gradient based methods for both linear predictors and NNs. Nacson et al. (2019); Ji & Telgarsky (2021); Ji et al. (2021) show fast convergence of GD-based methods to the max-margin predictor. For MLPs, early works study the implicit bias of GD/ /gradient flow using exponentially-taiiee points of the corresponding max-margin problem in finite (Ji & Telgarsky, 2020; Lyu & Li, 2020) and infinite width (Chizat & Bach, 2020). Additionally, works by Phuong & Lampert (2021); Frei et al. (2022b); Kou et al. (2023b) study the implicit bias of GD trained LeLUULLeay-ReLU networks on orthogonal data. Other works also study the implicit bias towards rank minimization with square loss in regression settings (Vardi & Shamir, 2021; Arora et al., 2019; Li et al., 2021). We encourage the reader to go through these works and a recent survey (Vardi, 2022) for a thorough understanding.
classification losses towards the KKT
Parametric convergence rates for NGD and PS in non-convex settings. We compare the rates obtained in our results to the corresponding results for linear logistic regression (sce Sec. 6). Contrary to these results, we prove the first parametric convergence for NGD and PS in non-convex settings. The work by Loizou et, al. (2021) shows loss convergence results for PS. However, the rate of directional convergence might be different from the rate of loss convergence. This has been seen in the series of works analyzing gradicnt-based methods for lincar logistic regression (Soudry et al., 2018; Nacson ct al., 2019; Ji & Telgarsky, 2019). If the problem had unique global minima, then loss convergence implies parameteric convergence. However, the problem of training self-attention does not have a unique minimum. For instance, using some other norm in place of l2 norm in the objective of (W-SVM) would yield a different solution. While the norm diverges, the convergence is in the direction. In addition to this, there has been some work showing the benefit of adaptive step sizes such as NGD and PS for transformer optimization (Kunstner et al., 2023; Schaipp et al., 2024), which is connected to Q3. We next discuss related work on transformer theory.
Transformers theory. Several studies, such as those by Baldi & Vershynin (2022); Dong et al. (2021); ))
Yun et al. (2020a;b); Sanford et al. (2023); Bietti et al. (2023) have delved into exploring the expressivity of allention, while Baldi & Vershynin (2022); Dong et al. (2021); Yun et al. (2020a;b); Mahdavi et al. (2023) have initiated an investigation of its memory capacity. To gain insights into the optimization aspects of training attention models, Sahiner et al. (2022); Ergen et al. (2022) have explored convex relaxations of attention. Furthermore, subdomain experiencing growing attention involves the theoretical exploration of in-context learning, as evidenced by recent studies (von Oswald et al., 2022; Akyürek et al., 2023; Zhang et al., 2023; Li et al., 2023c).
In this context, we discuss studies that seek to understand the optimization and generalization dynamics of transformers. Jelassi al. (2022) show that Vision Transformers (ViTs) learn spatially localized patterns in a binary classification task using gradient-base based methods. For a three-layer ViT starting from some structured initialization to mimic a pre-trained network, Li et al. (2023b) show sample complexity bounds to achieve zero generalizate anneralization and show attention maps sparsify as SGD training proceeds. As a step towards understanding the training dynamics of the closely related prompt-attention, Oymak et al. (2023) study the initial trajectory of GD for lne-layer attention. Further, there has also been work to obtain optimization and generalization guarantees of GD in multi-head attention models (Deora et al., 2023). Additionally, recent work by Tian et al. (2023a) attempts to understand SGD-dynamies for the task of next-token prediction for one-layer transformer with a linear decoder showing a similar sparsifying effect in the attention map. More recently, Tian et al. (2023b) extend this by analyzing the joint training dynamics of multi-layer transformer with an MLP. Other than this, there has been recent progress towards a theoretical understanding of next-token prediction dynamics under GD training (Li et al., 2024; Thrampoulddis, 2024). Furthermore, Makkuva et al. (2024) discussed the impact of the transformer architecture and distributional properties on the loss landscape when modeling the data as a Markov source, and Ildiz et al. (2024) linked the dynamics of self-attemtion to context-conditioned Markov chains.
46


Reviews:
[{"review_id": "KwLE2LpRQK", "reviewer": "Reviewer_2WpF", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The paper provides theoretical results in an important area: convergence rates of single layered transformers as well as their implicit bias to converge towards certain solutions\n- The results addresses several limitations of prior works like asymptotic results and local convergence results, and also shows that adaptive step sizes in gradient descent can provide much faster convergence\n- The paper also provides empirical results for convergence rates matching the intuitions obtained from the theoretical results\n\n\nWeaknesses:\n- I found several portions in the paper difficult to understand, especially, it seems that the author assumes that the reader is very familiar with prior works of Tarzanagh et al., which needs to be addressed. The writing should cover the background and problem setting better. E.g., “Our second technical assumption is similar to Tarzanagh et al. (2023a) with two key distinctions: Firstly, it applies to self-attention rather than their simplified attention model.” Please write in more detail what were the assumptions and results from previous work that you are improving upon. Also, in Assumption 2: what is “\\gamma_i”? When using a notation, please refer to the definition where they are defined, its very difficult to find the definitions with so many notations.\n- In eq 1, what is “x_1” and where is it defined? How does it fit into the attention equation? I am aware of self attention, but the way the equation is written confuses me. Can you please explain using dimensions of each matrix/vector and the outputs of the layers of the attention layer and the decoder layer.\n- Similarly, in page 4, in eq (W-SVM), please provide more details about the equation, and how is it obtained (it is OK to provide details in the appendix, but the paper should have all the necessary preliminaries for a reader to understand all the details)\n- How does the results translate to transformers with more layers? Are there any practical takeaways from this paper for training transformer models?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work provides new theoretical results on convergence guarantees of single layered transformers. The results in the paper addresses limitations in prior work in this area that provided convergence guarantees that were (a) local, (b) asymptotic. In response to these limitations, this paper shows that (a) under some assumptions, global convergence is possible, (b) also provides finite convergence rate. Further, this paper shows that using normalized SGD and Polyak step-size, faster convergence can be obtained. This paper provides these results for two problem setting: (a) when the linear decoder after the attention layer is frozen, (b) when both the attention layer and the decoder layers are trained simultaneously. They also provide empirical evidence illustrating the convergence rates for various algorithms on synthetic as well as small vision and language datasets.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "- Please address the points raised in the weaknesses.", "broader_impact_concerns": "N.A.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "XwlFQH9Q2r", "reviewer": "Reviewer_zaLY", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "*Strengths:*\n\n* The paper nicely extends similar results for logistic regression, but here additional technical difficulties arise, most importanly a non-convex loss function. \n\n* Adaptive step sizes are studied which are shown to converge faster (empirically) than constant step size GD.\n\n\n*Weaknesses / Questions:*\n\n* I could not understand the notation in Assumption 2. How is $\\gamma_i$ defined? Can you describe this assumption informally?\n\n* You assume that W-SVM is feasible. Is this assumption satisfiable under mild conditions on the data? Does it hold for the datasets you use in the experiments? For instance, matrices $W$ that have $x_{i,1}$ in their kernel are not feasible. For sufficiently many input examples which are linearly independent, would this imply that the solution has full rank? \n\n* On page 12 you claim that this is the first iterate-convergence for Polyak stepsizes and NGD in non-convex settings. I am not sure whether this is correct: e.g. in Theorem 3.6 of Loizou, 2021 (on stochastic Polyak step sizes), objective function convergence is shown for the PL-condition (possibly non-covex). If one assumes a unique minimum, then this would imply also iterate convergence. Further, the statement is somewhat misleading, as your result does not proof that $W_t$ actually converges to a point (as I understood, the norm of $W_t$ diverges).\n\n* I did not understand, why the W-SVM problem is using the first token always? Can you explain in more detail why this is chosen, and if/how this limits the results of the submission?\n\n* On the experiments on real-world data: did you verify the convergence to the max-margin solution on real data? It seems that all the real-data plots in the paper only show train and test curves. I think it would be interesting (or even mandatory) to verify if the theoretical results hold on real data, especially given that the paper makes quite a lot of assumptions on the data model. The insight that adaptive step sizes help for training (on real data) is nice, but not particularly new, as it has been shown in multiple papers that Polyak-type step sizes and normalized GD can work well on transformer models (e.g. https://arxiv.org/abs/2305.07583, https://arxiv.org/abs/2304.13960).\n\n\n*Minor:*\n\n* The definition of $\\bar W$ seems to be missing.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the convergence of gradient-based methods for self-attention modules. It shows that for normalized GD and Polyak step size GD, the normed iterates converge to the normed solution of a max-margin SVM problem. Empirical results show that these adaptive step sizes indeed achieve faster convergence than constant step sizes on synthetic data.\n\n\nFor transparency, I did not check the proofs of the submission in detail.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See weaknesses above, in particular the last bullet.", "broader_impact_concerns": "None.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "Vg2wT8ZBKO", "reviewer": "Reviewer_jEuR", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n1. Provides formal proofs for global convergence and finite-time convergence rates\n\n2. Demonstrates concrete advantages of adaptive step-size strategies over standard gradient descent\n\n3. Covers both theoretical foundations and some practical parts of optimization strategies\n\n\nWeaknesses::\n\n\n1. Analysis focuses only on binary classification with linear decoder, may not generalize to more complex tasks or multi-class settings\n\n\n2. Studies single self-attention layer rather than full transformer architecture\n\n\n3. Some theoretical results rely on specific conditions that may not hold in practical applications", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "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": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work advances the understanding of self-attention optimization by proving global convergence conditions and analyzing adaptive step-size strategies (normalized GD and Polyak step-size), demonstrating their finite-time convergence rates and attention map sparsification properties. The study shows these methods can accelerate parameter convergence compared to standard gradient descent in non-convex settings, while establishing stronger parallels between self-attention's implicit bias and linear logistic regression, despite self-attention's non-convex nature.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Provide more discussions on the impact of different initialization schemes on the convergence properties, as this may affect practical implementation\n\n\n2. Testing the validity of  Assumptions 1 and 2 in practical scenarios like in your experiments on language and vision datasets\n\n\n3. Provide more discussions on the limitation of binary and single attention layer setting.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: As all reviewers unanimously agreed on the great contributions of this paper, it should be accepted.
Decision: Accept (poster)