Max-Margin Token Selection in Attention Mechanism
Authors: Davoud Ataee Tarzanagh
Yingcong Li
Xuechen Zhang
Samet Oymak
Venue: NeurIPS 2023 spotlight (2023)
Keywords: attention mechanism
implicit bias
margin maximization
nonconvex optimization
prompt tuning
TL;DR: We prove that softmax-attention weights trained with gradient descent converge to a max-margin solution that separates locally-optimal tokens from non-optimal ones.

Abstract: Attention mechanism is a central component of the transformer architecture which led to the phenomenal success of large language models. However, the theoretical principles underlying the attention mechanism are poorly understood, especially its nonconvex optimization dynamics. In this work, we explore the seminal softmax-attention model $f(X)=\langle Xv, \texttt{softmax}(XWp)\rangle$, where $X$ is the token sequence and $(v,W,p)$ are trainable parameters. We prove that running gradient descent on $p$, or equivalently $W$, converges in direction to a max-margin solution that separates *locally-optimal* tokens from non-optimal ones. This clearly formalizes attention as an optimal token selection mechanism. Remarkably, our results are applicable to general data and precisely characterize *optimality* of tokens in terms of the value embeddings $Xv$ and problem geometry. We also provide a broader regularization path analysis that establishes the margin maximizing nature of attention even for nonlinear prediction heads. When optimizing $v$ and $p$ simultaneously with logistic loss, we identify conditions under which the regularization paths directionally converge to their respective hard-margin SVM solutions where $v$ separates the input features based on their labels. Interestingly, the SVM formulation of $p$ is influenced by the support vector geometry of $v$. Finally, we verify our theoretical findings via numerical experiments and provide insights.

Max-Margin Token Selection in Attention Mechanism
Davoud Ataee Tarzanagh University of Pennsylvania tarzanag@upenn. edu
Yingcong Li University of California, Riverside {yli692, x2)an394)juurcr edu
Xuechen Zhang
Samet Oymak University Michigan UC Riverside oymakθumi ch, edu
Abstract
Attention mechanism is central component of the transformer architecture which led to the phenomenal success of large language models. However, the theoretical principles underlying the attention mechanism are poorly understood, especially its nonconvex optimization dynamics. this work, we explore the seminal softmax- attention model f(X) (Xv, softmax(XXW) ()P))) p)), where X is the token sequence and (v, W.p) are trainable parameters. We prove that running gradient descent on p. equivalently W, converges in direction to max- mar-mnrgin solution that separates locally-optimal tokens from non-optimal ones. This clearly formalizes attention as optimal token selection mechanism. Remarkably, our results applicable general data and precisely characterize optimality of tokens in terms of the value embeddings Xv and problem geometry. We also provide broader regularization path analysis that establishes the margin maximizing nature of attention even for nonlinear prediction heads. When optimizing and p simultaneously with logistic loss, we identify conditions under which the regularization paths directionally converge to their respective hard-margin SVM solutions where v separates the input features based on their labels. Interestingly. the SVM formulation of is influenced by the support vector geometry of v. Finally, we verify our theoretical findings via numerical experiments and provide insights.
1 Introduction
Since its introduction in the seminal work [1], attention mechanism has played influential role in advancing natural language processing, and more recently, large language models [2, 3, 4, 5]. Initially introduced for encoder-decoder RNN architectures, attention allows the decoder to focus on the most relevant parts of the input sequence, instead of relying solely on a fixed-length hidden state. Attention mechanism has taken the center stage in the transformers [6], where the self-attention layer- which calculates softmax similarities between input tokens serves as the backbone of the architecture. Since their inception, transformers have revolutionized natural language processing, from models like BERT [7] to ChatGPT [8], and have also become the architecture of choice for foundation models [9] addressing diverse challenges in generative modeling [3, 10], computer vision [11, 12], and reinforcement learning [13, 14. 15].
The prominence of the attention mechanism motivates a fundamental theoretical understanding of its role in optimization and learning. While it is well-known that attention enables the model to focus on the relevant parts the input sequence, the precise mechanism by which this is achieved is far from clear. To this end, we ask
Q: What are the optimization dynamics and inductive biases of the attention mechanism?
We study this question using the fundamental attention model f(X) (Xv, S(XW p)). Here, X is the sequence of input tokens, v is the prediction head, W is the trainable key-query weights, and S
37th Conference on Neural Information Processing Systems (NeurIPS 2023).

---

