On the Optimization and Generalization of Multi-head Attention
Authors: Puneesh Deora
Rouzbeh Ghaderi
Hossein Taheri
Christos Thrampoulidis
Venue: Accepted by TMLR ()

Abstract: The training and generalization dynamics of the Transformer's core mechanism, namely the Attention mechanism, remain under-explored. Besides, existing analyses primarily focus on single-head attention. Inspired by the demonstrated benefits of overparameterization when training fully-connected networks, we investigate the potential optimization and generalization advantages of using multiple attention heads. Towards this goal, we derive convergence and generalization guarantees for gradient-descent training of a single-layer multi-head self-attention model, under a suitable realizability condition on the data. We then establish primitive conditions on the initialization that ensure realizability holds. Finally, we demonstrate that these conditions are satisfied for a simple tokenized-mixture model. We expect the analysis can be extended to various data-model and architecture variations.

Published in Transactions on Machine Learning Research (04/2024)
On the Optimization and Generalization of Multi-head Attention
Puneesh Deora"
puneeshdermaccc.uucc..
University of British Columbia
Rouzbeh Ghaderi*
gghaderiSece..uc..omm
University of British Columbia
Hossein Taheri"
hosscin@ucsb..cmm edu
University of California, Santa Barbara
Christos Thrampoulidis
ethrempo@ece.ube ubc.
University of British Columbia
Reviewed on OpenReview: /t/p::///operevvieet net/ forum? id- wTGjn7JoYK
Abstract
The training and generalization dynamics of the Transformer's core mechanism, namely the Attention mechanism, remain under-explored. Besides, existing analyses primarily focus on single-head attention. Inspired by the demonstrated benefits of overparameterization when training fully-connected networks, investigate the potential optimization and generalization advantages of using multiple attention heads. Towards this goal, we derive convergence and generalization guarantecs for gradient-descent training of single-layer multi-head self-attention model, under a suitable realizability condition on the data. We then establish primitive conditions on the initialization that ensure realizability holds. Finally, we demonstrate that these conditions are satisfied for a simple tokenized-mixtture model. We expect the analysis can be extended various data-model and architecture variations.
1 Introduction
Transformers have emerged as promising paradigm in deep learning, primarily attributable to their distinctive self-attention mechanism. Motivated by the model's state-of-the-art performance in natural language processing (Devlin et al., 2019 Brown et al., 2020; Raffel et al. 2020) and computer vision (Dosovitskiy al., 2021 Radford et al.] 2021 Touvron al. 2021). the theoretical study of the altention mechanism has seen notable surge in interest recently. Numerous studies have already explored the
expressivity of Attention, e.g. (Baldi & Vershynin 2022 Dong al. 2021 Yun et al. 2020a Sanford et al.
2023 Bietti et al. 2023). and initial findings regarding memory capacity have been very recently studied
in [Baldi & Vershynin 2022 Dong et al. 2021 Yun et al. 2020a b Mahdavi et al. 2023 In an attempt
to comprehend optimization aspects of training attention models, Sahiner et al. 2022 } Ergen et al. (0022
have investigated convex-relaxations while Tarzanagh et al. (2023a) investigates the model's implicit bias.
Additionally, Edelman et al. (2021) have presented capacity and Rademacher complexity-basee generalization
bounds for Self-Attention. However, the exploration of the finite-time optimization and generalization
dynamics of gradient-descent (GD) for training attention models largely remains an open question.
Recent contributions in this direction, which serve motivation for work, include the studies by Jelassi ct al. 2022) Li ct al. (023a)); Oymak et al. 20223.. These works concentrate on singlo-layer attention models with a single allention head. Furthermore, despite necessary simplifying assumptions made for the data, the
"These authors contributed equally. Alphabetical ordering used.

---

Published in Transactions on Machine Learning Research (04/2024)
analyses are rather intricate and appear highly specialized on the individual attention and data model. These direct and highly specialized analyses present certain challenges. First, it remains uncertain whether they can be encompassed within a broader framework that can potentially be extended to more complex attention architectures and diverse data models. Second, they appear disconnected from existing frameworks that have been flourishing in recent years for conventional fully-connected and convolutional neural networks e.g., (Jacot et al. 2018 Ji & Telgarsky 2020 Richards & Rabbat, 2021 Liu et al. 2020 Taheri & Thrampoulidis/ 2023 Consequently, it is also unclear how the introduction of attention alters the analysis landscape.
In this work, we study the optimization and generalization properties of multi-head attention mechanism trained by gradient methods. Our approach specifically leverages the use of multiple attention heads. Despite the operational differences between attention heads in an attention model and hidden nodes in an MLP, we demonstrate, from an analysis perspective, that this parallelism enables the exploitation of frameworks developed for the latter to study the former. Particularly for the generalization analysis, we leverage recent. advancements in the application of the algorithmic-stabliiy framework to overparameterized MLPs (Richards) & Kuzborskij 2021 Taheri & Thrampoulidis 2023).
Contributions. We study training and generalization of gradient descent optimization for multi-head attention (MHA) layer with H heads in a binary classification setting. For this setting, detailed in Section 2 we analyze training with logistic loss both the attention weights (parameterizing the soft logits), as well as, the linear decoder that turns output tokens to label prediction.
In Section 3. we characterize key properties of the empirical loss L, specifically establishing that it is self-bounded and satisfies key self-bounded weak-convexity property, i.e. >mii((²T.(0)) 2 VHH a parameter x. that depends only mildly on the parameter vector O. Establishing these properties (and also quantifying x) involves carefully computing and bounding the gradient and Hessian of the MHA layer, calculations that can be useful beyond the context of our paper.
_x_ T((0) for
we present our training and generalization bounds in their most general form. The bounds are given in terms of the empirical loss T.(0) and the distance |0 10-0ol to initialization Og of an appropriately chosen target vector 0. The distance to initialization also controls the minimum number of heads H 2 |00 - 0ol required for the bounds to hold. The choice of an appropriate parameter @ that makes the bounds tight is generically specific to the data setting and the chosen initialization. To guide such a choice, in Section 4.3 we formalize primitive and straightforward-to-ee
In Sections
conditions on the initialization Og that ensure it is possible to find an appropriate 0. In short, provided the model output at initialization is logarithmic on the train-set size n and the data are separable with respect to the neural- -tanett kernel (NTK) features of the MHA model with constant margin y, then Corollary [ shows that with step-size 17 - =((1 and O(n) gradient descent steps, the train loss and generalization gap is bounded by Ō(1/n) provided only a polylogarithmic number of heads H = 2(log" (n)). We remark that the aforementioned NTK separability assumption, although related to, differs from the standard NTK analysis. Besides, while this assumption is sufficient to apply our general bounds, it is not a necessary condition.
In Section 5 we investigate a tokenized mixture data model with (ireeeean tokens. We show that after one randomized gradient step from zero initialization, the NTK features of the MHA model separate the data with margin 7.. Thus applying our general analysis from Scction 4.1 establish training and generalization bounds described above, for a logarithmic number of heads. Towards assessing the optimality of these bounds, we demonstrate that MHA is expressive enough to achieve margin Vattm that is superior to y. The mechanism to reach Yattn involves selecting key-query weights of sufficiently large norm, which saturates the softmax nonlinearity by suppressing label-irrelevant tokens. We identify the large-norm requirement as a potential bottleneck in selecting those weights as target parameters in our theory framework and discuss open questions rogarding extending the analytical framework into this specific regime.
The remaining parts organised as follows. Proof sketches of our main training/generan generalizzatoon bounds are given in Section 0 The paper concludes in Section 7 with remarks on our findings' implications and open questions. Detailed proofs are in the appendix, where we also present synthetic numerical experiments.
Related work. We give a brief overview of the most relevant works on understanding optimiza- tion/generalization of self- Attennioon or its variants. Please see Section H for more detailed exposition.
2

---

Published in Transactions on Machine Learning Research (04/2024)
Oymak et al. (2023) diverges from traditional Attentioon by focusing a variant called prompt-Attention, aiming to gain understanding of prompt-tuning- Jelassi et al. (2022) shed light on how ViTs learn spatially localized patterns using gradient-basee methods. Li et al. (2023a) provides sample complexity bounds for achieving zero generalization error on training three-layer ViTs for classification tasks for a similar tokenized mixture data model as ours. Contemporaneeuus work Tian et al. (2023) presenis SGD-dynamics of single-layer attention for next-token prediction by re- re-parameterizing the original problem terms of the softmax and classification logit matrices, while Tarzanagh et al. (2023bla) study the implicit bias of training the softmax weights W with a fixed decoder U. All these works focus on single attention head; instead, we leverage the use of multiple heads to establish connections to the literature on GD training of overparameterized MT.Ps. Conceptually, Hron et al.] (2020) drew similar connections, linking multi-head attention to a Gaussian process in the limit as the number of heads approaches infinity. In contrast, we study the more practical regime of finite heads and obtain finite-time optimization and generalization bounds.
Among the extensive studics on training/generalizztion of overparameterized MLPs, our work closely aligns
with Nitanda et al. (2019); Ji & Telgarsky| (2020); Cao & Gu (011)); Chen et al. (2020); ; Telgarsky (2022);
Taheri & Thrampoulddis (2023) focusing classification with logistic loss. Conceptually, our findings extend
this rescarch to attention models. The use of algorithmic-ssablliy tools towards order-optimal generalization
bounds for overparameterized MLPs has been exploited recently by Richards & Kuzborskij (2021): : Richards
& Rabbat (0021) ; Taheri & Thrampoulidis (2023); Lei et al. (2022). To adapt these tools to the MHA
layer, we critically utilize the smoothness of the softmmx function and derive bounds on the growth of the
model's gradient/ Hessian, which establish a self- boffdoouded weak convexity property for the empirical risk (see
Corollary [.
Our approach also involves training both the classifier and attention weights, necessitating
several technical adjustments detailed in Section t and Appendix B.1
2 Preliminaries
Notation. (p-) RR + RR denotes the softmax map and ('(v) == Vip(v) = diag()p(v)) (()) its gradient at v E R² For E [T]. (v)) is the t-th entry of (vv) E R². For A E Raxm A,, is its i-th row 1 | و and particularly the and A. is its j-th column. Recall the induced matrix norm || p.7 | 0. = max |A, loo For simplicity, = max |A,. and
following: | 0,.00 = max,{[=] | A,:ll |All
|A |oll denote Euclidean norms and Amin ( (A) the minimum eigenvalue. We let a ^ = min(a, b} and a v ma concat denotes vector concatenation. All logarithms are natural logarithms (base e). We represent the line segment between W1, W2 E Rd as [Wi, w2] = (w w aw; + (1 - a)w2, a [0, 1]]. Finally, to simplify the exposition we use "2" or "s" notation to hide absolute constants. We also occasionally use standard notations O, Ω and Ō, to hide poly-log factors. Unless otherwise stated these order-wis. notations are with respect to the training-set size n. Whenever used, exact constants are specified in the appendix.
Single-hcad Sclf-attention. A single-layer self-attention head ATTN: : RTxd + RR''d with context size T and dimension d parameterized by key, query and value matrices Wo, WK e Raxde Wr e RRdxdd is given by:
ATTN(X; WK. Wv) = XXXXWWWXXXXXI
Here, X = =[ ...,, TT] e Rxd is the input token matrix and (p WoWWWW xxT) e RTxT" is the attention matrix. (Softmax applied to a matrix acts row-wisc.) To turn the Attention output in a prediction label, we compose ATTN with a linear projection head (aka decoder). Thus, the model's output i
& X: W, U) == (U, (XWX)x)
(1)
Note that we absorb the value weight matrix Wy into the projector U = [u...... ur RTd Also, we parameterize throughout the key-query product matrix as W := Wq K- Multi-head Self-attention. Our focus is on the multi-head attention (MHA) model with H heads:
Σ ATTN(X Wo WK.. Wvi)Won Re[
"While we focus on (i) Full-projection: trainable matrix U RTxd, our results also apply to (ii) Pooling: U = u.J with trainable v Rd, and (fi)) Last-token output: U= [0xx(=-1) u]" with trainable Rd.
3

---

Published in Transactions on Machine Learning Research (04/2024)
for output matrices Won e Ra.xd_ Absorbing Wyn Won into projection layer (similar to the single-head attention) and parameterizing Wn == Won WKy we arrive at the following MHA model:
W 0)
Enef 4(x; hn. Un)
V
Enet hee (U s (X )))
(2)
VH
VHI
parameterized by W concat (WWhhee and Ū concat ((U The 1/VH scaling is analogous
the normalization in MLP literature e.g. (Du et al. 2019 Ji & Telgarsky ensuring the model variance is of constant order when Un is initialized (1). Note that these relaxations sacrifice some generality since it is common practice to set dn and d, such that d, = d/H < d, thus imposing low-rank restrictions on matrices Wan Wkh WV, Won. We defer a treatment of these to future work.
2021 Richards & Kuzborskij 2021].
Throughout, we will 0₁ concal Wx) RRTT+ to denote the trainable parameters of the h-attention head and 0 = concat ((0x) (R"(dT+a)) for the trainable parameters of the overall model. More generally, we use the convention of applying =7" notation for quantities relating to the multi-head model. Finally, with some slight abuse of notation, we define: | 2,00 maX 10x]l.
Training. Given training set (X it[]) with IID samples, we minimize logistic-loss based empirical risk
T(00)
1 Σ (v, 0)) ie[m]
1 Σ log(1 +e (xx.))))) ve(m]
Our analysis exlends to any convex, smooth, Lipschitz and self-bounded loss
| The empirical risk is minimized
approximation of the test loss defined as I.(0) (xx)))(((((( ();0)) We consider standard gradient- descent (GD) applied to empirical risk L. Formally, initialized at Ø(0) and equipped with step-size > 0, at each iteration 2 O, GD performs the following update:
(k++1) k(k) 1/V (((*))
3 Gradient and Hessian bounds of soft-max attention
This section establishes bounds on the gradient and Hessian of the logistic empirical risk T.(.) evaluated on the multi-head attention model. To do this, we first derive bounds on the Euclidean norm and spectral-norm for the gradient and Hessian of the self-attention model. In order to simplify notations, state here the bounds for the single-head model (see App. A.1 for multi-head model): o ((x;0) & ((X:WU) (U, ((XW(")x) Lemma 1 (Gradient// Hessian formulas). For all a RT, b Rd the model's gradients satisfy:
HeadienntHHesiaan
Vvl(X:0) = and Vwo(X:0) Σ a.u, X 4' (xWTa:)) x:)
T V (X):0)) Σ a.a X x'XW'''' x,)X, and
T (w [(c a. 6'(W x.) xx
where diag( Xb) X Xu, X (xX æ.)- Xbu;X Xb,]xT<((XWW) (x(W"x,)
These calculations imply the following useful bounds.
Proposition 1 (Model Gradient Boadienn/HHessinnn Hessian bounds). The Euclidean norm of the gradient and the spectral norm of the Hessian of the single-head Attention model
1 are bounded as follows:
x
Σ
Xo + VTIX X
|| 2,,00
vi)((x0)) <d X -
X 1,00
Xu +oll==22VTTE + 2,00 X 1,,00
function 2: R-+ RR self-bounded C> such that
4

---

Published in Transactions on Machine Learning Research (04/2024)
Next, we focus the empirical loss L. To derive bounds on its gradient and Hessian, leverage the model's bounds from Proposition 1 and the fact that logistic loss is self-bounded, i.e., |e(()) s E(t). To provide concrete statements, we introduce first a mild boundedness assumption.
Assumption 1 (Bounded data). Data (X,y) RTxd x R satisfy the following conditions almost surely: y {+1}, and for some Rt 2 1, it holds for all te [T'] that ||ze|| s R.
Corollary 1 (Loss properties). Under Assumption
| the objective's gradient and Hessian satisfy the bounds
(1) |vī s (1)) L(θ).
A1 = /T R R2
|00 1,122,011
(2) |v*[(0)] s (0)
32(0) 1
VH B3 (0) (1)(
(8) Aminn (v²[(0))= 2 (0(0) III (())
Bs = 2VTAR (3 10 + +1)
The loss properties above are crucial for the training and generalization analysis. Property (1) establishes self-boudeedess of the empirical loss, which is used to analyze stability of GD updates for generalization. Property (2) is used to establish descent of gradient updates for appropriate choice of step-size n). Note that the smoothness upper bound is d-dependent, hence to show descent we need to also guarantee boundedness of the updates. Finally, property (3) establishes a self-bounded weak-convexity property of the loss, which is crucial to both the training and generalization analysis. Specifically, as the number heads H increases, the minimum eigenvalue becomes less negative, indicating an approach towards convex- like behavior.
4
Main results
In this section, we present our training and generalization bounds for multi-head attention.
4.1
Training bounds
We state our main result on train loss convergence in the following theorem. See App. for exact constants and the detailed proofs.
Theorem 1 (Training loss). Fix iteration horizon 2 1 and any Øe H(dTT++") and H satisfying
2H 1/2 R$ 00 0 - ps
(3)
I(oO) with p(0) s d T R 10 10 T2.) Then, the following bounds hold for the training loss and the weights' norm at iteration K of GD:
Fix step-size 1) 1 A 1/p(0) A
KL A
5 5 - 10 Σ I(O.) s 2T)
K
[(O(K)
s
(4)
K k=1
4n
o k 1. s 0 1.
Yielding a concrete train loss bound requires an appropriate set of target parameters Ø in the sense of minimizing the bound in 4 Hence, Ø should simultancously attain small loss (I(0) and distance to initialization (1@ (()) 1). This desiderata is formalized in Assumption 2 below. The distance to initialization, as well as 10 determine how many heads are required for our bounds to hold. Also, in view of the bound in (4) it is reasonable that appropriate choice for Ø attains I(θ) of same order 110 (()) ||/KK. Hence, the theorem's 's restriction on the step-size is governed by the inverse local-smoothness of the loss: 7 1/p(0).
4.2 Generalization bounds
Next we bound the expected generalization gap. Expectations are with respect to (w.r.t) randomness of the train set. See App. C for the detailed proof, which is based on algorithmic-stablig
In all the bounds in this paper involving |00|2, (2) possible to substitute this term with max U. However, for the sake of notation simplicity, we opt for slightly loaser bound ne) [U,]| max |00,| =: m
5

---

Published in Transactions on Machine Learning Research (04/2024)
Theorem 2 (Generalization loss). Fiz any K 2 ony 0 and H satisfying 3 and any step-size n satisfying
the condilions of Thm. 7 Then the expected generalization gop al ileration K saläsfies,
4 - [[2K - - s
-
(5)
4n
The condition on the number of heads is same up to constants to the corresponding condition in Theorem ! A Also, the generalization-gan bound translates to test-loss bound by combining with Thmn [ Finally, similar to Thm. [ we can get concrete bounds under the realizability assumption; see Cor. 4 in App. C.2 For the generalization analysis, we require that the realizability assumption holds almost surcly over all training sets sampled from the data distribution.
The bounds optimization and generalization are up to constants same as analogous bounds for logistic regression (Soudry et al.) 2018 Ji & Telgarsky 2018 Shamir 2021). Yet, for these bounds to be valid, we require sufficiently large number of heads as well as the existence of an appropriate set of target parameters 0, as stated the conditions of theorem. Namely, these conditions are related to the realizability condition, which guarantees small training error near initialization. The next assumption formalizes these conditions.
Assumption 2 (Realizability). There exist non-increasing functions g R. + R, and : R, + R4 such that Ve 0, there erists model parameters 0(e) e RH((T++))) for which: (i) the empirical loss over data samples salisfies (()( s s,, (ti) (0)=) 0 $ go()), and, (iii) (0(c) S g(s).
With this assumption,
specialize the result of Thms. above to specific data settings; see Cor. and in App. B.5 and C.2 In the next section we will further show how the realizability assumption is satisfied.
4.3 Primitive conditions for checking realizability
Here, we introduce a set of more primitive and straighhfforwarrdto-rhhhhh conditions on the data and initialization that ensure the realizability Assumption 2 holds.
Definition (Good initialization). We say g()) concat(O ((! a good initialization with respect to training data (X,,ye) provided the following three properties hold.
P1. Parameter T2 -bound: There erists parameter B₂ 2 such. that vh E [H] it holds |0( ll s B2.
P2. Model bound: There exists parameter B& 2 1 such that ViE [n] it holds (Xx; o(o)) S BB.
P3. NTK separability: There erists 0, (d(TT=d")) and y > 0 such that |0.1 /2 and Vi [m], it holds X (()) 0.)2 27
Prop. in the appendix shows that starting from good initialization we can always find σ.,) satisfying the realizability Assumption 2 provided large enough number of heads. Thus, given good initialization, we immediately apply Theorems and to get the following concrete bounds.
Corollary 2 (Gencral bounds under good initialization). Suppose good initialization g and let
VH 2 12 H³ B3 (3o(1/K))
1 - where go
2B + log K) K
Purther far sizp-size 1 1/p(K) A "log[I+e with p(K) 2 3/2 RaB go )² Then, it holds that
4B3
(2BB log(K))²
and - -((OK)) 17(2Bs + k))² 72 n
+ I(O(K)) s K 4 n K
Consider training loss after K GD steps: Assuming B3 - Ōx (1) and y 0x then choosing 1) Ō (1), the corollary guarantees train loss is xx( provided polylogarithmic number of heads H 22log (K)). Moreover, after K a GD steps the expected test. loss is (1/n).

---

Published in Transactions on Machine Learning Research (04/2024)
Remark 1. The last two conditions (P2 and P3) for good initialization are similar to the conditions needed in (Taheri & Thrampoulidis 2023 Ji & Telgarshy 2020; Nilanda et al.) 2019) for analysis of two-layer ML.Ps. Compared to (Ji & Telqarsky, 2020 Nitanda et al. 2019) which assume random Gaussian initialization g00, and simálar to (Taheri & Thrompoulidis, 2023) the NTK separability assumption (P3) can potentially accommodate delerminislie ( Condition (P1) appears because we allow Iraining both layers of the model.
Specifically the I2200 norm originates from the Hessian bounds in Corollary 7
5 Application to
model
We now demonstrate through an example how our results apply to specific data models.
Data model: An example. Consider M distinct patterns {1+, ,,, 22, .., PMb vMl, patterns M4 correspond to labels +1. The tokens are split into (i) a label-relevant set (R) and (īi) a label-irroeevant set (R [T].R). Conditioned on the label and R, the tokens æ,,e ee [7] are IID as follows
where discriminative
My
t R
(DM1)
Vy. + 21 3e
and te R
where 2, are noise vectors. Let denole the joint distribution induced by the described (X, pairs.
Assumption 3. The labels are equi-pree pri-ppobabbe and we further assume the following:
Orthogonal, equal-energy means: All palterns are orthogonal each other, i.e. μ4+ p.- Me 1 ve, V.E e [M]. Also, for all y {=1),EE [M] that | ||ve|| = S, where S denotes the signal strength. Sparsity level: The number of label-relevant tokens is [R] = =T; for sporsity level Ç (0,1).
sampled from distribution Dz, such that it holds almost surely for 2t ~ D= that (2t,, Myy] S Zp, y [+1} and (21, Zu/M, v& [M]. Moreover, | s Z. Overall, Assumption satisfied with R 52 Z 27./M.
Noise distribution: The noise tokens
The above assumptions can be relaxed, but without contributing new insights. We have chosen to present a model that representative and transparent in its analysis.
5.1 Finding good initialization
To apply the general Corollary 2 to the specific data model DM1 it suffices to find good initialization. While cannot directly show that (1)) = is good, g(), Thus, we consider training in two phases as follows.
show this is the case for first step of gradient descent
First phase: One step of GD as initialization. We use ni training samples to update the model parameters by running one-step of gradient descent starting from zero initialization. Specifically,
(U w( -0=)== =0 CAVH. Vo, (0 where
v [H).
Here, On denotes the step-size for head h [H] and the scaling by VH guarantees the update of each head is (1). The lemma below shows that at the end of this phase, have || Con Iru]]e 0(1/mm) where u, is the oracle classifier updated; the interesting aspect of the training lies in the second phase, which involves updating W.
μ4 - p_.. On the other hand, the attention weight-matrix does not get
Lemma 2 (First phase). After the first-gradient step as described above, we have U.!1) an 1r (su. +p")
and ( h 0. where with probability at least 1- 8 (0,1) over the randomness of labels there exists positive universal constant 0 such that
log( 1/5)
Ipll s C (2S Z) (V
=: P.
(6)
71
71

---

Published in Transactions on Machine Learning Research (04/2024)
Second phase: GD with constant step size. During the second phase, K gradient steps are performed nt new samples (distinet from those used in the first phase). Concretely, ((k+1) O(*) -n- V.T.. (O(k),, k = I....,K, K, with Ø(2) = concat (f0%) he(n) the step obtained by the first-phase update and v the step-size of the second phase. In order to analyze the second phase, during which both W and Ú get updated, employ i the general results of Section 4 To do so, we show that ((2) serves as good initialization as per Definition 1
Proposition 2. Consider the first-phase iterate (0( He and condition on the event Ip|| < P (depending only on the data randomness in the first phase) of Lemma & Suppose the step-size of the first phase is chosen IID a₁ ~ Unif h e [H]. Then, the initialization 0(1) = concat (o(") ,...,,((!) is good with respect to data sampled from DM1 and satisfying Assumption 3. Specifically, the three desired properties hold as follows.
Almost surely, P1 holds with B2= = =V=VT(S++ P).
With probability 1- (0, 1), P2 holds with BB = TR(S 2)oglo/((/))
**T(S+P) V2 Then, with probability 1 - e (0, 1), P3 holds with y n./2 where Suppose vH 2 7
7/x == T(1- T(-- - -755 - 127 16 T
Zu S SVT(s (s- 2(1 - 53
ST
PT Z)
v2
and := Zn Zy. The randomness is with respect to the sampling of an h [H].
The paramoter 7 in represents the NTK margin of the model at initialization Ø(1), By Corollary 2 larger margin translates to better reaiaiggneraliztiion bounds and smaller requirements on the number of heads. For a concrete example, suppose T v M (1) and Z Z O OS).. Then, provided first-phase sample size n1 & 5ddd so that P = O (1); it holds 7. "lin + (²( -))st, where Mia R((S) is the margin of a linear model for the same dataset (see App. Fr. Overall, applying Cor. 2 for K = and polylogarithmic
1 polylog(n) number of heads lcads to m3n
train loss and expected gencralization gap.
5.2 Proof sketch of P3: NTK separability
It is instructive see how P3 follows as sheds light on the choice of an appropriate target parameter 0 We choose
per Thms.
and]
W. p. p
Σ ve(u= +T) and U. = - = =1((4= - --) dc[.M]
and normalize parameters such that := O=: (U==UU-= TU.= U. sign(a) W. sign(a) [. W.). It is easy to see that U, is the optimal classifier for the label-relevant tokens. To gain intuition on the choice of W., note that W. WK,- WT with key-query matrices chosen as WK, [++ μ V /ar] Rdx(M+2) and Wo.- [u, μ._.. pl, + + ... ne, + p-] Rdx(M+2), With these choices, the relevance scores (aka softmax logits) of relevant tokens turn out to be strictly larger compared to the irrelevant tokens. Concretely, we show in App. D.2.3 that the t-th row (X; W.) XW. x, of the softmax-logit matrix satisfies the following:
(O($) (S(2) ,t' RS.
vt: [r.]
Thus, under this parameter choice, softmax can attend to label-relevant tokens and supresses noisy irrelevant tokens. In turn, this increases the signal-to-noise ratio for classification using U..
We now show how to compute Egcu VVVo o (x; 0.) for a single head. Recall 0. consists of U., W.. First,
since (11) 0, using Assumption 3 simple calculation shows y (( U.) SVT
1/ (c- (c- (1-1)))) 83
Second, to compute Ea- w (1) sign(a)W. W.) it follows from Lemma that
VWD(X;0()) ac Σ U. X x +a Σ xxp X x0)(
[c[T]
[e[T]
8

---

Published in Transactions on Machine Learning Research (04/2024)
Note the first term is dominant here since the second term can be controlled by making |p]p small as per Lemma Thus, ignoring here the second term (see Appendix D.2.3 for full calculation) y(V KVVw (()), W.) is governed by the following term: 2 (0( I 11 To simplify the exposition here, let focus the identity component and leave treatment of the the rank-one term to the detailed proof. The corresponding term then becomes
of Ex[T] yu ("))) (
2 of Ex(T) yu, (0) rr. Note that rz
oc Σ Σ Σ [c[T)ee[[] [T[T]]
(vulae)) ([r:)e)
class. logits softmax Ingits
which involves for each output token 1, the sum of products over all tokens e' [T] of softmax logits (i.e. relevant scores [rt]e) and corresponding classification logits (i.e. y yuīæe). Note that by choice of u, and W. both the classification and softmax logits large from label-relevant tokens, while being small for noise
tokens. Intuitively, this allows for positive margin % stated in Proposition 2 calculations to Appendix D.2.3
We defer the detailed
In the appendix, we also detail how to yield the computation for the MHA, which builds on the calculations for the single-head attention model above. In short, we simply choose multi-head parameter 0, as 0, = 0. (0 This guarantees that |0. v2 and maintains the multi-head NTK margin be least 7 in expectation. To complete the proof, it remains to get a high-probability version of this bound. To do this, notice that ( are IID, hence we can apply Hoeffding's inequality, which finally gives the desired bound on the NTK margin provided sufficient number heads H, which controls the degree of concentration when applying Hoeffding's incquality. See Lemmas 14 and for details.
VH concat (.
5.3 Is the NTK margin optimal?
Below, we discuss the optimality of the NTK margin 7.. First, define set of parameters Oopt := (Uopt, W Wopt):
U. Uppt == S/2T
and
W. W oppt == 52 /2 +
99)
normalized so that [|0opc I== Recall here the definitions of U. W. in the section above. As we already explained above, this choice of parameters guarantees that relevant tokens are assigned larger relevance and classification scores compared to irrelevant ones. Specifically about W., we saw in Eq. (8) that it ensures a gap of O($2) between relevance scores of label-relevant and label-irrelevant tokens. Thanks to this gap, it is possible for softmax to fully attend to the label-relevant tokens by saturating the softmax. To do this, it sulfices to scale-up W. by an amount oc 1/5². This is formalized in the proposition below.
Proposition 3 (Attention expressivity for tokenized mixture model). Consider single-head atlention model. Suppose the noise level is such that Zμ = Zu s $2/8. For any > 0, consider T. satisfying log E ( Then, the attention scores corresponding to weights D.- Wops satisfy
8/2(M+1)
T 35
Vt [T] 0 px( T.W. opt x²)= Σ (x) W. opt X²)SE.
SE.
(10)
Thus, abmost surely over data (X, .)) generated from data model with paramelers (U app Te. W o salisfies
DM1 the margin of single-head attention
VT V2s
yT (opsp Wopt) 2 Nattn (c)t(()
(1- -
(11)
From Eq. (10), note that as + 0 and I, - co, the softmax map saturates, i.e. it approaches a hard-max map that attends only to the label-relevant tokens (R) and suppress the rest (R"). As a consequence of this, Eq. (1]) shows that the achieved margin approaches Yattn == SVT//22 Note this independent of the
sparsity level Ç. In particular, Tattn 2 2%, Min and the gap increases with decreasing sparsity. See appendix for experiments and discussion regarding the margin achieved by GD for data model DMT
9

---

Published in Transactions on Machine Learning Research (04/2024)
Following Proposition 3 natural question arises: Is it possible to choose "good" parameters θ (Ü, W) based on the set of optimal parameters Oppt? This would then yield train-loss and expected generalization-gap bounds Ô (1/(T7 after (n)) steps of GD starting at 0 0. To investigate this question, dcfine the following parameters for each head, aligning with the aforementioned "good" directions of Proposition
Un := log(n)
C W == H1 Wapt:
Yaten H1/2 Uppt:
for some C> 0, p > O, and e [.]. To yield the margin Yatta of (10) we need that each Wn. has norm at least I. x 1/5². Thus, we need [| 2 st = S C HP Now, in order apply Thms. and the
requirement on the number of heads H in terms of distance of 0 to O(0) = 0 yields the following condition:
12 2
(12)
Note that |0||
Ing(n) (ng
W = C. 222-PP. Hence, in computing
|@] we distinguish two cases.
Tartn
s
First, assume that W
Ü which implies that S S log(n), H0-1/2 and @ W v |00| - C- C- Since
C1/2 HPl2 log(n) HP-1/2 1/2
S2
by using Eq.
12 we get the following conditions on H:
112 2 56 C². -3/2-399 39 2 C2/2, H2/2-012 v log (n) H2 HP-2 2 C and HP-1/4 5
C
log (n()
log(n) $ 1, a contradiction since n > 1. Thus, there are no possible choices for and C that satisfy both conditions. The case [W] s |Õ] can be treated similarly leading to the same conclusion; thus, is omitted for brevity.
Combining these two gives C'S (n)
Intuitively, this contradiction arises because of the large |Wall requirement to achieve margin Tastn Finally, one can ask if it is possible to resolve the contradiction by changing the scaling of normalization with respect to H in the MHA model Eq. (2) from 1/H to 1/H for e> 0. It can be shown via the same argument. that such value of exists for which @ constructed above satisfies the overparameterization requirement H 2 S5 @01. We thus conclude that the construction of weights in Proposition 3 docs not yield target parameter that sinultaneously achieves low empirical loss and allows choosing H large enough as per (3) This triggers interesting questions for future research: Does GD converge to weights attaining margin Yattn as in Proposition If so, under what conditions on initialization? See also the remarks Section 7
6 Proof Sketch of Section 4
Throughout this section we drop the Ø and (XX;; 0) as everything refers the full model. Moreover, Φ(K) and g(0) are denoted by 0k and 0g- Refer to Figure 1 in the App. for summary of the sketch.
6.1 Training analysis
The proof begins by showing sicp-wisc descent for any itcration k 0 of GD (sce Lemma 7 where step-size at each ileration 7i Ps. depends on the objective's local smoothness parameters Ps (Ox) v 32 (2(01)):
T(( xT( [[:))
(13)
Now, using Taylor's theorem we can link T.(0x) to T.(0) for any 0 as follows:
I(0) I(e.) (VI(O.), 0 0x) min Amin (v²I(On.)) 10 1,
(14)
Ona
10

---

Published in Transactions on Machine Learning Research (04/2024)
where 0k. == := + (1 0, a [0, 1]. We can plug this into
13 to relate the loss at iterates 0g and 0x-1- To continue, we need to lower bound mino., Amin ( (O,k,) For this, use the following property of the loss objective from Corollary 0A Amin (v²I 2 -x( I(O), where x(0) == vH
Ra(0) Note from the
definition of {-(-) that 101. 02 maxo<[a, ..] (0) (0)) (2)) Thus, the above property of the loss implies the following local self- bounded weak converity property the line [01;0s for arbitrary points 0,, 02:
Amin (v²T(0)) (0)) v (a(
A 10
min 0c(0;.02]
max T((). 0c[01.0s]
(15)
VH
Therefore, using Eq.
15 in Eq.
14 we can get:
1 (01) v (s(
((0) 2 L(Ow) + (VZ(O.). 0z)
max I(On..) 10-0&11
(16)
H
oc(o.1]
To apply the Descent Lemma in 13 at each iteration 17 S Th for all & < K. Then, combining with Eq. summation, we arrive the following:
we need to fix a step-size such that satisfies the condition of the Lemma
(16) and applying standard telescope
10 0ol
80(0) (0x)))
I(0))
max I(O,,) 10..)100000 021 oc(0,1]
(17)
K
2n
2K
VI
k=
Next, we use the following generalized local quasi-convexitty (GLQC) of the loss function.
Proposition (GLQC property: Slight variation Prop. of Taheri & Thrampoulidis (0023 ) Let 0 and 0g be two points that are sufficiently close each other, such that
2 (3)( Ba(O2)) 10 021 SVH.
(18)
Then, max T((0) 3 (T.(0)) v T((O2))
Using Proposition in Eq. 2 (Ba(0) v (0x)) 0 12 we can get the advertised regret bound in
17 and assuming sufficiently large heads H such that VH
2
In order to remove the dependence of И on iteration k, by an induction argument we can show bounded iterates-norm norm i.c. 10x 0|| S 3|0 - Bull (see Lemma 10 Using this and the definition of 83(-) we can control (0) v (0x) as (0) (ax))) & 110 Gol + ||0|2,000 to get the desired requirement of heads
VH 2 10 0 - 00 stated in Eq. 33.
The remaining piece to guarantee descent at each step is establishing a p(0) such that Pk S (00) for all & < K. To do this, we rocall that P = (0x)) v (2(0x+1). By definition of (2(+) in Corollary [ we can control (x() 32(0u.1) with controlling |kl|2,00 10x-1 12.00 as (1|Ox|1,,00 10x1012,00) $ |0 Using nterates-norm bound and setting p(0) (2VTaR" TR² (0) with a())
0+l 2,
VH 3 v R2? 33 R² (3 0oll =1012,0))+++0000)0))))0))))))))) +2VT n). satisfies the desired condition for the Descent Lemma completing the proof.
6.2 Generalization analysis
In order to bound the expected generalization gap, we leverage the algorithmic stability framework. To begin, consider the leave-one-out (loo) training loss --) == + Eje: (0) for i [n], where (, := E(y, (Xx;; 0)) denotes the j-th sample loss. With these, define the loo model updates GD on the loo loss for rj >
0 ++ == - (0T"). 1: 0,
0 O
The following lemma relates expected generalization loss to average model stability for any G- G-Lipschitz loss. Lemma 3 Lei & Ying 2020), Thm. 2). For G-Lipchhitz loss and for all iterates K, it holds that
E ([((0)) - 2G E 00x 0%4
11

---

Published in Transactions on Machine Learning Research (04/2024)
To bound the average model-stability on the r.h.s of the lemma's incquality, use GD expansiveness. Specifically applying (Theri & Thrampoulidis 2023 Lemma B.I.) to our setting, gives v0, 0::
((-)) - (0' -nVI((')) max {(1+ m7 (0.) I(O.)) v (82(0.))10000) 011, aefo,1]
(19)
VH
where, 0a =0 +1 a1-a)0', a [ 1]. Using this and gradient belf-boundedness from Corollary we get:
{(1 (0
km I (OEE) v n)2((E))) 110 0 n (0x), & (0x),
102+1- - s max nc(0.1]
(20)
VH
where 0 a0x + -1-00 a1-a)0p O) for [0.1]. Further using the bounded iterates- property from the training analysis, we can control (2(0 r
Kg s (0) and 3,(02z) S (e) making them independent of k (See
Lemma. 11 for the definitions of ,,,)),,)))) In order to invoke the Descent Lemma, we set the step-size
same as in the training analysis. Thus, (0 becomes:
0k+1 o il s ((1+ ())
max L T='()".) ||0 0,'| ()3(((s) e, (0x)
(21)
VH
ac[0.]]
n
As in the training analysis, we can control the loo empirical loss D for any point the line [0k, 0,²] of two sufficiently close points satisfying vH (Ba(Ox) v (s())r)) 10x 0-¹ 12. Using Prop.
Eq. 21
becomes
0z.11 0gill s (1 + an 10 0.'|| 1)((()) la (0x),
(22)
n
and Bi (0s) s 3 (0) similar to (-),Baa)) using bounded iterates- norm. Unrolling the iterates in (22 summing over i [n] and using training regret bounds, we have the following average model stability bound for any iterate K: E. |0x 0 2n²8(0) (2Kī(0)= 9) -10-0ol"
where Oggi == 3/H
4n/81(0)
"
4n
Combining this with application of Lemma 3 for our objective, which is G s 3 (0) Lipschitz from Corollary and using 7 S p())
1 S (3)())) we get the desired generalization gap stated in Thm.
7 Concluding remarks
We studied convergence and generalization of GD for training multi-head attention layer in classification task. Our training and generalization bounds hold under an appropriate realizability condition asking for the existence of an target model 0 achieving good train loss while being sufficiently close to initialization. In particular, from the condition on the number of heads H in (3, we need @ is at most Ô (2-1/37-116 R-5/3 f/
far from initialization (provided |,02,00 0, (1)). In Sec. 4.3 we showed that such a model exists if the initialization is chosen appropriately. Specifically it suffices that 1200 = (1), the model output at initialization is Ô (1)-bounded and that the data are linearly separable with margin y with respect to the NTK features of the model at initialization. Then, O (RPTRIO polylog(n)//{) number of heads guarantee that (n)) GD steps result in train and test loss bounds Ö (1)(m²n)). In Sec. 3 we applied our results to a tokenized-mixtune model. We showed that after one randomized gradient step from 0, the model satisfies the above conditions for good intialization. For this initialization, we computed the NTK margin % which in turn governs the guaranteed rate of convergence and generalization based on our general bounds. This opens several interesting questions for future work.
First, docs random initialization of attention weights satisfy NTK scparability, and if so, what is the corresponding margin? Second, are there other initialization strategies that guarantee the realizability conditions are satisfied? Here, note that our conditions for good initialization are only shown to be sufficient for realizability leaving room for improvements. Third, how suboptimal is the best NTK margin (among other potential natural initializations) compared to the model's global margin arg orgmax a m y/8 (xX;;;;;)) In Proposition we showed for the data model DMT that there exists single-head attention model 0opt
12

---

Published in Transactions on Machine Learning Research (04/2024)
v2 (1-c)-2eZu/3)) for all Te 2 0 and r. (for which the softmax map gets saturated and attends to tokens with highest relevance score) the achieved margin approaches Natt == S/T//22 which is independent of the sparsity level Ç. In particular, Nattn 2%74 2 Mi and the gap increases with decreasing sparsity. it possible to establish finite-time convergence bounds to models with margin & "Yttnn under appropriate initialization? How is the answer affected by the fact that the optimal attention weights in this case are diverging in norm (T, + oo)? Using our approach, we argued in Sec. 5.3 that the key challenge is the saturation of norm of Wopt (I+), which does not allow the appropriate realizability condition to hold (at least for 0 initialization). Finally, it is interesting to consider other data models for which multiple heads are necessary to interpolate the data.
(Uopt: Wopt) with l/0opz|| /2 such that yP(X; Wopt) SVT
1o)((5))-1)/e) and any (0, 1) (see App. E In particular,
8 Acknowledgements
This work is supported by NSERC Discovery Grant RGPIN-2021-03677, NSF Grant CCT-2009030, and a CIFAR AI Catalyst grant. The authors also acknowledge use of the Sockeye cluster by UBC Advanced Research Computing. PD thanks Bhavya Vasudeva for the helpful discussions.
References
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference Learning Representations, 2023. URL. https::/openreviews http://wwww....eeooo//////rrr/a/
Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- parameterization. In International Conference on Machine Learning, pp. 242-252 PMLR, 2019.
Sanjeev Arora, Simon Du, Wei Hu, Zhiyuan Li, and Ruosong Wang, Fine-grained analysis of optimization and generalization for overparameterized two-layer neural networks. In International Conference on Machine PMLR, 2019.
Learning, pp.
Pierre Baldi and Roman Vershynin. The quarks of atlention. arXiv preprint arXin:2202.0837 2022.
Arindam Banerjee, Pedro Cisneros- Cisneros-Velarde, Libin Zhu, and Mikhail Belkin. Restricted strong convexity of deep learning models with smooth activations. arXio preprinl ar arxi:22009.151 15100, 2022.
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of a transformer: A memory viewpoint. arXiv preprint arXiv:2306.000 00802, 2023.
Tom Brown, Benjamin Mann, Nick Ryder, Melanic Subbiah, Jared D 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 Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models few-shot learners. In H. Tarochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.). Advances in Neural Information Processing Systems, volume 33, pp. 1877 1901. Curran Associates, Inc., 2020. URL ttttps:/ppoceeddie 14a Saefff44-Paaaa Paperr....... .df
/tt.://///0...000000000000111110 seurips.cc/papee. http:/c//pppeerlle/ppppeeeeeeeee http://eesss//pe/222/000 /
Yuan Cao and Quanquan Gu. Generalization bounds of stochastic gradient descent for wide and deep neural networks. Advances in neural information processing systems, 32, 2019.
Zixiang Chen, Yuan Cao, Difan Zou, and Quanquan Gu. How much over-parameterizaaion is sufficient to learn deep relu networks? In International Conference on Learning Representations, 2020.
Jianpeng Cheng, Li Dong, and Mirella Lapata. Long short-term memory-networke for machine reading. In Proceedings of the 2016 Conference on Empirical Methods Natural Language Processing, pp. 551- 561. Association for Computational Linguistics, November 2016. doi: 10.18653/v1/DD D15-1053. URL https://cccaattaoooyyyor////1
13

---

Published in Transactions on Machine Learning Research (04/2024)
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. pp- 4171-4188 Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10. 1..1863//11NN1914222 URL httppps//acclanaoooooooo.../////
Yihe Dong, Jean- Tean-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. 2803. PMLR, 2021.
Alexey 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.
Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai. Gradient descent finds global minima of deep neural networks. In International conference machine learning, pp. 1675-1685. PMLR, 2019.
Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. arXiv preprint 2r 10. 10090, 2021.
Tolga Ergen, Behnam Neyshabur, and Harsh Mehta. Convexifying transformers: Improving optimization and understanding of transformer networks. arXiv preprint arXiv:2211.11 11052,, 2022.
Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nugp and ntk for deep atlention networks. In International Conference on Machine Learning, pp. 4376 4386. PMLR, 2020.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. Advances in neural information processing systems, 31, 2018
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 Nearal Information Processing Systems, 2022. URL http://ppparoooeeee.eeerrrrraaaa wttp://orumm.dddeaaaaaaaaaa.
Ziwei Ji and Matus Telgarsky. Risk and parameter convergence of logistic regression. arXiv preprint vrXii:1803.07300 2018.
Ziwei Ji and Matus Telgarsky. Polylogarithmic width suffices for gradient descent to achieve arbitrarily small test error with shallow relu networks. In International Conference on Learning Representations, 2020.
Ziwei Ji and Matus Telgarsky. Characterizing the implicit bias via a primal-dual analysis. In Algorithmic Learning Theory, pp. 772 804. PMLR, 2021.
Yunwen Lei and Yiming Ying. Fino-grained analysis of stability and generalization for stochastic gradient descent. In International Conference on Machine Learning, pp. 5809-5819. PMLR, 2020.
Yunwen Lei, Rong Jin, and Yiming Ying. Stability and generalization analysis of gradient methods for shallow neural networks. In Advonces in Neural Information Processing Systems, 2022.
Hongkang Li, Meng Weng, Sijia Liu, and Pin- Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. In International Conference on Learning Representations, 2023a.
Yingcong Li, M. Emrullah Iidiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning, 2023b.
Valerii Likhhsherstov, Krzysztof Choromanski, and Adrian Weller. On the expressive power of self-attention matrices, 2021.
Zhouhan Lin, Minwei Feng, Cicero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio. A structured self-attentive sentence embedding. In International Conference on Learning Representations, 2017. URT. httts:///pperrevier http://wwwww..eeeoooooooo//ao..d
14

---

Published in Transactions on Machine Learning Research (04/2024)
Chaoyue Liu, Libin Zhu, and Misha Belkin. On the linearity of large non-linear models: when and why the tangent kernel is constant. Advances in Neural Information Processing Systems, 33:15954-15996 2020.
Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer, and Veselin Stoyanov. Roberta: A robustly optimized bert pretraining approach, 2019.
Nya Loshchilov and Frank Hutter. Decoupled weight. decay regularization, 2019.
Haoye Lu, Yongyi Mao, and Amiya Nayak. On the dynamics of training attention models. In International Conference on Learning Representations, 2021. URL. htttps://operrevieieee// http://www....ooommodddoooooooaa
Sadegh Mahdavi, Renjie Liao, and Christos Thrampoulidis. Memorization capacity of multi-head attention in transformers. arXiv preprint ar Xiv:2300 02010, 2023.
Quynh Nguyen, Marco Mondelli, and Guido Montufar. Tight bounds on the smallest eigenvalue of the neural tangent kernel for deep relu networks. In Marina Meila. and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 8119 8129. PMLR, 18 18-24 Jul 2021.
Quynh N Nguyen and Marco Mondelli. Global convergence of deep networks with one wide layer followed by pyramidal topology. Advances in Neural Information Processing Systems, 33:1 1961-11972, 2020.
Konstantinos Nikolakakis, Farzin Haddadpour, Amin Karbasi, and Dionysios Kalogerias. Beyond Lipschitz: Sharp generalization and excess risk bounds for full-batch gd. arXiv preprint or X1x:2224..14410 2022.
Atsushi Nitanda, Geoffrey Chinot, and Taiji Suzuki. Gradient descent can learn less over-parameterized two-layer neural networks on classification problems. arXiv preprint arXiv:1905.09897 2019.
OpenAl. Openai: Introducing chatgpt, 2022. URL http:://peeaa..coooboog///aaggg 2022 2022
Ap Gpt-4 technical report, 2023.
Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparamettriizaionn Global convergence guarantces for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory, (1):84 105, 2020.
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of allention prompt-tuning. ICLR 2023 Workshop on Muthemalical and Empirical Understunding of Foundation Models, 2023.
Ankur Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit. decomposable allention model for natural language inference. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pp. 2249 2255, Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10.18653/v1/DD/ DD111D161111 1244. URL https://aclanthoolog wtw.aaatttoooooooooooo--/---- org/D16-1244
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 (ods.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 3748-8763. PMLR, 18-24 Jul 2021. URL https::/prroceedngss... .mr press/v139/radforddl 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 unified text-to-text transformer. 21(1), 2020. ISSN 1532-4435.
Dominic Richards and Πja Kuzborskij. Stability & generalisation of gradient descent for shallow neural networks without the neural tangent kernel. Advonces in Neural Information Processing Systems, 34: 8609 8621, 2021.
15

---

Published in Transactions on Machine Learning Research (04/2024)
Dominic Richards and Mike Rabbat. Learning with gradient descent and weakly convex losses. In International Conference on Arlificial Indelligence and Statistics, pp. 1990-1998. PMT.R, 2021.
Itay M Safran, Gilad Yehudai, and Ohad Shamir. The effects of mild over- paramttrieetinionnnn on the optimization landscape of shallow relu neural networks. In Mikhail Belkin and Samory Kpotufe (eds.), Proceedings of Thirty Pourth Conference on Learning Theery, volume 134 of Proceedings of Machine Learning Research, PMLR, 15-19 Aug 2021.
pp.
Arda Sahiner, Tolgs Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Unraveling attention via convex duality: Analysis and interpretations of vision transformers. International Conference on Machine Learning, 2022.
Clayton Sanford, Daniel Hsu, and Matus Telgarsky, Representational strengths and limitations of transformers. arXiv preprint ar Xix:29900. 02896, 2023.
Matan Schliserman and Tomer Koren. Stability vs implicit bias of gradient methods on separable data and beyond. In Po-Ling Loh and Maxim Raginsky (eds.). Proceedings of Thirty Fifth Conference on Learning Theory, volume 178 of Proceedings of Machine Learning Research, pp. 3380 3394. PMLR, 02 05 Jul 2022.
Ohad Shamir. Gradient methods never overfit on separable data. Journal of Machine Learning Research, 22 (85):1 20, 2021.
Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher D. Manning, Andrew Ng, and Christopher Potts. Recursive deep models for semantic compositionality over a sentiment treebank. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pp. 1633-1642, Seattle, Washington, USA, October 2013. Association for Computational Linguistics.
Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research, 19(1):2822 2878, 2018.
Hossein Taheri and Christos Thrampoulidis. Generalization and stability of interpolating neural networks with minimal width. aiv preprinl orXiv:2302.0022 09235,, 2023.
Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines, 2023a.
Davoud Alace Tarzanagh, Yingeong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism, 2023b.
Matus Telgarsky. Margins, shrinkage, and boosting. In International Conference on Machine Learning, pp- 307-315. PMLR, 2013.
Matus Telgarsky. Feature selection and low test error in shallow low-rotation relu networks. In The Eleventh International Conference on Learning Representations, 2022.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in l-layer transformer, 2023.
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 Reseurch, pp. 10347-10357 PMLR, 18 24 Jul 2021. URL html
https://proceeddngs. mmr.innnss........mm.m.m....m..m
Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, Aurclien Rodriguez, Armand Joulin, Edouard Grave, and Guillaune Lample. Llama: Open and elficient foundation language models, 2023.
Ashish Vaswani, Noam M. Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, 2017.
16

---

Published in Transactions on Machine Learning Research (04/2024)
Johannes von Oswald, Eyvind Niklasson, E. Randazzo, João Sacramento, Alexander Mordvintsev, Andrey Zh- moginov, and Max Vladymyrov. Transformers learn in-context by gradient descent. ArXiv, als/2212.076777, 2022.
Weihang Xu and Simon Du. Over-parameeteriaation exponentially slows down gradient descent for learning single neuron. In Gergely Neu and Lorenzo Rosasco (eds.), Proceedings of Thirty Sirth Conference Learning Theory, volume 195 of Proceedings of Machine Leurning Reseurch, pp. 1155 1198. PMLR, 2023. Are
Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. transformers universal approximators of sequeence-o--sequenee functions?, 2020a.
Chulhee Yun, Yin-Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawai, Sashank J. Reddi, and Sanjiv Kumar. o(n) connections are expressive enough: Universal approximability of sparse transformers, 2020b.
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in-context, 2023.
Mo Zhou, Rong Ge, and Chi Jin. local convergence theory for mildly over-parameterized two-layer neural network. In Mikhail Belkin and Samory Kpotufe (eds.), Proceedings of Thirty Fourth Conference on Learning Theory, volume 134 of Proceedings of Machine Learning Research, pp. 4577-4632. PMLR, 15-19 Aug 2021.
Zhenyu Zhu, Fanghui Liu, Grigorios Chrysos, Francesco Locatello, and Volkan Cevher. Benign overfitting in deep neural networks under lazy training. In International Conference on Machine Learning, pp. 43105 431288. PMLR, 2023.
17

---

Published in Transactions on Machine Learning Research (04/2024)
Contents
A Gradients and Hessian Calculations
18
A.1 Gradient Gradient/Hessian Hessian calculations for mullihead-atteetionn
19
A.2 Proof of Corollary
Objective's Gradient Hessian
25
B Training Analysis
26
Preliminaries
26
B.2
Proof of Proposition
26
B.3
Key Lemmas
27
B.4
Proof of Theorem
29
B.5
Corollary
Generalization Analysis
31
Proof of Theorem
34
C.2
Corollary
35
C.3 From good initialization to realizability
35
0.4
Proof of Corolarry 2
36
Proofs for Section
37
Uselul lacts
D.2
Proof of Proposition
37
D.2.1 Proof of P1: Bounded norm per head
38
D.2.2 Proof of P2: Bounded initialization
38
D.2.3 Proof P3: NTK separability
38
Proof of Lemma 2
43
E Optimal Model
44
E. Proof of Proposition
45
F Linear Model
46
G Experiments
46
Related work
48
A Gradients and Hessian Calculations
We define the following for convenience
('(0)= 1 Σ ((,,
1 Σ (x, 0))
and Σ"(
[e[m]
n [c[
18

---

Published in Transactions on Machine Learning Research (04/2024)
For logistic loss f(z) log +e"²):
1 and (0) 4
(0) s T((0) 1
(24)
We Á to denote the the vectorization of matrix A and denote the Hadamard product. Finally, we define e; (n)) as the i-th standard basis vector in R
A.1 Boadient/Hessian calculations for multihead-atteni
Lemma 4. Let the softmax allention model Φ 0)) in Eq. holds:
Then, for all veclors a E RR" and b, c Id a
(xO(X ((XWX'XX
Vw Σ x. x (XW'"x;;X. [=1
T
8. vwla.vo0(xx0)0)
Σ
a.aa'' x (XW'x,)))
Vw(cVw0((:000))
Σ c ni ae u X diag( diag(Xb) (xX Xb((,,)
-x((X x;) Xu, (Xb)7 4' xx)
Proof. For simplicity, we denote G= XWX and the rows of (G) (g))) = x(W x.), te [T]. Recall, for any v e Rd,
= -d)
(25)
i.c.,
(0 =p()) ((ins((())-pp(())))))))----
(26)
We start with the gradient with respect to U,
(27)
Next step to compute the gradient with respect to W, V ΣT, Vw 26).
Using Eq.
n X (XX(W a)))) = un T (+(()) + ((iig(((())) ((g))) XA xx.) + -(1
-u. x ((g))) tr X x' x X
a( A
Thus,
Vw x =x,4, x x,)
(28)
and
T Vw > x X (XW
(29)
1
19

---

Published in Transactions on Machine Learning Research (04/2024)
For the third statement of the lemma, have the following sequence of equalities,
Va (x;(( Vwla V (a' ((G) Xb)
T
-wir ((G))
Vw (a:b'X'
Σ aaa X ''WW x;)X.
(30)
where in the last equality used Eq. 28 For the Hcssian with respect to W, we have:
Vw(cVwo0(x0000) Σ Vw(( (' C'axux'' 6'((W'' xe) Xb
p
9:
Vwp,
P(g)) di (g)))
(31)
Termy
Termp
where we used the properly diag(ai) 02 diag (2) ai for vectors a,, ag IRT.
First, we compute the Termy above. Using Eq. (26):
pi (iag(q)))(X(WW + -A)xa xx) = pl diag( q:) (<(g))) + ('W' x,) x:)) +
= P. (()g(()))(()))) + tr(x,p) diag( ( ) (X)',, o(1)))))
(32)
Therefore,
Vwp. diag(q:)<((g.) x.p; diag (:)((())))))) 4:) (XWTx,;> (xxa(()))((XW''))
a. c¹ nz u. X diag( Xb) (XW'x::X xx)X
(33)
Similarly for Termm we have:
(X ())x))g(((() a)) (p]+((:)+ (p. x; XA w. + a((A|3))
x, XA a + o( A (2))
+l (())p'
(x +o(AA1)9)
(34)
Thus,
Vwrle((a)q =x(q'' p(g)pp" +p (XW"'xx)))) x
c² xu n ((g)) (Xu, (Xb)* + Xb(Xu) (''
(XW (WWx))x x.
(35)
Combining Eqs.
33
and
35
and plugging in Eq.
31 we conclude that
VwlcVVw4(X::)))) Σ
c x; x u X diag( - (p(g)) Xb(x;))
- Xuu (Xb))) (XX'',)) xx)
(36)
We restate Proposition
1
here for convenience.
20

---

Published in Transactions on Machine Learning Research (04/2024)
Proposition 5 (Restatement of Prop.
1
Let the softmar attention model ((X:0) in Eq. Then, it
holds:
1. (x:0( VT X
2,00
T
2. 1 52
Σ
X
Xu: lo0
T
|vo0(X:0) 52
Σ
X 00
VT X 12,00
4. |v)0(X:00 0) X 00 11.000 Σ
Xuelo + x 11,00
t=1
Moreover, if Assumption holds,
|Vo 5 | [1]++)).
0) (3VaR' U+| 1)
Proof. Using the first statement of Lemma
T
O(x0)1 =
Σ x (g:))
Σl 2 (
T max vT X 2,,00 [q[7]
(37)
where we usod that 1
Using Lemma 4 for the gradient with respect to W
T Vxx;)
a X
(38)
Therefore, by applying triangle inequality,
T x
W (x:
X (g:) Xuel.
In the next step we bound |X
| v'(W'' x;)
X
|X ((:((()))))))) (gr))) Xuull
(39)
Term
Termir
For Termr we do follows:
X diag(sp(g)))
max weRd and ||||1
(Xv) diag()(g:)) Xu:
max
[diag(<<(:))])= [Xu]] [Xe], |0|=1
max (gg:) ((Xu:) @ (Xv))
|o|-1
< max |o
(pg:))l (Xu:) @ (Xv) o0
< X Kurllo
max |||=1
X
= Xuul
X 2,
(40)
21

---

Published in Transactions on Machine Learning Research (04/2024)
where used Hölder's incquality in the first inequality.
Then, we compute Termi:
|x ((())))))())) Xuil
max weRe and ||||-1
(Xv) (X))))((())) (pg)))(())) Xu
max |
(g:))) Xu, oo
X
Xuell
max |
X0/
Xuull
|
2,00
(41)
where we used Hölder's inoquality in the first inoquality. Combining Eqs. 40
and
41 and plugging in Eq.
39
yields
((xx))) s 2| ||,,00
(42)
Therefore,
V =2 2,00 Σ
Xuel
(43)
1=1
Using Eqs.
37
and
43
conclude that
T Σ
Xueloo
(44)
2,00
By using Assumption to simplify Eq. 14
T vob(X:0) 22R Σ Xuilo +TR,
(45)
or
\vo((X;0) s S2VTR [ VTR.
(46)
We need to derive the Hessian to bound the greatest eigenvalue of it.
v3 Vzx(x:) WD O(X:0)
(xx;0)))
V V (x:0)
First, we compute the Hessian with respect to U,
vi ((x;0) X = OraxTa
(47)
In the next step, use the third statement of Lemma and set c and b of00) to compute the Hessian with respect to W and U,
V Vw([voo(X;0)) -xx
(48)
recall, X.j is the j-th column of X. Therefore,
V WO ((x0)) = Vw([vo
(49)
( ())m))))
Vw([vo((X:0))
Lastly, in order to compute the Hessian with respect to W, we use the last statement of Lemma 4 and set c= (d) and b= (.)
(1)
T Σ X [ X diag(X -xp.)) "x X,
(T(TT((()00)))
1=1
-(g( Xu: ) (
(50)
22

---

Published in Transactions on Machine Learning Research (04/2024)
Thus,
(v([v)0(X:0))1)
Vw((F
2
W 1 (x;)=
(51)
Vw([V (X:
Vw((Fw6(X;0))u)
To find the maximum eigenvalue of the Hessian we need upper-bound
max (w, (x)(0)) |=|=1
where
concat (PI....,P7 ....... qa) E Td
and
Rd t [T] and qi R² .je [d].
Then,
V2 1 x; 0)
X q xz (Xxu)) dia)( " (Xu:) -p(g:)" Xu,X., ('(::)
< max
1==1 1- j=1
Term
T
+2 max Σ |= --1
Σ
p x. X y 3 (g:)
(52)
Tormg
First, we bound Term;:
T N q. wz (Xu)² diag( X..) - :p(g:)) "x(g))"".,XXux)) (Xx)( - Xu:X! ('(g:)
max
1
j=1
$ X 1,00
X ||,00
max
Σ
qell |1 I(Xu,)
|-|-1
Term:
max Σ Σ
Σ
quil | (1))))))))))))))))))(( (Xu,) ('(:) (xx)))
|
1=1 j=1 e
Termi
d
T
max Σ Σ Σ
q |q (l)(())))))))))))))(( Xu, X. X|
(53)
1-1
Termia
For Termi, we have:
max ΣΣ
Σ
Σ
q
|1 (Xu:))
di) X
- 1=1 j-1 1=
d d T
2 max
Σ Σ
Σl
q:il [|g
|X |diagg( X Xuullo
|1|-1 s=1j=1 t=1
d T Σ Σ
T
2 max
Σ
qill q;
||
|X 0z
X /0
22d |
2,,00
X 1,00
Σl
X Xulloo
(54)
||==1 1=1=11 m
23

---

Published in Transactions on Machine Learning Research (04/2024)
where in the last inequality used Cauchy-Schwart incquality. Similarly for Termu write:
d d T
max |=|=1
Σ Σ
Σl
qill
(p(:))))()))))
(uu)) (':))))))))
2 max |0|
Σ
Σ
Σ
qill
| (
X 2,000 o S2d|X 11,00
Σlxuell (55)
j-1
where we used Hölder's inequality in the last inequality, And at the end, for Termii we have:
d
q.ll 1 (p())) Xu, X ('),) (
max |0|11
max
qill q (g(g)) Xui| |2, | s2d|
| 1,
Σ
X Xuelo
(56)
Then, we upper-bound Term2:
2 max |-|-| j=1t=l
Σ pi ve X. (T))) $2 max
Σ
Prll
X
((t)
|0|-1
j=1 t=1
T X
4
1
=
2.00
X 1,00
(57)
The last inequality comes from the calculation of the gradient concluded in Eq. conclude that
42
Pl Plugging in Eqs.
54
55
56
and
57 in Eq- 52
III
2,00
X 1,.00
(58)
Applying Assumption
we have;
T [(00()00 <6dR R
+ +2VTdR
(59)
(3)0(;001 $6dR
Σ max x u ddR
1=1 re[T]
$ 6d R
Σ max |x, |u|| + +2 TE[T]
<6VTRR² 1
(60)
where in the second and the last inequalities used Cauchy-Schwand incquality.
7
We now derive bounds for the multi-head head attention model.
Lemma 5. Recall the multihead model in Eq. 2
1
h
(
[U., (XW.
(61)
H
The following are true under Assumption
(x VTR 2R²
max Unll=+1) ke[H]
24

---

Published in Transactions on Machine Learning Research (04/2024)
2. (x
2 2VTRR3
3VdRR
max he[H]
| 11
VH
Proof. From Equation
46
for single-head attention,
JUAl= + VTR) 2
|V (x;
Σ 2 T
H
h-1
(2VTR
max he
UN| UN U
(62)
Denote v² = [a]
wH} (H(TT+E)) where | 1. Using Equation 60
(0, (v,vzá(X;0)))
1
(OVTaR*
Un| +2VT d 123
max hc[M
VHI
1
(VVTRR*
UNl + +2VTRR2) be[H]
(63)
VHI
A.2 Proof of Corollary 1 Objective's Gradient/Hessian
Corollary follows immediately by the result of Lemma below and using more relaxed bound ma U I |10|2,
Lemma 6 (Tight version of Corollary D. Let Assumption hold and we use logistic loss function. Then, the following are true for the loss gradient and Hessian:
1. |vī(O) VTR 2 R
max Unll + T.(0). he[H]
2 RVTRR |v\T(0)]
T TR
)² Unlll [Unll++]]
max 3VdR² he[
Un Unllo +1
R²
max he[H]
|
VH
4
3. Amin (v²T(()) 2
2VTRR
3van
Un|[++1 T((O). he[H]]
max
VH
Proof. The loss gradient is derived as follows.
1 R vi(0)
(x) vx((;
1-1
Recalling that y/ E {+]} we can write
[v[ = 1 & o(x;)) VV 1 - ((x;;))) V 0)
Thus, using Lemma to bound the norm of the model's gradient:
Vi((0) S VT
R²
max U I +1 I(0O. he[H]
Ux
(64)
For the Hessian of loss, note that
v'T(0)) 1 Σ (v: xxx ( 3 X 0))
(x
(65)
25

---

Published in Transactions on Machine Learning Research (04/2024)
It follows that
|v"
Σ " (xxx (())) V 0)
n
R
$ Σ
el (v: o x
(x,; 011 (1 ((x;;) 0)) 1V (xx
1-1
2 TTRR
maxI|Unlle 1) TR he]AA]
3
max
2 22
max he[H]
U
(66)
VH
To lower-bound the minimum eigenvalue of the Hessian of loss, note that {(-) is convex and thus ("(-) 0. Therefore, the first term in 65 is positive semi-definite and the second term can be lower-bounded as follows,
Amin (v (I(())
-1 Σ (v क( 0)) y (X, 0)
11
2 - Σ
(y: x (
V 1|
F=
2 d R
3 3VVRR
Un r T((O).
max he[H]
VH
B
Training Analysis
B.1 Preliminaries
Throughout this section we drop in 0 and ((x;;)) as everything refers to the full model. Moreover, g(k) and (()) are denoted by 0g and 0g.
The proof of both the training and generalization analysis follows the high-level steps outlined in Taheri & Thrampoulidis (2023). However, our analysis focuses on the self-attention model, which differs from the two-layer perceptron studied in the referenced work. Notably, train both the attention weights and the classifier head, whereas Taheri & Thrampoulddis (2023) assumees fixed outer layer weights. This introduces new challenge as the smoothness and curvature of the objective function at point O become dependent on O itself (see Corollary
D. Consequently, we make careful adjustments in the proof to account. for this challenge.
Our analysis critically uses the following property of the loss objective from Corollary ! VO: Amin (v²]())) ()) T(()), x=())= := (z(0) 3a (02). Thus, the above property of the loss implies the following local self-bounded weak convenily property on the line [0,, 02]:
vH Note from the definition of f-(-) that V0,, maxe<(0: .0)) (0 (3 v
min Amin (v*I(O)) 3a v (02)) L(0). 0c[0,.02] .2]
vO,, 02:
(67)
VH
In turn, Equation 67 can be used to prove that the loss satisfies a generalized local quasi-convexing (GLQC) property as formalized in the proposition below. The proposition is only a slight modification of Tahhri & Thrampoulidis 2023, Prop. 8). While,, direct application of their result is not possible since the self-bundedd
weak convexity property
67 hhlds only locally on the line [0,, 02], an inspection of their proof shows
that this is sufficient.
B.2 Proof of Proposition
First, we restate the proposition below for the reader's convenience. Proposition 6 (GLQC property). Let 01, 02 be points sufficiently close to each other, such that
(0)(0)) (s(02)) 00: 02 VH.
26

---

Published in Transactions on Machine Learning Research (04/2024)
Then, the following generalized local quasi- (GLQC) property holds:
4 max T((0)= < (T.(0.) v T((O2)). 0c(0; .02]
Proof. Although the loss L is not uniformly self-bounded weakly convex as assumed in (Thhri & Thram- poulidis 2023 Prop. 8), an inspection of their proof shows that for every 01, 02 it suffices that the loss is locally self-bounde weakly convex. With this observation, we can apply their proposition with K + k(0,) v k(02), which gives the desired.
B.3
Key Lemmas
The proof of Theorem | consists of several intermediate lemma,, which we state and prove in this section. Lemma (Descent Lemma). Let Assumplion hold. Then, for any iteration O have step-wise descent:
T((k+1) s T(a)) ! [[vT((
(68)
provided n where Pk is the objective's 's local smoothness porameter defined below,
P (2x) (2(0x+1).
Proof. By Taylor's expansion, there exists 0' [0xx 0x+1] such that,
L (0x-.) = I(0s) +VI Ok.1 0z) (0x-1 Os. V'I((' -0s) (T.(0))+ + (0.) ([(())+(TTT.0.))
Px+1 0z)
| - 0x11
max
0E..)
TPpp ||
s <0(
2
where the last step follows from Corollary ! and using maxgre(o,.0... ) S L. we conclude the claim.
32(00)==) 33 v 32(0)) = Pk. For
PE
Lemma 8. Assume 7) > 0 such that step-wise descent
68
holds for all [K 1J. Then, for any O, the
Sollowing holds:
10
0ol
K-1
Σ ((o))s s L(())
Σ
Tx| - 02 0,12
(69)
2
2K
k=0
1 where Th /H (8s(0) (s)) max.e[o.1] I(Ok.).
Proof. By Taylor's theorem, for Ox., == a0; (1 - a)), a [0. 1] we know that:
+())
0x11 0z) Ami (v2I(Ok.)) 10- 10=
Using
68
we gel:
I(0 I(()+ + (VZ((+) -2-0) Amin (V' ! VI())))
10- 0
|0 0 |
1
I(0)=
0 0x1/3
(70)
2n
2n
2
where the last step follows by completion of squares using O.11 0, =n) and also uses Corollary with Tv. max,c[(.1] Th, where Tk. (0(0,..)
)(( L(Ok.). Telescoping in 70 for ie O.... K get the VH
desired.
27

---

Published in Transactions on Machine Learning Research (04/2024)
Next, use the generalized local quasi-convexity property (Proposition to obtain explicit regret bound from Lemma 8
Lemma 9. Suppose the assumptions of Lemma hold. Moreover, assume for all [K - 1] it holds that v H 2 (0) 3s(0x)) |00 0xl Then
1 K
5 2I 3|10 110-0of
0ol2
I(0s)
4n + 2K
(71)
Proof.
We have minioe(0000)) Amin (²I((0)) (1(0)v/a(0.) maxgr:[0.0. T(O') from Corollary Thus, using
VH
Proposition we can control co..c (0kk)) Specifically, by assumption have for all & [ 1] that
(0)) |00 0xll² max 8z(On.) 10- 02112.
ae[0,1]
Then, by Proposition
we have
4 max) [(0+)++ 4 I()))
max I(Ok.) s ac[0,1]
Thus, applying this to 69
we have:
Σ I(Ox) s I(0) + 0 0ol
2
Σ
(3g(0) v (s(0,)) (I(O)) + +I(O)) -
K 1=1
2n
3K
VH
k=
s <I(0) + 0 0ol
K-1
(I(0x) I()))
2n
3K
k=
4 I()= + 10 00l $ 3
Σ I(Ox).
(72)
2
3K k=
where the condition on H in the second inequality. Rearranging terms above we conclude the claim of the lemma.
Lemma 10 (Iterates-norm bound). Suppose the descent property 68 holds vk [K 1], and let Assumption hold. Further, assume that
|0 10-0o1 22 max( mmx((nk((0),nn<((0)))
(73)
and
VH 2 36 / 3VaRR
3|0- 0ul ++ 2,00 +1] 110 00ll
+1002,,00
(74)
Then, for all is [K],
|| 0|| 3 6ul.
(75)
Proof. Denote A = = 0|. Start by recalling from Eq. (70 that for all k:
A A 2nI(()) 2n, n max TR.,) ) A
(76)
oc(0.]]
where recall that Th., Pa(O2, I(0...) We will prove the desired statement 75 VH
using induction. For k = 0,
Ag ||e 0ul. Thus, the assumption of induction holds. Now assume A, s 3 Ool, Vj [k- 1]. We will then prove it holds for k.
correct for j E [k 1], i..
28

---

Published in Transactions on Machine Learning Research (04/2024)
By induction hypothesis for all j [k 1], and all a [0, 1]:
VH2 36 TdR
3 VRRR
3 0 |0|12,c
+1] 10 0
24VTdR"
3 3v R2?
3/0 0ol 110
+ 1 100 0
[10- 0,112,00 + 10112,00
24VTRR*
+1 10 0,||²
24VTdR*
(3v 10 10 + 1) 10 0,1²= 28, 0,1²
Thus, by Proposition [jc[k-1
max T((Og)) xT(,)) + T.)). ec[0.1]
Using this
76 we find for all j [k- 1].
13.1 s A3 21)T((0) 20TT(00(11) + maxx<((.1) [0.1]] B3
Bs(O),) )43
+ (T)
44
+T
VH
3
3
2 (
2
$ A3 + (0) 27T((,,11) +
where in the second inequality we used to get
77
We proceed by telescoping the above display over j 0,1....,KK k-1
8
2 (TT((())) (T)
4 k- Σ T((O,) 2n
A A0 +
3
3
3"
j-1
$ A/ + 8
2
I (1((0)))
3
3
2
$ 0 1
0 0ol 12
0-001
s 9 - 101
(78)
where the second line follows by non-ncgativity of the loss. Thus, A 5310--ol1- This completes the induction and proves the lemma.
B.4 Proof of Theorem
We restate the theorem here for convenience, this time also including exact constants. Theorem 3 (Restatement of Thm. D Fix any 0 and H satisfying
VH 2 36 /dR' (3va1e 3 0o1 +11012,00 0z
)
+11 )+1] 110 10-0oll2
0oll
Purther, denole for convenience
a(0)==3vdr² ( BT 0ol +1002,00 + +2VTR
a(0)² Then, for any step-size n 1/p(0) and p(0)= - (200)R2 TR² bounds hold for the training loss and the weights" norm at iteration K of GD:
10-0.12 KI(0)
10-0o12 I(0,)
the following
/H 1
K E I(0.) 2[( 5 0ul k=1
[(0x) <
and |0 0oll 4 4 0ul
(79)
4nj
29

---

Published in Transactions on Machine Learning Research (04/2024)
Proof. Define
= =(0))=vVRR' ( (3)1010-00) doll +1012,=)) (V
and
(2 T
(,)
a,)
VHI
From Lemma recall that:
== (2 v (2(0x+1)
Now, recalling the definition of (0) from Corollary al step depends 1,0 10 ||,000; where:
we see that Pk. the objective's smoothness parameter
0 1,.00 v 110x-1112,00 s (10 0x1 2,0 ) v (10- 1x+1l + =0)
$ (10 0zl 2,0 ) v (|)0 0₂1 10k.1 - 0xl + 2,000
1x-0x11 + 10 0&l + |0|2,,
/ 1 + +1) + ||0 0μ |00 2,00
(2n/TR² + (10 -10--x11 110 |
TR- Ok.
where the second-last inequality follows from Corollary For each Pk. define corresponding Pry (r)())
2vTaR (3VdR20 +) TR² (2 R 0 + 1)) Consider Pa. it is easy to see that S Po, (G,) p. (0).
VH
Hence, the descent property of GD holds in first iteration as per Lemma 7 Since, for n) 1, (,) p
Thus, S 1 ((0) = イ p()) Pa(0) S (e.)
Moreover, note that the assumptions of Lemma 10
are satisfied. Thus, by induction over Lemmas 10 and noting that Pk: $ (Ox)) (,) for all k E [K
by using& similar argument as above, we oblain for any 1 p(O)
Vk [K]
10 01 s 3 0p0l
1 K
3 - 0ol
I(0o)
and
Σ ((0x)s2((0)
(80)
An
2K
Moreover, by assumptions of the theorem we immediately find that 10 0ol |10 0& + ||0 0 410-0ol:, This completes the proof.
1
I(Oo) 5
10-0.12 2n
We also have
2K
The training proof is summarized in Figure
B.5 Corollary 3
Corollary 3 (Training loss under realizability). Let Assumptions and 2 hold. Fir K 2 1. Assume any H such that
1 VH 36 3T (3VaR* (390 (3900(
o( (22 gu( )
(81)
K
K
Further, denote for convenience
+ R2' 330 K + g6 1 + 2 n]
a 3 3 3VT
and p(K) = 12/TTR" TR.)
a(K). Then, for any step-size S 1/p(K) go)2
go(1)²
I(o,) the following
VHI
bounds hold for the weights' norm and objective at iteration K of GD:
2 T(Ok)
5gol )
and |0 0ol s 4go K
(82)
K
4nK
30

---

Published in Transactions on Machine Learning Research (04/2024)
Proof by induction:
((0)) x ((0) + Em]
Σ
110 1 Gbll
and
||0x Abll s ||0 Poll, for 2 |10l| 110 6,11", 0(1).
n.K
Clearly, bounded iterates holds for K O.
Next, we see Despent L Lemma holds for & 0, Lo. t.0 s E(0)
78 (a)
wm
11 S 10
Easy show that: 16 110 110 B,ll + 1101 Hance Mo O())
110 22
Assume bounded herates til step
descent lemma holds lil step
Use second-order Taylors: Z(0) Z(e) + (VZ00)0 0) min Amca (12(1)) 10-1,1. and local awf-bounded wwai--coveutty of the loss
lower bound v0,0g
mim Assa (v
(A00)) (A()))
max Ze0 0010,00
VII
Combine with descent tomma and teloscape obtain regret boundh
K
110y s 2100 +
10 &.l
2q
8,00 1.00)
VH
max
L0 a.F
For suffitiently large # (1(66 16 8,1" / VA E [K
centrol
max (cil.]
[(o,)) E Le0 to got the desired regret bound.
Firish by using the regret bound show that for 2 1011 A,ll" Herstes are bounced for step K, Lo. 100g Hbl is 110 8ill
Figure 1: Training proof schema.
Proof. According to Assumption 2 for any sufficiently small e > O, there exists a ((c) such that I(((=)) SE and |0(=) 0ol (o(E). Pick K. Let the step-size rj > O, satisfy the assumption Descent Lemma Since T((0(1/K)) K. have:
10(1K) -012 Kī(O(1/K) 2 |((//K) 10001))0001 go! )
2 go (1)
and
10(11))-00)) - )²
To (1)²
T((0)) T(,o)
T(Oo)
Therefore, following our assumption on step-size 7, conclude that
go(1) (0(1/)) - [0(1K) -
n7 s gl))01)
(83)
I(Oo) Kī(OC/K) I(Oo)
where in the second inequality we used the fact that 90(-) is non-increasing function. The desired result is obtained by Theorem I
C Generalization Analysis
Throughout this scction we drop t in Ø and T(x;0)) everything refers to the full model. Moreover, ((K) and (() are denoted by 0x and 0g.
For the stability analysis below, recall the definition of the leave-one-out (loo) training loss for i [m] and note that by denoting (0) == & (X;;0)) to be the j-th sample loss: I- I=((0)= : Σe (0)) With these,
define the loo model updates of GD on the loo loss for some 77 >> 0:
0 k+ 0 ( 1 0,
0% On-
31

---

Published in Transactions on Machine Learning Research (04/2024)
Lemma 11. Assume the conditions of Theorem 1 hold and
VH 256 TTRR (3Van² (310 3 0ol 1011,,00) + 1) 10 10-0.12
(84)
Then, the on on-average model stability at iteration K of GD satisfies,
2n (VTR n² (3 0ol 10 2,00) ) + 1)) 22KT(0)= 9|10 -0000.0l)
0oll
Σ
Ox
0x
ORl
n
4n
Proof. First recall from Corollary that gradient and hessian's norm satisfy
[vT( < (O)T(()
Amin (VII())) 2 (0( I(0).
VH
Applying (Taheri & Thrampoulidis
2023
Lemma B.I.), two arbitrary points 0, 0' satisfy the following
GD-expansiveness inequality:
nBa (1. nBa(0a) (0a) T(O.)) ((0(O))(((0))) v n20(0.)) 0'l,
[1(0 ((0-~vV(())- - -O'- max
(85)
oe(0.1]
VH
where 0, a0 (1-a)0 a)0' denotes point parameterized by a [0, in the line segment between and 0'.
We aim to bound the on-average model stability in the r.h.s of the incquality in Lemma 3 Based on Eq. (85)
| 1 O ||
(0x 1 (0x)) (0r² II (0r"'((rr))
- 7 v (0.)
(Ox vVL (0k)) -0 nVV (0r')) " |7,, (0x)
T (0%)) (0 v (0r")) n)1(0x)) e, (0x)
(0x
n
max [e(00_]
[
n3 (Z.)) T. T.=()rt)) n)2((0))) ))10- 10:0-0p711
03'1
73 e, (0x).
(86)
VH
n
In the above we denoted OF == a&x (1 (1-a)00' We note that based on our guarantees for the weights' norm during training 75 it can be deduced that for all a e [0,
(s(O RV (3VRR² |00x (1 - =0)00 2, + +1
=2TTR
(3vdR² (10x112.co 10 1, +
$2VTOR (3VÆR* ((10) 01 + |0112,00) (100 -l| -0112,,00) +1)
22VTdR'
(3v (3]0 0nl + 10 +
+1) 1) B))
(87)
Similarly, we obtain
2 T
T R2
(2(0
(3VOR2 1 2,00 1)+
2R² 110m 12.00
2,,00 +1)
VH
4
d.
(3VÆR²
33100 00l
110 11012,z0) +
VH
T T
2
2 R
(3|
0ol
10 +1
=: (0)
(88)
4
32

---

Published in Transactions on Machine Learning Research (04/2024)
Hence, by the notation introduced above and noting that by our assumption on the step-size it holds s (0), we can rewrite Eq. (86) as follows,
n (1+
n8I(0s) L, (0x)
Ok+1 Okizl
max 1-)(
110 110s-00'000 0
VH
ac [c(0.1]
Assume
VH 2 128 (0) 10 0ol1² 2 4 (0) (110; 0ol12 +100²² 0oll²) 2 21 22) 10 -.'1²
2 2 (3a(0)) (a()))) |0: 0
where we used Theorem in the first inequality. We also have
min Amin Amin (v"L= 30(02) 8s(0,)) I (or)))
ad [c
VH
Thus, by applying Proposition
on the leavo-onc-out loss, it holds that for all a [0. 1].
max (O)( ) (L ac[0,1]
I (O')
Thus,
(1+
4n
+ T.-( ) (T.+
0x 0.'1l +|+ n11
0k+1 0hill
(0x))
(89)
3/H
n
In order to remove the dependence on k, note that
31(0w) VT R 2 10 +1
VT R
2 (10 -10001122,00
0x112,00 +
=1012,000
SVT R
2 (10 0xll +|0++11000 -1012,000 +
S 18
0oll + -10 + (.). 310 2
(90)
For simplicity of exposition, denote a,,; have,
10))(0) 3/H
(Z,)+I=(00))))
Then by unrolling the iterates we
1 1 (1 ax.) 0= - - ())
4 (0x)
(1 ax.))
e, (0x-1)+ n),
$ (1+ax,))1 (1+ + 10-11 - 0-4
l, (0))
(0) (00021) m133 (0) e, (0s)
J-1
n
n
mi (0,-1) o
l (0x)
$ Σ
exp(
+ n) (0)
an,1) m3 (0) (0
e, (0x)
<
Σ
exp((
Σ
n3 (0)
(91)
1=1
n
n
33

---

Published in Transactions on Machine Learning Research (04/2024)
where in the above used that 0 = O
in unrolling the iterates well as the fact that for 0 S
Note that by definition (0x) s T(Oa).. By training loss guarantees from Eq. 79 we have
(0)
&
Σ o1 $
Σ
(0. + (0,')
3
[=
4nj3a (0)
k
(T((0)) + -[(0T'))
3/H
1
(0)
5kL((0)+ + 5/10 - 0ol
3/H
2n
10,3g (0) 0oll
<
VH
1
$
10
where the last step stems from the condition on VH. Proceeding from Eq. 91 iterate
find that for the last
K-1 Σ
(0(--))) +m31
l: (0K-1)
0 0 s 2m)
k=
n
K-1
(0c))
$ 2nm3
Σ
(92)
k=Ù
It follows that the on-average model stability satisfies,
(0) K-1
Σ 0x
0 5
Σ
Σ
(0x)
n²
1=1
1= k=0
3 (0)
I(Ox).
K=0
Applying our training loss guarantees from Eq.
80
to the r.h.s. above yields,
1
0k 0 2m (0)
(2K[()) 910-0.12
n
4n
where here we used the assumption that, I(e.) s 0o0l011/// simplify the final result. This completes the proof.
C.1 Proof of Theorem 2
We restate the theorem here for convenience, this time also including exact constants.
Theorem 4 (Restatement of Thm.
2 Fix any O and H salisfying
H 256 R (3V 310 - 001 001+1012,00 + +11 10 -
10--011. 0012.
Further, denote for convenience
a(0) 3vdR 3V (310 0o1 1012,00 + +2
and p(0) - -2VTTRRR T 4 (0) Then, for any step-size n S 1 A 1/p(0)
10-0.12
10-0.12
the expected
Kī(0)
I(o.)
generalization gap at ileration K satisfies,
E(I(O(())- T((O(K)) -((()))))T((((K)))) - =[(2xT((0) 9 O)
(93)
4n
34

---

Published in Transactions on Machine Learning Research (04/2024)
Proof. Note that the assumptions of Lemma 11 are satisfied. Moreover, per Corollary the objective is
Lipschitz at all iterates with parameter 31 (0) since Ak E [K] : |10& 2,00 z 3 0nl | 2.00 Thus, by Lemma
and Lemma
I we have,
- E (3)(0))) (2KI + 9 - 0oll
}
(LOO))I((R))
(94)
4n
Recalling the condition on step-size 17 1 ,()) (3)(0))²
1 concludes the proof.
C.2 Corollary
Corollary Generalization loss under realizability). Let boundedness Assumplion hhold. Also let realizabilily assumption holds almost surely over the data distribution. Fir K 1. Assume any H such that
H 256 VR (3VaR² (3gul k o( ))
+ gol K )2
(95)
Let the step-size satisfy n 1/p(K) go(1)² ga(1)² I(o,) where p(K) as defined in Corollary expected generalization gop at iteration K of GD satisfies,
Then the
17 90 )²
[[L(0)-I((k))
(96)
17
Proof. According to Assumption 2 for any sufficiently small e O, there exists a ( such that I(O(e)) S and (0(c) 0oll = go()). Pick e Let the step-size, 1j 0 satisfies the assumption of Descent Lemma Since I(O(2/K)) s K. we have:
(0(1K)) 0ol3 KT(0U//)) 2 [10((/K) - -bull- = =go( x ) 2² 90)
and
[10(1/K) -0012 )² (o(2)
I(0o) I(Oo) I(Ou)
Therefore, we can conclude that
(0(1) I(Oo)) <
(0(11K)000) - ^ Kī(OC/K))
|0(1K) -
S go()
(97)
I(Oo)
where in the second inequality used the fact that go(.) is non-increasing function. The desired result is obtained by Theorem 2 and the fact that
k((O(VK)) 10(1K) -012 g0( )2
C.3 From good initialization to realizability
The proposition below shows that starting from a good initialization we can always find ((s) satisfying the realizability Assumption 2 provided the number of heads is large enough. Proposition 7 (From good initialization to realizabilityy. Suppose good initialization Og per Definition 1 Fir any 1 2 0 and let.
SVTdR B2
) 2 (1 (3vanz+) (2Bo (og(
2Br + log(1/e)
2
(98)
B
?
Then, the realizability Assumplion holds with go())= (2Bs log((1e)) and g(6) B2 + gc
35

---

Published in Transactions on Machine Learning Research (04/2024)
Proof. By Taylor expansion there exists 0' [00, 0o] such that,
3:0(X;;0) y; + y/s ( (X; Oo) v (0 0u (20((;;0) (0 (0--0))
(99)
Pick
2B log (1E) O.
0: 0(=) = 00+
Substituting this in 99 and using that 8g is a good initialization, we obtain for all i [n]]
(2Bg (og(/))) ] 1 x ||| 0 Onl 1
1 x 0) 2
2B + log 1/=) 2
B+ (2Bg log(1/e)
1/e)) (x;
(100)
To continue, we show in Lemma in the appendix that (X; 12 s (a(0') /V/ where recall 8₂(0) = dR (3vaR² 1012,00 1) defined in Corollary Now, note that
2B + log (//))
s ) v
B3(0.) ))= 23,
7
22 22/2 / B2 (3 Vd 11 ()+ 2B + (1/{
?
100T By (3vda=++ 1). ()1(1 2B + log(1/E)
where the first two inequalities follow by triangle inequality and the inequality after those follows because 10, 12, $ ||0, /2, 1 s B2 and also |,0122 B2 by good initialization assumption.
Plugging in this bound in 1O0 and using the assumption on H yields that Db x;; 0) 2 log(1/<) for all e [n]. This in turn implics that (0) == ((x: 0)) s log(1 E) $ and thus I(0) S E as desired.
Furthermore, note by definition of ((t) that go(e) 2fg (lg((e) and g(t) Bg + 90(+) For the latter, This completes the proof.
used triangle inequality and the rough bound 00. |2,-
7
C.4 Proof of Corollary 2
We restate the corollary here for convenience, this time also including exact constants. Suppose good initialization Oo and let
Corollary 5 (Restatement of Cor.
1 H 2 256 VTRR By 3VaR2 (480( +++ B2 B2) + B=)+1) gol )
K
where go ( 2ool K)
Purther, denote for convenience
a(K)= 3v 33TR R²
) B2 + 2 2
and p(K) = (2VTan² ())))((TTr))))) VH
TR a(K)² Then, for any fixed step-size
4B2
1
7 1/p(K)
72
log(1 eBs)
the following bounds hold:
2
(2Bo + log K))²
17(2Bo + log( K))²
ī((0))
and E[I(Ox) T((0))s
(101)
K
4-2 K
22 nn
36

---

Published in Transactions on Machine Learning Research (04/2024)
Proof. First, prove that the given assumption H satisfies the conditions of Proposition 7 for
VH 2 256 VT B2
3 3VRR²
Ago K
+ B2 + g0
)
)
K
1
1
1
256 /T B2
3 3v
3
+ g(
gu
K
1
R)( )
V 2 256 VTRR B2 (3van Ago (
)+
B2 + 90 )
K
25 B2 (3V + +1) go
+
K
It means that the condition of Proposition
on VH is satisfied. Moreover, s go k)
v H 256 VTRR B2 3v 4 ( B3
g0 _ )²
+
256 VTRR B2 122Va gol
}) so
=²
1
K
K
2 2VTdR B2 (3VdR² + +) go +))
and again the condition of Proposition 7 on H is satisfied. Then, can apply the results of Corollaries 3 and g for fixed K which satisfics K 2. Note that
2B γ
go (1)
and I(Ou) s log(1 BB)
482 7ー Ing(I=0"0) Thus, the condition on step-size simplifies 7) s 1/p(K)
This completes the proof.
D Proofs for Section 5
D.1 Useful facts
Fact 1. Let x E Rd be subgaussian vector with || S K. Then, for anyy de (0,1) and absolute constant C> it holds with probability at least 1-8 8 that (Vl(sC /log(1/5)) Fact 2. Suppose Xa.h [H] are IID realizations of random variable X for which F[X] μ and |X| B almost surely. Then, for any 8 e (0, 1), with probability at least 1 - it holds that
1 H
log(1/0)
Σ he[H]
Xv.
$ 2B 2BB
2H
D.2 Proof of Proposition 2
Recall
Tb (x:0)=
Σ (x W.. UN)= /H he[H] VH hc[
(U
where Ø concat 0, O,,..., ..., 0n) denotes the trainable parameters. After completing the first. phase of training, the initialization for the second phase is as follows for all he [H]]
(1) concat(UP (1( w(
W 0, Uf(i) an 2
S Sluu 1rp") 1rp')
(102)
37

---

Published in Transactions on Machine Learning Research (04/2024)
where recall that 04 parameter P is defined in 6 Onwards, we condition on this good event.
(1) and from Lemma it holds with probability 1- that ||p| s P where the
We prove the three properties P1, and P3 in the order stated below.
D.2.1 Proof of Pl: Bounded norm per head This is straightforward by noting that for all h [H]:
0 12= 12= $ S STTluu.l2 vTipl2s 5 S VTP.
D.2.2 Proof of P2: Bounded initialization
Lemma 12 (Initialization bound). Let any X sampled from DMI and satisfying Assumption 3 inilialization in 102 Jor any 5 e (0,1) 1) it holds with probabilily at least 1 5 that
Given the
TT
Proof.
क ((x:((1)
Σ ox(x:0(
Σ ( ( X()x).
VH hc[m]]
VH he[H]
Using the initialization
102 and recalling (p(0) 1r I/T, we have
1 Σ ( (xw x ) VH he[H]
1
Σ oh (U 11 1,1 VH he[H]
1
Σ Xn. he[H]
(103)
VH
Note for each he [M] that Xn, in 103 depends only on the random variable og- Recall that on, h [H] are IID Uni Thus, (eH) IID with 0 mean Ea O. Further, note that
|xx|
T (U 1r1,x)
< 1TuT 1rp 1-1,x)
S /T VTTPPl /TI te|?]
1 S
/ P).
$ T
Σ æ.1
Thus, using Hoeffding's inequality (see Fact 2 we have for some absolute constant 0, with probability at.
least 1 1- 8:
D.2.3 Proof of P3: NTK separability
We prove property P3 in two steps each stated a separate lemma below Lemma 13. Let
W. m.u. p E velm, + p.)
(104)
[[[M]
U. Lru..
(105)
Given the initialization ((1) and salisfying Assumplion
102 and (U., sign(a) (.) we have for any (X, v) sampled from DM1
F FOCD (XV (1)))) 0.) 7+7
38

---

Published in Transactions on Machine Learning Research (04/2024)
where the expectation is taken over Unif(+1) and
T(- ( 55* -775- - -1772 16 Z = 4/2(M+1) +1)
SVT (5-2(1- (()))))))))
2g
PS Z)²
5
V2
where 2 = Zp Zuu and U. W. denote normalized U. W., respeclicely.
Proof.
( ).0.) 0.) sign(a) (V W.) +.) U.).
Using ((1) concat(U()) 0) and U. U. 1
|0.|o SV2T Lr we have:
(VV D.),).)))))))),), = ( U
y (1, 11.11x111100 (
TS
SVT
(s 2c 1 -)) 7
8V2T u Σ
8t
(106)
V2
The gradient with respect W evaluated at ((1) concat (1)), .0) is
V ac Σ x R,X Σ x.p x
te[T]
te[T]
where
1 IT
R: Rg T
17
Vte
T
Thus,
v(V D (x;0(1) sign(a)W..
lal
W..XWW w.))
Σ
uu. R (X: Σ x.p X X R [T][T]
(107)
W.
[e[T]
Terma
Terma
where set for convenience:
Te X; W.) XW. RT
To compute re (X; W.W,, we consider two cases where row corresponds to signal relevant of noisy token. We denote the / E [T] entry of r. as [r:]e IR.
Case 1. Relevance scores of signal tokens: Consider signal token te R. so that py. Using orthogonality in Assumption 3 we can compute for all t' & [T]
(51
R'R
t R: [re]e
(108)
(HTTZ))))))
R
Therefore, again using Assumption
S
ERRR it R
t e R :rele
(109)
S S
Case 2. Relevance scores of noisy tokens: Similar to the calculations above, using Assumption 3 for parame- ters W as in 126 it holds for noisy tokens t R that
St S* ( (HI2)(HL)e) (2Izz)(1Iz))
1'
t R.: [re]==
t' R
(110)
(x( + Ee (v] (:) (MI + + + T 2y
39

---

Published in Transactions on Machine Learning Research (04/2024)
Now, we can start to bound each of the two terms in 1107] separately below.
Bounding Termy: Recall the data matrix complies with Assumption 3 hence.
1y8² 11/ - -e_
.t' E R
[Xu.]] le
(111)
,t' E R
Using this, we can compute:
Σ
x.
rx
(x;
= V
Σ
u
r. (X;
Σ u.XTT,(X:
T te|T]
T
l
Re
y
Σ
Xu 1:
(X; W.)r + Σ (xu,).,((X:W..)
T
R
MER
Σ
Σ [Xu.JoT(X::W.))
MER
y/
CT
CT y S*
Σ 2 -e_) S (11T2e))
T
MER
(Tys? (S S22 { Uj 21)
Σ
2 p-) ((II)) + (1L (HL22e)
Σ MeR
+ [(1 + [(1 + (p + )
de
J8
Further using the noise bounds in Assumption 3
and rx(X; W.) from
127
110
we have:
y Σ (X;
T [
2 <T
55 2(1 5) + +(1- ()T (S
- -022--2-2-)----2,00-002207 +2, +
2 T c50
+,,,++22200 -20---50000000000000--033--000-0 +2))
(112)
For the second part of Term::
Σ
u.
1y 11..
= y Σ
1 Xu. (X:W..
T2
T :t[T]
[c[
1
22²
Σ
(k - ( m
Σ PER
Σ 5*+ Σ 522T110
T
PER
RC
MER
MERR
Σ
Σ
52 x My + S² Σ Euje
Σ (=) (z)
R
MER
31
MERC
[Eu), 21 )1,) + {Σ [2),20)221111 S2 (1++ + p1)
51
3t
40

---

Published in Transactions on Machine Learning Research (04/2024)
Using Assumption to simplify the second term:
Σ
u 1 1T (x::)
T [[]]
2 T
- - + - (1-))
- -< -
- -21(111))
( (((1-))2
- 2Ç(1 - -201-99 - 2zo)
Therefore,
y Σ (''''Ir1(r((;;W.) - (|C_::/00000_-300000//0/33333--0 T te[2]
+ +222))),93--011--922,,,++,
(113)
Bounding Termi:
2( Σ x.p X w.)2
Σ | hpll|lp||11111 X tE[T]
[e[T]
2 - 2-C
2 2 ( 2)3). (114)
Combining 112
(113)
114
107
we get:
y( V h x; x20()), W.)
TC
2 a
2
TC
2 ((1 1001::/4442222200000000000000030 + +
(1--- (1-))((+++10(++++))))+ 2)9)
T(1 95
S55 - <(5Zg
+ -(55- 2
2)),)
a
2
4(1 (2 (Z + 2,)z2 (0M---)977770000-222233<4-----00 2)9)
T(1 5)
2 a
555 (52, + 24 4
+ 2.)2, 8(2, + 2,)
2
+ 115 (S+ (5- ++++))+++$+)))))) +S++2))
((1-C)C
555 72S1 12 - 11 2)3
(115)
2
41

---

Published in Transactions on Machine Learning Research (04/2024)
where Z = Z Zy. Using this, we get
sign(() W.
=: = Un-U
VwD(X;00)
|W.lFF
F ImUU(((=11
aal
[T(1- 5)
<55 <:55 -775 127 -1673
22/2(M+1))P)) 15 2)
WW
T(1 ()5
2d ( -12 16
(S+
(116)
4/ 4/2(M+1) +
5
Combining 106 and
116 concludes the proof.
Lemma 14 (NTK Separability). Assume ireitialization @(1) = concat (0 ,..., (!))) as in 102 and IID
Oh Un) for all e [H]. Recall 7, and 0. (.) from Lemma above. Sel
1 concat (0. (0(()) ... 0. (
0.
VHI
Let any (X,y) from DM1 Then, with probability at least over the randomness of an, h [H] it holds
2lo((1/8)
y (v [x:0 (x;00)) 2m -22===
H
Proof. We start by expanding the empirical margin:
0.. (0. Σ v(Vo (x: y u (x;;(1)) Ø.) H he[H]
1
Σ Nn. H hc[H]
(117)
Note that cach summand Xn, defined above depends only on an, [H]. Thus, [XAhheHH] are IID random variables because (an)eetH) are IID random variables. Moreover, X₁ is almost-surely bounded satisfying
Xxl s Vo (x: |1|0. (o V2T R n2 10 I (n+1) PP)+ VTR).
(118)
where the second inoquality follows by Proposition and by the assumption (0, )|12 = /2 for all he [H], and the last incquality follows because |0( (VT u. 1 VT P).
Finally, note that [[X,] v. from Lemma 13 Given these the desired claim follows by applying Hoeffding's
inequality (see Fact 2 to 117
Lemma 15 (Margin). Define 2 Zμ v Zy and
7. = T(1-C) -(SS - - 122
16 73
SVT 2g (P
(s- (5-)((1c) 5) 82
(119)
4v2((M+1)
5
/2
Suppose
VH 24. T) VTR
2 llg((n/))
7
Then, with probability 1- 8 10, P3 holds with 7 7./2.
Proof. The proof is straightforward by using union bound and plugging the condition of in the result of Lemma 14
42

---

Published in Transactions on Machine Learning Research (04/2024)
D.3 Proof of Lemma 2
The proof of the lemma follows directly by combining the two lemmas below and using Ç s 1.
Lemma 16. Fix any h [H]] Suppose zero initialization 0 and consider first gradient step as (102) It then holds that (( =0 and U,, RTxd has identical rows all equal to
Çan
Son ()
az
-
u. 2
E ic[m]]
yiply u.
Σ yi
Ev 2,,1)
(120)
2
2
T ie[
(1
5)T
LER
P1
Pa
where recall that u. = p4+ p__..
Proof. We start by computing
Vo, )
Σ
Vo, (o,
Σ (X::0 ))oo ((x:0
2t[m1]
(0)) Σ y;Vo, (x:00
(0)
Σ
VVo. ((::0
VH tt[ma]
[c(n1]
Σ V.Vo, 1/ (X:0% (0))
2/H H it[mg]
where we used that PA (x: ()) = (x: = 0 because U, = 0 and also (()) 1/2 for the logistic loss. Now, recall that 4 (X;0)) =
(X: Wi)). Hence, Ww, t ( ())) = 0, which gives us
w(1)) ( Also,
1 VU, 4r (X W ((0) -1117X.
Hence, the 7-th column of (()) becomes for all T [T]:
[ l.r
0hh Σ
Σ
aiit
2nT
[[[m]]
Sc[T]
5 Σ
yilty.
Ohh Σ <e[ng]
yi Σ (M,, + zi,)) teR;
2m
2nT
[e[ns]
The claim of the Jemma follows by rearranging the above.
Lemma 17. Suppose labels are IID and equal probable, i.e. y: ~ Rad(+1). Then for the two terms P1, P2 in 120 it holds with probability at least 25 and absolute constant C O over the randomness of labels that
d n
log 1/8)
1pil s CS(
+
(121)
and
122l CCS+ ()( Mi
log(1/0)
(122)
n1
Proof. For arbitrary ||a| = 1, let X. = (v,
1 n Exe(m.] yifly.). Then,
X, V/2 1 Σ vo Mee [e[ma]
CS
11 My
/ni
se[rc]]
43

---

Published in Transactions on Machine Learning Research (04/2024)
where in the second inequality used approximate rotation invariance of subgaussians and in the last step we used that for all i [m1], (1, M4y) S, thus they are subgaussians with parameter CS. Further note that Note that [[x,]= = (v, u.). Thus, by centering property of subgaussians (v,Pil is also subgaussian with same constant CS/ /ni. Since this holds for all v on the sphere, we conclude that p1 is Sigsaagssn From this, we can directly apply Fact for concentration of Euclidean norm of random vectors to arrive
121
We follow exactly same steps to prove unit norm
122 for p2- The only difference is noting that for all i [n,] and
1
Evv,, 2 v) s S Z.
y/ (1
C)T
zeRy
E Optimal Model
We first restate the optimal parameters Popt = (Uopt. W Wopt):
Unpt =
la(p - -
(123)
SVT
W Wopt" 52 22
1
[lrpee + pop + E velu ve((+++ +
(124)
&[[M]]
normalized that 0 |e V2.
The following lemma about saturation in softmax scores is used to prove Proposition 3
Lemma 18 (Softmax saturation). Let relevance-scores vector [ri...., rT] RRT be such that for some Le [T' and A, B R:
21 2 2
and
B max(rili- L ,.... T).
Further assume A 2B. Fix any and
T/I.
T2 log
A
Then, for the allention weights a [ai,..., uT) == (TTr)
RT holds that
T
051 Σ
04 Σ
a.
a.s..
(125)
[
i=L+1
Proof. For convenience denote D Ejery Pryy Note that D2 LA, Consider any i > Then,
TB
TB
0
D
Le
LeT(A-B)
Suppose 1'2 2 log (?), which ensures that
TA Cle
Setting C= T-L L 1, and combining the above two displays yields the desired:
i> T.
T-L'
Thus, Esl 0; S. The proof is complete by recalling that Eiery ai 1, hence Eiery 04
7
44

---

Published in Transactions on Machine Learning Research (04/2024)
E.1
Proof of Proposition 3
First, compute the altention matrix when the parameter W is set to the value below:
W p.p p. p-u.+ + E velu, + u_)" [c[M]
(126)
For convenience, use the notation r! x W RRT for the t-th row of matrix used to find attention scores. Similar to the proof of Lemma 13 we consider two cases where row corresponds to a signal relevant of noisy token. We denote the t' e [T] entry of re as [r:]e e R.
Case 1. Relevance scores of signal tokens: Consider signal token so that D, My. Then for weights W
in 12 and using orthogonality in Assumption 3 we can compute for all t' [2]
S4 t R: [re]e
t' RR I' e R
Therefore, again using Assumption
,t' RR Z I' ER.
t R: [re]e
(127)
Case 2. Relevance scores of noisy tokens: Similar to the calculations above, using Assumption for parame- ters W 126 holds for noisy tokens t R that
S* ( + 5
t' R
te RR: [[]]==
(nIz)((IZe) (112z
l' R
+ ()(r))))((zz))) + (v ( 5² + p
Therefore, using the noise bound assumptions, we have
22 St Zu)
t E [ri]e
(128)
22Z + 22 2Z 222
,t' R
Combining the above two cases, specifically Equations relevance-score vectors v2 are such that
127 and
128
we find that for all t e [T] the
5- 22 + Zn + SZ), = =
,t' R
t [T]:[er
(129)
t' Re
where we defined the parameters and B for convenience. Note from assumption that
52
3 -54 A
1.25 S > 2B.
Zp Z,
8
4
Thus, the conditions of Lemma that for
18 hold for L |R| CT. Applying the lemma we immediately conclude
8
(
11-11
5-11
r. 2 log
2 log
384
A
it holds:
VLE [T]: : 05 Σ [p(T.r:)p Σ [v(TT.n)) se.
(130)
PeR
PeR
45

---

Published in Transactions on Machine Learning Research (04/2024)
Now, recall that
W. w.=WW for T=1/(S T=I1 1)).
Thus, it holds for all & [T] that (0 w. ))= (((r)) Combining this with 130 and the proposition's assumption on I (satisfying I 2 T,), we have found that
[T) 0s Σ [s<x
Σ W. W.))] sc.
(131)
PeR
PER
In the rest of the proof, we use 131 to lower-bound the margin:
g - AUUN,,,
y
Σ E [o (.))), (M+ [e[T] Ne[T]
y/
Σ Ut.
SV2T
SV2T te[T]
where we defined tt, [T] for convenience. For any t [T], we have
4/
Σ [. W. ys²
Σ [o
W.x)], (1+ -L)' ze
PER
M'R
2 ys² Σ [d(z) W. (x 22p Σ [o(x,W.x"))
2 yys (1 (1- c) 2cZ
Putting the last two displays together and using y² completes the proof of the proposition.
F
Linear Model
To gain additional insights into the classification of the data model DM1 and also contrast our results to a simplified model, we examine here a linear classifier:
1n =UU)=(U,K)
For this lincar model, consider the oracle classifier
U. ST LTu,, with u. = p4 p..
and normalization such that |JU. 11=11111- 1. By using Assumption 3 almost surely for all examples (X,y) y) the margin of the oracle classifier is lower bounded by:
1
y (X:U.)=
[R] 2 y Σ (1+ - 21,
S/
tcR
1
[R|-
52 1- Z
SVT
Zu (5-
24)
Mim
(132)
S/2T
V2
52
G Experiments
In this section we provide some experiments discussing the role of number of heads H in the training dynamics on synthetic data models.
Data Model DM1 We set the mumber of tokens T 10 and sparsity level
0.1. We set
(Mt, 22,... UMs) fo the canonical basis vectors in Rd with d = 4, M = and signal strength $ = 2. Noisy tokens z are sampled from a Gaussian o0 with a = 0.1. We use = 100 training samples in each experiment and evaluate on a test set of size 300 (total trials). All models are initialized as 0(0) 0.
Figure D shows the effect of increasing the number of heads when running GD with constant step-size " = 1.0 and data generated from data model DMI Notice that rate of train loss decay reduces as we increase highlighting potential downside of overparameterization. A similar observation has been recently noted by
46

---

Published in Transactions on Machine Learning Research (04/2024)
Xu & Du 2023) when optimizing with GD to learn single neuron with ReLU activation under Gaussian input. We also observe that at least. for smaller H, GD indeed achieves margin Nattn- It is worth noting that these observations do not contradict our theoretical findings in Theorems [ and 2 which guarantee training convergence and generalization docay given sufficiently large mumber heads H, without making an explicit connection belween the rates of convergence as we increase H. We also Lest how these rates change if we scale step-size as 1 = (VH) for GD (Figure left) or optimize with Adam (Figure right) using constant. step-size 1) 0.06. It interesting to observe that in both these cases, the convergence speeds up for larger H, especially when optimizing with Adam, but somewhat strangely the margin attained by GD for larger H continues to fall away from Yattn- Further, note that our theory only covers step-size v O (1) and the trends observed in Figure with n- O (VH) for GD fall outside this regime. In essence, we believe that it would be interesting future work to see how well these observations generalize to different datasets and develop theory that explains the relation of overparameterization to rates of convergence.
GD with 17 (1)
#=1
= 100
Yeenn
#= 10
= 300
100
Train loss
10-1
Test Loss
10-1
80
10-3
10-3
60
10
10
10
10a
10
10a
15
1.00
1.0
Attn-score score
Avg-1-1 1-91-
10
Nargin
0.75
0.5
0.50
5-
0.0
0.25
0
10
10a
10
10a
10a
10a
Iterations
Iterations
Iterations
For data model DM1 Effect of number of heads H on convergence rates when trained with GD for constant step-size 17 = (1). The average ||-| illustrates 1/H and 1/VH average for W and U across heads, respectively. Attn-score denotes the softmax scores for the relevant tokens averaged across all train samples and heads. The average [W] indicates the saturation of softmax scores and consequently the token-selection (attn-score), and the average |U] controls the loss behaviour. Results demonstrate that overparameterization slows down GD with constant step-size. The circled area shows a O(1tt) trend similar to what our training and gencralization bounds predict.
Figure 2:
Planted data model Fix some W* Raddd U* E RTxd Entries within X are sampled IID X1 ~ (O, vie [T]. j e [d]. Given such an X e RTxx generate the label using W as the attention matrix and U as the projection classifier:
y = sign (";;"",, U*)) sign((U' (XXX (XX")))X)),
(DM2)
Data generated using model DM2 is used to train a multi-head self-attention model as given in equation 2 Such a teacher-student setting (train the student network to learn the ground truth parent) has been well explored in the past in the context of neural networks (Zhou al.) 2021: Safran et al. 2021). We set d 5, T 10. The train set contains n 1000 samples in each experiment and we evaluate on a test set of size 3000. Each result is averaged over 5 trials. For mumerical ease, while generating (example, label) pairs drop the samples for which joutput logit/ Nattor where call Nattn > 0 to be margin for the data model. We set Yattn = 0.2 in all the experiments. All models are initialized as g()) = O. From Figure 4 we observe that overparameterizall tokenized mixture model (Figure 3 left). Addition momentum significantly helps speeding up convergence (see Figure 5] left)) and so does optimizing with Adam (Figure 5 right). Interesting note that all the models reach the expected margin Yattn which was not the case for large for the tokenized mixture model. Further, observe that initializing at ((0) 0, all the optimizers find the planted model.
somewhat improves convergence speeds for GD with step-size O (VH), similar to
47

---

Published in Transactions on Machine Learning Research (04/2024)
GD with 1) = -0(
Adam with 1) (
H=1
100
M=1
H- 100
#=10
300
W
H- 10
H- 300
100
100
10
Thain lass
Train loss
Test ecc.
10
30
80
10
10
10-1
10
10a
10
10a
10
10a
10a
10a
10.
10a
10a
10a
1.00
20
1.00
Avg-1-1
100
rergin
Mttr- score
Aug-1-1
Nergin
0.5
0.50
0.50
0.25
0.25
10a
10a
10a
10a
Iterations
Ite attons
Ite rationson
Iterations
Iterations
Iterations
Effect of number of heads H on convergence rates when (left) trained with GD when scaling step-size as 1) O (vH); (right) trained with Adain with constant step-size 77 O (1). Quantities plotted same in Figure 2 Results demonstrate that overparameeeerization speeds-up with train and test loss convergence in both the scenarios.
Figure 3: For data model DM1
SST2 dataset We conduct an additional experiment simple real-world dataset. The SST2 dataset [Socher et al. 2013) consists of sentences, with each sentence having a associated binary label to classify the sentiment. We fine-tune RoBERTa based models with varying number of heads using AdamW (Loshchilov & Hutter 2019) optimizer with a learning rate 6. We train all the models for epochs, with the batch-size set 32. We use the Hugging Face pyterch-trame
rmeer implementation of the sabberttabsee model, with pretrained weights (Liu et al., 2019 In Figure t we see that incrcasing the number of heads speeds up the optimization and generalization. This behaviour is similarly observed for GD with momentum and Adam in Figure
GD with 7 (VV
#=1
#=100
#=
2000
W
H= 10
1000
Yatn
100
Train toss
109
Test loss
Test Acc.
10-1
10-2
10
60
10
10a
10a
10a
10a
10a
Avg. 1-1
Alignnent
Margin
10
101
10a
10a
10a
10a
103
Iterations
Iterations
Iterations
Figure 4: For data model DM2 Effect of number heads on convergence rates when trained with GD when scaling step-size as n= O (VH). See Figure caption for get more context on average I-||. Alignment of W with the planted- -ead W at any iteration k is given W*) by IWall w where W == concat ((W* he(m) contains W repeated H times. Alignment between Ú and Ō is computed similarly.
H
Related work
This section elaborates the paragraph related work in Section
48

---

Published in Transactions on Machine Learning Research (04/2024)
GD with 7/ O (VH) momentum
Adam with 17 = (
2000
H=1
H=100
2000
W
H 10
H- 1000
#=
#= #=1000
Vrin
Train loss
10a
Teet lacc
10"
Test Acc.
Truin Loss
10
loss
10a
Test Acc.
80
80
10
10-1
60
10
10
60
10
10
10
10a
10
10a
101
10
104
107
10a
10
1.0
Avg. 1-1
20
AL ALigmnent
Nargan
41 1ggsent
0.5
Margin
10
-2-
0.0
16'
102
10
10a
10a Iterations
10a
10
30
10
107
10a
10
Iterations
Iterations
Iteracions
Iterations
Iterations
Figure 5: For data model DM2 Effect of number of heads H on convergence rates when trained with (left) GD + momentun where step-size scales as 11 = (VHH) (right) Adam with constant step-size = 0 (1). Quantities plotted same as in Figure Results demonstrate
that overparamettrrization speeds up convergence in both scenarios.
H=1 HH4 HH16
train loss
train accuracy
100
0.6
0.4
60
0.2
20
Ste0
10
20
100
800
test loss
test accuracy
0.7
0.6
80
0.5
60
0.3
0.1
Step
Step
20
100
300
1
300
Figure 6: For SST-2 dataset (Socher et al., 2013). Effect of number of heads H convergence rates when trained with AdamW. Results demonstrate that increasing the number of heads speeds up the training and generalization dynamics.
Transformers and Self-attention. The landscape of NLP and machine translation was profoundly reshaped by the advent of Transformers, as pioneered by Vaswani et al. (2017) building upon earlier investigations into self-attention as explored in the works of Cheng et al. (2016): ; Lin et al. (2017) Parikh
et al. (016). More recent developments include the transformative success of large language models like, LLaMA (Touvron et al. 2023), ChatGPT (OpenAT 2022), and GPT4 AOpenAT
2023 Despite this, the
learning dynamics of the self-attention mechanism remain largely unknown. Some recent works have focused
understanding the expressive power (Baldi & Vershynin 2022, Dong et al. 2021 Yun al.) 2020a b
Sanford et al. 2023 Bietti et al.) 2023 and memory capacity of the attention mechanism Baldi & Vershynin
2022; Dong et al. 2021 Yun et al. 2020a Mahdavi et al. (022). Other aspects which are explored include
obtaining convex reformulations for the Training problem (Sahiner el al.] 2022 Ergen et al. (022), studying sparse function representations self-attention mechanism (Edelman et al., 2021: Likhosherstov et al.)
2021] and investigating the inductive bias of masked self-attention models. Additionally, a sub-arca gaining increasing popularity theoretical investigation of in-context learning, e.g. (von Oswald et al. 2022 Zhang
et al. 2023 Akyürek et. al.] 2023 Li et al. 2023b)
49

---

Published in Transactions on Machine Learning Research (04/2024)
Here, we discuss works that aim to understand the optimization and generalization dynamics of self-Attention or its variants. Oymak et al. (2023) diverges from traditional self-Attention by focusing a variant called prompt-Attention, aiming to gain understanding of prompt-tuning. Lu et al. (2021) show that for a bag of words model, an attention model optimized with gradient flow for a topic classification task discovers the "lopic" word as training proceeds. However, they don't provide finite-time optimizalion--generlllaaaz rates. Jelassi et al. (2022) shed light on how ViTs learn spatially localized patterns, even though this spatial structure no longer explicitly represented after the image is split into patches. Specifically, they show that for binary classification using gradient- based methods from random initialization, transformers implicitly prefer the solution that learns the spatial structure of the dataset. Li et al. (2023a) provided theoretical results on training three-layer ViTss for classification tasks for a similar data model as ours (tokenized mixture data). They provide sample complexity for achieving zero generalization error, and also examined the degree of sparsity in attention maps when trained using SGD. Contemporaneous works include (Tian et al.) 2023 Tarzanagh et al.) 2023a): The former presents SGD-dynamics of single-layer transformer for the task of next-token prediction by re-parameterizing the original problem in terms of the softmax and classification logit matrices, and analyzing their training dynamics instead. The latter studies the implicit bias of training the softmax weights W with fixed decoder U via a regularization path analysis. All these works focus on a single attention head. Instead, we leverage the use of multiple heads to establish connections to the literature on GD training of overparameterized neural networks. Conceptually similar connections have also been studied by Hron et al.] (2020) who connect multi-head attention to a limiting Gaussian process when the number of heads increase to infinity. In contrast, we study performance in the more practical regime of finite number of heads and obtain and obtain finite-time time optimization and generalization bounds.
There has been an abundance of literature that discusses convergence and generalization dynamics via NTK type approach, e.g. (Allen Zhu et al. 2019 Oymak
Overparameterized MLPs.
NN training
& Soltanolkotabi 2020 Arora et al., 2019 Nguyen & Mondelli 2020 Banerjee et al.] 2022. Nguyen et al,
2021 Zhu et al. 2023). However, most of these works focus on GD dynamics on regression problems using square loss and relate the training convergence and generalization to the minimum eigenvalue of the Hessian of the NTK. On the other hand, relatively fewer works focus on classification with logistic loss under an NTK separable data assumption, and (Nitanda et al. 2019 Ji & Telgarsky 2020 Cao Gu 2019 Chen et al., Telgarsky] 2013 Taheri & Thrampoulidis 2023) are most. relevant works that share overlapping ideas with work. We refer the reader to these for more thorough overview of the NTK-regime analysis for NNs.
2020
Other than these, Richards & Kuzborskij 2021 Richards & Rabbal 20211 Taheri & Thrampoulidis (023)
Lei et al. (2022) use algorithmic-stabilit based tools to understand the training dynamics and generalization of GD in shallow NNs. Lei et al. (2022); Richards & Kuzborskij (2021) establish generalization bounds with polynomial width (((polynn) requirement while minimizing square loss. Here, we make use of the tools developed by Taheri & Thrampoulidis (2023 who work with self-bounded Lipschitz loss functions, like logistic loss, similar to our analysis. The algorithmic stability arguments in the analysis of the above referenced papers are rooted on a technique to bound generalization gap by directly relating it to train loss based on the notion of average model stability introduced by Lei & Ying (020).. This technique has also been leveraged by Schliserman & Koren (2022) to study lincar logistic regression and Nikolakakis et al. (2022) who establish generalization--ris bounds for Lipschitz function optimization with bounded optimal set.
50


Reviews:
[{"review_id": "XnomLHKH4o", "reviewer": "Reviewer_LzY5", "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": "The paper is well written and presents the results in a very instructive fashion, showing how the presented results relate to previous works. The strengths and limitations of the derived bounds is discussed in length, in particular the authors do not shy away from the fact that the bounds cannot accomodate convergence towards the optimal margin of their data model. The proofs appear correct and the derived bounds seem useful and general enough.\nAltough not in the main paper and results, the experiments in appendix are somewhat underwhelming as they do not point towards one of the conclusions that I took from the paper, which is that realizability is achieved for a large enough number of attention heads; indeed, the model with one head already achieves the optimal margin and converges faster. I would be interested if there where settings in which the scaling induced by the derived bounds were more appearant.\nAll in all, I believe this is an excellent contribution to TMLR, and that the framework provided should advance the state of the theoretical study the convergence and generalization properties of attention transformers.", "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": "The authors derive training and generalization bounds for the multi-head attention layer of transformers under a binary classification setting with logistic loss under gradient descent optimization.\nAfter demonstrating some properties of the loss related to weak-convexity with respect to the model parameters, the authors first show general training bounds bounded by the distance of the initial parameters to a target vector satisfying some realizability condition. The authors then show generalization bounds after K iterations in terms of the empirical loss and the same distance of the target vector as in the training bounds. \nThe authors show in Sec. 4 that the realizability condition is realized under some initialization regimes and for a large enough number of attention heads H. In Sec. 5, they apply these results to a tokenized mixture model with orthogonal token vectors. They show that the required initialization condition can be found after one randomized initial gradient step from 0, and derive the generalization bounds in this setting given a polylogarithmic number of heads.\nThe authors leave questions to future work: studying convergence and margin achieved under random initialization; studying the convergence to the optimal margin for the tokenized model, which they show is not possible under this theoretical framework due to the constraints on the weights; and studying different data models.\nIn the appendix, the authors include all the proofs, as well as an experimental study of the convergence and generalization bounds under the data model.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I have only minor changes to suggest:\n* The cross footnote in page 3 is difficult to parse right alongside a mathematical equation. I would keep the remark in that footnote in the main text.\n* p.5: \"the minimum eigenvalue becomes less positive\" -> I believe this should be: the minimum eigenvalue becomes less negative?\n* p.7: patters->patterns in \"all patters are orthogonal\".", "broader_impact_concerns": "No broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "IxiXrIIVSs", "reviewer": "Reviewer_rp84", "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": "The paper explores convergence and generalization. It is well-written and outlines the contributions and steps in detail, making it easy to follow the overall structure.\n\nI am not an expert in this area by any means (not even an intermediate) and I could hardly follow the maths and derivations. As such, I can only acknowledge that it seems correct and likely interesting for experts in the area, given the related work and placement in the literature.\n\nGiven that is hard to verify for me: would it be possible to evaluate the bounds on an actual example of a simple MHA architecture? I have no idea what size of constants are involved. Then, it would be interesting to see how this compares to training such a network using finite precision and see how tight the bounds are in practice.", "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": "The paper presents a theoretical analysis of multi-head attention in transformer models, contributing new insights into its optimization, that is convergence, and generalization dynamics, with a theoretical application to a tokenized mixture model.\n\n---\n\n## Summary with GPT-4\n\n### Page 1-2\n- **Abstract and Introduction**: This paper delves into the Transformer's multi-head attention mechanism and the lack of analyses concerning its optimization and generalization dynamics. The authors establish conditions for a multi-head self-attention model to be stable under certain realizability conditions on the data. They focus on logistic loss optimization and derive bounds on the Hessian of the attention weights, concluding that their results can be generalized to other data models and architectures.\n\n- **Related Work**: The paper situates itself within the field of deep learning and specifically attention mechanisms, citing recent work in optimization and generalization dynamics of gradient descent for Transformers. The literature review spans various studies, including those on single-head self-attention and multi-head attention models, pointing out the research gap in finite-time optimization and generalization dynamics for the latter.\n\n### Page 3-4\n- **Preliminaries**: The authors provide notations for the softmax map and attention functions and explain the architecture of a multi-head self-attention model. They also mention the logistic regression loss used for training.\n\n- **Training and Gradient and Hessian bounds of soft-max attention**: Here, the authors discuss the empirical risk minimization framework for training, the gradient and Hessian bounds for softmax attention, and the model gradient/Hessian bounds. They introduce lemmas and propositions to establish the bounds, with detailed derivations provided for the bounds of the softmax attention model's gradients and Hessians.\n\n### Page 5-6\n- **Main results**: The authors present theorems detailing training and generalization bounds for multi-head attention. They provide theorems and corollaries that formalize the bounds on training loss and expected generalization gap, including assumptions and implications. These results connect the number of heads in the attention mechanism with the stability of optimization and generalization.\n\n- **Primitive conditions for checking realizability**: This section introduces conditions necessary for the realizability assumption. The conditions involve the data and initialization of the model, and bounds on the model parameters. Corollaries provide generalized bounds under good initialization.\n\n### Page 7-8\n- **Application to tokenized mixture model**: This section applies the theoretical findings to a specific data model, the tokenized mixture model. It outlines the model, assumptions, and how to find a good initialization. They also prove that under certain conditions, their approach can achieve NTK separability, which aids in suppressing noise and improving classification accuracy.\n\n### Page 9-10\n- **Concentration on maximum margin**: The authors discuss the optimality of the NTK margin in relation to the softmax attention scores, providing propositions that link attention scores to margin maximization. They show that as the number of heads grows, the model approaches a hard-max function, which is desirable for classification tasks.\n\n- **Gradient-based optimization and margin maximization**: Here, they discuss the application of their theoretical bounds to gradient-based optimization methods, specifically how these methods contribute to maximizing the NTK margin. This involves detailed mathematical analysis and lemmas to support their claims.\n\n### Page 11-12\n- **Algorithmic stability and expected generalization**: This section provides a framework for algorithmic stability and how it relates to generalization. The authors present lemmas and proofs to support their claims on stability and then tie this into the expected generalization gap.\n\n- **Concluding remarks**: The paper wraps up with a summary of the results and their implications for the training and generalization of multi-head attention layers. The authors note that random initialization of attention weights satisfies NTK separability, and they point to interesting open questions regarding different data models for future research.\n\nThis paper is quite technical, focusing on the mathematical aspects of multi-head attention mechanisms within deep learning models, specifically in the context of Transformers used for classification tasks. The authors contribute new theoretical insights into the optimization and generalization dynamics of these models, which are important for understanding the behavior of deep learning algorithms.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Maybe:\n\n1. the evaluation above so that it is possible to see the order of magnitude of the constants and the tightness of the bounds;\n2. adding a schema with the proof and result structure of the paper. (the outline is already quite comprehensive, but it might be nice to have that visualized?).\n\nAnd fixing the following typos:\n\n1. The derivative of the softmap map in §2 Preliminaries is wrong. It should be $$\\varphi'(v) := diag(\\varphi(v)) - \\varphi(v) \\varphi(v)^\\top;$$\n2. Just before (19) \"fives\" -> \"gives\".", "broader_impact_concerns": "None", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "rYNBHzRzVv", "reviewer": "Reviewer_HGup", "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 ## \n1. Although the model under study is simple, I think it is interesting to provide theoretical insights into this popular architecture. \n2. The work covers the multi-head cases, which I believe is non-trivial. \n\n## Weakness ##\n1. It appears the numerical study is not well-aligned with the theoretical insights, although the authors discuss why this does not contradict their theoretical findings. \n2. It would be interesting to see some empirical studies on some simple real-world datasets, e.g. SST2. [1]\n\n[1] Richard Socher, Alex Perelygin, Jean Wu, Jason Chuang, Christopher Manning, Andrew Ng, and Christopher Potts. Parsing With Compositional Vector Grammars. In Empirical Methods in Natural Language Processing (EMNLP), 2013.", "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 learning trajectory and generalization properties of a single-layer multi-head attention mechanism trained by gradient descent. Unlike the related work based on NTK, this work does not assume infinite width and infinite head of attention. Additionally, they apply their theoretical results to a bag-of-word synthetic dataset and show that the training involves picking key-query weights of sufficiently large norms and suppresses label irrelevant tokens to saturate the softmax nonlinearity.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I am not actively working in this area; thus, it would be challenging to provide useful comments on how to improve the work on the theoretical study. Overall, I am not aware of any critical issues. Authors may refer to the Weaknesses for suggestions to strengthen their work. Additionally,  Lu et al. [1] consider a related problem and characterize the training dynamics of attention models. Authors may consider discussing its relationship to this work. \n\n[1] On the Dynamics of Training Attention Models. Haoye Lu, Yongyi Mao, Amiya Nayak, ICLR 2021.", "broader_impact_concerns": "There aren't any major ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: All reviewers find the contributions of the paper strong and meaningful. The paper considers the challenging problem of analyzing training dynamics of multi head self attention in the non-asymptotic setting. The paper makes good progress in this direction and the results can be build upon in followup works. Overall I am happy to recommend acceptance of this draft.
Decision: Accept (poster)