10
10
(611111, optimal toben (y-11
1.3,, optimal toler Ey-
wine: Optimel tokkess 1y-1) Non opt tokkens (yw0.5)
(1.5) non opt taken (y-2.9)
(1,01, non opt taken (y-0,
(0.0), (ym0.9)
taken (ya0)
2
2-
-2
-2
2
4
10
-2-
2
10
12 -1
3
(a) Global convergence
(b) Global and local optimal directions
(c) Multiple inputs
Figure 1: The convergence behavior of the gradient descent on the attention weights using the logistic loss in (ERM). The arrows (->-) represent trajectories from different initializations. Here, (---) and (- --) denote the globally- and locally-optimal max-margin directions (GMM, LMM). y denotes the score of token per Definition 1. Discussion is provided under Theorems 2 and 3.
denotes the softmax nonlinearity. For transformers, corresponds to the [CLS] token or tunable prompt [16, 17, 18], whereas for RNN architectures [1], p corresponds to the hidden state. Given training data (Y, X with labels Y, E (-1, 1) and inputs X; E R we consider the empirical risk minimization with decreasing loss function ((-) :R
-
L(v. W) W)= cY f(X)), where f(X) = X/ p).
At a high-level, this work establishes fundamental equivalences between the optimization trajectories of (1) and hard-margin SVM problems. Our main contributions are as follows:
Optimization geometry of attention (Sec 2): We first show that gradient iterations of and W admit a one-to-one mapping. thus we focus on optimizing p without losing generality. In Theorem 3, we prove that, under proper initialization:
Gradient descent on p converges in direction to a max-margin solution - namely (ATT-SVM) - that separates locally-optimal tokens from non-optimal ones.
We call these Locally Max-Margin (LMM) directions and show that these thoroughly character- ize the viable convergence directions of attention when the norm of its weights grows infinity. We also identify conditions under which path converge to Globally-optimal Max-Margin (GMM) direction in Theorems and 2, respectively, A central feature of our results is precisely quantifying optimality in terms of token scores y: Y xx, where x, the (l token of the input sequence X. Locally-optimal tokens are those with higher scores than their nearest neighbors determined by the SVM solution. These are illustrated in Figure 1.
(nniittd--iddeeeddeet) regularization path and gradient descent
Optimize attention p and prediction-head v jointly (Sec 3): We study the joint problem under logistic loss function. We use regularization path analysis where (ERM) is solved under ridge constraints and we study the solution trajectory as the constraints are relaxed. Since the problem is linear in if the attention features x- = X S(X:WT p) are separable based on their labels Y, would implement max -argin classifier. Building on this, we prove that p and converges their respective max-margin solutions under proper geometric conditions (Theorem 5). Relaxing these conditions, we obtain a more general solution where margin constraints on p are relaxed on the inputs whose attention features are not support vectors of v (Theorem 6). Figure 3 illustrates these outcomes.
The next section introduces the preliminary concepts, Section 4 presents numerical experiments!, Section 5 discusses related literature, and Section 6 highlights limitations and future work.
1.1 Preliminaries
Notations. For any integer N 2 1, let [N] := (N. We use lower-case and upper-case bold letters (e.g, a and A) to represent vectors and matrices, respectively. The entries of are denoted a, We use (AA) to denote the maximum singular value A. We denote the minimum two numbers b as a b, and the maximum as aVb. Big-O notation O(-) hides the universal constants. Throughout, we will use L(p) and L(v, p) to denote Objective (1) with fixed (v, W) and W, respectively.
Optimization. Given an objective function L Rd R and an form bound R, define the regularized solution as
The code for experiments can be found https: hwtpda/hss///iitit..cc/////rrrre htt.th...ommuuccooooaaaaaaaammma ant.einiiiinnannniinn.....

---

B(R) arg min _(p).
(2)
UplisR
Regularization path - the evolution of (() as R grows - is known to capture the spirit of gradient descent the ridge constraint R provides a proxy for the number of gradient descent iterations. For instance, [19, 20, 21] study the implicit bias of logistic regression and rigorously connect the directional convergence of regularization path (i.e. limg-co /(R)R) R) and gradient descent. For gradient descent, we assume the objective L'(p) is smooth and describe the gradient descent process as
p(1 1) p()) n() L((())((())
where 77() is the stepsize time and v((((()) the gradient of L' p()-
Attention in Transformers. Next, we will discuss the connection between our model and the attention mechanism used in transformers. Our exposition borrows from [17], where the authors analyze the same attention model using gradient-base techniques on specific contextual datasets.
Self-attention is the core building block of transformers [6]. Given an input consisting of T tokens X = .y. E RR self-attention with key-query matrix W Rakd and value matrix V V Rdxy, the self-attention model is defined as follows:
fa(() = S(XWX")XV
Here, S(-) is the softmax nonlinearity that applies row-wise on the similarity matrix XWXT.
Tunable tokens: [CLS] and prompt-tuning, In practice, we append additional tokens to the raw input features X: For instance, a [CLS] token is used for classification purposes [7] and prompt vectors can be appended for adapting a pretrained model to new tasks [16, 18). Let p E R be the tunable token ([CLS] prompt vector) and concatenate it to X to obtain X, == [p X E (T- Consider the cross-attention features obtained from Xp and given by
(X) Sa,())
SXp"WX") XV. 3(XWXT)
= 3(X,WXT)XX X
The beauty of pross-attention that it isolates the contribution of p under the upper term fei(X) VTX S(XWT p) e R". In this work, we use the value weights for classification, thus we set 1, and denote v = V Rd. This brings us to our attention model interest:
f(X) x s(K)), where K XWW.
Here, (v. W. p) are the tunable model parameters and K is the key embeddings. Note that W and are playing the same role within softmax, thus, it is intuitive that they exhibit similar optimization dynamics. Confirming this, the next lemma shows that gradient iterations of p (after setting W - Identity) and W admit one-to-one mapping.
Lemma Fix E R& (0)) Let v Rd - R and R be differentiable functions. On the same training data (Y,, XY,,X)); define Σ(p) := 1/n Ei,,00 UXX, S(X,))) and L(W) := 1/n E (YY (X; S(X.WT u))). Consider the gradient descent iterations on p and W with initial values p0) and W(O) = up(0) /|luls and stepsizes and n//lu/p, respectively:
p(t + 1) p(t) (((()).)
" (L)(()).
W(r + 1) W(r)
We have that W() up(t) /|| for all 1 z O.
This lemma directly characterizes the optimization dynamics of W through the dynamics of p. allowing us to reconstruct W from using their gradient iterations. Therefore, we will fix W and concentrate on optimizing p in Section 2 and the joint optimization of (v, () in Section
Problem definition:
Throughout, (Y,, X),, denotes training dataset where Y, E 1-1, 11 and X, Txd. We denote the key embeddings of X, via K; X,WT and explore the training risk
L(. p) = Σ (V.--"X(((.p))))
(ERM)
Importantly, our results apply to general tuples (Y,, X,, Kj) and do not assume that (X,, K)) are tied via W. Finally, the nh tokens of X,, K; are denoted by Xis kn e Rd, respectively, for E (7).
3

---

The highly nonlinear and nonconvex nature of the softmax operation makes the training problem in (ERM) a challenging nonconvex optimization problem for p. even with a fixed v. In the next section, we will introduce set of assumptions to demonstrate the global and local convergence of gradient descent for margin maximization in the attention mechanism.
Global and Local Margin Maximization with Attention
In this section, we present the main results of this paper (Theorems 2 and 3) by examining the implicit bias of gradient descent learning p Rd given fixed choice of Rd Notably, our results apply to general decreasing loss functions without requiring convexity. This generality is attributed to margin maximization arising from the exponentially-taiie nature of softmax within attention, rather than C. We maintain the following assumption on the loss function throughout this section.
Assumption A (Well-hehaved Loss) Over any bounded interval: (1) e: R R. is strictly decreas- ing. (2) E' is Mo-Lipschitz continuous and (e' (u) M
Assumption includes many common loss functions, including the logistic loss {() = log (1 + +e e", exponential loss (u) e and correlation loss ((u) -u. Assumption A implies that & (p) is Lp-smooth (see Lemma 6 in Supplementary), where
1p 1 Σ (Mollo? 3M |Illl I /III)
(6)
We now introduce a convex hard-margin SVM problem that separates one token of the input sequence from the rest, jointly solved over all inputs. We will show that this problem captures the optimization properties of
Fix indices a = (a,) and consider
mim (a) arg min lplll subject to
min p (ki.. k2) 2 1, for all ssisn. (ATT-SVM) rew,
Note that existence of (a) implies the separability of tokens a from the others. Specifically, choosing direction n (a) will exactly select tokens (xir; at the attention output for each input sequence, that is, limg-so XT S( (a)) = Xia;- We are now ready to introduce our main results that characterize the global and local convergence of the attention weights p via (ATT-SVM).
2.1 Global convergence of the attention weights
We first identify the conditions that guarantee the global convergence of gradient descent for p. The intuition that, in order for attention to exhibit implicit bias, the softmax nonlinearity should be forced to select the optimal token within each input sequence. Fortunately, the optimal tokens that achieve the smallest training objective under decreasing loss function t(-) have clear definition.
Definition (Token Scores, Optimality & GMM) The score of loken X of input X; is defined as Ya Yp xax The optimal tokens for input X; are those tokens with highest scores given by
opt: E arg max Yi. KLT
Globally-optimal max-margin (GMM) direction is defined as the solution of (ATT-SVM) with optimal indices (opt,)= by p
is worth noting that score definition simply uses the value embeddings v Xu of the tokens. Note that multiple tokens within an input might attain the same score, thus opt; or pm may not be unique. The theorem below provides our regularization path guarantee on the global convergence of attention.
Theorem (Regularization Path) Suppose Assumption A on the loss function holds, and for all [n] and i # opt, the scores obey Yu < Yoopt Then, the regularization path ((R) = arg min L(p) converges to the GMM direction i.e. limg-se P(R)/R mune nmlp
Theorem shows that as the regularization strength R increases towards the ridgeless problem min, L(p), the optimal direction ((R) aligns more closely with the max-margin solution n Since this theorem allows for arbitrary token scores, it demonstrates that max-margin token separation an essential feature of the attention mechanism. In fact, it is corollary of Theorem 8, which applies to

---

the generalized model f(X) (XX S(XW") p)) and accommodates multiple optimal tokens per input. However, while regularization path analysis captures the global behavior, gradient descent lacks general global convergence guarantees. In Section 2.2. we show that due the nonconvex landscape and softmax nonlinearity, gradient descent often converges to local optima. We first establish that when (ERM) is trained with gradient descent, the norm of the parameters will diverge. For the restrictive setting of n 1, gradient descent also exhibits a global convergence guarantee.
Assumption B For all i e [n] and opt, the scores per Definition 1 obey y V/ir Viopt.
Theorem 2 (Global Convergence of Gradient Descent) Suppose Assumption A on the loss func- tion& E and Assumption B on the tokens' score hold. Then, the gradient descent iterates p(t + 1) = p(t) 17V 7V _(p(r)) on (ERM), with the stepsize ll L/ and any starting point p(0) satisfy lim,-wolp(()) lpp(Ol = co. n I, we also have lim p())llp(()) = pmmm mmlp
Theorem 2 shows that gradient descent will diverge in norm, and when n 1, the normalized predictor p())lpp(D|| converges towards pmmmi, the separator of the globally optimal token. While n1 = is a stringent condition, this requirement is in fact tight as discussed in Appendix E. To illustrate this theorem, we have conducted synthetic experiments. Let us first explain the setup used in Figure 1. We set d 3 as the dimension, with each token having three entries first two coordinates as key embeddings & = [x1, X2, 0] by setting W = diag([1, OJ). This is what we display in our figures as token positions. Finally, in order to assign scores to the tokens use the last coordinate by setting = [0, O. lj. This way score becomes Y v x Y X3, allowing us to assign any score (regardless of key embedding).
[x1, X2. We reserve the
In Figure I(a), the gray paths represent gradient descent trajectories from different initializations. The points (0, 0) and (1, 0) correspond to non-optimal tokens, while the point (-0.1, 1) represents the optimal token. Notably, gradient descent iterates with various starting points converge towards the direction of the max-margin solution prmm* (depicted --) Moreover, the iteration count r increases, the inner product (p(t)/lp(())l pmmm /||p"m**|] Il> consistently increases. Figure I(c) also depicts the directional convergence of gradient descent from various initializations multiple inputs, with the gray dotted line representing the separating hyperplane. These emphasize the gradual alignment between the evolving predictor and the max-margin solution throughout the optimization.
Lemma Suppose for all E [n]] and I + opt, Y; and yu Viopt, Also assume W Radd full-rank. Then mm exists i.e. (ATT-SVM) is feasible for optimal indices a + opt,
2.2 Local convergence of the attention weights p
Theorem on the global convergence of gradient descent serves as a prelude the general behavior of the optimization. Once we relax Assumption B by allowing for arbitrary token scores, we will show that p can converge (in direction) to a locally-optimal solution. However, this locally-optimal solution is still characterized in terms of (ATT-SVM) which separates locally-optimal tokens from the rest. Our theory builds on two new concepts: locally-optimal tokens and neighbors of these tokens.
Definition (SVM-Neighbor and Optally-Optimal Tokens) Fix token indices a = (a:);, for which (ATT-SVM) is feasible to obtain mm (a)). Consider tokens T, c [7] such that (Kin,- k p = for all I E Tú We refer to T; as SVM-neighhors of kin, Additionally, tokens with indices a = (a are called locally-optimal for all e [m] and t E Ti scores per Definition I obey Yia > Yu- Associated pmnm is called a locally-optimal max-margin (LMM) direction.
To provide a basis for discussing local convergence, we provide some preliminary definitions regarding cones. For a given & and a scalar > O, we define (q) as the set of vectors E Rd such that the correlation coefficient between p and q is at least
Cysimal Takenn
Locally Optinal Takon
P
mtm
Optiml Tokon
:
(()
(on == E Rd
P Illll
2
12114
Mj.
llqll
Given O, the intersection of (o)e,(()) and the set
Figure 2: Gradient descent initialization p(0) inside the cone containing the locally solution nm
(p E Rd1 |ppll 2 R) is denoted as Cr.R(q):
CuR(() (q) == (g) n Rd ||pl| 2 R
5

---

Next, we demonstrate the existence of parameters μ = p(a)) 0 and R 0 such that when R sufficiently large, there are no stationary points within Cuæ(pp"ing Further, the gradient descent initialized within (p. converges in direction to pmm [[]] refer to Figure 2 for a visualization.
Theorem 3 (Local Convergence of Gradient Descent) Suppose Assumption A on the loss function holds and assume a (a)))== are indices of locally-optimal tokens per Definition Then, there a constant pla) E (0, 1) and R > O such that C mp does not contain any stationary points. Further, for any starting point p(O) E C gradient descent iterates p(t 1) p(t)
- (p(t))
on (ERM) with stepsize S 1 L satisfies lim llp(r)) = co and lim p(t) p())l mm /llp ""]l.
To further illustrate Theorem 3, we can consider Figure 1(b) where n = and T = 3. In this figure, the point (0, 0) represents the non-optimal tokens, while (1, represents the locally optimal token. Additionally, the gray paths represent the trajectories of gradient descent initiated from different points. By observing the figure, we can see that gradient descent, when properly initialized, converges towards the direction of pmmm (depicted by -). This direction of convergence effectively separates the locally optimal tokens (1, 0) from the non-optimal token (0, 0).
2.3 Regularization paths can only converge to locally-optimal max-margin directions
An important question arises regarding whether our definition of LMM (Definition 2) encompasses all possible convergence paths of the attention mechanism when Ipll - co. To address this, we introduce the set of LMM directions as follows:
a is locally-optimal per Definition 2
(a)|l
The following theorem establishes the tightness of these directions: It demonstrates that for any candidate q & pamm, its local regularization path within an arbitrarily small neighborhood will provably not converge in the direction q.
Theorem 4 Fix q pamm with unit l2 norm. Assume that token scores are distinct (namely Yu # 7 for t # T) and key embeddings kú are in general position (see Theorem 7). Fix arbitrary e 0, > Define the local regularization path of q its (E, Ro)-conic neighborhood:
p(R)=
arg min L(p), where Cr, (q) cone, (p (pe ppERR Rd lpll 2 Ro]-
(9)
peC.n, (1)llSE
Then, either limg- [|p(R)))
lim ((R)) # q. In both scenarios
lim ((R) # q.
R-400
The result above nicely complements Theorem 3, which states that when gradient descent is initialized above a threshold (l))()) Ro) in an LMM direction, [Ip()]] diverges but the direction converges to LMM. In contrast, Theorem 4 shows that regardless of how small the cone is (in terms of angle and norm lower bound ||pl| 2 Ro), the optimal solution path will not converge along q parm.
Joint Convergence of Head y and Attention Weights
In this section, we extend the preceding results to the general case of joint optimization of head v and attention weights p using logistic loss function. To this aim, we focus regularization path analysis, which involves solving (ERM) under ridge constraints and examining the solution trajectory as the constraints are relaxed.
High-level intuition. Since the prediction is linear as a function of v, logistic regression in y can exhibit its own implicit bias to max-margin solution. Concretely, define the attention features - If this dataset D is linearly separable, then fixing p and optimizing only y will converge in the direction of the standard max-margin classifier
x X. p(K,p) and define the dataset De = (Y,, x/,, y
arg min ||||| vER
subject to Y,-vT,
for all 1 Isisn.
(SVM)
after setting inputs to the attention features c x [22]. This motivates a clear question:
Under what conditions, optimizing v. jointty will converge to their respective max-margin solutions?
We study this question in two steps. Loosely speaking: (1) We will first assume that, at the optimal

---

12
0.01
12
(8.5), (1.11.
1.0
0001, (2,11, mm-1
00.01. (9.5.1). Y=1
(0.9) 00.3 (6..1.5). Y=1
(0.5), (1.3)
0.9
Probabilities
0.8
4
0.6
0.5
Softmax prob.
Logistic prob.
0.4
2
10 12
-2
12
14
30
60
90
120
150
Iterations
(a) All inputs are support vectors (b) (0.5,1.5) is not a support vector
(c) Probability evolutions in (a)
Figure 3: (a) and (b) Joint convergence of attention weights p (->) and classifier head v
(->-)
to max-margin directions. (c) Averaged softmax probability evolution of optimal tokens and logistic
probability evolution of output in (a).
tokens Nar, i [n] selected by p. when solving (SVM) with Kia;, all of these tokens become support vectors of (SVM). (2) We will then relax this condition to uncover a more general implicit bias for that distinguish support non-support vectors. Throughout, we assume that the joint problem is separable and there exists (v, p) asymptotically achieving zero training loss.
3.1 When all attention features are support vectors
In (SVM), define label margin to I [[] Our first insight in quantifying the joint implicit bias that, optimal tokens admit natural definition: Those that maximize the downstream label margin when selected. This formalized below where we assume that: (1) Selecting the token indices (a));1 from each input data achieves the largest label margin. (2) The optimality of the choice is strict in the sense that mixing other tokens will shrink the label margin in (SVM).
Assumption C (Optimal Tokens) Ler > 0 be the label margin when solving (SVM) with r; < Xia- There exists 0 such that for all p. solving (SVM) with + x? results in label margin of most y maxjeln - Sta.) where s; SK,p)
Example: To gain intuition, let us fix a E Rd and consider the dataset obeying Xn = Y; a and |xx |aal| for all f 2 and all E [n]. For this dataset, can choose 0; 1, mm a /|| r= |/
= ||a|| and v = ||a|| suPretajre2 ||x
Theorem Consider the ridge ridge-constrained solutions (v,, PR) of (ERM) defined as
(vr, PR) = arg min L(v. p)) i
Suppose there are token indices a (ar) for which (a))l exists (ATT-SVM is feasible) and Assumption C holds for some r, > 0. Then, limg- PR/R = pm "IIp" where pmum the solution of (ATT-SVM); and oo v,/r rr- pm mm where umm is the solution of (SVM) with r; Xin;
As further discussion, consider Figure 3(a) where set 3,7 d 2 and W Identity. All three inputs share the point (0, 0) which corresponds to their non-optimal tokens. The optimal tokens (denoted by *) are all support vectors of the (SVM) since -m = [0, 1] is the optimal classifier direction (depicted by -- Because of this, p will separate optimal * tokens from tokens at the (0. 0) coordinate via (ATT-SVM) and its direction is dictated by yellow and teal colored which the support vectors.
3.2 General solution when selecting one token per input
Can relax Assumption C, and if so, what is the resulting behavior? Consider the scenario where the optimal p diverges to oo and ends up selecting one token per input. Suppose this selects some coordinates a (a,)² Let S c [n] be the set indices where the associated token X)a, is support vector when solving (SVM). Set S [n] S. Our intuition is as follows: Even if we slightly perturb this p choice and mix other tokens # a, over the input set S c [n]] since S is not support vector for (SVM), we can preserve the label margin (by only preserving the support vectors S). This means that
p may not have to enforce max-margin constraint over inputs i e S, instead, it suffices just select
7

---

1.0
20
Softmax probabilit)
Attention norm ||p]]
Normalized GD
0.8
15
(Constant
nonopttocm Oyed
12 antima lohen Cyeis
0.6
10-
0.4
Normalized
0.2
GD (Constant
Cx2
100
150
200
50
100
150
200
C+139
Iterations
Iterations
-2
0
2
8
10
(a) Evolution of softmax probability
(b) Evolution of attention weights
Figure 4: Evolution of softmax probability and attention weights when training with normalized gradient descent or constant step size 17 respectively.
Figure 5: Trajectories of p with different loss functions and scores in Theorem 2.
these tokens (asymplotically). This results in the following relaxed SVM problem:
for r * a; iES
p pelax arg min lipll such that p (kin.
ki) 2 for all I # a, iE
(10)
Here, p (kin k2) 2 0 corresponds the selection idea. Building on this intuition, the following theorem captures the generalized behavior of the joint regularization path.
Theorem Consider the same (ERM) problem as discussed Theorem Suppose S(K,pR),, - 1, i.e., the tokens (==);= are asymptotically selected. Let yum be the solution of (SVM) with r; = Xin, and be its set of support vector indices. Suppose Assumption C holds over i.e. having Sia, < shrinks the margin when (SVM) is only solved over S C [m]. Then, lim,- v,/r mm /|ly and limg- PR/R- relax /Ilp relx where p nelax is the solution of (10) with (a:);= choices.
To illustrate this numerically, consider Figure 3(b) which modifies Figure 3(a) by pushing the yellow the northern position (0.5, 1.5). We still have pmm [0. 1] however the yellow is longer support vector of (SVM). Thus, p solves the relaxed problem (10) which separates green and teal *'s by enforcing the max-margin constraint on p (which is the red direction). Instead, yellow only needs to achieve positive correlation with p (unlike Figure 3(a) where it dictates the direction). We also display the direction of p qm using a gray dashed line.
We further investigate the evolution of softmax and logistic output probabilities throughout the training process of Figure 3(a), and the results illustrated in Figure 3(c). The averaged softmax probability of optimal tokens represented by the red curve and is calculated as Σ max S S(K,)), An achievement of for this probability indicates that the attention mechanism successfully selects the optimal tokens. On the other hand, the logistic probability of the output is represented by the blue curve and determined by 1/n Σx 1/(1 + e e-(-/0)).. This probability also reaches value of 1, suggesting that the inputs are correctly classified.
Experiments
Sparsity of softmax and evolution of attention weights. It is well known that, in practice, attention maps often exhibit sparsity and highlight salient tokens that aid inference. Our results provide a formal explanation of this when tokens are separable: Since attention selects locally-optimal token within the input sequence and suppresses the rest, the associated attention map S(X p) will (eventually) be sparse vector. Additionally, the sparsity should arise in tandem with the increasing norm of attention weights. We provide empirical evidence to support these findings.
Synthetic experiments. Figures 4(a) and 4(b) show the evolution of the largest softmax probability and attention weights over time when using either normalized gradient or a fixed stepsize 1j for training. The dataset model follows Figure 1(c). The softmax probability shown in Figure 4(a) is defined as ; Σn max S S(K,p), When this average probability reaches the value of 1, it means attention selects only a single token per input. The attention norm in Figure 4(b), is simply equal to llpll-
The red curves in both figures represent the normalized gradient method, which updates the model parameters p using p(t+ 1) p(t) 1/V (( with 17 0.1. The blue curves correspond
8

---

0.9
0.8
IW1. (K100)
and the logistic loss 0 0.7
0.6
0.5
Input image
(b) Epoch 0
(c) Epoch 100
0.4
1
0.2
100
200
300
400
Epoch
Figure 7: Red curve is the sparsity level nnz(s)/ T of the average attention map which takes val- ues [0.1]. A sparser vector implies that few key tokens receive significantly higher atten- tion, while the majority of the tokens receive minimal attention. Blue curve is the Frobenius norm of attention weights ||W]l of the final layer. We display their evolutions over epochs.
(d) Epoch 200 (e) Epoch 300
(I) Epoch 400
Figure 6: Illustration of the progressive change in attention weights of the [CLS] token during training in the transformer model, using specific input image shown in Figure 6(a).
gradient descent with constant learning rate given by p(1 1) p()) - v/V L(p(t)) with 7 1. Observe that the normalized gradient method achieves a softmax probability of quicker as vanilla GD suffers from vanishing gradients. This is visible in Figure 4(b) where blue norm curve levels off.
Real experiments. To study softmax sparsity and the evolution of attention weights throughout training, train vision transformer (ViT-base) model [23] from scratch, utilizing the CIFAR- 10 dataset [24] for 400 epochs with fixed learning rate x 10-3. ViT tokenizes an image into 16x 16 patches, thus, its softmax attention maps be easily visualized. We examine the average attention map- associated with the [CLS] token computed from all 12 attention heads within the model. Figure provides visual representation of the resulting attention weights (16 x 16 grids) corresponding to the original patch locations within the image. During the initial epochs of training, the attention weights are randomly distributed and exhibit dense pattern. However, as the training progresses, the attention map gradually becomes sparser and the attention mechanism begins concentrate on fewer salient patches within the image that possess distinct features that aid classification. This illustrates the evolution of attention from a random initial state to more focused and sparse representation. These salient patches highlighted by attention conceptually corresponds to the optimal tokens within our theory.
We quantify the sparsity of the attention map via a soft-sparsity measure, denoted by nnz(s) where is the softmax probability vector. The soft-sparsity is computed as the ratio of the l-norm to the squared no-nom,, defined as nnz((s) = [|s|| /| nnZ(s) takes values between 1 to T 256 and a smaller value indicates sparser vector. Also note that ||s||| Ex1 s, 1. Together with sparsity, Figure also displays the Frobenius norm of the combined key-query matrix of the last attention layer over epochs. The theory suggests that the increase in sparsity is associated with the growth of attention weights- which converge directionally. The results in Figure/ 7 align with the theory, demonstrating the progressive sparsification of the attention map as ||W]| grows.
Transient optimization dynamics and the influence of the loss function. Theorem 2 shows that the asymptotic direction of gradient descent is determined by m However, it is worth noting that transient dynamics can exhibit bias towards certain input examples and their associated optimal tokens. We illustrate this idea in Fig 5(a), which displays the trajectories of the gradients for different scores and loss functions. We consider two optimal tokens (*) with scores 71 = and y2 C, where C varies. For our analysis, we examine the correlation loss ((x) {(x) log(1 + *))
In essence, as C. increases, we can observe that the correlation loss f(x) -x exhibits a bias towards the token with a high score, while the logistic loss biased towards the token with a low score. The underlying reason for this behavior can be observed from the gradients of individual inputs: vL)(p) E. KJ Xr, where S'(-) represents the derivative the softmax function and E/ == CY, X s)).)) Assuming that (approximately) selects the optimal tokens, this

---

yt- With the correlation loss, [("| resulting in [Plll x Yr meaning that a larger score induces a larger gradient. On the other hand, the logistic loss behaves similarly to the exponential loss under separable data, i.c., |e'| = e e=1(11 + + = e=x Consequentyy, [[V[l & yee = e e indicating that a smaller score leads to a larger gradient. These observations explain the empirical behavior we observe.
simplifies f (()) and Lipill x (((()))
5
Related Work
Implicit Regularization. The implicit bias of gradient descent in classification tasks involving separable data has been extensively examined by [22, 25, 26, 27, 28, 29]. These works typically losses to make connections to margin
use logistic loss or, more generally, exponentially-taine maximization. These results are also extended to non-separable data by [30, 31, 21J. Furthermore, there have been notable investigations into the implicit bias in regression problems/losses utilizing techniques such as mirror descent [32, 25, 33, 34, 35, 36]. In addition, several papers have explored the implicit bias of stochastic gradient descent [37, 38, 39, 40, 41, 42], as well as adaptive and momentum- based methods [43, 44, 45, 46]. Although there are similarities between optimization approach for y and existing works, the optimization of stands out as significantly different. Firstly, our optimization problem nonconvex, introducing new challenges and complexities. Secondly, it necessitates the introduction of novel concepts such as locally-optimal tokens and requires a fresh analysis specifically tailored to the cones surrounding them.
Attention Mechanism. Transformers, introduced by [6], revolutionized the field of NLP and machine translation, with earlier works on self-attention by [47, 48, 49, 50J. Self-attention differs from traditional models like MLPs and CNNs by leveraging global interactions for feature representations, showing exceptional empirical performance. However, the underlying mechanisms and learning processes of the attention layer remain unknown. Recent studies such as [51, 52, 53, 54, 23] have focused on specific aspects like representing sparse functions, convex-relaxations, and expressive power. In contrast to our nonconvex (ERM), [52] studies self-attention with linear activation instead of softmax, while [53] approximates softmax using a linear operation with unit simplex constraints. Their main objective is to derive convex reformulations for ERM-based training problem. [55, 56] have developed initial results characterize the optimization and generalization dynamics of attention. [17] is another closely related work where the authors analyze the same attention model (ERM) as us. Specifically, they jointly optimize V. for three gradient iterations for a contextual dataset model. However, all of these works make stringent assumptions on the data, namely, tokens are tightly clusterable or can be clearly split into clear relevant and irrelevant sets. Additionally [56] requires assumptions on initialization and [55] considers a simplified attention structure where the attention matrix not directly parameterized with respect to the input. Our work links attention models to hard-margin SVM problems and pioneers the study of gradient descent's ''s implicit bias in these models.
Discussion
We have provided a thorough aprimization-theoreee characterization of the fundamental attention model f(X) rT X S(XW by formally connecting it to max-margin problems. We first established the convergence of gradient descent on (or equivalently W) in isolation. We also explored joint convergence of (v.p) p) via regularization path which revealed surprising implicit biases such (10). These findings motivate several exciting avenues for future research. An immediate open problem is characterizing the (local) convergence of gradient descent for joint optimization of (v, p). Another major direction is to extend similar analysis to study self-attention layer (4) or to allow for multiple tunable tokens (where p becomes a matrix). Either setting will enrich the problem by allowing the attention to discover multiple hyperplanes to separate tokens. While our convergence guarantees apply when tokens are separable, would be interesting characterize the non-separable geometry by leveraging results developed for logistic regression analysis [31, 22]. Ideas from such earlier results can also useful for characterizing the non-asymptooic//ransieee descent aligns with the max-margin direction. Overall, we believe that max-margin token selection is a fundamental characteristic of attention mechanism and the theory developed this work lays the groundwork of these future extensions.
dynamics of how gradient
10

---

Acknowledgements
This work was supported by the NSF grants CCF-2046816 16 and CCF-2212426, Google Research Scholar award, and Army Research Office grant W911NF2110312.
The authors express their gratitude for the valuable feedback provided by the anonymous reviewers and Christos Thrampoulidis, which has significantly improved this paper.
References
[1] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. The International Conference on Learning Representa- lions, 2015.
[2] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, and al. Language models are few-shot learners. In Advances neural information processing systems, vol- ume 33, pages 1877-1901, 2020.
[3] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, et al. Evaluating large language models trained on code. arXiv preprint (rXiv:2107.03374. 2021.
[4] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 1(8);9, 2019.
[5] Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311, 2022.
[6] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems, volume 30, 2017.
[7] Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapler of the Association for Computational Linguis- lics: Human Language Technologies, Volume I (Long and Short Papers), pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational I Linguistics.
[8] OpenAL. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 2023.
[9] Rishi Bommasani, Drew Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, et al. On the opportunities and risks of foundation models. arXiv preprint 1rXiv:2108.07558. 2021.
[10] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alee Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference Machine Learning, pages 9821-8831. PMLR, 2021.
[11] 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. International Conference on Learning Representations, 2021.
[12] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agar- wal, Girish Sastry. Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine PMLR, 2021.
learning, pages
[13] Danny Driess, Fei Xia, Mehdi SM Sajadii, Corey Lynch, Aakanksha Chowdhery, Brian Ichter, Ayzaan Wahid, Jonathan Tompson, Quan Vuong, Tianhe Yu, et al. Palm-e: An embodied multimodal language model. arXiv preprint
2023.
11

---

[14] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems, volume 34, pages 15084-15097, 2021.
[15] Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXXv:2205.06117 26750.1..., 2022.
[16] Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-efficien prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing, pages 3045-3059, 2021.
[17] Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. In International Conference on Machine Learning, 2023.
[18] Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint 2 101.00190, 2021.
[19] Saharon Rosset, Ji Zhu, and Trevor Hastie. Margin maximizing loss functions. Advances in neural information processing systems, 16, 2003.
[20] Arun Suggala, Adarsh Prasad, and Pradeep K Ravikumar. Connecting optimization and regularization paths. Advances in Neural Information Processing Systems, 31, 2018.
[21] Ziwei Ji, Miroslav Dudík, Robert E Schapire, and Matus Telgarsky. Gradient descent follows the regularization path for general losses. In Conference on Learning Theory, pages 2109-2136. PMLR, 2020.
[22] 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-2877 2018
[23] Yihe Dong, Jean-Baptien Cordonnier, and Andreas Loukas. Attention not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pages 2793 2793-2803. PMLR, 2021.
[24] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton. The cifar-10 dataset. online: http://www. cs. toronto. edu/kriz/ifar: html, 55(5), 2014.
[25] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro. Characterizing implicit bias in terms of optimization geometry. In International Conference on Machine Learning, pages 1832-1841, PMLR, 2018.
[26] Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry. Convergence of gradient descent on separable data. In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3420-3428. PMLR, 2019.
[27] Ziwei and Matus Telgarsky. Characterizing the implicit bias via a primal-dual analysis. In Algorithmic Learning Theory, pages 772-804. PMLR, 2021.
[28] Edward Moroshko, Blake E Woodworth, Suriya Gunasekar, Jason D Lee, Nati Srebro, and Daniel Soudry. Implicit bias in deep linear classification: Initialization scale vs training accuracy. Advances neural information processing systems, 33:22182-22193 2020.
[29] Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balean, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 17176-17186 Curran Associates, Inc., 2020.
[30] Ziwei and Matus Telgarsky. Risk and parameter convergence of logistic regression. arXiv preprint 1803.v1803.77.0 2018
12

---

[31] Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent nonseparable data. In Conference on Learning Theory, pages 1772-1798. PMLR, 2019.
[32] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro. Kernel and rich regimes in overparametrized models. In Conference on Learning Theory, pages 3635-3673, PMLR, 2020.
[33] Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view implicit bias in training linear neural networks, arXiv preprint arXiv:2010.0250/ 2020.
[34] Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery. Advances in Neural Information Processing Systems, 32:2972-2988. 2019.
[35] Ehsan Amid and Manfred K Warmuth. Winnowing with gradient descent. In Conference on Learning Theory, pages 163-182. PMLR, 2020.
[36] Ehsan Amid and Manfred KK Warmuth. Reparameterizing mirror descent as gradient descent. Advances in Neural Information Processing Systems, 33:8430-8439, 2020.
[37] Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization effect of initial large learning rate in training neural networks. arXiv preprint arXiv: 2019.
[38] Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenten like process. In Conference on learning theory, pages 483-513. PMLR, 2020.
[39] Jeff Z HaoChen, Colin Wei, Jason D Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. arXiv preprint arXiv:2006.0886 2020.
[40] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss?
a mathematical framework. In International Conference on Learning Representations, 2022.
[41] Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers. arXiv preprint wrXiv:210600653 2021.
[42] Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, Dean P Foster, and Sham Kakade. The benefits of implicit regularization from sgd least squares problems, Advances in Neural Information Processing Systems, 34:5456-5468, 2021.
[43] Qian Qian and Xiaoyuan Qian. The implicit bias of adagrad on separable data. Advances in Neural Information Processing Systems, 32, 2019.
[44] Bohan Wang, Qi Meng, Huishuai Zhang, Ruoyu Sun, Wei Chen, and Zhi-Ming Ma. Momentum doesn't it change the implicit bias. arXiv preprint arXiv:2/10.0333 2021.
[45] Bohan Wang, Qi Meng, Wei Chen, and Yan Liu. The implicit bias for adaptive optimization algorithms on homogeneous neural networks. In International Conference on Machine Learning, pages 10849-10858. PMLR, 2021.
[46] Ziwei Ji, Nathan Srebro, and Matus Telgarsky. Fast margin maximization via dual acceleration. International Conference on Machine Learning, pages 4860 4869. PMLR, 2021.
[47] Jianpeng Cheng, Li Dong, and Mirella Lapata. Long short-term memory-networf for machine reading. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 551-561, Austin, Texas, November 2016. Association for Computational Linguistics.
[48] Ankur Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit. A decomposable attention model for natural language inference. In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 2249-2255, Austin, Texas, November 2016. Association for Computational Linguistics.
13

---

[49] Romain Paulus, Caiming Xiong, and Richard Socher. A deep reinforced model for abstractive summarization. In International Conference on Learning Representations, 2018.
[50] 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.
[51] Benjamin L. Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang. Inductive biases and variable creation in self-attention mechanisms. In International Conference on Machine Learning, pages 5793-5831. PMLR, 2022.
[52] Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci. Unraveling attention via convex duality: Analysis and interpretations of vision transformers. In International Conference Machine Learning, pages 19050-19088, PMLR, 2022.
[53] Tolga Ergen, Behnam Neyshabur, and Harsh Mehta. Convexifying transformers: Improving optimization and understanding of transformer networks. arXii:2221 2022.
[54] Pierre Baldi and Roman Vershynin. The quarks of attention. arXiv:2202.0883 2022.
[55] 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, editors, Advances in Neural Information Processing Systems, 2022.
[56] Hongkang Li, Meng Wang, Sijia Liu, and Pin- Yu Chen. theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015, 2023.
[57] Simon Du, Jason Lee, Haochuan Li, Liwei Wang, and Xiyu Zhai Gradient descent finds global minima of deep neural networks. In International Conference on Machine Learning, pages 1675-1685. PMLR, 2019.
[58] Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization neural networks. arXiv preprint arXiv: 07572, 2018
[59] Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterization: Global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas Information Theory, 1(1):84-105, 2020.
[60] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over- wer-parameteriizaton. In International Conference on Machine Learning, pages 242-252. PMLR, 2019.
[61] Vladimir Vapnik. Estimation of dependences based on empirical data. Springer Science & Business Media, 2006.
[62] Peter Bartlett. For valid generalization the size of the weights is more important than the size the network. Advances in neural information processing systems, 9, 1996.
[63] Albert Novikoff. On convergence proofs for perceptrons. Technical report, STANFORD RESEARCH INST MENLO PARK CA, 1963.
[64] Peter Bartlett, Yoav Freund, Wee Sun Lee, and Robert Schapire. Boosting the margin: new explanation for the effectiveness of voting methods. The annals of statistics, 2665):6651 1998.
[65] Tong Zhang and Bin Yu. Boosting with early stopping: Convergence and consistency. Annals of Statistics, page 1538, 2005.
[66] Matus Telgarsky. Margins, shrinkage, and boosting. In International Conference on Machine Learning, pages 307-315. PMLR, 2013.
14

---

[67] Ganesh Ramachandra Kini, Orestis Paraskevas, Samet Oymak, and Christos Thrampoulidis. Label-imbalanced and group-sensitive classification under overparameterization. Advances in Neural Information Processing Systems, 34:18970-18983, 2021.
[68] Mahdi Soltanolkotabi, Dominik Stöger, and Changzhi Xie. Implicit balancing and regulariza- tion: Generalization and convergence guarantees for overparameterized asymmetric matrix sensing. arXiv:2303. 2.N.:::333../244.. 2023.
[69] Hossein Taheri and Christos Thrampoulidis. On generalization of decentralized learning with separable data. In International Conference on Artificial Intelligence and Statistics, pages 4917-4945. PMLR, 2023.
[70] Samet Oymak and Mahdi Soltanolkotabi. Overparameterizeed nonlinear learning: Gradient descent takes the shortest path? In International Conference on Machine Learning, pages 4951-4960. PMLR, 2019.
[71] Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. arXiv preprint urXiv:
2018.
[72] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019.
[73] Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arviv preprint arXiv: 190.000000 2019.
[74] Lenaic Chizat and Francis Bach, Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss. In Conference on Learning Theorry pages 1305-1338. PMLR, 2020.
[75] Spencer Frei, Gal Vardi, Peter L Bartlett, and Nathan Srebro. Benign overfitting in linear classifiers and leaky relu networks from kkt conditions for margin maximization. arXiv e-prints, pages arXiv-2303, 2023.
[76] Gal Vardi, Ohad Shamir, and Nati Srebro. On margin maximization in linear and relu networks Advances in Neural Information Processing Systems, 0.1-77224-7736 2022.
[77] Navid Azizan, Sahin Lale, and Babak Hassibi. Stochastic mirror descent on overparameterized nonlinear models. IEEE Transactions Neural Networks and Learning Systems, 33(12):7717- 7727, 2021.
[78] Navid Azizan and Babak Hassibi. Stochastic gradient/mirror descent: Minimax optimality and implicit regularization. In International Conference on Learning Representations.
[79] Guorui Zhou, Xiaoqiang Zhu, Chenru Song, Ying Fan, Han Zhu, Xiao Ma. Yanghui Yan, Junqi Jin, Han Li, and Kun Gai. Deep interest network for Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, pages 1059-1068, 2018.
rate prediction. In
[80] Qiwei Chen, Huan Zhao Wei Li, Pipei Huang, and Wenwu Ou. Behavior sequence transformer for e-commerce recommendation in alibaba. In Proceedings of the 1st International Workshop on Deep Learning Practice for High- Dingn--immsionall Sparse Data, pages 1 1-4, 2019.
[81] Fei Sun, Jun Liu, Jian Wu, Changhua Pei, Xiao Lin, Wenwu Ou, and Peng Jiang. Bert4rec: Sequential recommendation with bidirectional encoder representations from transformer. Proceedings of the 28th ACM International Conference on Information and Knowledge Managemeni, pages 1441-1450 2019.
[82] Mia Xu Chen, Orhan Firat, Ankur Bapna, Melvin Johnson, Wolfgang Macherey, George Foster, Llion Jones, Mike Schuster, Noam Shazeer, Niki Parmar, Ashish Vaswani, Jakob Uszkoreit, Lukasz Kaiser, Zhifeng Chen, Yonghui Wu, and Macduff Hughes. The best of both worlds: Combining recent advances in neural machine translation. In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 76-86, Melbourne, Australia, July 2018. Association for Computational Linguistics.
15

---

[83] Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement learning as big sequence modeling problem. In ICML. 2021 Workshop on Unsupervised Reinforcement Learning, 2021.
[84] Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In Proceedings the 39th International Conference Machine Learning. 2022.
[85] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In Interaational Conference on Learning Representations, 2020.
[86] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In International Conference on Machine Learning, pages 10347-10357. PMLR, 2021.
[87] Zi-Hang Jiang, Qibin Hou, Li Yuan, Daquan Zhou, Yujun Shi, Xiaojie Jin, Anran Wang, and Jiashi Feng. All tokens matter. Token labeling for training better vision transformers. In Advances in Neural Information Processing Systems, volume 34, pages 18590-18602 18602, 2021.
[88] Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant self-attention. In International Conference Machine Learning, pages 5562-5571. PMLR, 2021.
[89] Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Negp and ntk for deep attention networks. International Conference on Machine Learning, pages 4376-4386. PMLR, 2020.
[90] Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006. 14548, 2020.
[91] Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations, 2018.
[92] Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal upproximators of peuuence-t--eqqeeeee functions? In International Conference on Learning Representations, 2019.
[93] Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. arXiv:2301.13196 2023.
[94] Yoav Levine, Noam Wies, Or Sharir, Hofit Bata, and Amnon Shashua. Limits depth effi- ciencies of self-attention. In Advances in Neural Information Processing Systems, volume 33, pages 22640-22651, 2020.
[95] Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprini
2021.
[96] Jason Wei. Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le. Finetuned language models are zero-shot learners. arXiv preprint arXiv: 2 109.01652, 2021.
[97] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. arXiv:2211.15661, 2022.
[98] Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances Neural Information Processing Systems, 35:30583-30598 2022.
[99] Yingcong Li, M Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in-context learning. International Conference on Machine Learning, 2023.
16

---

[100] Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoe Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903, 2022.
[101] Guhao Feng, Yuntian Gu, Bohang Zhang, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: theoretical perspective. arXiv preprini arXiv:2305. 15408, 2023.
[102] Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting thain-of-thought: A study on compositional in-context learning of mlps. arXiv preprint arXiv:2305. 18869 2023.
[103] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. arXiv:2305. /6380, 2023.
[104] Tan Minh Nguyen, Tam Minh Nguyen, Nhat Ho, Andrea L Bertozzi, Richard Baraniuk, and Stanley Osher. A primal-dual framework for transformers and neural networks. In The Eleventh International Conference on Learning Representations, 2023.
[105] Davoud Atace Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Trans- formers as support vector machines. arXiv preprint arXiv:2308. 16898, 2023.
17

---

Roadmap. The appendix is organized as follows: Section A provides basic facts about the training risk. Section B presents the proof of local and global gradient descent and regularized path for learning E R with fixed v E Rd choice. Section C provides the proof of regularized path applied the general case of joint optimization of head and attention weights p using logistic loss function. Section D presents the regularized path applied to more general model f(X) d(X S(WW p)) with a nonlinear head v. Section provides implementation details. Finally, Section discusses additional related work on implicit bias and self-attention.
Table of Contents
A Addendum to Section
18
A.1 Preliminaries on the Training Risk
18
A.2 Proof of Lemma
20
B Addendum to Section
21
B.1 Descent and Gradient Correlation Conditions
21
B.2 Proof of Theorem
................................
29
Proof Theorem
................................
30
B.4 Proof of Theorem
30
B.5 Proof of Theorem 4
33
B.6 Proof of Lemma 2
39
C Addendum to Section
39
C.I Proof of Theorem 5
39
C.2 Proof of Theorem
41
D Regularization Path of Attention with Nonlinear Head
43
D.1 Proof of Theorem 8 Application Linearly-miss mixed Labels
44
45
E Implementation Details and Additional Experiments
46
F Addendum Section
48
F1 Related Work on Implicit Regularization
48
F2 Related Work on Attention Mechanism
48
A Addendum to Section 1
A.1 Preliminaries on the Training Risk
By our assumption Rd R and E R + R are differentiable functions. Recall the objective
1 Σ (Y, S(K,)(()))))
L(p. W)
(11)
n
with the generic prediction model v(X ((XTS((p))) and K XWT.
18

---

Here, write down the gradients of W and p in (11) highlight the connection. Set q WT p. z(X) = S(P),, and a(X) == Kp. Given X and using K = XW we have that
V. W) = (a (x(x)),
(12a)
V pp W) WV W).
(12b)
Vwu(p. W) pV W).
W).
(12c)
where
(a(a())) = diag(S(a - 3(a(x))(((()) E RRT
Setting (zz) = = z for linear head, we obtain
(p W) = X (a(xx)
(13a)
(p W) = WV W) = K
(13b)
Vp W) = pV. W) = pv x''xx''(aX
(13c)
Recalling (12b) and (12c), and defining e =- CY) ((z(x())) E R. we have that
1 n
Σ e Y, WV,
DlL(p, W)
(14a)
n
Vw
Σ
Gi Y, WV)
(14b)
i=1
Setting (zz) = v z for linear head and y, = Yi X,y, obtain
Vp W) Σ
t K. S' (a(()))))>
(15a)
t=1
VwL(p, W) p
1
Σ
"
S S(axxxxx
(15b)
n
im1
Lemma (Key Lemma) For any p. E R
let a = Kq. = S(Kp). and y = Xv. Set
r= sup 1 - Yzl and A = sup ||&,|| 1lgll. [Tt[7]
We have that
a diag(s)y a ss y
Σ 122
(a: - a,), - -72) s 2[A(1 -
Proof. y= Σi= V,5, 7s,, We have
T
7 7 Y1
(1 7,)s, and 1y 31 I(I -1).
Then,
a diag(()) a ss ss Σ a,y,s, Σ
a,s, 1:S:
m=1
m=1
mml
T
=ais((y) - -)
Σ
a -i).
(16)
122
Since
T s AT - 3
Σ
a,s, vi) Σ
a,s,(y: - 7i)
P22
19

---

obtain
a' d'diag(s)y
a'ss"y a
aisi(yy) -V) Σ a,si(y1 -1) #A(( -s1)²
I
a1s1 (1 -7)s, a;si(y1 - Xi) 7) AT(1 si)²
(a)si a,)8,(11 - 71) AT(1 -1))
(a: - a;), y,) 2AT(1 - si)P. Σ
122
Here, the right handside uses the fact that
T
(a1s) a1), mic - (1-5s))A
Σ
Σ
8 (1 AT)
122
122
A.2 Proof of Lemma
Proof. Let us prove the result for general step size sequence (1))20- On the same training data X recall the objectives L(p) = E ((Y ((XT and L(W) Σ ((Y; MXXT(XXWTTI) u))) Suppose claim is true till iteration t. For iteration I, using W()) = p(). define and observe that
= S'XX n) 5(Xi(()..
s; 3(X,W()) S(X,()).
z|X,) x7 S(X,p(()) (7)((X u).
for all E [it]. Thus, using (14), we have that
vi(p())
Σ
G Y. X,S, ('(((
1=1
VL((()) 1
Σ
C Y, S,S, ((((A)
i=1
Consequently, we found that gradient rank-1 with left singular space equal to u, i.c.,
VL(W()) = uV L(P()).
Since W(1's left singular space is guaranteed to be (including W(O) by initialization), we only need to study the right singular vector. Using the induction till 1, this yields
W(1 1)TH = W(I) u nllla| P L((())
p())) 17 1 uV L(p())
p(t 1).
This concludes the induction.
FFor simplicity, the right hand side denote the upper and lower bounds.
20

---

B Addendum to Section 2
B.1 Descent and Gradient Correlation Conditions
The lemma below identifies conditions under which pr is global descent direction for L'(p).
Lemma Suppose (()) is a strictly decreasing differentiate loss function and Assumption B holds. Then, for all p Rd, the training loss (ERM) obeys (V (p)), pmmnit O.
Proof. Set
y = - K,p, a, = K,p and e l' (== (r]S(Kip).
(17)
Let us recall the gradient evaluated p which is given by
V((()
π Σ G
sa)
(18)
This implies that
1 a
(V pm
Σ c (ā. saam)))
(19)
To proceed, we will prove that individual summands are all strictly negative. To show that, without losing generality, let focus the first input and drop the subscript for cleaner notation. This yields
(a. 3' (a)y)> a" diag(S(a))y a S(a)S(a)
(20)
Without losing generality, assume optimal token is the first one and y, is a constant for all / 2 2.
To proceed, will prove the following: Suppose y = Yr2 is constant, 71, a the largest indices of y, a. Then, for any s obeying Every 8, = 1, 2 O, we have that ā diag(s)y a aTss" y > 0. To see this, write
à diag(s)y - ā ss'
Σ
a,y,8 Σ
a,S; 7:8;
1=1
1=1
1=1
0:7181 +
Σ
ā,S;
- -(7181 + + (s)(--ss)
āisi +
ā,s,
222
122
T Σ
=ā1771 -))s(1 - -1))
(y- 0-11111 + x(1 -s
a,s;
(21)
R22
2 (1 - 51) (y: y)s:
a,s;
12?
Σ/32 a,s, Σ S
(y1 - -)) (1 s1)s|
ā
To proceed, let gap 71 and ggpp = ā, max a, With these, we obtain
ā" diag(s)y ā_ ss 2 agnpY (1--)
(22)
Note that
a 2 inf (kiopt, kw)) mmm 2 1, gap soopt,
y gap
inf Yiapt, - Yi > 0, reopt;
Sn sn) > 0. On the other hand, by our assumption & < O. Hence, infimum 'ing (22) over all inputs, multiplying by & and using (19) give the desired result.
21

---

Lemma 5 (Gradient Correlation Conditions) Consider and let pm p be (ATT-SVM) solution separating a = opt from remaining tokens of input X. Suppose ((-) is a strictly decreasing differentiate loss function and Assumption B holds. For any choice of π > O, there exists R == Rg such that, for any p with lplll 2 R, we have
2 (1 + ( [(p).
L(p)
lpll
Above, observe that as R - co, we eventually get to set π 0. 0.
The proof is similar to Lemma 4 at a high-level. However, we also need to account for the impact of p besides pT pm in the gradient correlation. The main goal is showing that p pmu is the near-optimal descent direction, thus, P cannot significantly outperform it.
Proof.
To proceed, let أ ||p""||pp/ppp M sup, ||k| @ 1/ S(Kp), Kp. ā Kpnm. Without losing generality assume opt = 1. Set y = 7r22. Repeating the proof of Lemma 4 yields
Σ/ a,s,
Σ 8 (V (VP))) mmm) = l' (71 y) s1)8i
Σ a,S;
(VL(())))) = (y1 y)(1 si)si
a1
Σ 8,
Given æ, for sufficiently large R, wish to show that
Σ a,s, a:
Σ a,s;
S (1+ (T)
(23)
Σ S
Σ7 8t
We consider two scenarios.
Scenario 11p -n s E (2M) In this scenario, for any token, we find that la, - ā = kk, (p )m MIi -='m|| Me.
Consequently, obtain
ΣV a,s,
2 Σ a,s, Σ/ 5 2M
Σ a,s,
ā Σx2 8,
Σ 8,
Also noticing ā Scenario 11p
2% a.s.
2 (thanks to nm satisfying margin), this implies (23).
E/2 5;
2 e = (2M) In this scenario, for some y= r(e) and r 2 2, we have that p (k: k.) = a: a. 2y.
Here r arg max pd k; denotes the nearest point to k Recall that = 3(Ra) where =|ppl//lp==m] To proceed, split the tokens into groups: Let N be the group of tokens obeying p" (k: k,) 1-v for E N and [T] N be the rest. Observe that
EEITI-N St
ΣIeTTI-S 8;
enR
ST
Te
Σ S
ST
e2CR
Set M M/O and note that ||all "Ip"""| plugging in the above bound, we obtain
||kil M. Using (K k,) S 1 1111-y00 over E N and
ET a;)s;
Eien(a a,),,
EMdTI-N(a -a)s;
Σ/2 S,
E/ 5,
E/ 8,
-1- 21-e+2MTe 2MTe-R-
Using the fact that ā
Eg a.s; E 3
1, the above implies (23) with x' 2M vv To proceed,
choose R₂ 1 log(2MMT/4 /x) to ensure π' S
The following lemma states the descent property of gradient descent for L(p) under Assumption A. important note that although the infimum of the optimization problem is L", it not achieved at any finite Additionally, there are finite critical points p.
22

---

Lemma 6 Under Assumption A, the function L(p) is L-p-smooth, where
Lp := A Σ (MolloP + 3M,llell [IIIIXI)1")
(24)
Furthermore, if n & 1/ppp- then, for any initialization p(O), with the GD sequence p(r + 1) p(t) - ((p((()))) we have
L(p(t + 1)) - L(p(r)) s s-2 |(_((((
(25)
for all 2 O. This implies that
Σ
2 < co, and lim [| 0.
(26)
1000
Proof. Recall that defined y; V/: X,y and a; K.p. The gradient of L(p) is given by
[() Σ e' (xTS(Kpp)) s'))))
Note that for any p Rd, the Jacobian of s(Kp) is given by
dS(K,p) (K ((dn))-((iaaggaaa (K)))) s(kp)) K;
(27)
Op
The Jacobian (27) together with the definition the softmax function S() implics that (0S(K,p)/0pl ||K Hence, for any p. Rd we have [|((K - p)l = ||]]|| llp Pll
(28a)
and
S'(k,p)))(((pp))))) E diigg(((K))) diag(S(K,))) + ()S(Kp))((Kpp) ((K,))S(k,)))
$ Kil ||p pll-
(28b)
Here, the last inequality uses the fact that lab- - cdl s Idlla cl + lallb d|.
Next, for any p. P Rd, have
[V>((p) - <(PII S Σ l (y S(K,))))) K. S'KK ( (x]S(K,p)) K. s'Kpp')
K. S' S''KK'Kp''''
(K,p)ill (> (xTs(Kpp))
( S(K,p)) |KTS'(K,p>) S,
2 | [S(K, S(K,p))
Σ Millxil |K]|| - S'(Kip)||
(29)
where the second inequality follows from the fact that lab call s |d/lue - cl lallb d) and the third inequality uses Assumption A.
Substituting (28a) and (28b) into (29), get
[V_[(p) - v.(l S
[oll + 3M lpp Pll
22
(MaIMMPIIIPIIIXI 3MIII// Pll
n 1=1
< 1, pll.
23

---

where Lp is defined in (24).
The remaining proof follows standard gradient descent analysis (see e.g. [22, Lemma 10]). Since L (p) is smooth,, get
L(p(1 1)) L(p()) VE(DO))(((())))) (p(+)))))))111))111))))))))))11) L 1p 1)
=(p(0) (p)) + 4n (p))(
= =(((()) 7/ 4p
$ {(((() - (
where the last inequality follows from assumption the stepsize. The above inequality implies that
Σ S 2 (C(((()) - L"),
(30)
where the right hand side is upper bounded by a finite constant. This is because, by Assumption A, (O)) and L s (t)()), where denotes the minimum objective.
Finally, (30) yields the expression (26).
In the following lemma, demonstrate the existence of parameters μuaa)) > O and R, > 0 such that when R, is sufficiently large, there are no stationary points within Cμ.R, (p" (p""). Additionally, we provide the local gradient correlation condition.
Lemma 7 (Local Gradient Condition) Suppose Assumption A on the loss function & holds. Let (a:) be indices of locally-opting tokens per Definition 2.
L1. There exists positive scalar pla)) such that for sufficiently large Ry, stationary point exists within C,,,, (p where Cy,R, is defined in (8).
L2. For all q, E cone, mm) (P" with || = | and l|pll 2 R2 with same R2 choice as (LI.), there exist dataset dependent constants C, such that
Σ (1 s(K,p)a)) (V cc Σ -1-8((()).)) 0. [T
C.
(31a)
i6[A]
|VL(1)) 3 AC Σ
S(K,)))))..)) ACTe R2,9/2
(31b)
M0[M]
-19
vL(p)
c
> 0,
(31c)
[V)iI
CA
Here, Ā maxeta) llk
killl and = mm
L3. For any there exists Rę such that R₂ R2 and all p e Cy.R, (m
obeys
(L(())
P Iplll
2 (1 n. V((p))
mm
mm
Proof. Let mm (ar) be the solution of (ATT-SVM), Recall
C,.R, (p mm cone, mm
(p) |Ppll Ry).

---

Let (Tr) be the sets of all SVM-neighhors per Definition Let = T, [a;) be the set of
non-SVM-neighbor tokens, i E [n]. Let
@ 1/1pp"'!!
5 0.5 min min (kur - kit) " pmm. rejn) Re.,re
A = max
(32)
le[r]_e[7]
# (()) miin0..5,0)
A
When T; = 0 for all E [n] (i.e. globally-optimal indices), we set ò = oo as all non-neighbor related
terms will disappear. Since pmm the max max-margin model ensuring (kia, 10M,,-k;,, k2)) p 2 1 for all [n]. the following inequalities hold for all q cone, || || and all i e Tú:
(k,, kx)) q 2 8200 0,
(kin; k;-) q 2 + 6,
(33)
3/2 (km.. k q 2 1/2.
Here, we used l|q - /Il z 2μ which implies llq - |r 2 v 24//0.
LI. and L2.. Now that the choice of local cone is determined, we need to prove the main claims. We will lower bound -q [((p)) and establish its strict positivity for || R. where R Rg. This will show that there is no stationary point as a by product.
Consider any g Rd satisfying liqll = ||p" To proceed, we write the gradient correlation following (18) and (21)
(V q) - Σ E (a, S'(d)ri)
(34)
where we denoted f; - - XI 8 ((K a, k,g, aj K,p. = S(Kp)) Using (33), for all / E Tú E F, for all E ( we have that d',, d 2 RO(1 8), and a d', ROS. Consequently, we can bound the softmax probabilities s; S(K,p) as follows: For all e [n].
S
Σ
Sir S Sin, Σ S S Te R012 Siz R012
re7
Q:
Σ
Sir Te RoN s4 I Te ROMS, Si Vt, E Tí
(35)
ref:
Recall scores Yn Y, v xg.
Define the score gaps over neighbors:
/:t0 Vio.
max and 7/ Vio. IE
min 7n-
It follows from (32) that
A = max | 2 max ||a = ||kn jaIn)aT
Define the a-dependent global scalar I' suPeejol.rrr(T] 17 yirl Let focus on fixed datapoint
ie [n], assume (without losing generality) a; and drop subscripts that is, == d' = 1, X X;
Y : Y, K K,, a' Kp, Kq. s S(Kp). y Y. Xr, yau ; zap seap 7. Q 2. and
S = S, Directly applying Lemma 3, we obtain
la diag(s)y a T
Σ (a) a,), - 2[A(1 2 22
To proceed, let us decouple the non-neighbors within Etz a.)s.(yi 7 via
Σ
(a--a))s y.) 2QTA.
25

---

Aggregating these, found
a t diaiggs)y a' ss E(ar - (a)-a)) 71) s 2r s1)² + Q) 167;
(36)
To proceed, let us upper/lower bound the gradient correlation. We use two bounds depending on q E (p (Case 1) or general q R (Case 2).
Case I: E cone,(p"m)) Since 1.5 2 a1 a, 2 0.5 following (33), we find
S. gea Σ (a: (7) y,) 0.5 S yeu..
Next we claim that dominates ((1 - 51) Q) for large R. Specifically, we wish for
y /4 / 4 (1) Q) S z 16 LA
max((1 5:)
(37)
yeap
Now choose R 2 8-1 log(7)/O to ensure S since Q s T' RO Consequently
s1)² (Q S)² 45 Te R012
Combining these, what we wish is ensured by guaranteeing
AA yeap
S 2
max -R0/2 Te--ee-Roo (R-R5)))
(38)
This in turn is ensured for all inputs i E [n] by choosing
6
64TTA
R
log
(39)
pap
where yop = minjajat yi" is the global scalar which is the worst case score gap over all inputs. With the above choice of we guaranteed
S y
(1
syyye
2(1 - 81) yp 2 S 7p a diag(s)y a=ss"" y2
4 2 8
Since this holds over all inputs, going back to the gradient correlation (34) and averaging above over all inputs e [n] and plugging back the indices i, we obtain the advertised bound by setting = 1 Sia, (where set = I above without losing generality)
2 Σ nI rEIN|
G.
7/00 2 (V 2
Σ ME|n)
yi
gap
(40)
8n
Let -E'mi max be the min/max values negative loss derivative admits over the ball [-B, B] for B = ||||| max;, [|xael and note that maXie(w] 7. O and minjetu] y. 0 are dataset dependent constants.
Then, we declare the constants C -2f m maXie[u] g0u obtain the bound
0, ((1/))( min 7 0 to
C Σ qi -VC(p)
2
4i-
(41)
K[A)
[0[m]
which is the desired statement in (31a).
Case 2: E Ied and || mI Define Ā = max:e]n]are(7) likr use the fact that
ki For any l|gll = |m]] we
A S
||a - a,ll s ||,
k,ll
Note that by definition
1. To proceed, we can upper bound
S
_pp >
Σ IET
(a1
a,)s -D).
(42)
θ
26

---

By choosing the same R in (39) ensure S dominates ((1 -1) Q) and since & 1, guaranteed
2A
spap 2 a diag(s)y - a ss"
Going back to the gradient correlation (34) and averaging above over all inputs E [n]. with the same definition of C > O, we obtain
AC
2
(43)
On He[p]
To proceed, since (43) holds for any q R² ||""||
and ll
||
observe that when choosing
q | (p), this implies that
AC
v (VL())) = [(D)l l|p
||
1
Σ
qi-
On
Simplifying O 1/Ilp"M||| on both sides yields (31b). Incorporating (35) in the bound above provides exponential upper bound that decay with R.
Combining this with (41), we obtain that for all q. g.pe E cone, and |ل 2 R₂
v.((p) ||lll [V
2 c
CA
This gives the desired result in (31c). L3.: Establishing gradient correlation.
Our final goal is establishing gradient comparison between p, pr for the same choice of provided in (32). Define " = ||p||| to be the normalized vector. Set notations a: K,P. a; K,p"' and y; Y, X,v.
To establish the result, using (34), we will prove that, for any > there sufficiently large Rs such that for any Cp.(( )m)
P 1pll
v(()).
G (a. S'(K,))r))
π
Σ e ('K,''(K,p)))) = (1+ ( -V[(p), i=l
nm
(44)
I
Following (36), for all [n], for all cone, mm) with l| = ||p a' Kq and s S(Kp).
we have found
aa diag(s)) - atTTT s.s. yi
E d,)s m s 20A sn1)² 2.).
(45)
TeT,
Plugging in a, ā, in the bound above and assuming æ S (w.l.o.g.). (44) is implied by the following stronger inequality
G
6LA((1
sa) + Σ (a, a=)8a(71 yiu)
.TT
<
(ān az)Sa(7n - 7a)
n
[=1
sni7n - 7u).
n
First, claim that 0.5 Σiey su(yn 7n) 2 6TA((1 - 2 Q.) for all i [n]. The proof of this
claim directly follows the earlier argument, namely, following (37), (38) and (39) which leads to the choice
max
Co TTA
R
log
(46)
ry gap
27

---

for constant Co Here, choose sufficienlly large Co 64x ensure R= R₂ Ry. Following this control over the perturbation term 6TA((1 - Sa )2 + Qi), to conclude with the result, what remains proving the comparison
Σ 1=1
G
Σ
(an
aa)s yu) S
1 + 0.5x
Σ
s 77d)
(47)
IET;
IE
To proceed, we split the problem into two scenarios. Scenario l1p
|l
e
R for some E O. In this scenario, for any token, we find that 4л6
la ā.l kk (p p=' - A@E 7/4.
Consequently, we obtain
a: ag S ān ã. + 2 0.5%.
Similarly, 4, 4u 1 0.5x 2 0.5. Since all terms an aa. Siz Yn Ya in (47) are nonnegative and (a,1 (7n) - 7u) & (1 0,5 Yu), above implies the desired result (47).
Scenario 2: ||ø - -| 2 =
JAG Since ý not (locally) max-margin, in this scenario, for some ie [n], r(e) and Ti, have that
_
(kk kx)) an ai - 2v.
Here arg ma p" kú denotes the nearest point to ką (along the direction). Note that non- neighbor E Ti cannot be nearest because P E , and (33) holds. Recall that s; = S(Ra,) where R = O 2 RO. To proceed, let a, == min ai ag.
= i E [n] a S 1- 2r)
- I [ [m] 2y a
For ill i [n] I,
Σ IET
a1 - a,)5 7n) (1 0.5z))
Σ s417111 - 7
IET
s (2A - (1 + 0.5m)) r
Σ
8it
(48)
rT,, du 2221-221
(2A - (1 O.5 ITe
R(1+g)
2ATTe R0+4)
For all i E I, split the tokens into two groups: Let N; be the group of tokens obeying an au S
and T, N; be the rest of the neighbors. Observe that
ΣieT,-N Su
exR
Te
Ene,, Sur
22vR
Thus, using lan aul 2A and recalling the definition of gap
observe that
2rATe Rv
(an an)ss(rn 7u)
Σ
solyn yi).
gap mmn
T0T,
Thus,
Σ
(an an)si((n yu)
Σ
(an an)sn((a 7id Σ
(an an)si(yn Vn)
Te7.
IeN.
IET.NN
2TATe -iy
Σ
(1 v)sa(ynn - -aa) +
Σ
SilYa ya)
y
ieN.
IET.
2TATe Rv
Σ
+
SulYa Y
Sap
reT:
2TTTe Rv
Σ
5
+
Su(Ya - -i)).
20p
IET
28

---

Hence, choosing
1
8TAY
R
log
(49)
v0
gaap π
results in that
aa an)sh(yy yn) (1+ mh
Σ Sir(yn yn)
2 IET,
IET
2TATe RR
Σ
5 - -Yi)
Eup min
2
IET
(50)
π
Sirya Va)
4 IET;
π
y m
R(1-2v)
47
Here, the last inequality follows from the fact that Exer, S4 2 max Sp 2 Σl, Ma R(1-2/1 IT.
Ri 20
From Assumption A, we have Cmin S -l" S Cmax for some positive constants Cmin and Cruxx It follows from (48) and (50) that
Σ
C--
Σ NET
(an
an)sa(yn - Ya) - Σ 0.5r)) yn)
IET
$ Cmax 2ATTTe R((+
Cmin nT
TT
R(1-2v)
e
4
s 0.
Combing with (49), this is guaranteed by choosing
8ГAT
8nT Cmax zap π
R 2 max
log
log
y
Umnn"
(2v x/2)
Cmin min
where v(AAg) depends only on and global problem variables.
Combining this with the prior choice (46) (by taking maximum), we conclude with the statement.
B.2
Proof of Theorem
This proof is a direct corollary of Lemma 14 which itself is special case of the nonlinear head Theorem Let us verify that f(X) = v S(Xp) satisfies the assumptions of Lemma where replace the nonlinear head with linear To this, set the optimal sets to be the singletons
Proof.
O = (opt,). Given (X, Y), let s; = S(K,p) and qi of Ersopt, Sy Recalling score definition
y: = Y, X,y and setting y; Viopt; and Z; == Enoopt, VuSar particular prediction can be written
y. 8; Yiopt, (1 - qi) + Σ VuS
reopt;
=v - - Z4-
To proceed, we demonstrate the choices for C. 0. Let := (o Yu A 0 and gmax maxiefs] qi- Note that Z; Ex YuSu min 2 -Cqmax Now, using strict score optimality of opt,'s for [n], we set
Σrsopt, PuSu
sup Yu
e: sup
sup ietnl
0.
iEinl
viqi
Yiopt,
Empops Yo
We conclude by observing Z, s vigi
S Vigie as desired.
29

---

B.3 Proof of Theorem 2
Proof. We first show that lim |p (f)ll = co. From Lemma 4, we have
(V mume
C(y, v X S(K,))) (K,p" s'(aa))).
where y; Y, X,v and K.p-
It follows from Lemma 4 that (V [(p), p < 0 for all p R Hence, for any finite p. (V [(p). pmm*) cannot be equal to zero, as a sum of negative terms. Therefore, there are no fi- nite critical points p. for which V [(() O. On the other hand, Lemma states vL(((()) 0 which implies that || (f)l - 00.
Next, we provide the directional convergence for the setting 1. Let consider arbitrary value of E E (0, 1) and set æ = E/(1 e). As lim llp(r)] = co, we can select specific & such that for all r fe, it holds that I|p(D)) 2 R. V 1/2 for any choice of R.- To proceed, we choose R, based on Lemma 5 so that for any 2 fes we have that
P mmm
p(t)
v((((t))
(1
|
||
Multiplying both sides by the stepsize 77 and using the gradient descent update, we get
man
p(()
p(( 1) -p())
2 p(r+ 1) 11-pippppp -(()), ||p(0||
p
(1 6)
(IL((I 1)11 - |Ip(0l1 11p(( + 1) ((O1I)
211p(011
(51)
21lp(Dll (llp(t+ + 1) - -11(()11))) - 11plr + 1) p()1))) -
2 (1 c)
2 (1 (lp(( 1)ll - |1(())l - 11 P(O113)
2 (1 (IIp(r + 1)ll ||p(()) 2n/ (L(p()) L(p(1 1))) (1)))+))))))
Here, the second inequality is obtained from ||p(r)] 2 1/2; the third inequality follows since for any a, b > O, we have (a² b2)/(2b
2 0; and last inequality uses Lemma
Summing the above inequality over 2 l gives
p()
mmm
C(e, 17) [|((r)]
2 1 +
PpOOll
mmm
for some finite constant C(e, 17) defined as
man
C(e, = p(t.)
(1 O)pp(())l 2n(11)) (L)(((t)) -)),
(52)
mmm
where L" denotes the minimum objective. Since lpp (f)ll + oo, we get
p(r)
mms
lim inf 1-000
21-E
||p(Oll
Given that we can choose any value of e E (0, 1), we have p(t)/lp(())) + mm= /|lp mm |]
B.4 Proof of Theorem' 3
Proof. Following the proof of Lemma 7, let (TD),, denote the sets of SVM-neighbors defined in
Definition 2. We define T, [7] -T, [ar;) as the tokens that are non-SVM neighbors. Additionally, let y defined in (32). Let us denote the initialization lower bound as R,. := R, where given
in the Theorem 3's statement. Consider arbitrary value of e (0, 4/2) and let 1/(1 + œ) -
30

---

We additionally denote R R2 V 1/2 where R defined in Lemma 7(I.3.). At initializaaton p(0), we set e 2u/2 to obtain R Rp/z and provide the proof four steps:
Step 1: There are no stationary points within C,,R, (p ". We begin by proving that there are no stationary points within C, Hugg (p Then, since R 2 R2 per Lemma 7, we can apply (L2.) to find that: For all q. p E cone,, with # 0 and ||pll 2 Ry. we have that -q" -q" VL((() is strictly positive.
Step 2: It follows from Lemma 7(.3.) that, for E (0. all E CuR. (p satisfy
- L())
||pm|] 2 -V- (1-
mm
P
(53)
Ipll
The argument above applies a general E (0. p/2). However, initialization p(O), we set c 4/2 to obtain our earlier Ry, choice. To proceed, for any (0/), we will show that after gradient descent enters the conic set CH,R, (p"mm) for the first time, it will never leave the set. Let & be the first time gradient descent enters Cp.RR (p" In Step 4, we will prove that such & is guaranteed to exist. Additionally, for μ/2, note that 4, 0 i.e. the point of initialization.
Step 3: Updates remain inside the cone (p (p" By leveraging the results from Step 1 and Step 2, we demonstrate that the gradient iterates, with an appropriate constant step size, starting from p(te) Cy.R, (p""),, remain within this cone.
We proceed by induction. Suppose that the claim holds up to iteration s 2 t.. This implies that p(t) E C,.RR (p" Hence, recalling cone definition, for > 0 and R,, we have and [Ip(r)] 2 Re. Let
pt0
2
mm
p()) == 1-e
VL((())
Note that p()) due Step 1.
This together with the gradient descent update rule gives
p(1 1)
p())
mmm T
V L(P()),
[|p(||]
p
2
vL(((())
(54a)
mp(r)(1 E)
1 -
||p(t))
Note that from Lemma 7, we have <V-((p())), p()) 0 which implies that I|p(r + 1)Il (|l(t)))- This together with R, definition and (Ip(()) 2 1/2 implies that
II + 1)1l S 2|[p (II(( + 1) (IP(OlF)
211p(0ll (21le(011 - 2n (VL((()) P()) +
(54b)
= |1p(0)l [|(())l (VL(p(()), P(t) + +
Hence, using (53)
[|p( 1)1l |p(()]l
P())
+ II LCPOLI
1 ||p(1||
VL(((()).
[|((0||
I|p(OlI
1
17
VL((())
T7 IVP
(54c)
(1 - e))PPOOII
|
mp(t)
I(p(OII
=: C)(p()).))
[p((0l
Here, the second inequality follows from (53).
31

---

Now, it follows from (54a) and (54c) that
p(t 1)
mm
1
1-
()((n(1 -c)
Ipp(t
DlI
Ci(p().)
||p(1)I
1 ++ Ci(p(t),))
(1-x)((-C((pt))
11p(011
(M- 1)(-PP)) p()) "IVVLOPO
P001- -
(55)
1 P Ci(p(0.
1IP(O1I
Ipp(O))
p( 77(1 -1)
1 / Ci(p(0.
IIp(OlI
Ip((I
2 11-44 μ.
where the last inequality uses our choice of stepsize S L/ in Theorem 3's statement. Specifically, need 17 be small to ensure the last inequality. We will guarantee this by choosing proper R in Lemma 7. Specifically, Lemma leaves the choice of Co in R, lower bound (46) open (it can always be chosen larger). Here, by choosing Co 1/1, will ensure 7) 1/L, works well.
To proceed, we have that
(1-C) p()) P-6 C 1 1 L(p((II
K,022
1 H C C ACT
(56)
c O
10122
e
2(1 #)(1 - 4) ACT
Here, the second inequality uses our choice of (0,)12) (see Step 2), and the first inequality is obtained from Lemma 7 since
p())
VL((())
O
IV L(POII
[(p(til
m
T-ECA
IV L(P(t
AC Σ (1 ) ÁCTe-RO01
for some data dependent constants and C, Ā maxggln)are(7) [ik, kzlll and 4 //
Next, we will demonstrate that the choice of 17 in (56) does indeed meet our step size condition as stated in the theorem, i.e., 17 S 1/Lp. Recall that 1/(1 + πT which implies that π e/(1 -). Combining this with (46), we obtain:
R max(2,8
CoTTA sap
log
where Co 2 64T,
max 2-)
log
CoTAA Ex
where
Co 2 64 E 1-6
64
On the other hand, at the initialization, we have 11/ which implies that
o2x log
(2 7T
R
where Co 2 64
(57)
(2- 4))
In the following, we will determine a lower bound on Co such that our step size condition in Theorem 3's statement, i.e., 7 S 1/L., is satisfied. Note that for the choice of n (56) to meet the condition 175 1/Lp, the following condition must hold:
# (2 1) C.T
10022
R
2 θ
log
(2 1)
C.T
(58)
Lp
Lp
r
where C₁ (1 4) FC 6c
This together with (57) implies that for sufficiently large
max(2,
(2 - H)C3T
CaTA
C2
RC
log
where
C
#
y eap
Lp
32

---

the step size bound in (56) ensures that 17 S 1/ guarantees (55). Hence, p(t 1) remains within the cone, i.e,, p(t + 1) C,.R(p"mn)) (pmm))
Step 4: The correlation of p(t) and pnn increases over The remainder is similar to the proof of Theorem 2. From Step 3, we have that all iterates remain within the initial conic set i.e. p(t) e Cuee (p for all 2 O. Note that it follows from Lemma 7 that (V [(), phmm// /Ilp=m]]) < 0, for any finite p Cu.ey(pmim) Hence, there are finite critical points p Cu (p""), for which [(p) 0. Now, based on Lemma 6, which guarantees that V V(((r)) - O, this implies that Up (?)l - co. Consequently, for any choice of e e (/22) there is a time I such that, for all 2 4. p(t) E Cu.RR (p Once within Cp.RR (p following similar steps in (51) and (52), for any fe.
p())
P mm
C2(E.T)
-+
p(t) E CH.., (p
IIp(Oll
||p
Ilp(r))
for some finite constant C2(E, m) Consequently,
p())
pm
lim inf 1-:00
21 where p(1) E Cp.., (p
||
Since the choice of e E (0, 2/2) is arbitrary, we obtain p(r)/1lp(rii - mm /1pp |m
B.5 Proof of Theorem 4
B.5.1 Supporting Lemma
We present a lemma that will aid in simplifying our analysis. We begin with a definition.
Definition 3 -Selected-tokens, Neighbors, Margins, and Neighbor-optimality of a direction) Let q E (0) and (Y, K, Xi); be our dataset. We define the (possibly non-unique) selected-tokens of as follows:"
a; arg max k..q.
(59)
2E[T]
Next, we define the margin and pirretioonllneighhbn for q as the minimum margin tokens to the delected-tokens, i.e.,
rg = min (kja, kn) q.
(60)
ie(o],taa
M₂ (i. (kia, ku) ( TT))
(61)
Finally, we say that q is neighbor-optimal if the scores of its directionalnneighoor are strictly less than the corresponding selected-token. Concretely, for all (i. Ma. require that
Y; x < Yia. Y;- x
Lemma 8 (When does one direction dominate another?) Suppose q.p E Rd be two unit Eu- clidean norm vectors with identical selected tokens. Specifically, for each i [n], there exists unique a; E [7'] such that a;; = arg max k,,9 arg max K.P.. Suppose directional margins obey I p and set ör = T
Suppose and p are both neighbor-optimal. Then, for some R(ōr) and all R > R(6r), we have that L(R p) L(R q)
Suppose q has a unique directional neighhor and not nnighbor-optined (i.e. this neighbor has higher score). Let òg be the margin difference between unique directional-neighbor and the second most minimum-mman marignn neighbor (i.e. the one after the unique one, see (65)) of q. Then, for some R(ör A 6g) and all R > R(ôT 6g), we have that L(R q) L(R p)-
Proof.
We prove these two statements in order. First define the directional risk baseline induced by letting R - co and purely selecting the tokens (ari)) This is given by
L. Σ e( x
'If' is unique for all i [m]. let call it, unique selected tokens.
33

---

We evaluate q. p with respect to L. To proceed, let S(RK,q). Define =y=k k Note that, the smallest value for a; is achieved for g. For sufficiently large R z O(log(()) /T g). observe that, for f # d;
e
e RTy 2 S/1 Everyy
e 2 0.5e-
(62)
Recalling the score definition and let M.. be the upper and lower bounds on -(" over its bounded domain that scores fall on, respectively. Note that, for some intermediate M, 2 M, 2 M_ values, we have
L(Rq) L. =
Σ 1=1
E(
s
((
re[7]
M,
Σ
Sui(Yon - 7u).
Now, using (62) for a refreshed M, 2 Ma 2 0.5M values, we can write
[(Rq) L. Σ
Σ
Mae RT (y Yu).
(63)
t=l
The same bound also applies with some My multipliers
[(Rp) - --= -=
Σ
Mre-RTi
Mge RT (T Yvd).
(64)
2m1 FRCE
We can now proceed with the proof.
Case 1: q and are both neighbor-opting This means that Via; O for all i [m], or Let K2 > K_ be upper and lower bounds on Via: Yu values. We can now upper bound the right hand side of (63) via
1 M.K.Te RI, L(Rq) L. M_K "
2n
Consequently, L(Rq) > L(Rp) as soon as bn _M___e" > M,K, M,K,Te-*i; Since M,, K., n, T are global constants, this happens under the stated condition the margin gap - T..
Case has a unique directional-neighhoor and not neighbor-optimal. In this scenario, L(Rq)= L. is actually negative for large R. To proceed, define the maximum score difference K, = suptreg: Vin Yul- Also let (j.B) be the unique directional neighbor achieving the minimum margin rT Then, 6g the margin difference between unique directional-neighhee and the second minimum-margin neighbor (i.e. the one after the unique one) of 4 is defined as
8y min
I 9
g-
(65)
iE(n], FRa,, (())A(j()
To proceed, we can write
L(Rp) L. 2 -M,K,Te RT,
On the other hand, setting = y ju;; O, we can bound
[(Rg) L. ge
Σ
Mie RI 2 - 74)
n
n
ielnl. MAs. (LAA())))
M_e RTax M.K.Te +6),
n
Consequently, have found that L(Rp) > L(Rq) as soon
M_e RT, < 2 M,K, RIT (64) RT,)
This happens when R 2 -, (up to logarithmic terms) establishing the desired statement.
34

---

B.5.2 Proof of Theorem
Define the locally-optimal unit directions p p (Ip""(())) (a'il
a is a locally-optimal set of indices
The theorem below shows that cone-restricted regularization paths can only directionally converge to an element of this set.
Theorem (Non-LOMM Regularization Paths Fail) Fix unit Euclidean norm vector E I& such that q & spmm. Assume that the token scores are distinct (i.e., Yu # Yir for & T) and the key embeddings k are in general position. Specifically, we require the following conditions to hold" +
When m d, all matrices K Rnnd where each row of K has the form ku- -ii, o a unique (i,a;, t + 0;) tuple, are full-rank.
When m = d 1, the vector of all ones is not in the range space of any such K matrix.
Fix arbitrary e 0, Ro O. Define the local regularization path of q as its (c, coDic neighborhood:
P(R) arg min L(p), where Ce.Ro co)e. n E Rd1 lpll Ro/. (olllll
(66)
Then, either limg- [|)(
lim ((R)) q. In both scenarios R-400
lim ((R) R # q-
R-100
We will prove the result by dividing the problem into distinct cases. In each case, we will construct an alternative direction that achieves strictly better objective than some 8 = 8(e) neighborhood of thereby demonstrating the suboptimality of the q direction. Let's define the 8 neighborhood follows:
Proof.
s and lPlll 2 Ro
N1 Ipll
(67)
Now, let's recall few definitions based Definition 3. First, the tokens selected by q are given by (59). To proceed, let's initially consider the scenario where is unique for all E [n]. meaning that as we let co, c q will choose single token per input. Later, we will revisit the setting when arg max is not a singleton, and is allowed to select multiple tokens.
Additionally, it's important to note that [p)() is non-decreasing by definition. Suppose it has finite upper bound [()( M for all R < In that scenario, we have limg (( 0 # q.
(A) q selects a single token per input: Given that the indices a (ar)) defined in (59) are uniquely determined, we conclude that the q direction eventually selects tokens kia, Recall the definition of the margin I from (60) and the set directional neighbors, which is defined as the indices that achieve I q. shown in (61). Let us refer to q as neighbor-optimal if ya Y in; for all (i,1) My-
We will consider two cases for this scenario: when q is neighbor-optimal and when is not neighbor- optimal.
(A1) is neighhor-optimal In this case, we will argue that max-margin direction from := mm (m)) can be used construct strictly better objective than q. Note that (a) exists because is already viable separating direction for tokens a. Specifically, consider the direction gmm Observe that, q' lies within cg).
1-6
1-
I+e
We now argue that, there exists 5 6, > 0 such that for all R.
min L(r q') R.SeRR
min PEN. Slplis
L(p).
(68)
*hii requirement holds for general data because it guaranteed by adding arbitrarily small independent gaussian perturbations to keys kg. As result, let prove the result for 2e without losing generality.
35

---

To prove this, we study the margin y induced by q' and the maximum margin I's induced within E Ns- Concretely, we will show that lg > I' and directly apply the first statement of Lemma 8 to conclude with (68).
Let = (/llp""(((r) [Ir be the margin induced by frmm. Note that I > rg by the optimality of (a) and the fact that q (a)) Consequently, we can lower and upper bound the margins via
T min min(kia - -ka)T Ty ee
r - -r 2
ranl 19a;
l+E
To max min min(kia. P.N, (e[s] tan.
ku) p/llpl
max min min Iriss (e[]] Fea.
kn) (q ôr)
r Mб,
where M = max [|k, - kirll- Consequently, setting 6 Tm (I rq). find that
E r 2 r3 4 -r)
Equipped with this inequality, we apply the first statement of Lemma 8 which concludes that for some R₂ = RC (g)) and all R > Re, (68) holds. This in turn implies that, within C., the optimal solution
either upper bounded by R. in l2 norm (i.e. limg- [[))] < 0o) or
P(R) at least 8 = 5(c) > 0 away from q after fz-normalization i.e. [|([R]]
qil 2 6.
In either scenario, we have proven that lim
PR) R-000 R
q.
(A2) is not neighbor-optimal. In this scenario, lim P(K)RR = 0 # To start, assume that conic neighborhood e of q is small enough so that selected-tokens remain unchanged within C.- This is without generality because if directional convergence fails in a small neighborhood of q, it will fail in the larger neighborhood as well. Secondly, if lom [|][
will prove that (R) is finite obtain
and P(R) Cz, since softmax will eventually perfectly select (i.e. assigning probability I on token indices (i, ())), we would have
Σ (l
lim L((() L. R
Note that, this simply by selection of a and regardless of P(R) directionally converges to q. This means that, if there exists finite p C. such that L(p) < L. (i.e. outperforming the training loss of [|p(R)]| + oo), then [|)(R)]] < co. This would conclude the proof.
Thus, we will simply find such a obeying L(p) < L.. To this aim, we first prove the following lemma.
Lemma Given a fixed unit Euclidean norm vector p. if all directional neighbors of p consistently have higher scores for their associated selected tokens, iee, Via; Yup for all (i, a;) and directional neighbor (i,B), then there exists R such that for all R > R,
L(R (RRp))= < L. lim L(R p).
(69)
R-s00
Proof. Define the maximum score difference K. supreintJza, 17:8 Yul. Also let Mp be the set of directional neighbors achieving the minimum margin rp (61). Define 1 K p- k., p. Define ö, to be the margin difference between the directional- neighbors and the second-most minimum-margin neighbors defined as
6p
min ieln), L.0M,,
I Tp-
(70)
'Here, we apply this lemma to compare q' against all p E N₂. We can do this uniform comparison because the R requirement in Lemma 8 only depends on the margin difference and global problem variables and not the particular choice of e No.
36

---

To proceed, setting min,gAEM, 7 Yju; > 0 and using (63), bound
[(Rp) S Σ
M e RT k n jelu),so,iLLMMM
Σ
My e --R RI " - Yu)
$ M RI M, K. RU, R(T,+P)
Consequently, have found that L(Rp) L. as soon
M_e .k K. Re
This happens when R 2 1/6g (up to logarithmic terms) establishing the desired statement.
Based this lemma, what need is constructing perturbation to modify q's directional neighbors and make sure all of them have strictly better scores than their associated selected-tokens. Note that, construct new candidate (say 4o = q perturbation), all sufficiently large scalings of go will achieve L(R qo) < L.. Thus, we can find a strictly better solution than L. for any norm lower bound Rg which enforced within the definition of C..
Lemma There are at most d directional neighbors i.e. [Mgl S d.
Proof. Directional neighbors are indices (i, obeying the inequality
(ki. kú TT
Declare to the matrix with rows obtained by these key differences k; kan. D E RRMa where M MMg.. We then obtain Dq -r where IM is the all ones vector. If M d, the equality Dq = -T IM cannot be satisfied because 1M not in the range space D by our assumption of general key embedding positions.
To proceed, |Myl= & and let D be as defined the lemma above. D also full-rank by our assumption of general key positions. We D to construct perturbation follows. Let M; Mg be the set of directional neighbor with strictly higher scores than their associated selected-tokens. In other words, all (j.p) Mt obeys
Y j/a""
Define the score difference x min 7 - -Y,,, > O. We know x > 0 because a is not neighbor- optimal. Finally, define the indicator vector of 1, with same dimension as cardinality [Mgl. 1, is for the rows of D corresponding to M and is 0 otherwise. Finally, set the perturbation as
q D'
where used the full-rankness of D during escudd-invvessien To proceed, for small &) O. consider the candidate direction do g + t We pick 6 = Ole) sufficienlly small to ensure qo C.. To finalize, let consider the margins of the tokens within gu- Similar Lemma set 6y = miM., I - r Let é = l|qull 1 = llg cug || - 1. Using definition q we have that
For (i,1) M2 we achieve margin of
(kin, - k.) + (/(1 6n) = r 6g
1 + to
For (i, Mt we achieve margin of Ie
T+m
For (i,1) Mg. setting
I1 I| sup Ilkin
kill. we achieve margin of most
r 8g K
(kin, - kk)) eog //11 + En)
37

---

In short, since én O(En), setting eD sufficiently small guarantees that M; is the set of directional neighbors of qu. Since M has strictly higher scores than their associated selected-tokens, applying Lemma 9 on qo shows that, _(R qo) L. for sufficienlly large R implying [(ş()))) < co.
(B) q selects multiple tokens for some inputs i E [n]] In this setting, we will again construct perturbation to create a scenario where qo = q + Eq selects a single token for each input i [n]. We will then employ margin analysis (first statement of Lemma 8) to conclude that go outperforms 6 < e neighborhood of q.
Let I C [n]] be the set of inputs for which selects multiple tokens. Specifically, for each i I, there T, [T] such that [T,] 2 and for any I and 0 Ti,
k k arg max kg.
x477
From these multiply-selected token indices let select the highest score one, namely, Bi arg maXger, Yie for E I. Now, define the unique optimal tokens for each input as a e R" where a, B, for E I and a, = arg maX K q for i e I. Define L. = > (( as earlier.
Secondly, construct perturbation q show that qu qu To see this, define the matrix D where each (unique) row is given by kin - ka where E T, 0 * a; E I. Now note that, (kin, kia) q 0 for all positions, this implies that D has at most d - I rows and, thus, its rows are linearly independent. Consequently, choose q D"1 where t denotes pseudo-inverse and is the all ones vector, Also let rg be the margin of directional margin of that is
select tokens asymptotically.
Tú 7. ar i I. Since keys are in general
rg
min (kin, -ky) q. jEInl.eTT
With this choice and setting K = sup llkia, kirll: we have that
For 0 T # aj: (kia, k qo = (kia, kw) qq' ED. For all other (i. with # a;; (kin. kp) qo z -g- lleo
Choosing 4 < /[g (1 + Kllq II), together, these imply that,
(a:))__ is the litected-tokenss of qu- achieves a directional margin of
60
6D
go
qoll
1 + collq
qo is neighbor-optimal because directional neighbors are E Tú, A # a; and Y < Yia,-
Note that, these conditions lay the groundwork for applying the first statement of Lemma 8 with p - qo- We next explore the optimal directions within No and show that qo strictly outperform them in terms of training loss.
To proceed, given small & cp. let us study Ls min Rslpl L(p). Here, recall that p has 5-small directional perturbation around q which can modify the token selections by breaking the ties between the multiply-selecte assumption, for large R < IIpll, the optimal E N₂ is guaranteed to (uniquely) select indices a; E Ti Because all other p directions which, asymptotically. either select other tokens 0 Ti, # a or split the probabilities equally across subset of T, achieve a larger loss. For instance, q will split the probabilities equally across T; to achieve an asymptotic loss of
token indices by g. However, thanks to the distinct token score
L == lim L(R 9)
Σ ((ri.)) ieT
Σ
1 FET
e
Σ
V
BET,
Thus, L. > L. because ( Eeer, y l (y ((ymr) where a; has the highest score i.e. Yin,
i7 EoeT, Via Set P(R) = arg min peisllllll L(p). Consequently, there are two scenarios are:
lom [|]( is finite. This already proves the statement of the theorem as P(R)/R - within 5 ncighborhood of q.
38

---

For sufficiently large R, the selected-tokens of ((R) (a=))
Proceeding with the second (remaining scenario), we study the directional margin of p(R). More broadly, for any p E Ns and p = p/|lpll with selected-tokens a, using the fact that | qii s 8 & 60. we can bound the directional margin as
For 0 Ti 1. # aj: (kin, - p (kio; - k (p- - q) S Kö.
For all other (i.1) with 1 # a;: (kir. kx) p lg Kő.
This means that, any such p E N₂ achieves a directional margin of most
rp 5 K6.
Applying Lemma and setting ö O(ev). this implies that for
R
0 60
1tailg Kō
gn
we have that L(R qo) -i, L(p). Since this holds for all R. (68) holds (similar to Case
(A1)) and we conclude that whenever [[p(R)) + oo, it doesn't directionally converge within N (i.e. 8 0 neighborhood of proving the advertised result.
B.6 Proof of Lemma
We prove a slightly general restatement where we require v E ran) instead of full-rank W.
Lemma 11 Suppose for all i E [n] and i # opt, Y; = 1 and yu Yiopt, Also suppose v range(WTT). Then, p mam exists i.e. (ATT-SVM) is feasible for optimal indices a;; t- opt,
Proof. To establish the existence of pmm=, we only need to find a direction that demonstrates the feasibility of (ATT-SVM), i.. we need to find that satisfies the margin constraints. To begin, let's define the minimum score difference:
min Yiopt, reppt,
Yu-
We then set p = yWT) where f denotes pseudo-inverse. By assumption WT = v To conclude, observe that is a feasible solution since k,, Wxx, and for all i E [n] and # opt, have that
(kopt, kn) p = (Xiopt, xxx W p = y (Niopt, - xn) - T WT WWWT))t
(Kiopt, Nv 1,
which together with the constraints in (ATT-SVM) completes the proof.
C
Addendum to Section 3
C.1 Proof of Theorem 5
Proof. Suppose the claim is incorrect and either PR/R or fails to converge as R, r grows. Set E 1I = 1/ pmm REp" and gmm mm The proof strategy is obtaining
(R),P)) contradiction by proving that (fmm. p is strictly better solution compared to (v,, pr) for large R. Without losing generality,
will set a, = I for all e [n] as the problem is invariant tokens' permutation. Define of = S to be the amount of non-optimality (cumulative probability of non-first tokens) where x.. = S(K,p) is the softmax probabilities.
Case 1: P does not converge. Under this scenario there exists 6, y >(0)== > O such that we can find arbitrarily large R with ||pp/R - pum /R|| 2 and margin induced by PR/ R is at most E(1 - y) (from strong convexity of (ATT-SVM)). Following of definition above, set qmax suppe(n) g to be
39

---

worst non-optimality in pe and 4max suPieint " to be the same for pmm.. Repeating the identical argument in Theorem 8 (specifically (84)), we can bound the non-optimality amount af of pmm as
Lnno, exp(k x, m
E exp(k,, exp(k inc Bumm S eep((-RE).
t ExT exp(k; pmmm)
(71)
Thus, gmax = maXi q S TT ee Next without losing generality, assume first margin constraint y-violated by Pn and min (K1a k1)) pr ER(1 y). Denoting the amount of non-optimality of the first input as oh",, we find
Eren exp(k
1 Eitan exp(k (., PR)
9tth
T exp(-(1- R)))).
(72)
EmerTy exp(k
We similarly have gmmx 2 T-1 (xp((RE) to find that
log(smax) 2 (1 y)ER log T,
-ER log s log( S -ER + log
(73)
In words, pml contains exponentially less non-optimality compared to Pn as R grows. The remainder of the proof differs from Theorem as we need to upper/lower bound the logistic loss of (pom prim) and (v,, PR) respectively conclude with the contradiction.
First, let us upper bound the logistic loss of (f pmm",, (i""., Set r, = S,K, Observe that ||r: - s E, we have that w satisfies the SVM constraints on ri with Y, r y 2 e/T. Consequently, setting 6max = suPreinl E, y achieves a label-margin of r Emax on the dataset (Y, Tiliate- With this, we upper bound the logistic loss of (5mm, pmm) as follows. Let M suPiet]Aars(T] [| - Xirll. Let us recall the fact (73) that worst-case perturbation is
Emax S M Mexp(-ER + log 7) = MT exp(-ER).
This implies that
[(
m max log(1 + exp(-exxi Yir (mmm))m))) fmim)).
max exp(-Y/ r gmin) (D))
exp(-rl + r6max)
e MT Exp(-ER) e-/
(74)
Conversely, we obtain a lower bound for (v,, PR). Set X, 3(Kpp)) Using Assumption C, we find that solving (SVM) on (Y, rieeml achieves at most I ve (1- /T margin. Consequently, have
E(,,, PR) 2 max log(1= + cxp((Y,rr.
1
expeeepppxXp ,,, v,) log
2n le(n
2 cxp( ve 11 11 (T)) A log
2n
d1/Thcxxp-----EEM o 2.
(75)
2n
Observe that, this lower bound dominates the previous upper bound when large, namely, when (ignoring the multiplier 1/2n for brevity)
MT
-1-WER MTe- ER
=
R
Ry
:=
log y
Thus, we indeed obtain the desired contradiction since such large R is guaranteed to exist when Pa/R pm
Case 2: v./r does not converge. This is the simpler scenario: There exists > 0 such that we
can find arbitrarily large obeying [|v,/r mm Milll 2 5. If [ellpgRR m 0, then "Case
40

---

1" applies. Otherwise, have | E 0, thus can assume ||pa/R =p for arbitrary choice of E > O.
On the other hand, due to the strong convexity of (SVM), for some y= y(0) 0. achieves margin of at most (1 - -r)rrr on the dataset (Y, x))em-- Additionally, since |lpR/ R - =p 1 E, pR
strictly separates all optimal tokens (for small enough > 0) and @ma := f(e) 0 as R co.
Consequently, setting r; = x p(K) for sufficienlly large R 0 setting M = suPee[)).e(7] ||xا||| we have that
min Yv, r. min Y1r, xm + suplr, (n xa)) reln
(E(m]
ic(nl
y) Mf(e)r
(1 y/2)rr.
(76)
This in turn implies that logistic loss is lower bounded by (following (75)),
-T e A log2.
1
L(v, PR)
2n
Going back to (74), this exponentially dominates the upper bound of (pum, pam) whenever rMT exp(-ER) ryl (that is, whenever R, are sufficiently large), again concluding the proof.
C.2 Proof of Theorem
We will prove this result in two steps. Our first claim restricts the optimization to the particular quadrant induced by mingea (Kin, k₂) PR 2 O under the theorem's condition p(K,p)),,
Lemma 12 Suppose SRK, - 1. Then, there exists Ro such that for all R 2 Ro, we have that, min (km.. T
Pk 2 O, for all iE [n].
(77)
Proof.
Suppose the claim does not hold. Set s = 3(K,p)). Fix Ro such that s 2 0.9 for all R 2 Ro-
On the other hand, there exists arbitrarily large R for which (kin; ki²) pR < O for some # a; E [7]
and E [n]. At this (R, i, choices, we have that su 2 g Since s + s s 1, we find s < 0.5 in
which contradicts with s 0.9.
Let Q be the set of p satisfying the quadrant constraint (77) i.e. indices (a:); are selected. Let hg be the solution of regularization path of (v. p) subject to the constraint Q. From Lemma 12, know that, for some Ro and all 2 Ro. hR PR. Thus, if the limit exists, have that limg-co hg/R~ R limg-oa PR/R.
To proceed, will prove that limg-ss hg/R R exists and is equal to plaxx mlayy and simultancously establish v,/r > ymmm///y /Ily""mi
Lemma 13 limg- Rg/R- pllax clap and lim vr/r pmu mllm
Proof. The proof will be similar to that of Theorem 5, As usual, we aim to show that SVM-solutions constitute the most competitive direction. Set rlar
Case 1: hg/ R does not converge. Under this scenario there exists d = y(0) > 0 such that we find arbitrarily large R with |/h//R rlax]] 8. This implies that margin induced by hg/R at most E(1 ((1-y) y) over the support vectors (from strong convexity of (10).. The reason is that, hg satisfies (k,, k2) 0 for all a, by construction hg Q. Thus, constraint over the support vectors have to be violated (when normalized the same 62 norm as |laxy 1/3).
As usual, we will construct a solution strictly superior to hg and contradicts with its optimality.
Construction of competitor: Rather than using pllax direction, we will choose slightly deviating direction that ensures the selection of the correct tokens over non-supports S. Specifically, consider the solution of (10) where we tighten the non-support constraints by arbitrarily small > 0.
e-tx arg min ||l|| such that p" pkii.. kn) 2 { for all 1 # a, i S
for # ar i S
(78)
41

---

Let pmm be the solution of (ATT-SVM) with (a:) (which assumed to be separable). Observe that pt Epmn + (1 p satisfies the constraints of (78). Additionally, pt would achieve margin of (1-o/E+</A A+dE-A)
AE where A 1/ Using optimality of pf this implies that the reduced margin E. = 1/ (by enforcing e over non-support) over the support vectors is a Lipschitz function of c. That is E. 2 E cM for some M 2 0. To proceed, choose an > 0 such that, it is strictly superior to margin induced by hg, that is,
E 2 3 - -
To proceed, set pf = RE.p r/s Let us recall the following notation from the proof of Theorem s = S(K:p) and o. = - Sin, Set Qmux max q," to be worst non-optimality of he over support set. Similarly, define Fmax = maxis sfttt to the same for pr Repeating the identical arguments to (71), (72), (73), and using the fact that p achieves a margin E(1 Y) s E, S E. we end up with the lines
log ER log T s log(gmax) - 0.5y) + log T.
log(gmax)
(79a)
(79b)
In what follows, we will prove that prank achieves strictly smaller logistic loss contradicting with the optimality of Pn (whenever ||/// R - | 2 6).
Upper bounding logistic loss. Let us now upper bound the logistic loss of (i pe-1t)) where pmm = mmm with ym being the solution of (SV with r, + Xaa and I //]l Set r, = X S 5(K )r Set v= = minges Y, Y p I to be the additional margin buffer that non-support vectors have access to. Also set = suPie(a)ere(r] ||x - Xirll. Observe that we can write
Σ St((inn Xzx) |x rell q.M.
Xio; r
reo;
Non-supports achieve strong largln Using above and (78) for all i e S and # a; have that S Rsa e
-1(1--//2/R s/a, Consequently, whenever R 2 Ro := (e3(1 - (E2)) log( Tv
TM
Ei Su
Tv 5 M
s Te
Sia.
This implies that, on i
Y r.r.ymm y mm +I+++Y- xxa,) T + M 2 1.
(80)
In words: Above fixed RS that only depends on y y(8), features r; induced by all non-support indices i E S achieve margin at least 1. What remains is analyzing the margin shrinkage over the support vectors as Theorem 5.
Controlling support margin and combining bounds: Over S, suppose pm satisfies the SVM con-
straints r; with Y, r. pm 1 E/TT. Consequently, setting Enux = sup E, y achieves label-margin of r 6mux on the dataset (Y,, r:)ietn- Next, we recall the fact (79b) that worst-case perturbation is 6mux M exp 0.5y) log 7') MT exTxxxx( 0.5 With this and (80), we upper bound the logistic loss of (pmm F as follows.
[(prmm m max log(1 + exp(- pmnn)).
< max cxp(- Y pmm)) (an)
exp( + TEmax)
(81) Conversely, we obtain a lower bound for (r,, ht).. Set r, = X (K,,)) Recall the lower bound (79a) over the support vector set S. Combining this with our Assumption C over the support vectors
MT -1p((=1----055
42

---

of (SVM) implies that, solving (SVM) (Y, ri)epp achieves most I ve-th-WER/7 /7 margin. Consequently, we have
L(v, 2 max log(1 + xxx,(-Yrr,
iE|nl
2
max exp(- v,) A log2
2n
exp(-r(TT ve (1-)ER/ (T)) log
2n
10/T)exx((--1EER)) e- log
(82)
2n
Observe that, this lower bound dominates the previous upper bound when R large, namely, when (ignoring the multiplier 1/2n for brevity)
2 y
MT
(r/T)e MTe 511--51
R2
Rn
log
Thus, we obtain the desired contradiction since pe-ix is a strictly better solution compared to pr = hg (once R is sufficiently large).
Case 2: v,/r does not converge. This is the simpler scenario: There exists 5 > such that we can find arbitrarily large obeying |/v,/r ymm /||y!!||||| 2 6. First, note that, due to the strong convexity of (SVM), for some == y(0) > 0, v, achieves a margin of at most (I - y)r on the dataset (Y, xn hetal provided that px),,
By theorem's 's condition, choice of c = y/3 > 0, above some sufficiently large (7o, Ro), we have that ||x - rill S e. Following (81), this implies that, choosing fum = my mmm achieves a logistic loss of at most e'1B Again using [|x rill s e, for sufficiently large R) have that
This immediately implies that, for any
min Yiu r; S min Y,v, xa + suply, (r: xa) eE(s]
(e(m]
- y)+
(T 2y/3)r.
This in turn implies that logistic loss is lower bounded by (following (82)),
1 _ 2013 e TT log 2.
L( p&)
2n
This dominates the above upper bound e/v/3 e- of pam whenever L >1 = r> 1 log(2n).
(that is, when is sufficiently large), again concluding the proof.
D Regularization Path of Attention with Nonlinear Head
So far discussion has focused on the attention model with linear head. However, the conceptual ideas on optimal token selection via margin maximization also extends to a general nonlinear model under mild assumptions. The aim of this section showcasing this generalization. Specifically, we consider the prediction model f(X) ((X (OXTS((K) ((p)) where v/(-) Rd - R generalizes the linear head r of our attention model. For instance, following exposition in Section 1. vb(-) can represent multilayer transformer with p being tunable prompt at the input layer. Recall that (X, K,, YD is the dataset of the input-key-label tuples. We consider the training risk
1 L(p) = n Σ EYY. V(XT where S = =(K.)))RR
(83)
The challenge with nonlinear v/(-) that, we lack clear score function (Def. 1) unlike the previous sections. The assumption below introduces generic condition that splits the tokens of each X, into an optimal set O, and non-optimal set O = [7] O, words, non-optimal tokens are those that strictly increase the training risk L(p) if they are not fully suppressed by attention probabilities s..
43

---

Assumption D (Mixing non-optimal tokens hurt) There exists sets (OiYi [T] as follows. Let df Lmo li be the sum of softmax similarities over the non-optimal set for p. Set q maxxe[[] J.- For any A > 0, there exists p < O such that:
For all p. p' E R., if log(4ma) s A) log( A p. then L(p) L(p').
This assumption is titled mixing hurts because the attention output X; s. is mixing the tokens of X, and our condition is that, to achieve optimal risk, this mixture should not contain any non-optimal tokens. In particular, we require that, model p that contains exponentially less non-optimality (quantified via log(4max)) compared to p' is strictly preferable. As we discuss in the supplementary material, Theorem 1 in fact a concrete instance (with linear head v) satisfying this condition.
Before stating our generic theorem, we need to introduce the max-margin separator towards which regularization path of attention will converge. This is a slightly general version of Section 2's (ATT-SVM) problem where we allow for a set of optimal tokens O; for each input.
mm arg minl subject to max min p"kki. - kg) 2 1, for all iE [m]. (ATT-SVM')
oeO; BeO,
Unlike (ATT-SVM), this problem is not necessarily convex when the optimal set O; is not singleton. To see this, imagine = d = and T = 3: Set the two optimal tokens as K = and k = = and the non-optimal token as K3 = O. The solution set of (ATT-SVM') pmm 1-1, 1) whereas their convex combination = O violates the constraints. To proceed, our final result establishes the convergence of regularization path to the solution set of (ATT-SVM') under Assumption D.
Theorem Let Gtnm be the set of global minima of (ATT-SVM'). Suppose its margin := 1 1/Ilp"''|| > O and Assumption D holds. Let dist(:.-) denote the Ei-dittance between vector and
set. Following (83), define P(R) arg pin L(()). We have that lim dist
((R
Gmm)= 0.
ER
We note that Theorem is a corollary of this result where O,'s and GTLm are singleton. Based on this result, with multiple optimal tokens, Theorem I would gracefully generalize to solve (ATT-SVM').
D.1 Proof of Theorem 8
Proof.
The key idea is showing that, thanks to the exponential tail of roftmax-attentien (harmful) contribution of the non-optimal token with the minimum margin can dominate the contribution of all other tokens R co. This high-level approach is similar earlier works on implicit bias of gradient descent with logistic loss [31, 22].
Pick pTT E gmm and set P
=
T.uuu This will be the baseline model that pa has to compete against.
pm
Also let ğR ER. Now suppose dist (PR. Gumm + 0 as R - Then, there exists ō > 0 such that, we can always find arbitrarily large R obeying dist (pR. Gumm) 2 5.
Since px is 8 > 0 bounded away from Gmm, and || = II=""I,, px strictly violates at least one of the inequality constraints in (ATT-SVM'). Otherwise, we would have PR @ Without losing generality, suppose Pn violates the first margin constraint, that is, for some == y()) > O, max minged, P3 (k1,- - kıØ)= s 1 y. Now, we will argue that this will lead to a contradiction as since we will show that [(p))) < L(px) for sufficiently large
First, let us control L(pt). We study s = 3(K,p)) and let a; E O, be the index a in (ATT-SVM') for which margin, max mingoo, (ki.- kia) pauni 2 I is attained. Then, we bound the non-optimality amount o; of P as
tt Σiari exp(k,,ps, S Lnco, exp(k,,pp)
Eno, exp(k (T Px)
S T exp(-ER).
exp(k ic, Pf)
Thus, G = max of T exp(-ER). Secondly, we wish to control L(px) by lower bounding the non-optimality in PR. Focusing on the first margin constraint, let E O be the index in (ATT-SVM') for which margin, y is attained. Denoting the amount of non-optimality of the first input as Q1. we find?
Enso exp(k),pt) 1 ErEo, exp(kT,PR) T l exp(pEE(( -)).
Φ Eeiry exp(k T exp(k (Pp)
Here, assumed margin is non-negative i.e. &T. supRo, &R. Otherwise, suP &T.px is attained in O which implies T-1, Thus, can still use the identical inequality (84) with the choice y 1.
44

---

We similarly have 9mxx exp(-ERR In conclusion, for PR. Pk= denoting maximum optimality by 4max 2 91 and m we respectively obtained
log( 2 (1-7) log T,
(ER) log T' I log( -ERR log
(84)
The above inequalities satisfy Assumption D follows where p P and p' < Pg: Set Ro 3y log 7' so that log 7
Secondly, set po = - - log T. This way, 2 log(dmax)
3
implies R Ry and log 7' y
Using the latter inequality, bound the log 7' terms to obtain
1og(4max) 2 - (1 - 2y/ (ER),, and log(4mm>> 5 -1 - (ER)
To proceed, we pick 1+A 1-y/3
T-2y/3 implying A := 3'Zz. Finally, for this A, there exists p(A) which we need ensure log(4max S p(A). This can be guaranteed by picking sufficiently large R that ensures log(4ma)) S -(1 - (ER) S p(A) to satisfy all conditions of Assumption D. Since such large exists by initial assumption dist (Px. Gmmm + O, Assumption D in turn implies that L(pp) < L((n)
contradicting with the optimality of PR in (83).
D.2 Application Linearly-mixed Labels
The following example shows that if non-optimal tokens result in reduced score (in terms of the alignment of prediction and label), Assumption D holds. The high-level idea behind this lemma is that, if the optimal risk is achieved by setting & O, then, Assumption D will hold.
Lemma 14 (Linear label mixing) Recall of = End, sy from Assumption D. Suppose Y; () and
Y, (X) 5t) v/(1 ,,) + Zi,
for some (v;),_1 > O. Here Z4 some C, > and for all E R4, assume
Z(p) is the contribution of non-optinne con-optimal tokens to prediction. For
-C $ Z, S (1 e)v.dg.
(85)
Then, Assumption D holds for L(p) E. ((Y, (X7 s²)) when E(-) is strictly decreasing loss function with continuous derivative.
Proof. Recall the assumption Y, (X7 s v1(1 - Zi with Z, obeying (85). Let also write the loss function
[(p) Σ (vv(1 - 5t) Z,).
Define PPnx supssjn G.P. Let M be the maximum absolute value of score over tokens. Let
B max (x) A min -x' 0. ISM
R
Through Taylor's Theorem (integral remainder), we have that Blq.v; Z4) 2 ((v((1 d.) Z4) ((v)) 2 Alc.vi 2,) 2 eAvig."
Set L. = +122" ET ((v)). Set C. = B(C + max v,) and C_ = n Ae min v,. This also implies
C. 2 Σ
Blg.:: - 7)) L(p) L. 2
AlgP1: - 75) Felm)
Σ
Fe(m)
Σ
eAvigf 2 C- C-q
Thus, to prove L(p') L(p), simply need to establish the stronger statement C- C. Grans
45

---

Going back to the condition of Assumption D, any log(<max) (1 A) log(4 obeys Lmax
(max))+s i.e. Lmax 2 ( )(++)) Following above, we wish to ensure gmnxx > Og for such (p. p¹)
pairs where O = C+/C- > 1. This is guaranteed by
A
(amax (1+a) = 1+A A log(dmax) -log(O).
The above is satisfied by choosing p(A) := -2(1 + A log(O) in Assumption D. Thus, all p.
with log(d s p = p()) satisfies the condition of Assumption finishing the proof.
E Implementation Details and Additional Experiments
In this section, we provide implementation details and additional experiments.
We build one attention layer using PyTorch. During training, we SGD optimizer with learning rate 0.1 and train the model for 1000 iterations. To better visualize the convergence path, we normalize the gradient of p (and v) at each iteration.
Next, given the gradient solution p. we determine locally-optimal indices to be those with the highest softmax scores. Using these optimal indices, we utilize python package cvxopt to build and solve (ATT-SVM), and then get solution pmm After obtaining prlm, we also verify that these indices satisfy our local-optimal definition. The examples use in the paper are all trivial to verify (by construction).
In Figures 3()) and 3(b), m (blue dashed) is solved using python package skl earn. svm via (SVM) based on the given label information, and red dashed line represents p pelax direction instead, which is the solution of (10). Note that in both figures, ym /|| = [0. 1J. Therefore, in Figure 3(a) all optimal tokens are support vectors and p pelax pmm Whereas in Figure 3(b), yellow * is not a support vector and only needs to satisfy positive correlation with p. Gray dashed line displays the plt! direction.
Failure of gradient descent's global convergence when π 2 2 (refer to Theorem 2). Figure 8 provides a counter-cxample demonstrating that the n = restriction is indeed necessary and tight to guarantee global convergence of the gradient descent iterates p(r + 1) = p(t) L(p(() on (ERM).
For this example, we use logistic loss in (ERM). We set 7 = d = 2, implying that there is only one non- optimal token, thus Assumption B is satisfied. The red and blue lines represent GMM and LMM solutions, respec- tively. We note that the green star and teal square indicate the locally-optimal tokens. Specifically, referring to the local optimality definition (Definition 2), for LMM solution (p"")) represented by the blue line, the square teal token does not have any SVM--nighbors. The arrows indicate the two trajectories originating from different initializa- tions. This demonstrates that the gradient descent iterates p(t + 1)= p(1) L(p()((() (ERM) with two different initializations converge to two different SVM solutions (GMM and LMM). Results validate the necessity of и 1 in Theorem to provide the gradient descent convergence to mun from any initialization.
5
Optimal tokens
Nor-opt tokens
3
2
0-
-1
-2
-4
-2
2
Figure 8: The convergence behavior of the gradient descent on the atten- tion weights p using the logistic loss in (ERM) with n = T = =d= 2.
The convergence behavior of gradient descent under over-parameterization. To illustrate Theorems 3 & 4, we have investigated the convergence behavior of p(t) generated by gradient descent in Figure 9(a), using n = 4, T 6, and conducted 1,000 random trials for varying e 5, 10, 100, 300, 500). These experiments use normalized gradient descent with learning rate I for 1000 iterations. Inputs X2 and the linear head y are uniformly sampled from the unit sphere, while Y, is uniformly + 1, and W is set to I.
The bar plot Figure 9()) distinguishes between non-saturated softmax (red bars) and saturated softmax (other bars). Saturation is defined as average softmax probability over tokens selected by gradient descent are at least 10-5 and implies that attention selects one token per input. Note that, whenever the norm of gradient descent solution is finite, softmax will be naturatu.. For
46

---

100
((Kp(t)) <1 p(t) mm
d= 10
80
Percentage
p()) - p
Probability
0.3
d= 100
60
0.2
40
0.1
20
0.0
2
5
10
100
300
500
-0.2 0.0 0.2 0.4 0.6
1.0 1.2
Min score gap to neighbors (2)
(a) Perc. different convergence scenarios for p())
(b) Prob. of == min Vic. 7
Figure 9: Convergence analysis of p(r) trained with random data using gradient descent. (a) shows three scenarios: (1) attention failing to select one token per input (i.e. softmax is not saturated); (2) p converging towards pmmm, and (3) p equating to pmmm with red, blue, and green bars, respectively. Considering saturated softmax instances where p() selects one token o; per-input, (b) presents histogram the minimal score gap between o; and its corresponding neighbors Tj
small (e.g., 2). problem has small degrees of freedom to separate optimal tokens from the rest (i.e. SVM solution for LMM directions) especially due to label randomness. This results in a tall red bar capturing the finite-norm solutions. However, for larger d, we observe that softmax saturates (i.e. [|p()]] - oo) and we observe that the selected tokens almost always converges to an LMM direction (blue bar)- this is in line with Theorems
We also study the convergence to the oobaill--ptimall GMM which is represented by the green bar: GMM is strict subset of LMM however as d increases, we observe that the probability of GMM convergence increases as well. This behavior is in line with what one would expect from over- perrparaetttrzee deep learning theory [57, 58, 59, 60] and motivates future research. The average correlation coefficient between p()) and its associated LMM/GMM direction is 0.997, suggesting that, whenever softmax saturates, gradient descent indeed directionally converges to a LMM solution E pmmm, confirming Theorem 4.
Furthermore, we found that there exist problem instances, with saturated softmax and l!p())) - co, that do not converge to either LMM or GMM. We analyzed this phenomenon using the minimum score gap, y == minje(n)re7; 7 iz: Yu: where T,i e [m]. represents the sets of SVM-neighbor tokens. Figure 9(b) provides the probability distribution of y (with bins of width 0.01) and demonstrates the rarity of such cases. Specifically,
found this happens less than 1 1% of the problems, that is, Prob(y 0) < 0.01. Figure 9() also reveals that, in these scenarios, even if y < 0, it is typically close to zero i.e. even if there exists a SVM-neighbor with a higher score, is only slightly so. This not surprising since when token scores are close, we need a large number of gradient iterations to distinguish them. For all practical purposes, the optimization will treat both tokens equally and rather than solving (ATT-SVM), the more refined formulation (ATT-SVM') developed in Section D will be a better proxy. Confirming this intuition, we have verified that, over the instances 0, gradient descent solution is still > 0.99 correlated with the max- max-margin solution in average.
In Figure 9, again applied normalized gradient descent with learning rate equal to 1. Each trial involved randomly generated data and train- ing for 1000 iterations discussed in Theorem 4. The
probability
1.0
tokens selected by were denoted as (ar)) where an arg max,e(r] s(X,p). The averaged softmax prob- abilities were calculated as f == Σ" S(X,p),; (same Figure 3(c)). The red bars in Figure 9(b) represent the values of P55 < 1 10-5) for each choice of d. Figure 10 displays the cumulative probability distribution of y from Figure 9b). with the gray dashed line indicating y = 0. From this figure, we observe that the minimal score gap exhibit a sharp transition at zero (<1% of the instances have y < O), demonstrating that, in most random problem instances with ||pl| - co (5 1). problem directionally
0.8
0.6
Gumulative
0.4
0.2
10
100
0.0
Figure 10: Cumulative prob. of the gap
y: minietnee,, Via, - 7i in Figure 9(b).
47

---

converges LMM i.e. p() / We believe the occurrence of negative score gap is due to small score differences (so that optimal token is not clearly distinguished) and finite number of gradient iterations we run. Interestingly, even in the negative score gap scenarios, gradient descent is aligned with (a) (even if (a) is not LMM) which can be predicted from our Section D which handles the scenario where there are multiple optimal tokens per input.
F Addendum to Section 5
We provide an overview of the current literature on implicit regularization and attention mechanism.
F1 Related Work on Implicit Regularization
The introduction of Support Vector Machines (SVM), which utilize explicit regularization choose maximum margin classifiers, represents onc of the carliest relevant literature in this field [61]. The concept of maximizing the margin was later connected generalization performance [62]. From a practical perspective, exponential losses with decaying regularization exhibit asymptotic behavior similar to SVMs, as demonstrated in [22]. While the analysis of the pereeptron [63] originally introduced the concept of margins, the method itself does not possess an inherent bias it terminates with zero classification error. However, establishing a meaningful lower bound for the attained margin is not possible. Initial empirical investigations highlighting the implicit bias of descent methods focused on C reguularizztion revealing that coordinate descent, when combined with the exponential loss, exhibits inherent inclination towards { rregularized solutions [64].
This work draws extensively from the literature on implicit bias and regularization, which has provided valuable techniques and inspiration. A common observation in these studies is the convergence to a specific optimal solution over the training set. This phenomenon has been observed in various approaches, including coordinate descent [65, 66], gradient descent [30, 67, 25, 68, 69, 22, 70], deep linear networks [71, 72], ReLU networks [73, 74, 29, 75, 76], mirror descent [77, 78, 33, 36), and many others. The implicit bias of gradient descent in classification tasks involving separable data has been extensively examined by [22, 25, 26, 27, 28, 29J. The works on classification typically utilize logistic loss or exponentialll-ttaled losses to establish connections to margin maximization. The results have also been extended to non-separe separable data by [30, 31, 21]. Additionally, several papers have explored the implicit bias of stochastic gradient descent [37, 38, 41, 42], as well adaptive and momentum-bant based methods [43, 44, 45, 46].
While there are some similarities between our optimization approach for and existing works, the optimization of P presents notable differences. Firstly, our optimization problem is nonconvex and involves composition of loss and softmax, which introduces new challenges and complexities. The presence of softmax adds a nonlinearity to the problem, requiring specialized techniques for analysis and optimization. Secondly, our analysis introduces the concept of locally-optimal tokens, which refers to tokens that achieve locally optimal solutions their respective attention cones. This concept is crucial for understanding the behavior of the attention mechanism and its convergence properties. By focusing on the cones surrounding locally-optimll tokens, we provide a tailored analysis that captures the unique characteristics of the attention model. Overall, our work offers novel insights into the optimization of attention-baa -based models and sheds light on the behavior of the attention mechanism during training.
F.2 Related Work Attention Mechanism
As the backbone of Transformers [6], the self-attention mechanism [47, 48, 49, 50] plays a crucial role in computing feature representations by globally modeling long-range interactions within the input. Transformers have achieved remarkable empirical success in various domains, including natural language processing [4, 2]. recommendation systems [79, 80, 81], and reinforcement learning [82, 83, 84]. With the introduction of Vision Transformer (ViT) [85], Transformer- based models [86, 87] have become strong alternative convolutional neural networks (CNN) and become prevalent in vision tasks.
However, the theoretical foundation of Transformers and self-attention mechanisms has remained largely unexplored. Some studies have established important results, including the Lipschitz constant of self-attention [88], properties of the neural tangent kernel [89, 90], and the expressive power and
48

---

Turing-commpeteness. of Transformers [91, 92, 93, 51, 23, 94] with statistical guarantees [95, 96]. There is also a growing effort towards a theoretical understanding of emergent abilities of language models such as in-context learning [97, 98, 99] and inherently related to the models ability to attend to the relevant information within the input sequence.
[100, 101, 102] which are
Focusing on the self-attention component, Edelman et al. [51] theoretically shows that single self-attention head can represent a sparse function of the input with a sample complexity for the generalization gap between the training loss and the test loss. However, they did not delve into the algorithmic aspects of training Transformers to achieve desirable loss. Sahiner et al. [52] and Ergen et al. [53] further explored the analysis of convex relaxations for self-attention, investigating potential optimization techniques and properties. The former work applies to self-attention with linear activation (rather than softmax) whereas the latter work attempts approximate softmax via linear operation with unit simplex constraints. In contrast, we directly study softmax and characterize its non-convex geometry. In terms of expressive ability, Baldi and Vershynin [54] investigated the capacity of attention layers to capture complex patterns and information, while Dong et al. [23] illustrates the propensity of attention networks degenerate during the training process, with the result often being an output that is approximately rank-1 1 matrix.
Recent works have made progress in characterizing the optimization and generalization dynamics of attention [55, 56, 103, 17, 104]. Jelassi et al. [55] studied gradient-based methods from random initialization and provided theoretical analysis of the empirical finding that Vision Transformers learn position embeddings that recapitulate the spatial structure of the training data, even though this spatial structure is no longer explicitly represented after the image is split into patches. Li et al. [56] provided theoretical results on training three-layer ViTs for classification tasks. They quantified the importance of self-attention in terms of sample complexity for achieving zero generalization error, as well as the sparsity of attention maps when trained by stochastic gradient descent (SGD). In another related work, Nguyen et al. [104] proposed a primal-dual optimization framework that focuses on deriving attention as the dual expansion of a primal neural network layer. By solving a support vector regression problem. they gained a deeper understanding and explanation of various attention mechanisms. This framework also enables the creation of novel attention mechanisms, offering flexibility and customization in designing attention- based models. In another closely related work, Oymak et al. [17] analyzed the same attention model as ours, denoted by (ERM). Specifically, they jointly optimize v. p for three gradient iterations for contextual dataset model. This is in contrast to our emphasis on infinite-iteration behavior of p-only optimization. However, it is important to note that all of these works make certain assumptions about the data. Specifically, they assume that tokens are tightly clusterable or can be clearly split into relevant and irrelevant sets. Additionally, Li et al. [56] require specific assumptions on the initialization of the model, while Jelassi et al. [55] consider simplified attention structure where the attention matrix is not directly parameterized with respect to the input.
In contrast, our work offers a comprehensive opeimization-thhoreiicc analysis of the attention model, establishing formal connection to max- max-margin problems. While comparable works make assump- tions on the dataset model, our results apply under minimal assumptions for general data and realistic conditions. Our analysis based on moxx-margin-equivaleee of the optimization geometry of attention and its behavior during the training process. As articulated in our experiments, our results lead to novel insights even for n = 1,2 samples, = 2,3 tokens and d 2, 3 dimensions (in contrast to [55, 56, 103, 17J). Notably, theoretical understanding of the implicit bias exhibited by gradient descent methods in the context of the attention model. We remark that recent work [105] expands the theory presented in this work to 1-layer transformers. By uncovering the underlying optimization principles and thoroughly character- izing the directional convergence of attention, we provide valuable insights into the dynamics and generalization properties of attention- based models opening the path for future research.
allows us to gain a deeper understanding
work also presents the first
49


Reviews:
[{"review_id": "8J7IypPMGj", "reviewer": "Reviewer_sy1C", "summary": "The paper focusses on the optimization dynamics of attention mechanism. The authors analyze a softmax-attention model and demonstrate that running gradient descent on its parameters leads to a max margin solution, separating optimal tokens from non-optimal ones. The authors also present a regularization path analysis, demonstrating the convergence of solutions for nonlinear classifier heads. Overall, the paper aims to enhance the understanding of attention mechanisms and their optimization dynamics in large language models. ", "questions": "1. Can you provide more concrete examples illustrating the application of the findings to real-world LLM tasks?\n\n2. How do the revealed implicit biases in the joint parameter analysis affect the interpretability and generalization capabilities of the attention model?\n\n3. Can you elaborate on how the optimization-theoretic characterization and convergence insights presented in your work can be practically leveraged to enhance the training and fine-tuning of large language models?", "limitations": "The work is highly theoretical and does not present any potential negative societal impacts to the best of my knowledge.", "rating": 7, "confidence": 1, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Comprehensive Characterization: The paper analyzes the fundamental attention model and its connection to max-margin problems. Analysis of this connection is quite original to the best of my knowledge. It overall advances the understanding of the attention mechanism from another theoretical perspective.\n\n2. Convergence Insights: The paper looks into the convergence characterstics of gradient descent for tuning the token/prompt. This analysis can be used to drive futher improvements in optimization of large language models.\n\n3. Joint Parameter Analysis: Through the analysis of the regularization paths, this work highlights the implicit biases and interactions between parameters (v, p), and describing their joint convergence.\n\n4. Implications for Future Research: The work suggests promising avenues for future studies, such as exploring similar analysis for self-attention layers and multiple tunable tokens.\n\n5. Real-World Relevance: Exhaustively understanding attention mechanisms in large language models is crucial for enhancing their performance in natural language processing tasks.\n\n6. Numerical Validation: The authors provide empirical evidence supporting their theoretical findings through numerical experiments.", "weaknesses": "1. Lack of Concrete Examples: The paper could definitely benefit from providing more examples to illustrate the concepts. The findings of the paper are very abstract and make it hard for readers to grasp the implications.\n\n2. Limited support from other works: While the paper thoroughly analyzes the attention mechanism from a max-margin problem perspective, it does not highlight directly if the claims align with other relvant theoretical analysis of attention mechanism.\n\n3. Complexity of Analysis: The optimization-theoretic characterization might be challenging for readers without a strong background in the subject, making it less accessible to a broader audience.\n\n4. Lack of application to Real-World Data: The paper leaves uncertainty about the applicability of the findings in practical attention based model design and applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nThank you for your positive feedback and helpful suggestions.\n\n> **Q1:** Can you provide more concrete examples/applications on real-world tasks? \n\n**R:** Thank you for this suggestion. We have conducted new experiments using real data, demonstrating how our theory successfully predicts two important empirical phenomena related to the optimization dynamics of transformers and attention mechanisms:\n\n   - The attention map (i.e. softmax output) becomes more sparse over time by focusing on the most informative tokens.\n   - This is achieved by the norm of the attention weights ($W$) growing over time and leading to a \"saturating\" effect on softmax, resulting in a sparse pattern.\n\nOur experiments in Figure 2 verify both of these predictions in line with our gradient descent convergence in Theorems 1&3. We train a vision transformer (ViT-base) model from scratch with the CIFAR-10 dataset for 400 epochs with a fixed learning rate $3\\times 10^{-3}$.\n\n   - In Figure 2 (left), we present the progressive change in attention weights of the [CLS] token (which corresponds to our $p$ parameter) during training, computed from all attention heads within the model.\n  - In Figure 2 (right), we display the norm of attention weights and the sparsity level of attention maps averaged over all layers. We used $($L1norm/L2norm$)^2$ of the attention maps as a soft-sparsity measure, where a smaller value indicates a sparser vector.\n\nInitially, during the early epochs of training, the attention weights are randomly distributed, leading to a dense pattern. However, as training progresses, the weights grow, causing the attention map to gradually become sparser. Consequently, the attention map starts to focus on fewer salient patches within the image that possess distinct features that aid in classification.\n\n\n\n> **Q2:** How do the revealed implicit biases in the joint parameter analysis affect the interpretability and generalization capabilities of the attention model? Can you elaborate on how the optimization-theoretic characterization and convergence insights presented in your work can be practically leveraged to enhance the training and fine-tuning of large language models?\n\n\n**R:**  Our joint analysis in Section 3 has a surprisingly interpretable message: $p$ and $v$ (essentially) converge to their respective max-margin solutions, thus, optimization dynamics of “classification” (v) and “attention” (p) can be decoupled. Second, as also pointed out by Reviewer 8XUp, understanding implicit bias is essential across various domains, including interpretability, fairness, optimizer choice, and generalization bounds, because it connects complex optimization dynamics to amenable problems (like our attention SVM). Future works can study various aspects of transformers (TF) through the SVM lens. Below, we provide detailed discussion in the context of language models and how our optimization-theoretic characterization can enhance the training and fine-tuning of LLMs.\n\n* **Fairness and Bias Mitigation**: Gradient descent is a fundamental optimization algorithm widely used in the training and fine-tuning of large language models.  Language models trained using gradient descent can inherit biases present in the training data [[GPT-4] (https://arxiv.org/pdf/2303.08774.pdf)]. Understanding the implicit bias of gradient descent in this context allows researchers to identify and mitigate biases, ensuring fairness and ethical use of language models.\n\n\n* **Generalization Bounds**: The implicit bias of gradient descent is closely tied to the generalization capabilities of trained models [[Vardi23](https://arxiv.org/abs/2208.12591)]. Understanding this relationship helps in establishing theoretical bounds on a model's generalization performance. Specifically, the implicit bias of gradient descent influences how well language models can apply their learned knowledge to new language tasks. Understanding this relationship helps in determining how effectively language models generalize to various language-related challenges.\n\n* **Robustness and Regularization**: Implicit bias can influence the regularization properties of gradient descent. By understanding how the algorithm tends to favor certain solutions, we can develop regularization techniques that encourage better model generalization and robustness against noise and overfitting.\n\n* **Algorithmic Choices**: Knowledge of the implicit bias of gradient descent helps in selecting appropriate optimization methods when training language models. Different algorithms exhibit varying biases, and understanding these nuances can guide the choice of optimization approach based on the desired behavior of the language model. Finally, it is also possible that one can develop new training algorithms: For instance, can we literally solve an SVM during training to accelerate TF optimization (e.g. after identifying which tokens to separate with SVM)?\n\nIn conclusion, characterizing implicit bias is paramount for multiple aspects and, by comprehending the interplay between optimization and these aspects, we can potentially enhance training and fine-tuning processes, leading to more principled, efficient, and trustworthy language models.\n"}, {"review_id": "3LKougFpBq", "reviewer": "Reviewer_SE93", "summary": "This paper proposes to give a mathematical explanation and analysis for widely used attention mechanism. They formulate normal attention, self-attention and prompt tuning into one single formulation. And they connected attention to max-margin problems.", "questions": "Please address the questions in weakness part.", "limitations": "As indicated in weakness part.", "rating": 6, "confidence": 2, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. This paper proposes a mathematical analysis for attention mechanims, which deepen our understanding of the operator.\n2. The authors formulate several kinds of attentions and prompt-tuning into one single formulate, which is practical and novel.", "weaknesses": "1. To be honest, I am not an expert for this area. I appreciate the author's effort on the mathmatical part for deep learning. \n    However, I think it would be better to give conclusions and design guidances based on the authors' observations.\n    For example, can we link and explain some phenomena or pains during attention-based model training? Can we improve or accelerate training by improving network structure or losses?\n\n2. The authors discuss attention modules based on isolated simple operators. I think it is very helpful.\n    However, I wonder if we extend to real large-scale attention based models, do the conclusions remain the same?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time and helpful suggestions.\n\n> **W1:**  To be honest, ... improving network structure or losses?\n\n**R:** Thank you for your questions. In response to reviewer’s concern, under **W2**, we provide and discuss real data experiments which demonstrate that our theory successfully predicts important empirical phenomena related to the optimization dynamics of transformers and attention mechanisms. To provide a more general perspective, as highlighted by Reviewer 8XUp, the recognition of implicit bias is essential across various avenues, including but not limited to fairness, behavior optimization, and generalization bounds. We provide further discussions on these aspects as follows: \n\n- **Fairness and Bias Mitigation**: Gradient descent is a fundamental optimization algorithm widely used in the training and fine-tuning of large language models.  Language models trained using gradient descent can inherit biases present in the training data. Understanding the implicit bias of gradient descent in this context allows researchers to identify and mitigate biases, ensuring fairness and ethical use of language models.\n\n- **Generalization Bounds**: The implicit bias of gradient descent is closely tied to the generalization capabilities of trained models. Understanding this relationship helps in establishing theoretical bounds on a model's generalization performance. Specifically,  the implicit bias of gradient descent influences how well language models can apply their learned knowledge to new language tasks. Understanding this relationship helps in determining how effectively language models generalize to various language-related challenges.\n\n - **Robustness and Regularization**: Implicit bias can influence the regularization properties of gradient descent. By understanding how the algorithm tends to favor certain solutions, we can develop regularization techniques that encourage better model generalization and robustness against noise and overfitting.\n\n- **Algorithmic Choices**: Knowledge of the implicit bias of gradient descent helps in selecting appropriate optimization methods when training language models. Different algorithms exhibit varying biases, and understanding these nuances can guide the choice of optimization approach based on the desired behavior of the language model.\n\nIn summary, acknowledging implicit bias in gradient descent is crucial for exploring the fairness, optimization behavior, and generalization bounds of training LLM algorithms. Our theory takes the initial steps in this direction.\n\n> **W2:** The authors ... the conclusions remain the same?\n\n**R:** Thank you for your suggestion. We have carried out additional experiments to demonstrate the expansion of our findings to real large-scale attention-based models. These experiments also serve to illustrate how our theory proficiently anticipates two essential empirical phenomena associated with the optimization dynamics of transformers and attention mechanisms:\n\n   - The attention map (i.e. softmax output) becomes more sparse over time by focusing on the most informative tokens.\n   - This is achieved by the norm of the attention weights ($W$) growing over time and leading to a \"saturating\" effect on softmax, resulting in a sparse pattern.\n\nOur experiments in Figure 2 verify both of these predictions. We train a vision transformer (ViT-base) model from scratch with the CIFAR-10 dataset for 400 epochs with a fixed learning rate $3\\times 10^{-3}$.  \n\n   - In Figure 2 (left), we present the progressive change in attention weights of the [CLS] token (which corresponds to our $p$ parameter) during training, computed from all attention heads within the model.\n  - In Figure 2 (right), we display the norm of attention weights and the sparsity level of attention maps averaged over all layers. We used $($L1norm/L2norm$)^2$ of the attention maps as a soft-sparsity measure, where a smaller value indicates a sparser vector.\n\nInitially, during the early epochs of training, the attention weights are randomly distributed, leading to a dense pattern. However, as training progresses, the weights grow, causing the attention map to gradually become sparser. Consequently, the attention map starts to focus on fewer salient patches within the image that possess distinct features that aid in classification.\n\nIn light of the concerns raised by the reviewer regarding the framework investigated in this study, we find it crucial to highlight the practical significance of our attention model. Within transformers, the parameter $p$ corresponds to a trainable prompt [Lester et al. EMNLP’21] or the '[CLS]' token [Devlin et al. NAACL’19], as mentioned in [Oymak et al. ICML’23]. To comprehensively address the reviewer's concerns, we also provide new experiments that showcase the extensibility of our work. To this aim, we studied optimization dynamics of self-attention, as illustrated in the attached Figure 3, using the following objective \n\n$\n\\quad L(W)=\\frac{1}{n}\\sum_{i=1}^n \\ell(Y_i\\cdot v^\\top X^\\top S(X_i W^\\top x_{i1}))  \\quad  $ (SA-ERM)\n\nThis corresponds to running linear classification on the first token output of a self-attention layer ($x_{i1}$). We recently discovered a slightly modified ATT-SVM can predict the implicit bias of Self-Attention\n\n$\\min_{W} \\left\\Vert W \\right\\Vert_F \\quad     \\text{subject to} \\quad ( x_{i\\alpha_i}-x_{it})^\\top W x_{i1} \\geq 1 \\quad  \\text{for all } \\quad t\\neq \\alpha_i, i\\in[n] \\quad$  (S-ATT-SVM)\n\nFig 3 shows **self-attention solutions directionally align** with (S-ATT-SVM). Empirically, we observed that optimizing $W_{prod}$ biases gradient descent towards (S-ATT-SVM) with the Frobenius norm objective, while optimizing $(W_k, W_q)$ separately biases it towards (S-ATT-SVM) with the nuclear norm objective. In short, while self-attention introduces different behavior and deserves separate investigation, we believe the attention<->SVM connection introduced by our work is fundamental and remains valid.\n"}, {"review_id": "OQdldSe6M3", "reviewer": "Reviewer_5i57", "summary": "This work studies the mechanism for relevant token selection in the attention model by drawing connections with the implicit bias literature and max-margin SVM formulation. The authors consider the prompt attention model $f(X)=v^TX^T\\text{softmax}(XW^Tp)$, with tokenized input $X$, value weights $v$, key-query weights $W$, and tunable token/prompt $p$. They show that running gradient descent (GD) on $p$ and $W$ (for fixed $W$ and $p$, respectively) is equivalent, so they consider optimizing only $p$ for fixed $W$. \n\nThey consider a data setting where the quality of $t^{th}$ token of input $X$ is determined by its scores $Yv^Tx_t$, with the globally optimal tokens being the ones with the highest scores.\n\nFirst, they consider optimizing $p$ when $v$ is fixed using a loss that is decreasing and smooth. In this setting, they show the following results:\n\n- Under an assumption on the token scores, $p$ converges (in direction) to the global max-margin solution, which separates the globally optimal tokens from the rest.\n- They also develop a regularization path analysis to show global convergence by relaxing the said assumption.\n- The main result shows that with an appropriate initialization and a small enough step size, the GD iterates of $p$ converge (in direction) to a local max-margin solution that separates the locally-optimal tokens from the rest. Here, locally optimal tokens are defined as the ones which have higher scores than their SVM neighbors.\n\nNext, they consider the joint optimization of $p$ and $v$ using logistic loss. In this case, for a given $p$, if the resulting features are separable, $v$ has an implicit bias to converge to the max-margin solution as the problem is linear in $v$. Here, optimal tokens as the ones that maximize the downstream label margin. They show that:\n\n- When the attention features (for the max-margin $p$) are all support vectors (for the respective max-margin $v$), both $v$ and $p$ converge to their respective max-margin solutions.\n- When this is not the case (i.e. the attention features are not all support vectors), $p$ asymptotically selects one token per input, and it suffices to select tokens with the highest scores while also mixing other tokens. This does not impact the margin of $v$, which still converges to the max-margin solution.\n\nThroughout, the authors illustrate these results through numerical experiments.", "questions": "1. Experiments:\n- The numerical experiments illustrate the theoretical results well. However, it would be good to include additional experiments on semi-synthetic or real datasets, such as those considered in [1], [2].\n- Fig. 1(b) illustrates the local convergence of $p$ in the red or blue direction, depending on the initialization. It would be helpful to see if there are cases when the GD iterates do not converge to either of the two solutions shown in Fig. 1(b). This would give some insight into the gap between the empirical observations and the theoretical result for this case.\n\n2. Discussion/clarification:\n- Points of discussion:\n    - Generally, margin maximization refers to separating samples/features from two classes, whereas in this work, the main contribution is to show that the tunable prompt converges to a solution that separates the globally/locally optimal tokens from non-optimal ones. I suggest including more discussion on this part in the introduction. Relatedly, it would help to clarify lines 36-37.\n    - In Section 2, choice of $v$ determines the scores, and hence the solution learned by $p$. Some discussion on this would be nice.\n    - The data setting is interesting, but since there are a lot of cases, it would be helpful to show some connections between (some of) the synthetic data setting(s) and some datasets that are commonly used in practice.\n    - In the discussion on convergence to the local max-margin soution, and the description of Fig. 1(b) in Section 2.2, it would be helpful to clarify that depending on the initialization, the GD iterates will converge to either the global max-margin (when initialized in the cone associated with that solution) or the local max-margin.\n\n- Figures:\n    - In Fig. 1(a), it is unclear what the role of the blue line (local max-margin solution) is, since both the non-opt tokens take the same value.\n    - Fig. 1(c) description needs some clarification. The separating hyperplane in. the figure looks like it's the max-margin solution for the teal datapoint, but not the max-margin across all three colors (margin with green is small).\n    - In Figs. 2(b) and (c), it would be helpful to have the legend for red and blue curves associated with $p$ and $v$, respectively.\n\n- Other:\n    - There are some typos/inconsistencies in the proof of Lemma 1 that should be corrected.\n    - Theorem 2 shows global convergence of $p$ via regularization path analysis. It can be moved above assumption B to improve flow, as it is more general.\n    - In Section 2.2, it would be helpful to use some specific notation for the solution that separates the locally optimal tokens from the rest ($p^{mm}(\\alpha)$), such as $p^{mm}_l$.\n    - Fig. 2(a) comparing the transient dynamics for correlation and logistic loss is interesting. However, it needs a minor clarification in the description. It is stated that when $p$ selects the optimal token, the gradient norm $\\propto \\gamma_i$ for correlation loss, and $\\propto \\gamma_ie^{-\\gamma_i}$ for logistic loss, and we can compare loss for tokens with different scores. However, if $p$ selects the optimal token, the score would be fixed. Maybe, it can be rephrased to “if $p$ selects the token with score $\\gamma_i$”.\n    - In Section 3.1, label margin is defined as $\\frac{1}{||v_{mm}||}$, but in the example, label margin $\\gamma$ is defined as $||v_*||$. It would be helpful to mention what $v_*$ is for clarity.\n\nReferences:\n\n[1] Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models, 2023.\n\n[2] Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity. arXiv preprint arXiv:2302.06015, 2023.", "limitations": "N/A", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "1. This work gives interesting theoretical insights into the mechanism for relevant token selection in the attention model.\n\n2. It lays the groundwork to analyze attention mechanism using the lens of implicit bias and motivates several interesting directions for future work.\n\n3. Overall, it is an interesting paper, with clear exposition and well-connected ideas. It makes several meaningful contributions (as listed in the summary).", "weaknesses": "1. The experimental results seem limited and the paper would benefit from the inclusion of more experiments.\n\n2. Some points of discussion can enhance intuition, and certain aspects regarding figures need some clarification.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thorough review and helpful suggestions, they will definitely help improve the paper.\n\n> **Q1:** The numerical experiments such as those considered in [1], [2].\n\n**R:** Following your suggestion, we conducted additional experiments using real-world datasets to further substantiate our hypothesis concerning the optimization dynamics in transformers and attention. Our theory successfully anticipates two crucial empirical phenomena:\n\n   - The attention map (i.e. softmax output) becomes more sparse over time by focusing on the most informative tokens.\n   - This is achieved by the norm of the attention weights ($W$) growing over time and leading to a \"saturating\" effect on softmax, resulting in a sparse pattern.\n\nOur experiments in Figure 2 verify both of these predictions. We train a vision transformer (ViT-base) model from scratch with the CIFAR-10 dataset for 400 epochs with fixed learning rate $3\\times 10^{-3}$.  \n\n   - In Figure 2 (left), we present the progressive change in attention weights of the [CLS] token (which corresponds to our $p$ parameter) during training, computed from all attention heads within the model.\n  - In Figure 2 (right), we display the norm of attention weights and the sparsity level of attention maps averaged over all layers. We used $($L1norm/L2norm$)^2$ of the attention maps as a soft-sparsity measure, where a smaller value indicates a sparser vector.\n\nInitially, during the early epochs of training, the attention weights are randomly distributed, leading to a dense pattern. However, as training progresses, the weights grow, causing the attention map to gradually become sparser in line with our Thm 1-3.\n\n> **Q2:** Scenarios where GD doesn’t converge to locally-optimal directions\n\n**R:**  Thank you for the comment.  We've examined the gradient descent-generated behavior of $p(t)$ in Figure 4 (attached) with random initialization. We ran experiments with varying $d$ values for 1000 iterations.  We find that, **(1, red bar)** for small $d\\in[2,5]$,  $p(t)$ does not have to saturate softmax i.e. may not converge to a local ($p^{mm}$) or global ($p^{mm\\star}$) max-margin direction. We suspect this is because the **ATT-SVM is not feasible for small $d$**. **(2, blue bar)** For larger $d$’s,  $p(t)$  indeed converges to a locally-optimal direction. **(3, green bar)** As $d$ gets even larger, $p(t)$  converges more frequently to the globally-optimal direction; please also see table below.\n\n| d   | Finite p(t) | % of $p(t) \\rightarrow p^{mm}$ | % of   $p(t) \\rightarrow p^{mm \\star}$ |\n|-----|--------|-------|--------|\n| 2   | 69.1   | 30    | 10.4   |\n| 5   | 5.9    | 92.8  | 17.9   |\n| 10  | 0      | 99.7  | 38.3   |\n| 100 | 0      | 99.5  | 92     |\n| 300 | 0      | 100   | 99.1   |\n| 500 | 0      | 100   | 99.8   |\n\n> **Q3:** Clarify the main contribution (nature of ATT-SVM) and lines 36-37\n\n**R:** Thank you for the great suggestion. We will incorporate the following discussion in the main text:\n\nGradient descent on logistic loss and separable datasets converges to the hard margin SVM solution for linear classification   [Soudry et al. JMLR’18, Rosset et al. NeurIPS’03, Telgarsky ICML’13]. Similarly, the attention layer in neural networks, utilizing softmax nonlinearity, exhibits behavior resembling margin-maximizing solutions. However, attention mechanism operates on input tokens rather than performing direct classification. Thus, it aims to separate tokens within input sequences, favoring SVM-like solutions, represented by (ATT-SVM). Formalizing this intuition is challenging due to the highly nonconvex nature of the optimization landscape caused by the softmax operation.\n\n> **Q4:** Discussion on choice of $v$ and scores\n\n**R:** Agreed, we will elaborate on the score definition and how the choice of $v$ impacts the solution learned by $p$. \n\n\n\n> **Q5:** The data setting is  … used in practice.\n\n**R:** Please refer to our response to your **Q1**. We also emphasize that Def 2 and Thm 3 notably apply to general datasets. We only need ATT-SVM to be feasible.\n\n> **Q6:** In the discussion on convergence,  clarify dependence of GD on initialization\n\n**R:** Thank you for the suggestion. We will highlight that convergence of attention weights depends on the point of initialization and GD can potentially converge to any of the locally-optimal directions per Def 2.\n\n> **Q7:** In Fig. 1(a), the role of the blue line is unclear\n\n**R:** Agreed with the reviewer. In Fig 1(a), the blue line is unnecessary and will be removed. We will also clarify Fig 1(c) by providing additional explanatory notes. In Fig 2(b) and (c), we will include a legend to denote the red curve associated with $p$ and the blue curve associated with $v$, providing clear identification for each curve.\n\n> **Q8:** typos/inconsistencies in the proof of Lemma 1 that should be corrected.\n\n**R:** Thanks for catching this. We also noticed it after submission. We replaced all $q$ terms with $p$ and we are now consistently using $\\cal{L}_p$ and $\\cal{L}_W$ to distinguish the optimization objectives of $p$ and $W$.\n\n> **Q9:** Theorem 2 is more general.\n\n**R:** Totally agreed. We'll include your suggestion in the final paper version.\n\n> **Q10:**  In Section 2.2, it would be helpful … such as $p^{mm}_l$. \n\n**R:**  Thanks, good point. We may include this in the final paper.\n\n> **Q11:**  Clarify Fig. 2(a) … “if  $p$ selects the token with score $\\gamma_i$.\n\n**R:**  We will clarify that $n=2$ and there are two optimal tokens with scores $\\gamma_1=1,\\gamma_2=C$ (non-opt score is 0). Also the sentence “$p$ *approximately* selects optimal tokens” refers to softmax output assigning high-probability to the optimal tokens. This means $\\approx 1$ probability (which indeed occurs as we run GD longer) but not necessarily $=1$.\n\n> **Q12:**  Clarify $v_*$ and its margin.\n\n**R:**  Thank you, we will clarify that $v^{mm}=v^*/||v^*||^2$ resulting in $1/||v^{mm}||=||v^*||$.\n"}, {"review_id": "gAJpyORYZc", "reviewer": "Reviewer_8XUp", "summary": "The paper focuses on the attention mechanism which is commonly used in transformer architectures. In particular, the authors introduce a certain attention model and investigate its optimization dynamics and inductive biases under various assumptions on token's scores. In particular, the setting is a single-head attention mechanism trained by gradient descent and with decreasing losses, such as logistic or linear loss for binary classification. The results mainly hold for attention with linear head and fixed classifier head however some of the results hold without these limitations. The first main contribution of the paper is proving (under the assumption that all non-optimal tokens have the same score value) directional convergence of tunable prompt (denoted by $p$) to a certain max-margin solution which separate one token from the rest of tokens for each input. However, the imposed assumption on score values can be limiting; thus the authors assert that proving local convergence is possible if the initialization is within a cone of the final solution. The paper also studies some extensions such as the joint optimization of the classifier head and trainable parameters and demonstrate that under a specific label margin conditions, the classifier head and trainable parameters converge to their respective max-margin solutions. Some numerical results on synthetic data are designed to validate the claims of the paper.\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n\n ", "questions": "please see the section above. ", "limitations": "The authors adequately discuss the limitations throughout the paper. I do not see any potential negative impact with this work. ", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "The paper is the first work on the implicit bias behavior of GD for the attention mechanism. Understanding implicit bias is crucial for several directions such as fairness, optimization behavior and generalization bounds. Some assumptions of this work are idealized, but that is fine given that it one of the first solid works in this direction. The most important aspect of the paper is giving a formalized understanding of the attention mechanism as a token-selection mechanism and providing sufficient conditions for convergence to a solution which favors optimal tokens. By connecting the attention mechanism to the implicit bias literature and the max-margin SVM formulation (which for the attention mechanism takes a new and interesting form), the study establishes a solid foundation for future research. The required work for obtaining the results is non-trivial and sufficiently distinct from previous works in implicit-bias literature. The paper is also well-written in most parts, although some parts of the paper become very technical with little intuition and thus it can be difficult to clearly understand the results such as the results in section 2.3.", "weaknesses": "some questions and minor weaknesses and suggestions:  \n\n- Regarding theorem 1, the fact that any initialization leads to convergence is rather counter-intuitive. Does the assumptions of Theorem 1 imply that the problem convex or is there any other reason?\n\n- Also related to theorem 1, it is not trivial why parameters norm ($|| p_t ||$) is diverging since $p_t$ is inside the soft-max. can the authors please explicitly specify the behavior of $|| p_t ||$ in the statement of theorem? We know for GD on ReLU neural networks that parameters norm diverges, but that is due to using ReLU non-linearities which implies that the loss prefers large first-layer weights. \n\n- The contributions section can be more specific by providing more details for contributions; for example in line 52 the authors can be more specific in explaining their contributions for the model with non-linear head and specifing the key distinctions with previous parts. \n \n- How are the values of $\\mu$ and $R$ determined in theorem 3? This can be insightful as they are the key parameters in specifying the cone.  ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank the reviewer for their thorough feedback and helpful suggestions. \n\n>**W1:**  Is the problem in theorem 1 convex, any other reason?\n\n**R:** Thank you for the great question. First, let us clarify that the problem is not convex even under Assumption B. One reason is that Assumption A actually allows for very general nonconvex loss functions. However, even for the simplest setting (we can come up with), the problem is not convex: Concretely, let us set\n\n- $\\ell(x)=-x$\n- Pick n=1 samples and $T=2$ tokens. Make tokens unit $\\ell_2$ norm and orthogonal. \n- Set $v=x_1$, $p=cx_1$ with $c\\in \\mathbb{R}$. This way $x_2^\\top v=x_2^\\top p=0$.\n\nFollowing this, the training objective takes the form $max_p L(p):= \\mathbb{S}([c, 0])_1 = \\frac{e^c}{1+e^c}$. This is the standard logistic function and is not convex or concave. We will add this example to the paper.\n\nRather than convexity, our proof relies on establishing favorable **gradient correlation** presented in Lemmas 3 and 5. For instance, in Lemma 5, we establish that for any choice of $\\pi$, there exists $R_{\\pi}$ such that: \n\n$ \\langle \\nabla \\cal{L}(p),  \\frac{p}{|| p||}\\rangle \\geq (1+\\pi) \\langle \\nabla \\mathcal{L}(p), \\frac{p^{mm}}{|| p^{mm}||} \\rangle.$  \n\nIf $\\mathcal{L}(p)$ was convex, it would actually help with establishing the above. Instead, our analysis could be perceived as a directional convergence version of the restricted secant inequality [Karimi et al. ECML’16] where the gradient behaves nicely towards the $p^{mm}$ direction. Via gradient correlation, we are also establishing the **weak convexity** of $\\mathcal{L}$: As $p$ converges in direction to $p^{mm}$, it can be demonstrated from Lemma 4 that $\\nabla \\mathcal{L}((1+\\pi) ||p||p^{mm}) \\rightarrow 0$, which implies the weak monotonicity of the gradient and the weak convexity of $\\mathcal{L}$.\n\n> **W2:** It is not trivial why parameters norm ($p(t)$) is diverging since $p(t)$ is inside softmax. Can the authors please explicitly specify the behavior of $||p(t)||$?\n\n**R:** We will provide a better discussion in the final manuscript. The intuition is as follows: Softmax output is a probability vector, thus the attention output $f(X)=\\left<Xv, \\mathbb{S}(Xp)\\right>$ creates a convex combination of token scores $\\gamma=Xv\\in \\mathbb{R}^T$ (here we simply set label $Y=1$). When $v$  is fixed, $\\gamma$ is fixed and, since the loss function is decreasing, the smallest training loss the attention model can achieve is by **assigning all softmax probability to the tokens with the highest score**. Otherwise we are strictly worse off when the convex combination contains some non-optimal tokens. Thus, we want probability 1 for optimal tokens and 0 for others. On the other hand, softmax with finite weights cannot accomplish this because softmax output is strictly positive. This is precisely why norms go to $\\infty$: Softmax asymptotically sets token probabilities to $1$ and $0$. \n\nTo provide a GD-specific intuition, Lemma 4 in the supplementary shows: $\\langle \\nabla \\mathcal{L}(p), p^{mm} \\rangle < 0$ for all finite $p \\in \\mathbb{R}^d$. Consequently, there are no finite critical points $p$ for which $\\nabla \\mathcal{L}(p) = \\mathbf{0}$. This implies that $|| p\\left(t\\right) ||\\rightarrow \\infty$. In our Thms 1 and 3, we prove that $||p(t)||$ diverges and aligns with the SVM solution, and attention maps $\\mathbb{S}(Xp)$ select optimal tokens as $t$ grows. However, the precise quantification of behavior of $||p(t)||$ remains open.\n\n\n> **W3:** Contributions section can be more specific, specifically for nonlinear head\n\n**R:** Thank you for the suggestion. We will incorporate the following in the main text:\n\nTo establish the margin maximizing nature of attention under broader conditions, we study the general model $f(\\mathbf{X})=\\psi(\\mathbf{X}^\\top \\mathbb{S}(\\mathbf{X}\\mathbf{W}^\\top p))$ where $\\psi$ is a nonlinear head. This setting poses challenges as we lack a clear score function, unlike the previous sections. To address this, we introduce a generic condition that splits the tokens of each input into an optimal set and a non-optimal set. Non-optimal tokens are those that strictly increase the training risk if they are not fully suppressed by attention probabilities $\\mathbb{S}(X_iW^\\top p)$.\n\n\n> **W4:** How are the $\\mu$ and $R$ determined in Thm 3?\n\n**R:**  As stated in Step 1 of the proof sketch on page 6, we $\\mu$ is a function of the margin of the entire dataset. Specifically:\n$\\delta := \\frac{1}{2}\\min_{i\\in[n]}\\min_{t\\in \\cal{T}i,\\tau\\in \\bar{\\cal{T}}i}(\\mathbf{k}{it}-\\mathbf{k}{i\\tau})^\\top p^{mm}, \\quad A := \\max_{i\\in[n],t\\in[T]} ||\\mathbf{k}_{it}|| \\cdot  || p^{mm}||, \\quad  \\mu = \\frac{1}{8}\\left(\\frac{\\min(0.5,\\delta)}{A}\\right)^2.$\n\nFurthermore, Lemma 3 reveals that $R$ is inversely dependent on both $\\mu$ and the score gap:\n\n$ R \\geq \\frac{\\max(2,\\delta^{-1})}{||p^{mm} ||}\\log\\left(\\frac{64T\\Gamma A}{\\boldsymbol{\\gamma}}\\right),$ \n\nwhere $\\boldsymbol{\\gamma} = \\min_{i\\in[n]}\\boldsymbol{\\gamma}_i$ represents the worst-case score gap across all inputs.\n\nThe definition of $\\mu$ offers valuable insights into the gradient descent initialization process. When tokens are more separable, $\\mu$ increases, leading to a reduction in correlation ($1-\\mu$) and an expansion of the initialization cone. Similarly, $R$ exhibits an inverse relationship with both $\\mu$ and the score gap $\\boldsymbol{\\gamma}$. When the data's score gap and $\\mu$ are large, gradient descent can commence with a small norm, allowing for initialization within a wider cone around $p^{mm}$.\n\nWe will integrate the aforementioned insights into the discussion of Theorem 3, emphasizing the significant impact of considering the interplay between $\\mu$, the score gap, and the properties of tokens on the gradient descent initialization process.\n"}, {"review_id": "xLlE7H63Qp", "reviewer": "Reviewer_iE2T", "summary": "The paper is clear and well-written. \nUnderstanding the e optimization dynamics and implicit bias is a significant theoretical issue, especially for morden neural network models.\nThis paper provides a preliminary theoretical analysis of the margin maximization bias of attention-like models. \nTheoretically, the authors provide some global convergence and local convergence results to characterize the implicit bias of gradient descent for attention-like models.", "questions": "My primary concern is whether the training dynamics of standard attention models $Atten(X)=W_V X \\mathbb{S}(X^\\top W_K^\\top W_Q X)$ closely resemble the training dynamics of the attention-like model in this study. \nSpecifically, I find it difficult to discern the training behavior of $W_K$ and $W_Q$, which are crucial components in standard attention, from the results presented by the authors. \nIf the authors could provide further clarification on this matter, I could change my perspective.", "limitations": "As previously mentioned, the main limitation of the article lies in the oversimplification of the attention model, and it remains uncertain whether the training dynamics of the attention model differ fundamentally from the training dynamics of the attention-like model proposed by the authors.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "The authors present a comprehensive theoretical analysis of the implicit bias associated with margin maximization in attention-like models. \nSpeficically, they provide theories about global convergence, local convergence, and regularization paths in various training scenarios. \nThis work extends lots of theoretical results on the implicit bias for linear models to the context of attention-like models.", "weaknesses": "In this article, the authors investigate a model that bears resemblance to attention but involves a significant degree of simplification compared to standard attention models. \nThe model employed combines $W_K$ and $W_Q$ into a single matrix $W$, while substituting one of the $X$'s with $p$. \nWhile attention models are still relatively underexplored in the field of optimization theory, I believe that the level of simplification adopted in this study is overly simplistic and may even be somewhat irrelevant. \nWith such a simplification by the authors, especially for the optimization about $W$ or $p$, it becomes almost a one-layer neural network optimization problem that is largely irrelevant to attention.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their feedback and suggestions. Below, we respond to their concerns point by point. We would be happy to respond to future concerns they may have during the discussion period.\n\n> **W1:* This study is overly simplistic... With such a simplification.. it becomes almost a one-layer neural network problem that is largely irrelevant to attention.\n\n**Response:** We respectfully disagree with this assessment for the following reasons:\n\n1. Our attention model is practically relevant: In transformers $p$ corresponds to a tunable prompt or [CLS] token [Oymak et al. ICML’23]. When sample size is $n=1$, setting $p=x_1$ (i.e. first token) and optimizing $W$ (via Lemma 1’s equivalence), our theory specializes to establish the implicit bias of a 1-layer self-attention model.\n\n2. The problem does not become a one-layer neural network optimization. In fact, it is very different for the following reasons:\nSoftmax nonlinearity is different from applying $T$ nonlinear activation functions individually because it couples the $T$ nonlinearities to induce a probability distribution. This makes it the standard choice in attention/transformer layers and as a loss function (cross-entropy). Crucially, softmax also induces sparsity which is precisely what happens in real attention maps (e.g. see attached Fig 1) or when attention selects the optimal token within our theory.\nFeedforward neural nets only multiply weights and features. In contrast, our model $f(X)=v^TX^TS(Xp)$ as well as self-attention multiplies features with each other (i.e. the $X$ term appears twice or more).\n\n3. Finally, we understand the concern that self-attention or transformers may be more practically relevant. On the other hand, we firmly believe our SVM-equivalence framework is fundamental and extensible. To provide a concrete example, we recently discovered that a slight variation of our ATT-SVM seems to predict the implicit bias of self-attention (building on the aforementioned $n=1$ observation). This will be discussed under the next response.\n\n> **Q1:*  My primary concern is whether the training dynamics of standard attention models resemble the training dynamics of the attention-like model in this study. I find it difficult to discern the training behavior of  $W_K$ and $W_Q$, which are crucial components in standard attention. If the authors could provide further clarification on this matter, I could change my perspective.\n\n**Response:** We acknowledge this concern and respond it in two fronts:\n(1) Training behavior of $W_k$ and $W_q$ under our paper’s setting still follows max-margin directions,\n(2) We empirically demonstrate that our Attention<->SVM connection is extensible to self-attention.\n\n**(1) Training behavior of $W_k$ and $W_q$:** Standard self-attention calculates $S(XW_qW_k^\\top X^\\top)$ where $W_q,W_k$ are size $d\\times m$. Clearly, when $W_q,W_k$ are full dimensional ($m=d$), we don’t lose any expressivity by merging them into $W_{prod}=W_qW_k^\\top$. On the other hand, we acknowledge that the optimization behavior might be different. Fortunately, for the problem $L(W_k,W_q)=\\sum_{i=1}^n v^TX_i^TS(X_iW_qW_k^\\top a)$:\nWe can prove a version of Lemma 1 that creates a mapping between $W_q,W_k$ iterations and $p$ iterations for any $d\\geq m\\geq 1$.\nNumerically, we found that $W_{prod}(t)=W_q(t)W_k(t)^\\top$ still converges to max-margin direction.\n\nOur experiments are shown in Fig 1. Fig 1(left) is the outcome of our Lemma 1 (iterations on $W_{prod}$ and associated $p$ iterations) whereas Fig 1(right) is the joint $W_k,W_q$ iterations (akin to transformers) and associated $p$ iterations. $W_k,W_q$ still align with the max-margin direction albeit with a slightly different trajectory. To formalize this, we can map the joint gradient updates $W_k(t+1)=W_k(t)-\\eta \\nabla_{W_k} L(W_k,W_q)$, $W_q(t+1)=W_q(t)-\\eta \\nabla_{W_q} L(W_k,W_q)$ to the following $p(t)$ iterations on the $L(p)$ objective: Starting with proper $p(0)$ choice and scalar $\\nu_0=1$, run\n\n$\\nu_{t+1}=\\nu_t-\\eta\\nu_t^{-1}p(t)^\\top \\nabla L(p(t))$\n\n$p_{t+1}=(\\nu_{t+1}/\\nu_t)[p(t)-\\eta\\nu_t^2\\nabla L(p(t))]$\n\nThis mapping is not as simple as Lemma 1. Regardless, it strongly suggests that our max-margin theory on $p$ should extend to $W_k,W_q$. We emphasize that Lemma 1 is enabled by the vector $a$ being fixed: This way the gradient updates on $W_{prod}$ are rank-1 and stay along $a$ direction. Below, we empirically show that, situation is similar but more intricate for self-attention.\n\n\n**(2) Extensibility** In Fig 3, we study the self-attention objective \n\n$\n\\qquad L(W)=\\frac{1}{n}\\sum_{i=1}^n \\ell(Y_i\\cdot v^\\top X^\\top S(X_i W^\\top x_{i1}))  \\qquad  $ (SA-ERM)\n\nThis corresponds to running linear classification on the first token output of a self-attention layer ($x_{i1}$). We consider a slightly modified ATT-SVM to capture the inductive bias of this\n\n$\n\\qquad \\min_{W} || W||_F \\quad  \\text{subject to} \\quad (x_{i\\alpha_i}-x_{it})^\\top W x_{i1}\\geq 1 \\quad  \\text{for all }\\quad t\\neq \\alpha_i, i\\in[n]   \\qquad \n$  (S-ATT-SVM)\n\nThe intuition is as follows: In our ATT-SVM (with Lemma 1), $a$ is fixed whereas here $a\\gets x_{i1}$ is changing for each training example. Fig 3 shows **self-attention solutions directionally align** with (S-ATT-SVM). Empirically, we found that directly optimizing $W_{prod}$ biases the gradient descent towards (S-ATT-SVM) with the Frobenius norm objective, while optimizing $(W_k, W_q)\\in\\\\mathbb{R}^{d\\times d}$ separately biases it towards (S-ATT-SVM) with the nuclear norm objective. In our paper's setting (Lemma 1), both objectives coincide because the solution is rank-1 due to fixed $a$. Finally, if $(W_k,W_q)$ are $d\\times m$ with $m<d$, we suspect a low-rank constraint in (S-ATT-SVM) is needed. \n\nTo sum up, we agree with the reviewer that $(W_k,W_q)$ or self-attention introduce unique behavior, however, based on empirical evidence, the Attention<->SVM connection introduced by our paper remains valid.\n"}, {"review_id": "CU4k98liVj", "reviewer": "Reviewer_CU4X", "summary": "This paper aims to provide an optimization-theoretic characterization of the softmax attention model $f(X)=v^{\\top}X^{\\top}{\\rm softmax}(XW^{\\top}p)$ by linking it to max-margin problems. The authors established the convergence of gradient decent on $p$ for a fixed $v$ choice, and further explored the joint convergence of $(v,p)$ via regularization path analysis. They also showed that the idea on selecting the optimal token via max-margin can be extended to a general nonlinear model. Their results are verified through numerical studies.", "questions": "- Why Lemma 1 implies line 92-94?", "limitations": "See weaknesses.", "rating": 6, "confidence": 1, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "- Interesting and important result. Attention has played an important role in large language model. However, the theoretical understanding is relatively lack. This paper links the regularized solution and the gradient decent process to the max-margin solution, which may motivate several possible directions for the further research.\n- Clear presentation with adequate explanation.\n", "weaknesses": "- Some assumptions are relatively strong. In Assumption B, they assume all non-optimal tokens have equal scores, which may not be true in practice.\n- Lack the convergence of the gradient decent when jointly optimizing $(v,p)$.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and helpful comments.\n\n> **W1:** Some assumptions are relatively strong. In Assumption B, they assume all non-optimal tokens have equal scores, which may not be true in practice.\n\n**R:** Thanks for raising this. We agree that this assumption is fairly strong, however, it is only used within Theorem 1 and nowhere else. Theorem 1 provides a global convergence guarantee for gradient descent and serves as a prelude to the general behavior of the attention’s implicit bias. However, as you can see, our Theorem 3 on the local convergence of gradient descent and all other theorems do not require this assumption. Theorem 3 also clearly establishes that, in general (i.e. when Assumption B does not hold and non-optimal scores are different), the global convergence (to the optimal direction $p^{mm\\star}$) can fail due to the existence of locally-optimal directions.\n\nThus, the only way global convergence (from any initialization) can happen is if $p^{mm\\star}$ is the unique locally-optimal direction (per Definition 2). Assumption B is one condition that guarantees this. We believe another condition that guarantees this might be ensuring that “all tokens involved in ATT-SVM are support vectors”. This way, for any choice of non-globally-optimal tokens $\\alpha$, there will be some optimal token $opt_i\\neq \\alpha_i$ that is its SVM-neighbor. Then by Definition 2, $\\alpha$ cannot be locally-optimal because $opt_i$ has a higher score than $\\alpha_i$. Thus, $p^{mm\\star}$ becomes the unique direction satisfying Def 2. While this condition (i.e. all SVM constraints being support vectors) sounds technical, for classical SVM problems, it holds when the embedding dimension $d$ is large [Muthukumar et al. JMLR’21, Hsu et al. AISTATS’21]. \n\nBased on this intuition, we provide a new experiment in Figure 4. We solve many random instances of the attention problem for various values of $d$ for 1000 iterations and  investigate the convergence behavior of $p(t)$ generated by gradient descent. Experiments were conducted with various values of $d$ for 1000 iterations. The bar plot in Figure 4 distinguishes between non-local convergence (red bars), local convergence (blue bars), and global convergence (green bars). Global convergence is a strict subset of local convergence. In short, in line with our hypothesis, as $d$ grows, we observe global convergence with probability approaching $1$. While we do not have a proof of this, it certainly makes an interesting discussion for our paper and we hope to incorporate it. \n\n \n> **W2:** Lack the convergence of the gradient descent when jointly optimizing (v,p).\n\n**Response:** While this may appear to be a weakness, we emphasize that the contributions we make are novel and challenging even for $p$-only optimization. For joint optimization, we also provide a regularization path theory that successfully predicts the implicit bias of gradient descent (see Figure 2(b,c)). For joint optimization, there is no remotely similar result in the literature, and we have a surprisingly powerful message (see Sec 3.1): $p$ and $v$ (essentially) converge to their respective max-margin solutions, thus, optimization dynamics of “classification” (v) and “attention” (p) can be decoupled. We genuinely hope that this novel message (and other contributions) will spur interest in the community and invite future research to solve these open problems.\n\nWe believe that conducting a separate analysis of attention weights  $p$ can offer a clearer and more comprehensive understanding of the implicit bias ingrained in gradient descent for attention mechanisms. To accomplish this, our approach involves the introduction of concepts like token scores and locally-optimal tokens, each of which demands a more comprehensive and detailed explanation.  Additionally, we undertake an extensive convergence analysis that aims to capture the optimization dynamics through the lens of local SVM geometry and the conic initialization centered around the max-margin solution. \n\nWhile we acknowledge the potential benefits of a joint implicit bias analysis involving both $(v,p)$, our experiments showcased in Figure 2 of the paper provide evidence for the feasibility of this approach. However, it's important to note that a comprehensive treatment of these intricate technical details might require the incorporation of several novel techniques. Condensing all these details into a single paper could potentially be overwhelming and hinder the clarity of our main findings.   \n\n\n\n> **Q1:** Why Lemma 1 implies line 92-94?\n\n**Response:**  Thanks for asking this. Recall that, in our attention model $f(p,W)=v^\\top X^\\top S(XW^\\top p)$ in Line 30 , we can either optimize $W$ or optimize $p$. Throughout the paper, we optimize $p$ because Lemma 1 says that optimizing $W$ can be mapped to optimizing $p$. Lemma 1 creates this mapping as follows: Since we will also use $p$ as a variable, fix a vector $a$ and consider the $W$ optimization with $p\\gets a$, $f(W)=v^\\top X^\\top S(XW^\\top a)$, and associated loss function $L(W)$ as defined in Lemma 1. \n\nWe map these to the function $f_p(p)=v^\\top X^\\top S(X p)$ and associated loss function $L_p(p)$ where the idea is viewing the combined $W^\\top a$ as a single $p$ variable. We prove that **running gradient descent on $p(t)$ with learning rate $\\eta$** is same as **running gradient descent on $W(t)$ with learning rate $\\eta||a||^{-2}$** and $W(t)$ iterations precisely track $p(t)$ iterations via $W(t)=||a||^{-2}ap(t)^\\top$. Note that, for this to happen, initializations $p(0),W(0)$ should match (and that is all!). In general, for any $W$ initialization, there exists a matching $p$ initialization. Lemma 1 states this for a rank-1 $W$ initialization to avoid verbosity. We are happy to address further questions if anything is unclear."}]

Author rebuttal:
We sincerely appreciate the time and efforts of the reviewers. We highlight the **main contributions (C1-C3)** of the paper and present **new experiments (E1-E4)** along with explanations for the corresponding **attached figures (Figs 1-4)**. We would be grateful to respond to any reviewer inquiries during the discussion period.

**C1:**  As Reviewer 8XUp kindly states: "The paper is the **first work** on the implicit bias behavior of GD for the attention mechanism". Implicit regularization is extensively studied for linear models and standard neural net architectures [Soudry et al. JMLR’18; Gunasekar et al. NeurIPS'18; Arora et al. NeurIPS’19; Li, Wang et al. NeurIPS’22; Frei et al. ICLR’23 and more]. However, the attention mechanism remains an important unexplored topic. We investigate attention's optimization landscape and analyze its implicit bias, shedding light on the role of softmax nonlinearity. Soudry et al. and others connect logistic regression to standard SVM, which separates inputs based on their labels. Instead, we show attention is biased towards ATT-SVM which separates and selects optimal tokens within the input sequences.

**C2:** We make innovative and nontrivial theoretical contributions:
  - Defs. 1 and 2 introduce novel concepts of  **token scores** and **locally-optimal tokens**, and Lems 2-5 present innovative theories integrating these concepts with the softmax nonlinearity's special structure.
  - Gradient analysis (specifically Thm 3) requires novel proof ideas that capture the optimization dynamics in terms of the local SVM geometry and conic initialization around the max-margin solution.
  - Our regularization path analysis in Secs. 3 and 4 yield nontrivial findings, remarkably predicting the implicit bias of gradient descent when jointly optimizing $(v,p)$.

**C3.**  Reviewers sy1C, 5i57, and 8XUp acknowledge the potential of **our work to open new research avenues**.  ​​We see the connection between attention and ATT-SVM as fundamental, providing a general framework to comprehend complex architectures and generalization dynamics,  akin to the logistic regression and classical SVM connection in deep learning theory  [Soudry et al. JMLR’18, Rosset et al. NeurIPS’03, Telgarsky ICML’13 and more]. See **E3** below for supporting evidence.

**Supporting Experiments**

**E1:**  Reviewer iE2T remarks that transformers use separate key and query weights $(W_k, W_q)$, while our approach uses the combined matrix $W_{prod}:=W_q W_k^\top$. In Fig 1 (attached), we demonstrate that regardless of whether we optimize $W_k$ and $W_q$ separately or as $W_{prod}$, the resulting **trajectories align in direction with the ATT-SVM solution**. As predicted by Lemma 1, $W_{prod}$ converges to a rank-1 matrix of the form $a {p^{mm}}^\top$, where $p^{mm}$ is the parameter obtained by running gradient descent on the $p$ parameter. We use the objective 

$
\qquad   L(W_k,W_q)=\frac{1}{n}\sum_{i=1}^n \ell(Y_i\cdot v^\top X_i^\top \mathbb{S}(X_i W_k W_q^\top a)) \qquad  $  (QK-ERM)

with a fixed vector $a$, in accordance with Lemma 1. 

    More in the response to Reviewer iE2T.

**E2:** Our theory predicts crucial empirical phenomena in transformer optimization dynamics:

- The attention map (i.e. softmax output) becomes more sparse over time by focusing on the most informative tokens.
- To do so, the norm of the attention weights $W$ should grow over time and “saturate” softmax towards a sparse pattern.

In support of this, Fig 2 (left) displays the evolving attention map of the [CLS] token during training, aggregated from all attention heads of a vision transformer. Fig 2 (right) shows the average norm of attention weights and sparsity level of attention maps across all layers. Initially dense and random, the attention weights gradually grow, leading to a sparser attention map.
   
     More in the response to Reviewer sy1C.

**E3:** In Fig 3, we consider the **self-attention** objective 

$
\qquad L(W)=\frac{1}{n}\sum_{i=1}^n \ell(Y_i\cdot v^\top X^\top S(X_i W^\top x_{i1}))  \qquad  $ (SA-ERM)

This corresponds to running linear classification on the first token output of a self-attention layer ($x_{i1}$). We recently discovered a slightly modified ATT-SVM can predict the implicit bias of Self-Attention

$\min_{W} \left\Vert W \right\Vert_F \quad     \text{subject to} \quad ( x_{i\alpha_i}-x_{it})^\top W x_{i1} \geq 1 \quad  \text{for all } \quad t\neq \alpha_i, i\in[n]  \qquad $  (S-ATT-SVM)

Fig 3 shows self-attention solutions directionally align with (S-ATT-SVM). Empirically, we found that optimizing $W_{prod}$ biases gradient descent towards (S-ATT-SVM) with the Frobenius norm objective, while optimizing $(W_k, W_q)$ separately biases it towards (S-ATT-SVM) with the nuclear norm objective. In our paper's setting (Lemma 1 and **E1** above), this distinction vanishes as we fix $x_{i1}\gets a$, leading to a rank-1 matrix solution for S-ATT-SVM.

     More in the response to Reviewer iE2T.

**E4:** We study the convergence of attention with random problem instances with $n=4, T=6$ and varying dimension $d = 2,5,10,100,300,500$. We find that, 
- **(1, red bar)** for small $d=2,5$, $p(t)$ generated by gradient descent does not have to saturate softmax i.e. may not converge to a local ($p^{mm}$) or global ($p^{mm\star}$) max-margin direction. This is likely because the ATT-SVM is not feasible for small $d$.
-  **(2, blue bar)** For larger $d$’s,  $p(t)$  indeed converges to a $p^{mm}$. 
- **(3, green bar)** As $d$ gets even larger, $p(t)$  converges more frequently to the $p^{mm\star}$ (for $d\geq 300$, 99% of the problem instances). We believe over-parameterization (large $d$) enables global convergence because as $d$ grows, the global direction $p^{mm\star}$ becomes the only direction that satisfies our Definition 2 (local optimality). Thus, local and global convergence coincide.

.

    More in the response to Reviewers CU4X & 5i57.


Meta-review: The paper studies the optimization dynamics of softmax attention mechanism, which can potentially have a great impact as the softmax attention is widely used in different applications. The authors showed that  softmax-attention weights trained with gradient descent converge to a max-margin solution that separates locally-optimal tokens from non-optimal ones. Reviewers generally liked the idea and appreciated the contribution. This paper would be of broad interest of the community that are working on improving softmax attention models.
Decision: Accept (spotlight)