Title: Max-Margin Token Selection in Attention Mechanism
Venue: NeurIPS 2023
Accepted: yes
Keywords: attention mechanism

implicit bias

margin maximization

nonconvex optimization

prompt tuning
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

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 softmaxattention model f ( X ) = ⟨ Xv , softmax ( XWp ) ⟩ , 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.

1 Introduction

Since its introduction in the seminal work [1], attention mechanism has played an 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 of 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 ) = GLYPH<10> Xv , S ( XW ⊤ p ) GLYPH<11> . 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).

4

2

0

2

4

6

8

10

2

0

2

4

6

8

10

(-0.1,1), optimal token ( =1)

(1,0), non-opt token (

=0.9)

(0,0), non-opt token (

=0.9)

(a) Global convergence

2

0

2

4

6

8

10

12

2

0

2

4

6

8

10

(-0.1,1), optimal token ( =1)

(1,0), non-opt token (

=0.9)

(0,0), non-opt token ( =0)

(b) Global and local optimal directions

2

1

0

1

2

3

4

5

6

2

1

0

1

2

3

4

Optimal tokens ( =1)

Non-opt tokens (

=0.5)

(c) Multiple inputs

denotes the softmax nonlinearity. For transformers, p 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 ( Yi , X i ) n i = 1 with labels Yi ∈ {-1 , 1 } and inputs X i ∈ R T × d , we consider the empirical risk minimization with a decreasing loss function ℓ ( · ) : R → R ,

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 p 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 L ocally-optimal M axM argin ( LMM ) directions and show that these thoroughly characterize the viable convergence directions of attention when the norm of its weights grows to infinity. We also identify conditions under which (algorithm-independent) regularization path and gradient descent path converge to G lobally-optimal M axM argin ( GMM ) direction in Theorems 1 and 2, respectively. A central feature of our results is precisely quantifying optimality in terms of token scores γ t = Y · v ⊤ x t where x t is the t th 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.

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 v , if the attention features x att i = X ⊤ i S ( X i W ⊤ p ) are separable based on their labels Yi , v would implement a max-margin classifier. Building on this, we prove that p and v converges to 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 1 , Section 5 discusses related literature, and Section 6 highlights limitations and future work.

1.1 Preliminaries

Notations. For any integer N ≥ 1, let [ N ] : = { 1 , . . . , N } . We use lower-case and upper-case bold letters (e.g. a and A ) to represent vectors and matrices, respectively. The entries of a are denoted as a i . We use ¯ σ ( A ) to denote the maximum singular value of A . We denote the minimum of two numbers a , b as a ∧ b , and the maximum as a ∨ b . 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 : R d → R and an ℓ 2-norm bound R , define the regularized solution as

Regularization path - the evolution of ¯ p ( R ) as R grows - is known to capture the spirit of gradient descent as 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. lim R →∞ ¯ p ( R ) / R ) and gradient descent. For gradient descent, we assume the objective L ( p ) is smooth and describe the gradient descent process as

where η ( t ) is the stepsize at time t and ∇L ( p ( t )) is the gradient of L at p ( t ).

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-based techniques on specific contextual datasets.

Self-attention is the core building block of transformers [6]. Given an input consisting of T tokens X = [ x 1 , . . . , x T ] ⊤ ∈ R T × d , self-attention with key-query matrix W ∈ R d × d , and value matrix V ∈ R d × v , the self-attention model is defined as follows:

Here, S ( · ) is the softmax nonlinearity that applies row-wise on the similarity matrix XWX ⊤ .

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 ∈ R d be the tunable token ( [CLS] or prompt vector) and concatenate it to X to obtain Xp : = [ p X ⊤ ] ⊤ ∈ R ( T + 1) × d . Consider the cross-attention features obtained from Xp and X given by

The beauty of cross-attention is that it isolates the contribution of p under the upper term f cls ( X ) = V ⊤ X ⊤ S ( XW ⊤ p ) ∈ R v . In this work, we use the value weights for classification, thus we set v = 1, and denote v = V ∈ R d . This brings us to our attention model of interest:

Here, ( v , W , p ) are the tunable model parameters and K is the key embeddings. Note that W and p 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 a one-to-one mapping.

Lemma 1 Fix u ∈ R d \ { 0 } . Let ψ : R d → R and ℓ : R → R be di ff erentiable functions. On the same training data ( Yi , X i ) n i = 1 , define ˜ L ( p ) : = 1 / n P n i = 1 ℓ ( Yi · ψ ( X ⊤ i S ( X i p ))) and L ( W ) : = 1 / n P n i = 1 ℓ ( Yi · ψ ( X ⊤ i S ( X i W ⊤ u ))) . Consider the gradient descent iterations on p and W with initial values p (0) and W (0) = up (0) ⊤ / ∥ u ∥ 2 and stepsizes η and η/ ∥ u ∥ 2 , respectively:

We have that W ( t ) = up ( t ) ⊤ / ∥ u ∥ 2 for all t ≥ 0 .

This lemma directly characterizes the optimization dynamics of W through the dynamics of p , allowing us to reconstruct W from p using their gradient iterations. Therefore, we will fix W and concentrate on optimizing p in Section 2 and the joint optimization of ( v , p ) in Section 3.

Problem definition: Throughout, ( Yi , X i ) n i = 1 denotes training dataset where Yi ∈ {-1 , 1 } and X i ∈ R T × d . We denote the key embeddings of X i via K i = X i W ⊤ and explore the training risk

Importantly, our results apply to general tuples ( Yi , X i , K i ) and do not assume that ( X i , K i ) are tied via W . Finally, the t th tokens of X i , K i are denoted by x it , k it ∈ R d , respectively, for t ∈ [ T ].

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 a set of assumptions to demonstrate the global and local convergence of gradient descent for margin maximization in the attention mechanism.

2 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 on learning p ∈ R d given a fixed choice of v ∈ R d . Notably, our results apply to general decreasing loss functions without requiring convexity . This generality is attributed to margin maximization arising from the exponentially-tailed nature of softmax within attention, rather than ℓ . We maintain the following assumption on the loss function throughout this section.

Assumption A (Well-behaved Loss) Over any bounded interval: (1) ℓ : R → R is strictly decreasing. (2) ℓ ′ is M 0 -Lipschitz continuous and | ℓ ′ ( u ) | ≤ M 1 .

Assumption A includes many common loss functions, including the logistic loss ℓ ( u ) = log (1 + e -u ), exponential loss ℓ ( u ) = e -u , and correlation loss ℓ ( u ) = -u . Assumption A implies that L ( p ) is Lp -smooth (see Lemma 6 in Supplementary), where

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 softmax-attention. Fix indices α = ( α i ) n i = 1 and consider

Note that existence of p mm ( α ) implies the separability of tokens α from the others. Specifically, choosing direction p mm ( α ) will exactly select tokens ( x i α i ) n i = 1 at the attention output for each input sequence, that is, lim R →∞ X ⊤ i S ( R · K i p mm ( α )) = x i α i . 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 p

We first identify the conditions that guarantee the global convergence of gradient descent for p . The intuition is 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 ℓ ( · ) have a clear definition.

Definition 1 (Token Scores, Optimality & GMM ) The score of token x it of input X i is defined as γ it : = Yi · v ⊤ x it. The optimal tokens for input X i are those tokens with highest scores given by

Globally-optimal max-margin ( GMM ) direction is defined as the solution of (ATT-SVM) with optimal indices ( opt i ) n i = 1 by p mm ⋆ .

It is worth noting that score definition simply uses the value embeddings v ⊤ x it of the tokens. Note that multiple tokens within an input might attain the same score, thus opt i or p mm ⋆ may not be unique. The theorem below provides our regularization path guarantee on the global convergence of attention.

Theorem 1 (Regularization Path) Suppose Assumption A on the loss function holds, and for all i ∈ [ n ] and t , opt i , the scores obey γ it < γ i opt i . Then, the regularization path ¯ p ( R ) = arg min ∥ p ∥≤ R L ( p ) converges to the GMM direction i.e. lim R →∞ ¯ p ( R ) / R = p mm ⋆ / ∥ p mm ⋆ ∥ .

Theorem 1 shows that as the regularization strength R increases towards the ridgeless problem min p L ( p ), the optimal direction ¯ p ( R ) aligns more closely with the max-margin solution p mm ⋆ . Since this theorem allows for arbitrary token scores, it demonstrates that max-margin token separation is an essential feature of the attention mechanism . In fact, it is a corollary of Theorem 8, which applies to

the generalized model f ( X ) = ψ ( X ⊤ 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 to 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 ∈ [ n ] and t , τ , opt i , the scores per Definition 1 obey γ it = γ i τ < γ i opt i .

Theorem 2 (Global Convergence of Gradient Descent) Suppose Assumption A on the loss function ℓ and Assumption B on the tokens' score hold. Then, the gradient descent iterates p ( t + 1) = p ( t ) -η ∇L ( p ( t )) on (ERM) , with the stepsize η ≤ 1 / Lp and any starting point p (0) satisfy lim t →∞ ∥ p ( t ) ∥ = ∞ . If n = 1 , we also have lim t →∞ p ( t ) / ∥ p ( t ) ∥ = p mm ⋆ / ∥ p mm ⋆ ∥ .

Theorem 2 shows that gradient descent will diverge in norm, and when n = 1, the normalized predictor p ( t ) / ∥ p ( t ) ∥ converges towards p mm ⋆ , the separator of the globally optimal token. While n = 1 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 x = [ x 1 , x 2 , x 3]. We reserve the first two coordinates as key embeddings k = [ x 1 , x 2 , 0] by setting W = diag([1 , 1 , 0]). This is what we display in our figures as token positions. Finally, in order to assign scores to the tokens we use the last coordinate by setting v = [0 , 0 , 1]. This way score becomes Y · v ⊤ x = Y · x 3, allowing us to assign any score (regardless of key embedding).

In Figure 1(a), the gray paths represent gradient descent trajectories from di ff erent 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 p mm ⋆ (depicted by - - -). Moreover, as the iteration count t increases, the inner product GLYPH<10> p ( t ) / ∥ p ( t ) ∥ , p mm ⋆ / ∥ p mm ⋆ ∥ GLYPH<11> consistently increases. Figure 1(c) also depicts the directional convergence of gradient descent from various initializations on 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 2 Suppose for all i ∈ [ n ] and t , opt i , Yi = 1 and γ it < γ i opt i . Also assume W ∈ R d × d is full-rank. Then p mm ⋆ exists - i.e. (ATT-SVM) is feasible for optimal indices α i ← opt i .

2.2 Local convergence of the attention weights p

Theorem 2 on the global convergence of gradient descent serves as a prelude to 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 2 (SVM-Neighbor and Locally-Optimal Tokens) Fix token indices α = ( α i ) n i = 1 for which (ATT-SVM) is feasible to obtain p mm = p mm ( α ) . Consider tokens T i ⊂ [ T ] such that ( k i α i -k it ) ⊤ p mm = 1 for all t ∈ T i. We refer to T i as SVM-neighbors of k i α i . Additionally, tokens with indices α = ( α i ) n i = 1 are called locally-optimal if for all i ∈ [ n ] and t ∈ T i scores per Definition 1 obey γ i α i > γ it. Associated p mm 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 q and a scalar µ > 0, we define cone µ ( q ) as the set of vectors p ∈ R d such that the correlation coe ffi cient between p and q is at least 1 -µ :

Given R > 0, the intersection of cone µ ( q ) and the set { p ∈ R d | ∥ p ∥ ≥ R } is denoted as C µ, R ( q ):

Optimal Token

Locally-Optimal Token

Non-Optimal Token

Next, we demonstrate the existence of parameters µ = µ ( α ) > 0 and R > 0 such that when R is su ffi ciently large, there are no stationary points within C µ, R ( p mm ). Further, the gradient descent initialized within C µ, R ( p mm ) converges in direction to p mm / ∥ p mm ∥ ; refer to Figure 2 for a visualization.

Theorem 3 (Local Convergence of Gradient Descent) Suppose Assumption A on the loss function ℓ holds and assume α = ( α i ) n i = 1 are indices of locally-optimal tokens per Definition 2. Then, there is a constant µ = µ ( α ) ∈ (0 , 1) and R > 0 such that C µ, R ( p mm ) does not contain any stationary points. Further, for any starting point p (0) ∈ C µ, R ( p mm ) , gradient descent iterates p ( t + 1) = p ( t ) -η ∇L ( p ( t )) on (ERM) with stepsize η ≤ 1 / Lp satisfies lim t →∞ ∥ p ( t ) ∥ = ∞ and lim t →∞ p ( t ) / ∥ p ( t ) ∥ = p mm / ∥ p mm ∥ .

To further illustrate Theorem 3, we can consider Figure 1(b) where n = 1 and T = 3. In this figure, the point (0 , 0) represents the non-optimal tokens, while (1 , 0) represents the locally optimal token. Additionally, the gray paths represent the trajectories of gradient descent initiated from di ff erent points. By observing the figure, we can see that gradient descent, when properly initialized, converges towards the direction of p mm (depicted by - - -). This direction of convergence e ff ectively 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 ∥ p ∥ → ∞ . To address this, we introduce the set of LMM directions as follows:

The following theorem establishes the tightness of these directions: It demonstrates that for any candidate q < P mm , its local regularization path within an arbitrarily small neighborhood will provably not converge in the direction of q .

Theorem 4 Fix q < P mm with unit ℓ 2 norm. Assume that token scores are distinct (namely γ it , γ i τ for t , τ ) and key embeddings k it are in general position (see Theorem 7). Fix arbitrary ϵ > 0 , R 0 > 0 . Define the local regularization path of q as its ( ϵ, R 0 ) -conic neighborhood:

Then, either lim R →∞ ∥ ¯ p ( R ) ∥ < ∞ or lim R →∞ ¯ p ( R ) / ∥ ¯ p ( R ) ∥ , q . In both scenarios lim R →∞ ¯ p ( R ) / R , q .

The result above nicely complements Theorem 3, which states that when gradient descent is initialized above a threshold ( ∥ p (0) ∥ ≥ R 0) in an LMM direction, ∥ p ( t ) ∥ 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 ∥ p ∥ ≥ R 0), the optimal solution path will not converge along q < P mm .

3 Joint Convergence of Head v and Attention Weights p

In this section, we extend the preceding results to the general case of joint optimization of head v and attention weights p using a logistic loss function. To this aim, we focus on 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 v can exhibit its own implicit bias to a max-margin solution. Concretely, define the attention features x p i = X ⊤ i S ( K i p ) and define the dataset D p = ( Yi , x p i ) n i = 1 . If this dataset D p is linearly separable, then fixing p and optimizing only v will converge in the direction of the standard max-margin classifier

after setting inputs to the attention features r i ← x p i [22]. This motivates a clear question:

Under what conditions, optimizing v , p jointly 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

2

0

2

4

6

8

10

12

14

2

0

2

4

6

8

10

12

(0,0), (1,1), Y=1

(0,0), (1,-1), Y=-1

(0,0), (0.5,1), Y=1

2

0

2

4

6

8

10

12

14

2

0

2

4

6

8

10

12

(0,0), (1,1), Y=1

(0,0), (1,-1), Y=-1

(0,0), (0.5,1.5), Y=1

0

30

60

90

120

150

0.4

0.5

0.6

0.7

0.8

0.9

1.0

Probabilities

Softmax prob.

Logistic prob.

Iterations

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 x i α i , i ∈ [ n ] selected by p , when solving (SVM) with r i ← x i α i , all of these tokens become support vectors of (SVM). (2) We will then relax this condition to uncover a more general implicit bias for p that distinguish support vs 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 be 1 / ∥ v mm ∥ . Our first insight in quantifying the joint implicit bias is that, optimal tokens admit a natural definition: Those that maximize the downstream label margin when selected. This is formalized below where we assume that: (1) Selecting the token indices α = ( α i ) n i = 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) Let Γ > 0 be the label margin when solving (SVM) with r i ← x i α i . There exists ν > 0 such that for all p , solving (SVM) with r i ← x p i results in a label margin of at most Γ -ν · max i ∈ [ n ](1 -s i α i ) where s i = S ( K i p ) .

Example: To gain intuition, let us fix a ∈ R d and consider the dataset obeying x i 1 = Yi · a and ∥ x it ∥ < ∥ a ∥ for all t ≥ 2 and all i ∈ [ n ]. For this dataset, we can choose α i = 1, v mm = a / ∥ a ∥ 2 , Γ = 1 / ∥ v mm ∥ = ∥ a ∥ and ν = ∥ a ∥ -sup i ∈ [ n ] , t ≥ 2 ∥ x it ∥ .

Theorem 5 Consider the ridge-constrained solutions ( v r , p R ) of (ERM) defined as

Suppose there are token indices α = ( α i ) n i = 1 for which ∥ p mm ( α ) ∥ exists (ATT-SVM is feasible) and Assumption C holds for some Γ , ν > 0 . Then, lim R →∞ p R / R = p mm / ∥ p mm ∥ , where p mm is the solution of (ATT-SVM) ; and lim r →∞ v r / r = v mm / ∥ v mm ∥ , where v mm is the solution of (SVM) with r i = x i α i .

As further discussion, consider Figure 3(a) where we set n = 3 , T = 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 v mm = [0 , 1] is the optimal classifier direction (depicted by - - -). Because of this, p mm will separate optimal ⋆ tokens from tokens at the (0 , 0) coordinate via (ATT-SVM) and its direction is dictated by yellow and teal colored ⋆ s which are the support vectors.

3.2 General solution when selecting one token per input

Can we relax Assumption C, and if so, what is the resulting behavior? Consider the scenario where the optimal p diverges to ∞ and ends up selecting one token per input. Suppose this p selects some coordinates α = ( α i ) n i = 1 . Let S ⊂ [ n ] be the set of indices where the associated token x i α i is a 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 t , α i over the input set ¯ S ⊂ [ 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 ∈ ¯ S , instead, it su ffi ces to just select

0

50

100

150

200

Iterations

0.2

0.4

0.6

0.8

1.0

Softmax probability

Normalized GD

GD (Constant

)

0

50

100

150

200

Iterations

0

5

10

15

20

Attention norm ||

p

||

Normalized GD

GD (Constant

)

2

0

2

4

6

8

10

12

6

4

2

0

2

4

6

C=1

C=2~9

C=10

C=100

(

x

) =

x

(

x

) =

log

(1 +

e

x

)

(0,0), non-opt token ( =0)

(1,1), optimal token ( =C)

(0,0), non-opt token ( =0)

(1,-1), optimal token ( =1)

(a) Evolution of softmax probability

(b) Evolution of attention weights

these tokens (asymptotically). This results in the following relaxed SVM problem:

Here, p ⊤ ( k i α i -k it ) ≥ 0 corresponds to the selection idea. Building on this intuition, the following theorem captures the generalized behavior of the joint regularization path.

Theorem 6 Consider the same (ERM) problem as discussed in Theorem 5. Suppose S ( K i p R ) α i → 1 , i.e., the tokens ( α i ) n i = 1 are asymptotically selected. Let v mm be the solution of (SVM) with r i = x i α i and S be its set of support vector indices. Suppose Assumption C holds over S i.e. having s i α i < 1 shrinks the margin when (SVM) is only solved over S ⊂ [ n ] . Then, lim r →∞ v r / r = v mm / ∥ v mm ∥ and lim R →∞ p R / R = p relax / ∥ p relax ∥ , where p relax is the solution of (10) with ( α i ) n i = 1 choices.

To illustrate this numerically, consider Figure 3(b) which modifies Figure 3(a) by pushing the yellow ⋆ to the northern position (0 . 5 , 1 . 5). We still have v mm = [0 , 1] however the yellow ⋆ is no longer a 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 mm using a gray dashed line.

Wefurther investigate the evolution of softmax and logistic output probabilities throughout the training process of Figure 3(a), and the results are illustrated in Figure 3(c). The averaged softmax probability of optimal tokens is represented by the red curve and is calculated as 1 n P n i = 1 max t ∈ [ T ] S ( K i p ) t . An achievement of 1 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 is determined by 1 / n P n i = 1 1 / (1 + e -Yi · f ( X i ) ). This probability also reaches a value of 1, suggesting that the inputs are correctly classified.

4 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 a locally-optimal token within the input sequence and suppresses the rest, the associated attention map S ( Xp ) will (eventually) be a 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 η for training. The dataset model follows Figure 1(c). The softmax probability shown in Figure 4(a) is defined as 1 n P n i = 1 max t ∈ [ T ] S ( K i p ) t . 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 ∥ p ∥ .

The red curves in both figures represent the normalized gradient method, which updates the model parameters p using p ( t + 1) = p ( t ) -η ∇L ( p ( t )) / ∥∇L ( p ( t )) ∥ with η = 0 . 1. The blue curves correspond

(a) Input image

(b) Epoch 0

(c) Epoch 100

(d) Epoch 200

0

100

200

300

400

Epoch

0.2

0.3

0.4

0.5

0.6

0.7

0.8

0.9

1

Sparsity

0

1

2

3

4

5

6

7

W F

(×10 3 )

to gradient descent with constant learning rate given by p ( t + 1) = p ( t ) -η ∇L ( p ( t )) with η = 1. Observe that the normalized gradient method achieves a softmax probability of 1 quicker as vanilla GD su ff ers from vanishing gradients. This is visible in Figure 4(b) where blue norm curve levels o ff .

Real experiments. To study softmax sparsity and the evolution of attention weights throughout training, we train a vision transformer (ViT-base) model [23] from scratch, utilizing the CIFAR10 dataset [24] for 400 epochs with fixed learning rate 3 × 10 -3 . ViT tokenizes an image into 16 × 16 patches, thus, its softmax attention maps can be easily visualized. We examine the average attention map - associated with the [CLS] token - computed from all 12 attention heads within the model. Figure 6 provides a visual representation of the resulting attention weights (16 × 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 a dense pattern. However, as the training progresses, the attention map gradually becomes sparser and the attention mechanism begins to 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 a 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 d nnz( s ) where s is the softmax probability vector. The soft-sparsity is computed as the ratio of the ℓ 1-norm to the squared ℓ 2-norm, defined as d nnz( s ) = ∥ s ∥ 1 / ∥ s ∥ 2 . d nnz( s ) takes values between 1 to T = 256 and a smaller value indicates a sparser vector. Also note that ∥ s ∥ 1 = P T t = 1 s t = 1. Together with sparsity, Figure 7 also displays the Frobenius norm of the combined key-query matrix W 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 ∥ F grows.

Transient optimization dynamics and the influence of the loss function. Theorem 2 shows that the asymptotic direction of gradient descent is determined by p mm ⋆ . 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 di ff erent scores and loss functions. We consider two optimal tokens ( ⋆ ) with scores γ 1 = 1 and γ 2 = C , where C varies. For our analysis, we examine the correlation loss ℓ ( x ) = -x and the logistic loss ℓ ( x ) = log(1 + e -x ).

In essence, as C increases, we can observe that the correlation loss ℓ ( x ) = -x exhibits a bias towards the token with a high score, while the logistic loss is biased towards the token with a low score. The underlying reason for this behavior can be observed from the gradients of individual inputs: ∇L i ( p ) = ℓ ′ i · K ⊤ i S ′ ( Xp ) Xv , where S ′ ( · ) represents the derivative of the softmax function and ℓ ′ i : = ℓ ′ ( Yi · v ⊤ X ⊤ i S ( X i p )). Assuming that p (approximately) selects the optimal tokens, this

simplifies to ℓ ′ i ≈ ℓ ′ ( γ i ) and ∥∇L i ( p ) ∥ ∝ | ℓ ′ ( γ i ) | · γ i . With the correlation loss, | ℓ ′ | = 1, resulting in ∥∇L i ( p ) ∥ ∝ γ i , 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.e., | ℓ ′ | = e -x / (1 + e -x ) ≈ e -x . Consequently, ∥∇L i ( p ) ∥ ∝ γ ie -γ i ≈ e -γ i , indicating that a smaller score leads to a larger gradient. These observations explain the empirical behavior we observe.

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 use logistic loss or, more generally, exponentially-tailed losses to make connections to margin maximization. These results are also extended to non-separable data by [30, 31, 21]. 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 our optimization approach for v and existing works, the optimization of p stands out as significantly di ff erent. Firstly, our optimization problem is 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, 50]. Self-attention di ff ers 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 to 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 , p 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 is 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 implicit bias in these models.

6 Discussion

We have provided a thorough optimization-theoretic characterization of the fundamental attention model f ( X ) = v ⊤ X ⊤ S ( XWp ) by formally connecting it to max-margin problems. We first established the convergence of gradient descent on p (or equivalently W ) in isolation. We also explored joint convergence of ( v , p ) via regularization path which revealed surprising implicit biases such as (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, it would be interesting to characterize the non-separable geometry by leveraging results developed for logistic regression analysis [31, 22]. Ideas from such earlier results can also be useful for characterizing the non-asymptotic / transient dynamics of how gradient 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 in this work lays the groundwork of these future extensions.

References

Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio. Neural machine translation by jointly learning to align and translate. The International Conference on Learning Representations , 2015.

Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, and et al. Language models are few-shot learners. In Advances in neural information processing systems , volume 33, pages 1877-1901, 2020.

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 arXiv:2107.03374 , 2021.

Alec Radford, Je ff rey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAI blog , 1(8):9, 2019.

Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. arXiv preprint arXiv:2204.02311 , 2022.

Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in neural information processing systems , volume 30, 2017.

Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers) , pages 4171-4186, Minneapolis, Minnesota, June 2019. Association for Computational Linguistics.

OpenAI. Gpt-4 technical report. arXiv preprint arXiv:2303.08774 , 2023.

Rishi Bommasani, Drew A 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 arXiv:2108.07258 , 2021.

Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning , pages 8821-8831. PMLR, 2021.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations , 2021.

Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International conference on machine learning , pages 8748-8763. PMLR, 2021.

Danny Driess, Fei Xia, Mehdi SM Sajjadi, 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 arXiv:2303.03378 , 2023.

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.

Scott Reed, Konrad Zolna, Emilio Parisotto, Sergio Gomez Colmenarejo, Alexander Novikov, Gabriel Barth-Maron, Mai Gimenez, Yury Sulsky, Jackie Kay, Jost Tobias Springenberg, et al. A generalist agent. arXiv preprint arXiv:2205.06175 , 2022.

Brian Lester, Rami Al-Rfou, and Noah Constant. The power of scale for parameter-e ffi cient prompt tuning. In Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing , pages 3045-3059, 2021.

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.

Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. arXiv preprint arXiv:2101.00190 , 2021.

Saharon Rosset, Ji Zhu, and Trevor Hastie. Margin maximizing loss functions. Advances in neural information processing systems , 16, 2003.

Arun Suggala, Adarsh Prasad, and Pradeep K Ravikumar. Connecting optimization and regularization paths. Advances in Neural Information Processing Systems , 31, 2018.

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.

Daniel Soudry, Elad Ho ff er, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro. The implicit bias of gradient descent on separable data. The Journal of Machine Learning Research , 19(1):2822-2878, 2018.

Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas. Attention is not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning , pages 2793-2803. PMLR, 2021.

Alex Krizhevsky, Vinod Nair, and Geo ff rey Hinton. The cifar-10 dataset. online: http: // www. cs. toronto. edu / kriz / cifar. html , 55(5), 2014.

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.

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.

Ziwei Ji and Matus Telgarsky. Characterizing the implicit bias via a primal-dual analysis. In Algorithmic Learning Theory , pages 772-804. PMLR, 2021.

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 in neural information processing systems , 33:22182-22193, 2020.

Ziwei Ji and Matus Telgarsky. Directional convergence and alignment in deep learning. In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems , volume 33, pages 17176-17186. Curran Associates, Inc., 2020.

Ziwei Ji and Matus Telgarsky. Risk and parameter convergence of logistic regression. arXiv preprint arXiv:1803.07300 , 2018.

Ziwei Ji and Matus Telgarsky. The implicit bias of gradient descent on nonseparable data. In Conference on Learning Theory , pages 1772-1798. PMLR, 2019.

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.

Chulhee Yun, Shankar Krishnan, and Hossein Mobahi. A unifying view on implicit bias in training linear neural networks. arXiv preprint arXiv:2010.02501 , 2020.

Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini. Implicit regularization for optimal sparse recovery. Advances in Neural Information Processing Systems , 32:2972-2983, 2019.

Ehsan Amid and Manfred K Warmuth. Winnowing with gradient descent. In Conference on Learning Theory , pages 163-182. PMLR, 2020.

Ehsan Amid and Manfred KK Warmuth. Reparameterizing mirror descent as gradient descent. Advances in Neural Information Processing Systems , 33:8430-8439, 2020.

Yuanzhi Li, Colin Wei, and Tengyu Ma. Towards explaining the regularization e ff ect of initial large learning rate in training neural networks. arXiv preprint arXiv:1907.04595 , 2019.

Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant. Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process. In Conference on learning theory , pages 483-513. PMLR, 2020.

Je ff Z HaoChen, Colin Wei, Jason D Lee, and Tengyu Ma. Shape matters: Understanding the implicit bias of the noise covariance. arXiv preprint arXiv:2006.08680 , 2020.

Zhiyuan Li, Tianhao Wang, and Sanjeev Arora. What happens after SGD reaches zero loss? -a mathematical framework. In International Conference on Learning Representations , 2022.

Alex Damian, Tengyu Ma, and Jason Lee. Label noise sgd provably prefers flat global minimizers. arXiv preprint arXiv:2106.06530 , 2021.

Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, Dean P Foster, and Sham Kakade. The benefits of implicit regularization from sgd in least squares problems. Advances in Neural Information Processing Systems , 34:5456-5468, 2021.

Qian Qian and Xiaoyuan Qian. The implicit bias of adagrad on separable data. Advances in Neural Information Processing Systems , 32, 2019.

Bohan Wang, Qi Meng, Huishuai Zhang, Ruoyu Sun, Wei Chen, and Zhi-Ming Ma. Momentum doesn't change the implicit bias. arXiv preprint arXiv:2110.03891 , 2021.

Bohan Wang, Qi Meng, Wei Chen, and Tie-Yan Liu. The implicit bias for adaptive optimization algorithms on homogeneous neural networks. In International Conference on Machine Learning , pages 10849-10858. PMLR, 2021.

Ziwei Ji, Nathan Srebro, and Matus Telgarsky. Fast margin maximization via dual acceleration. In International Conference on Machine Learning , pages 4860-4869. PMLR, 2021.

Jianpeng Cheng, Li Dong, and Mirella Lapata. Long short-term memory-networks 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.

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.

Romain Paulus, Caiming Xiong, and Richard Socher. A deep reinforced model for abstractive summarization. In International Conference on Learning Representations , 2018.

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.

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.

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 on Machine Learning , pages 19050-19088. PMLR, 2022.

Tolga Ergen, Behnam Neyshabur, and Harsh Mehta. Convexifying transformers: Improving optimization and understanding of transformer networks. arXiv:2211.11052 , 2022.

Pierre Baldi and Roman Vershynin. The quarks of attention. arXiv:2202.08371 , 2022.

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.

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.

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.

Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and generalization in neural networks. arXiv preprint arXiv:1806.07572 , 2018.

Samet Oymak and Mahdi Soltanolkotabi. Toward moderate overparameterization: Global convergence guarantees for training shallow neural networks. IEEE Journal on Selected Areas in Information Theory , 1(1):84-105, 2020.

Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song. A convergence theory for deep learning via over-parameterization. In International Conference on Machine Learning , pages 242-252. PMLR, 2019.

Vladimir Vapnik. Estimation of dependences based on empirical data . Springer Science & Business Media, 2006.

Peter Bartlett. For valid generalization the size of the weights is more important than the size of the network. Advances in neural information processing systems , 9, 1996.

Albert B Noviko ff . On convergence proofs for perceptrons. Technical report, STANFORD RESEARCH INST MENLO PARK CA, 1963.

Peter Bartlett, Yoav Freund, Wee Sun Lee, and Robert E Schapire. Boosting the margin: A new explanation for the e ff ectiveness of voting methods. The annals of statistics , 26(5):1651-1686, 1998.

Tong Zhang and Bin Yu. Boosting with early stopping: Convergence and consistency. Annals of Statistics , page 1538, 2005.

Matus Telgarsky. Margins, shrinkage, and boosting. In International Conference on Machine Learning , pages 307-315. PMLR, 2013.

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.

Mahdi Soltanolkotabi, Dominik Stöger, and Changzhi Xie. Implicit balancing and regularization: Generalization and convergence guarantees for overparameterized asymmetric matrix sensing. arXiv:2303.14244 , 2023.

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.

Samet Oymak and Mahdi Soltanolkotabi. Overparameterized nonlinear learning: Gradient descent takes the shortest path? In International Conference on Machine Learning , pages 4951-4960. PMLR, 2019.

Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. arXiv preprint arXiv:1810.02032 , 2018.

Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems , 32, 2019.

Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. arXiv preprint arXiv:1906.05890 , 2019.

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 Theory , pages 1305-1338. PMLR, 2020.

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.

Gal Vardi, Ohad Shamir, and Nati Srebro. On margin maximization in linear and relu networks. Advances in Neural Information Processing Systems , 35:37024-37036, 2022.

Navid Azizan, Sahin Lale, and Babak Hassibi. Stochastic mirror descent on overparameterized nonlinear models. IEEE Transactions on Neural Networks and Learning Systems , 33(12):77177727, 2021.

Navid Azizan and Babak Hassibi. Stochastic gradient / mirror descent: Minimax optimality and implicit regularization. In International Conference on Learning Representations .

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 click-through rate prediction. In Proceedings of the 24th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining , pages 1059-1068, 2018.

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-Dimensional Sparse Data , pages 1-4, 2019.

Fei Sun, Jun Liu, Jian Wu, Changhua Pei, Xiao Lin, Wenwu Ou, and Peng Jiang. Bert4rec: Sequential recommendation with bidirectional encoder representations from transformer. In Proceedings of the 28th ACM International Conference on Information and Knowledge Management , pages 1441-1450, 2019.

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 Macdu ff 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.

Michael Janner, Qiyang Li, and Sergey Levine. Reinforcement learning as one big sequence modeling problem. In ICML 2021 Workshop on Unsupervised Reinforcement Learning , 2021.

Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In Proceedings of the 39th International Conference on Machine Learning , 2022.

Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations , 2020.

Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-e ffi cient image transformers & distillation through attention. In International Conference on Machine Learning , pages 10347-10357. PMLR, 2021.

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, 2021.

Hyunjik Kim, George Papamakarios, and Andriy Mnih. The lipschitz constant of self-attention. In International Conference on Machine Learning , pages 5562-5571. PMLR, 2021.

Jiri Hron, Yasaman Bahri, Jascha Sohl-Dickstein, and Roman Novak. Infinite attention: Nngp and ntk for deep attention networks. In International Conference on Machine Learning , pages 4376-4386. PMLR, 2020.

Greg Yang. Tensor programs ii: Neural tangent kernel for any architecture. arXiv preprint arXiv:2006.14548 , 2020.

Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, and Lukasz Kaiser. Universal transformers. In International Conference on Learning Representations , 2018.

Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar. Are transformers universal approximators of sequence-to-sequence functions? In International Conference on Learning Representations , 2019.

Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos. Looped transformers as programmable computers. arXiv:2301.13196 , 2023.

Yoav Levine, Noam Wies, Or Sharir, Hofit Bata, and Amnon Shashua. Limits to depth e ffi -ciencies of self-attention. In Advances in Neural Information Processing Systems , volume 33, pages 22640-22651, 2020.

Charlie Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. arXiv preprint arXiv:2103.07601 , 2021.

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:2109.01652 , 2021.

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.

Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant. What can transformers learn in-context? a case study of simple function classes. Advances in Neural Information Processing Systems , 35:30583-30598, 2022.

Yingcong Li, M Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference on Machine Learning , 2023.

Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Ed Chi, Quoc Le, and Denny Zhou. Chain of thought prompting elicits reasoning in large language models. arXiv preprint arXiv:2201.11903 , 2022.

Guhao Feng, Yuntian Gu, Bohang Zhang, Haotian Ye, Di He, and Liwei Wang. Towards revealing the mystery behind chain of thought: a theoretical perspective. arXiv preprint arXiv:2305.15408 , 2023.

Yingcong Li, Kartik Sreenivasan, Angeliki Giannou, Dimitris Papailiopoulos, and Samet Oymak. Dissecting chain-of-thought: A study on compositional in-context learning of mlps. arXiv preprint arXiv:2305.18869 , 2023.

Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in 1-layer transformer. arXiv:2305.16380 , 2023.

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.

Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308.16898 , 2023.

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 p ∈ R d with a fixed v ∈ R d choice. Section C provides the proof of regularized path applied to the general case of joint optimization of head v and attention weights p using a logistic loss function. Section D presents the regularized path applied to a more general model f ( X ) = ψ ( X ⊤ S ( XW ⊤ p )) with a nonlinear head ψ . Section E provides implementation details. Finally, Section F discusses additional related work on implicit bias and self-attention.

Table of Contents

A Addendum to Section 1

A.1 Preliminaries on the Training Risk

By our assumption ψ : R d → R and ℓ : R → R are di ff erentiable functions. Recall the objective

with the generic prediction model ψ ( X ⊤ S ( Kp )) and K = XW ⊤ .

Here, we write down the gradients of W and p in (11) to highlight the connection. Set q : = W ⊤ p , z { X } : = X ⊤ S ( Kp ), and a { X } : = Kp . Given X and using K = XW ⊤ , we have that

where

Setting ψ ( z ) = v ⊤ z for linear head, we obtain

Recalling (12b) and (12c), and defining ℓ ′ i : = ℓ ′ ( Yi · ψ ( z { X i } )) ∈ R , we have that

Setting ψ ( z ) = v ⊤ z for linear head and γ i = Yi · X i v , we obtain

Lemma 3 (Key Lemma) For any p , q ∈ R d , let a = Kq , s = S ( Kp ) , and γ = Xv . Set

We have that

Proof. Set ¯ γ = P T t = 1 γ t s t . We have

Then,

Since

we obtain 2

Here, ± on the right handside uses the fact that

A.2 Proof of Lemma 1

Proof. Let us prove the result for a general step size sequence ( η t ) t ≥ 0 . On the same training data ( Yi , X i ) n i = 1 , recall the objectives ˜ L ( p ) = 1 n P n i = 1 ℓ ( Yi · ψ ( X ⊤ i S ( X i p ))) and L ( W ) = 1 n P n i = 1 ℓ ( Yi · ψ ( X ⊤ i S ( X i W ⊤ u ))). Suppose claim is true till iteration t . For iteration t + 1, using W ( t ) ⊤ u = p ( t ), define and observe that

for all i ∈ [ n ].

Thus, using (14), we have that

Consequently, we found that gradient is rank-1 with left singular space equal to u , i.e.,

Since W ( t )'s left singular space is guaranteed to be in u (including W (0) by initialization), we only need to study the right singular vector. Using the induction till t , this yields

This concludes the induction.

B Addendum to Section 2

B.1 Descent and Gradient Correlation Conditions

The lemma below identifies conditions under which p mm ⋆ is a global descent direction for L ( p ).

Lemma 4 Suppose ℓ ( · ) is a strictly decreasing di ff erentiate loss function and Assumption B holds. Then, for all p ∈ R d , the training loss (ERM) obeys GLYPH<10> ∇L ( p ) , p mm ⋆ GLYPH<11> < 0 .

Proof. Set

Let us recall the gradient evaluated at p which is given by

To proceed, we will prove that individual summands are all strictly negative. To show that, without losing generality, let us focus on the first input and drop the subscript i for cleaner notation. This yields

Without losing generality, assume optimal token is the first one and γ t is a constant for all t ≥ 2.

To proceed, we will prove the following: Suppose γ = γ t ≥ 2 is constant, γ 1 , ¯ a 1 are the largest indices of γ , ¯ a . Then, for any s obeying P t ∈ [ T ] s t = 1 , s t ≥ 0, we have that ¯ a ⊤ diag( s ) γ -¯ a ⊤ ss ⊤ γ > 0. To see this, we write

To proceed, let γ gap = γ 1 -γ and a gap = ¯ a 1 -max t ≥ 2 a t . With these, we obtain

On the other hand, by our assumption ℓ ′ i < 0. Hence, infimum'ing (22) over all inputs, multiplying by ℓ ′ i and using (19) give the desired result.

This implies that

Note that

Lemma 5 (Gradient Correlation Conditions) Consider n = 1 and let p mm = p mm ⋆ be (ATT-SVM) solution separating α = opt from remaining tokens of input X . Suppose ℓ ( · ) is a strictly decreasing di ff erentiate loss function and Assumption B holds. For any choice of π > 0 , there exists R : = R π such that, for any p with ∥ p ∥ ≥ R, we have

Above, observe that as R →∞ , we eventually get to set π = 0.

Proof. The proof is similar to Lemma 4 at a high-level. However, we also need to account for the impact of p besides p mm in the gradient correlation. The main goal is showing that p mm is the near-optimal descent direction, thus, p cannot significantly outperform it.

To proceed, let ¯ p = ∥ p mm ∥ p / ∥ p ∥ , M = sup t ∥ k t ∥ , Θ = 1 / ∥ p mm ∥ , s = S ( Kp ), a = K ¯ p , ¯ a = Kp mm . Without losing generality assume opt = 1. Set γ = γ t ≥ 2. Repeating the proof of Lemma 4 yields

Given π , for su ffi ciently large R , we wish to show that

We consider two scenarios.

Scenario 1: ∥ ¯ p -p mm ∥ ≤ ϵ : = π/ (2 M ). In this scenario, for any token, we find that

Consequently, we obtain

Also noticing ¯ a 1 -P T t ≥ 2 ¯ a t s t P T t ≥ 2 s t ≥ 1 (thanks to p mm satisfying ≥ 1 margin), this implies (23).

Scenario 2: ∥ ¯ p -p mm ∥ ≥ ϵ : = π/ (2 M ). In this scenario, for some ν = ν ( ϵ ) and τ ≥ 2, we have that

Here τ = arg max t ≥ 2 ¯ p ⊤ k t denotes the nearest point to k 1. Recall that s = S ( ¯ R a ) where ¯ R = ∥ p ∥ / ∥ p mm ∥ . To proceed, split the tokens into two groups: Let N be the group of tokens obeying ¯ p ⊤ ( k 1 -k t ) ≤ 1 -ν for t ∈ N and [ T ] - N be the rest. Observe that

Set ¯ M = M / Θ and note that ∥ a t ∥ ≤ ∥ p mm ∥ · ∥ k t ∥ ≤ ¯ M . Using ¯ p ⊤ ( k 1 -k t ) ≤ 1 -ν over t ∈ N and plugging in the above bound, we obtain

Using the fact that ¯ a 1 -P T t ≥ 2 ¯ a t s t P T t ≥ 2 s t ≥ 1, the above implies (23) with π ′ = 2 ¯ MTe -¯ R ν -ν . To proceed, choose R π = ν -1 Θ -1 log(2 ¯ MT /π ) to ensure π ′ ≤ π .

The following lemma states the descent property of gradient descent for L ( p ) under Assumption A. It is important to note that although the infimum of the optimization problem is L ∗ , it is not achieved at any finite p . Additionally, there are no finite critical points p .

Lemma 6 Under Assumption A, the function L ( p ) is Lp-smooth, where

Furthermore, if η ≤ 1 / Lp, then, for any initialization p (0) , with the GD sequence p ( t + 1) = p ( t ) -η ∇L ( p ( t )) , we have

for all t ≥ 0 . This implies that

Proof. Recall that we defined γ i = Yi · X i v and a i = K i p . The gradient of L ( p ) is given by

Note that for any p ∈ R d , the Jacobian of S ( K i p ) is given by

The Jacobian (27) together with the definition of the softmax function S ( · ) implies that ∥ ∂ S ( K i p ) /∂ p ∥ ≤ ∥ K i ∥ . Hence, for any p , ˙ p ∈ R d , we have

and

Here, the last inequality uses the fact that | ab -cd | ≤ | d || a -c | + | a || b -d | .

Next, for any p , ˙ p ∈ R d , we have

where the second inequality follows from the fact that | ab -cd | ≤ | d || a -c | + | a || b -d | and the third inequality uses Assumption A.

Substituting (28a) and (28b) into (29), we get

where Lp is defined in (24).

The remaining proof follows standard gradient descent analysis (see e.g. [22, Lemma 10]). Since L ( p ) is Lp -smooth, we get

where the last inequality follows from our assumption on the stepsize.

The above inequality implies that

where the right hand side is upper bounded by a finite constant. This is because, by Assumption A, L ( p (0)) < ∞ and L ∗ ≤ L ( p ( t )), where L ∗ denotes the minimum objective.

Finally, (30) yields the expression (26).

In the following lemma, we demonstrate the existence of parameters µ = µ ( α ) > 0 and R µ > 0 such that when R µ is su ffi ciently large, there are no stationary points within C µ, R µ ( p mm ). Additionally, we provide the local gradient correlation condition.

Lemma 7 (Local Gradient Condition) Suppose Assumption A on the loss function ℓ holds. Let α = ( α i ) n i = 1 be indices of locally-optimal tokens per Definition 2.

L1. There exists a positive scalar µ = µ ( α ) > 0 such that for su ffi ciently large ¯ R µ , no stationary point exists within C µ, ¯ R µ ( p mm ) , where C µ, ¯ R µ is defined in (8) .

L2. For all q , p ∈ cone µ ( p mm ) with ∥ q ∥ = ∥ p mm ∥ and ∥ p ∥ ≥ ¯ R µ with same ¯ R µ choice as ( L1. ) , there exist dataset dependent constants C , c > 0 such that

Here, ¯ A = max i ∈ [ n ] , t ,τ ∈ [ T ] ∥ k it -k i τ ∥ and Θ = 1 / ∥ p mm ∥ .

L3. For any π > 0 , there exists R π such that R π ≥ ¯ R µ and all p ∈ C µ, R π ( p mm ) obeys

Proof. Let p mm = p mm ( α ) be the solution of (ATT-SVM). Recall

Let ( T i ) n i = 1 be the sets of all SVM-neighbors per Definition 2. Let ¯ T i = [ T ] - T i - { α i } be the set of non-SVM-neighbor tokens, i ∈ [ n ]. Let

When ¯ T i = ∅ for all i ∈ [ n ] (i.e. globally-optimal indices), we set δ = ∞ as all non-neighbor related terms will disappear. Since p mm is the max-margin model ensuring ( k i α i -k it ) ⊤ p mm ≥ 1 for all i ∈ [ n ], the following inequalities hold for all q ∈ cone µ ( p mm ) , ∥ q ∥ = ∥ p mm ∥ and all i ∈ [ n ] , t ∈ T i , τ ∈ ¯ T i :

Here, we used ∥ q -p mm ∥ 2 / ∥ p mm ∥ 2 ≤ 2 µ which implies ∥ q -p mm ∥ ≤ p 2 µ/ Θ .

L1. and L2. . Now that the choice of local cone is determined, we need to prove the main claims. We will lower bound -q ⊤ ∇L ( p ) and establish its strict positivity for ∥ p ∥ ≥ R , where R = ¯ R µ . This will show that there is no stationary point as a by product.

Consider any q ∈ R d satisfying ∥ q ∥ = ∥ p mm ∥ . To proceed, we write the gradient correlation following (18) and (21)

where we denoted ℓ ′ i = ℓ ′ ( Yi · v ⊤ X ⊤ i S ( K i p )), a i = K i q , a ′ i = K i p , s i = S ( K i p ).

Using (33), for all t ∈ T i , τ ∈ ¯ T i , for all p ∈ C µ, R ( p mm ), we have that

Consequently, we can bound the softmax probabilities s i = S ( K i p ) as follows: For all i ∈ [ n ],

Recall scores γ it = Yi · v ⊤ x it . Define the score gaps over neighbors:

Define the α -dependent global scalar Γ = sup i ∈ [ n ] , t ,τ ∈ [ T ] | γ it -γ i τ | . Let us focus on a fixed datapoint i ∈ [ n ], assume (without losing generality) α i = 1, and drop subscripts i , that is, α : = α i = 1, X : = X i , Y : = Yi , K : = K i , a ′ = Kp , a = Kq , s = S ( Kp ), γ = Y · Xv , γ gap : = γ gap i , ¯ γ gap : = ¯ γ gap i , Q : = Qi , and S : = S i . Directly applying Lemma 3, we obtain

To proceed, let us decouple the non-neighbors within P T t ≥ 2 ( a 1 -a t ) s t ( γ 1 -γ t ) via

It follows from (32) that

Aggregating these, we found

To proceed, let us upper / lower bound the gradient correlation. We use two bounds depending on q ∈ cone µ ( p mm ) (Case 1) or general q ∈ R d (Case 2).

Case 1: q ∈ cone µ ( p mm ). Since 1 . 5 ≥ a 1 -a t ≥ 0 . 5 following (33), we find

Next we claim that S dominates ((1 -s 1 ) 2 + Q ) for large R . Specifically, we wish for

Now choose R ≥ δ -1 log( T ) / Θ to ensure Q ≤ S since Q ≤ Te -R Θ δ S . Consequently

Combining these, what we wish is ensured by guaranteeing

This in turn is ensured for all inputs i ∈ [ n ] by choosing

where γ gap min = min i ∈ [ n ] γ gap i is the global scalar which is the worst case score gap over all inputs. With the above choice of R we guaranteed

Since this holds over all inputs, going back to the gradient correlation (34) and averaging above over all inputs i ∈ [ n ] and plugging back the indices i , we obtain the advertised bound by setting qi = 1 -s i α i (where we set α i = 1 above without losing generality)

Let -ℓ ′ min / max be the min / max values negative loss derivative admits over the ball [ -B , B ] for B = ∥ v ∥ · max i , t ∥ x it ∥ and note that max i ∈ [ n ] ¯ γ gap i > 0 and min i ∈ [ n ] γ gap i > 0 are dataset dependent constants. Then, we declare the constants C = -2 ℓ ′ max · max i ∈ [ n ] ¯ γ gap i > 0 , c = -(1 / 8) ℓ ′ min · min i ∈ [ n ] γ gap i > 0 to obtain the bound

which is the desired statement in (31a).

Case 2: q ∈ R d and ∥ q ∥ = ∥ p mm ∥ . Define ¯ A = max i ∈ [ n ] , t ,τ ∈ [ T ] ∥ k it -k i τ ∥ . For any ∥ q ∥ = ∥ p mm ∥ , we use the fact that

Note that by definition ¯ A Θ ≥ 1. To proceed, we can upper bound

By choosing the same R as in (39) to ensure S dominates ((1 -s 1 ) 2 + Q ) and since ¯ A Θ ≥ 1, we guaranteed

Going back to the gradient correlation (34) and averaging above over all inputs i ∈ [ n ], with the same definition of C > 0, we obtain

To proceed, since (43) holds for any q ∈ R d and ∥ q ∥ = ∥ p mm ∥ , we observe that when choosing q = ∥ p mm ∥ ∥∇L ( p ) ∥ · ∇L ( p ), this implies that

Simplifying Θ = 1 / ∥ p mm ∥ on both sides yields (31b). Incorporating (35) in the bound above provides the exponential upper bound that decay with R .

Combining this with (41), we obtain that for all q , p ∈ cone µ ( p mm ) and ∥ q ∥ ≥ ¯ R µ

This gives the desired result in (31c).

L3.: Establishing gradient correlation.

Our final goal is establishing gradient comparison between p , p mm for the same choice of µ > 0 provided in (32). Define ¯ p = ∥ p mm ∥ p / ∥ p ∥ to be the normalized vector. Set notations a i = K i ¯ p , ¯ a i = K i p mm , and γ i = Yi · X i v .

To establish the result, using (34), we will prove that, for any π > 0, there is su ffi ciently large R = R π such that for any p ∈ C µ, R ( p mm ):

Following (36), for all i ∈ [ n ], for all q ∈ cone µ ( p mm ) with ∥ q ∥ = ∥ p mm ∥ , a ′ = Kq and s = S ( Kp ), we have found

Plugging in a i , ¯ a i in the bound above and assuming π ≤ 1 (w.l.o.g.), (44) is implied by the following stronger inequality

First, we claim that 0 . 5 π P t ∈T i s it ( γ i 1 -γ it ) ≥ 6 Γ A ((1 -s i 1 ) 2 + Qi ) 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

for some constant C 0 > 0. Here, we choose su ffi ciently large C 0 ≥ 64 π to ensure R = R π ≥ ¯ R µ . Following this control over the perturbation term 6 Γ A ((1 -s i 1 ) 2 + Qi ), to conclude with the result, what remains is proving the comparison

To proceed, we split the problem into two scenarios.

Scenario 1: ∥ ¯ p -p mm ∥ ≤ ϵ = π 4 A Θ for some ϵ > 0. In this scenario, for any token, we find that ⊤ mm

Consequently, we obtain

Similarly, a i 1 -a it ≥ 1 -0 . 5 π ≥ 0 . 5. Since all terms a i 1 -a it , s it , γ i 1 -γ it in (47) are nonnegative and ( a i 1 -a it ) s it ( γ i 1 -γ it ) ≤ (1 + 0 . 5 π ) s it ( γ i 1 -γ it ), above implies the desired result (47).

Scenario 2: ∥ ¯ p -p mm ∥ ≥ ϵ = π 4 A Θ . Since ¯ p is not (locally) max-margin, in this scenario, for some i ∈ [ n ], ν = ν ( ϵ ) > 0, and τ ∈ T i , we have that

Here τ = arg max t ∈T i ¯ p ⊤ k it denotes the nearest point to k i 1 (along the ¯ p direction). Note that a nonneighbor t ∈ ¯ T i cannot be nearest because ¯ p ∈ cone µ ( p mm ) and (33) holds. Recall that s i = S ( ¯ R a i ) where ¯ R = ∥ p ∥ Θ ≥ R Θ . To proceed, let a i : = min t ∈T i a i 1 -a it ,

For all i ∈ I , split the tokens into two groups: Let N i be the group of tokens obeying a i 1 -a it ≤ 1 -ν and T i - N i be the rest of the neighbors. Observe that

Thus, using | a i 1 -a it | ≤ 2 A and recalling the definition of γ gap min , observe that

Hence, choosing results in that

Here, the last inequality follows from the fact that P t ∈T i s it ≥ max t ∈T i sit ≥ e -¯ R (1 -2 ν ) P T t = 1 e -¯ R ( a i 1 -a it ) ≥ e -¯ R (1 -2 ν ) / T .

From Assumption A, we have c min ≤ -ℓ ′ ≤ c max for some positive constants c min and c max. It follows from (48) and (50) that

Combing with (49), this is guaranteed by choosing

where ν = ν ( π 4 A Θ ) depends only on π and global problem variables.

Combining this with the prior R choice (46) (by taking maximum), we conclude with the statement.

B.2 Proof of Theorem 1

Proof. This proof is a direct corollary of Lemma 14 which itself is a special case of the nonlinear head Theorem 8. Let us verify that f ( X ) = v ⊤ X ⊤ S ( Xp ) satisfies the assumptions of Lemma 14 where we replace the nonlinear head with linear v . To see this, set the optimal sets to be the singletons O i = { opt i } . Given ( X i , Yi ), let s i = S ( K i p ) and qi = q p i = P t , opt i s it . Recalling score definition γ i = Yi · X i v and setting ν i : = γ i opt i and Zi : = P t , opt i γ it s it , a particular prediction can be written as

To proceed, we demonstrate the choices for C , ϵ > 0. Let C : = -min i ∈ [ n ] , t ∈ [ T ] γ it ∧ 0 and q max = max i ∈ [ n ] qi . Note that Zi ≥ P t , opt i γ it s it ≥ qi γ min ≥ -Cq max. Now, using strict score optimality of opt i 's for all i ∈ [ n ], we set

We conclude by observing Zi ≤ ν i qi P t , opt i γ it s it ν ≤ ν i qi ϵ

iqi as desired.

B.3 Proof of Theorem 2

Proof. We first show that lim t →∞ ∥ p ( t ) ∥ = ∞ . From Lemma 4, we have

where γ i = Yi · X i v and a i = K i p .

It follows from Lemma 4 that GLYPH<10> ∇L ( p ) , p mm ⋆ GLYPH<11> < 0 for all p ∈ R d . Hence, for any finite p , GLYPH<10> ∇L ( p ) , p mm ⋆ GLYPH<11> cannot be equal to zero, as a sum of negative terms. Therefore, there are no finite critical points p , for which ∇L ( p ) = 0. On the other hand, Lemma 6 states ∇L ( p ( t )) → 0 which implies that ∥ p ( t ) ∥ → ∞ .

Next, we provide the directional convergence for the setting n = 1. Let us consider an arbitrary value of ϵ ∈ (0 , 1) and set π = ϵ/ (1 -ϵ ). As lim t →∞ ∥ p ( t ) ∥ = ∞ , we can select a specific t ϵ such that for all t ≥ t ϵ , it holds that ∥ p ( t ) ∥ ≥ R ϵ ∨ 1 / 2 for any choice of R ϵ . To proceed, we choose R ϵ based on Lemma 5 so that for any t ≥ t ϵ , we have that

Multiplying both sides by the stepsize η and using the gradient descent update, we get

Here, the second inequality is obtained from ∥ p ( t ) ∥ ≥ 1 / 2; the third inequality follows since for any a , b > 0, we have ( a 2 -b 2 ) / (2 b ) -( a -b ) ≥ 0; and the last inequality uses Lemma 6.

Summing the above inequality over t ≥ t ϵ gives

for some finite constant C ( ϵ, η ) defined as

where L ∗ denotes the minimum objective.

Since ∥ p ( t ) ∥ → ∞ , we get

Given that we can choose any value of ϵ ∈ (0 , 1), we have p ( t ) / ∥ p ( t ) ∥ → p mm ⋆ / ∥ p mm ⋆ ∥ .

B.4 Proof of Theorem 3

Proof. Following the proof of Lemma 7, let ( T i ) n i = 1 denote the sets of SVM-neighbors as defined in Definition 2. We define ¯ T i = [ T ] - T i - { α i } as the tokens that are non-SVM neighbors. Additionally, let µ be defined as in (32). Let us denote the initialization lower bound as R 0 µ : = R , where R is given in the Theorem 3's statement. Consider an arbitrary value of ϵ ∈ (0 , µ/ 2) and let 1 / (1 + π ) = 1 -ϵ .

We additionally denote R ϵ ← R π ∨ 1 / 2 where R π was defined in Lemma 7( L3. ). At initialization p (0), we set ϵ = µ/ 2 to obtain R 0 = R µ/ 2 and provide the proof in four steps:

µ Step 1: There are no stationary points within C µ, R 0 µ ( p mm ) . We begin by proving that there are no stationary points within C µ, R 0 µ ( p mm ). Then, since R 0 µ ≥ ¯ R µ per Lemma 7, we can apply ( L2. ) to find that: For all q , p ∈ cone µ ( p mm ) with q , 0 and ∥ p ∥ ≥ R 0 µ , we have that -q ⊤ ∇L ( p ) is strictly positive.

Step 2: It follows from Lemma 7( L3. ) that, for all ϵ ∈ (0 , µ/ 2), all p ∈ C µ, R ϵ ( p mm ) satisfy

The argument above applies to a general ϵ ∈ (0 , µ/ 2). However, at initialization p (0), we set ϵ = µ/ 2 to obtain our earlier R 0 µ choice. To proceed, for any ϵ ∈ (0 , µ/ 2), we will show that after gradient descent enters the conic set C µ, R ϵ ( p mm ) for the first time, it will never leave the set. Let t ϵ be the first time gradient descent enters C µ, R ϵ ( p mm ). In Step 4 , we will prove that such t ϵ is guaranteed to exist. Additionally, for ϵ ← µ/ 2, note that t ϵ = 0 i.e. the point of initialization.

Step 3: Updates remain inside the cone C µ, R ϵ ( p mm ) . 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 ( t ϵ ) ∈ C µ, R ( p mm ), remain within this cone.

ϵ

We proceed by induction. Suppose that the claim holds up to iteration t ≥ t ϵ . This implies that p ( t ) ∈ C µ, R ϵ ( p mm ). Hence, recalling cone definition, for µ > 0 and R ϵ , we have D p ( t ) ∥ p ( t ) ∥ , p mm ∥ p mm ∥ E ≥ 1 -µ and ∥ p ( t ) ∥ ≥ R ϵ . Let

Note that ρ ( t ) > 0 due to Step 1 . This together with the gradient descent update rule gives

Note that from Lemma 7, we have ⟨∇L ( p ( t )) , p ( t ) ⟩ < 0 which implies that ∥ p ( t + 1) ∥ ≥ ∥ p ( t ) ∥ . This together with R ϵ definition and ∥ p ( t ) ∥ ≥ 1 / 2 implies that

Hence, using (53)

Here, the second inequality follows from (53).

Now, it follows from (54a) and (54c) that

where the last inequality uses our choice of stepsize η ≤ 1 / Lp in Theorem 3's statement. Specifically, we need η to be small to ensure the last inequality. We will guarantee this by choosing a proper R ϵ in Lemma 7. Specifically, Lemma 7 leaves the choice of C 0 in R ϵ lower bound of (46) open (it can always be chosen larger). Here, by choosing C 0 ≳ 1 / Lp will ensure η ≤ 1 / Lp works well.

To proceed, we have that

Here, the second inequality uses our choice of ϵ ∈ (0 , µ/ 2) (see Step 2 ), and the first inequality is obtained from Lemma 7 since

for some data dependent constants c and C , ¯ A = max i ∈ [ n ] , t ,τ ∈ [ T ] ∥ k it -k i τ ∥ , and Θ = 1 / ∥ p mm ∥ .

Next, we will demonstrate that the choice of η in (56) does indeed meet our step size condition as stated in the theorem, i.e., η ≤ 1 / Lp . Recall that 1 / (1 + π ) = 1 -ϵ , which implies that π = ϵ/ (1 -ϵ ). Combining this with (46), we obtain:

On the other hand, at the initialization, we have ϵ = µ/ 2 which implies that

In the following, we will determine a lower bound on C 0 such that our step size condition in Theorem 3's statement, i.e., η ≤ 1 / Lp , is satisfied. Note that for the choice of η in (56) to meet the condition η ≤ 1 / Lp , the following condition must hold:

where C 2 = (1 -µ ) ¯ A 2 C 2 Θ c .

This together with (57) implies that for su ffi ciently large

the step size bound in (56) ensures that η ≤ 1 / Lp guarantees (55). Hence, p ( t + 1) remains within the cone, i.e., p ( t + 1) ∈ C µ, R ϵ ( p mm ).

Step 4: The correlation of p ( t ) and p mm increases over t . 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 ) ∈ C µ, R 0 µ ( p mm ) for all t ≥ 0. Note that it follows from Lemma 7 that ⟨∇L ( p ) , p mm / ∥ p mm ∥⟩ < 0, for any finite p ∈ C µ, R 0 µ ( p mm ). Hence, there are no finite critical points p ∈ C µ, R 0 µ ( p mm ), for which ∇L ( p ) = 0. Now, based on Lemma 6, which guarantees that ∇L ( p ( t )) → 0, this implies that ∥ p ( t ) ∥ → ∞ . Consequently, for any choice of ϵ ∈ (0 , µ/ 2) there is a time t ϵ such that, for all t ≥ t ϵ , p ( t ) ∈ C µ, R ϵ ( p mm ). Once within C µ, R ϵ ( p mm ), following similar steps in (51) and (52), for any t ≥ t ϵ ,

for some finite constant C 2 ( ϵ, η ). Consequently,

Since the choice of ϵ ∈ (0 , µ/ 2) is arbitrary, we obtain p ( t ) / ∥ p ( t ) ∥ → p mm / ∥ p mm ∥ .

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 ∈ R d -{ 0 } and ( Yi , K i , X i ) n i = 1 be our dataset. We define the (possibly non-unique) selected-tokens of q as follows: 3

Next, we define the margin and directional-neighbors for q as the minimum margin tokens to the selected-tokens, i.e.,

Finally, we say that q is neighbor-optimal if the scores of its directional-neighbors are strictly less than the corresponding selected-token. Concretely, for all ( i , t ) ∈ M q , we require that

Lemma 8 (When does one direction dominate another?) Suppose q , p ∈ R d be two unit Euclidean norm vectors with identical selected tokens. Specifically, for each i ∈ [ n ] , there exists unique α i ∈ [ T ] such that α i = arg max t ∈ [ T ] k ⊤ it q = arg max t ∈ [ T ] k ⊤ it p . Suppose directional margins obey Γ q < Γ p and set δ Γ = Γ p -Γ q .

Suppose q and p are both neighbor-optimal. Then, for some R ( δ Γ ) and all R > R ( δ Γ ) , we have that L ( R · p ) < L ( R · q ) .

Suppose q has a unique directional-neighbor and is not neighbor-optimal (i.e. this neighbor has higher score). Let δ q be the margin di ff erence between unique directional-neighbor and the second-most minimum-margin neighbor (i.e. the one after the unique one, see (65) ) of q . Then, for some R ( δ Γ ∧ δ q ) and all R > R ( δ Γ ∧ δ q ) , 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 →∞ and purely selecting the tokens α = ( α i ) n i = 1 . This is given by

We evaluate q , p with respect to L ⋆ . To proceed, let s i = S ( R K i q ). Define Γ it q = k ⊤ i α i q -k ⊤ it q . Note that, the smallest value for t , α i is achieved for Γ q . For su ffi ciently large R ≳ O (log( T ) / Γ q ), observe that, for t , α i

Recalling the score definition and let M + , M -be the upper and lower bounds on -ℓ ′ over its bounded domain that scores fall on, respectively. Note that, for some intermediate M + ≥ Mi ≥ M -values, we have

Now, using (62) for a refreshed M + ≥ Mit ≥ 0 . 5 M -values, we can write

The same bound also applies to p with some M ′ it , multipliers

We can now proceed with the proof.

Case 1: q and p are both neighbor-optimal. This means that γ i α i -γ it > 0 for all i ∈ [ n ] , t , α i . Let K + > K -> 0 be upper and lower bounds on γ i α i -γ it values. We can now upper bound the right hand side of (63) via

Consequently, L ( R q ) > L ( R p ) as soon as 1 2 n M -K -e -R Γ q > M + K + Te -R Γ p . Since M + , K + , n , T are global constants, this happens under the stated condition on the margin gap Γ p -Γ q .

Case 2: q has a unique directional-neighbor and is not neighbor-optimal. In this scenario, L ( R q ) - L ⋆ is actually negative for large R . To proceed, define the maximum score di ff erence K + = sup i , t , α i | γ i α i -γ it | . Also let ( j , β ) be the unique directional neighbor achieving the minimum margin Γ q . Then, δ q - the margin di ff erence between unique directional-neighbor and the second minimum-margin neighbor (i.e. the one after the unique one) of q - is defined as

On the other hand, setting κ = γ j β -γ j α j > 0, we can bound

Consequently, we have found that L ( R p ) > L ( R q ) as soon as

This happens when R ≳ 1 δ q ∧ ( Γ p -Γ q ) (up to logarithmic terms) establishing the desired statement.

To proceed, we can write

B.5.2 Proof of Theorem 4

Define the locally-optimal unit directions

The theorem below shows that cone-restricted regularization paths can only directionally converge to an element of this set.

Theorem 7 (Non-LOMM Regularization Paths Fail) Fix a unit Euclidean norm vector q ∈ R d such that q < P mm . Assume that the token scores are distinct (i.e., γ it , γ i τ for t , τ ) and the key embeddings k it are in general position. Specifically, we require the following conditions to hold 4 :

When m = d, all matrices ¯ K ∈ R m × d where each row of ¯ K has the form k it -k i α i for a unique ( i , α i , t , α i ) 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 ϵ > 0 , R 0 > 0 . Define the local regularization path of q as its ( ϵ, R 0 ) -conic neighborhood:

Then, either lim R →∞ ∥ ¯ p ( R ) ∥ < ∞ or lim R →∞ ¯ p ( R ) / ∥ ¯ p ( R ) ∥ , q . In both scenarios lim R →∞ ¯ p ( R ) / R , q .

Proof. We will prove the result by dividing the problem into distinct cases. In each case, we will construct an alternative direction that achieves a strictly better objective than some δ = δ ( ϵ ) > 0 neighborhood of q , thereby demonstrating the suboptimality of the q direction. Let's define the δ neighborhood as follows:

Now, let's recall a few more definitions based on Definition 3. First, the tokens selected by q are given by (59). To proceed, let's initially consider the scenario where α i is unique for all i ∈ [ n ], meaning that as we let c →∞ , c · q will choose a single token per input. Later, we will revisit the setting when arg max is not a singleton, and q is allowed to select multiple tokens.

Additionally, it's important to note that ∥ ¯ p ( R ) ∥ is non-decreasing by definition. Suppose it has a finite upper bound ∥ ¯ p ( R ) ∥ ≤ M for all R < ∞ . In that scenario, we have lim R →∞ ¯ p ( R ) R = 0 , q .

(A) q selects a single token per input: Given that the indices α = ( α i ) n i = 1 defined in (59) are uniquely determined, we can conclude that the q direction eventually selects tokens k i α i . Recall the definition of the margin Γ q from (60) and the set of directional neighbors, which is defined as the indices that achieve Γ q , as shown in (61). Let us refer to q as neighbor-optimal if γ it < γ i α i for all ( i , t ) ∈ M q .

We will consider two cases for this scenario: when q is neighbor-optimal and when q is not neighboroptimal.

⋄ (A1) q is neighbor-optimal. In this case, we will argue that max-margin direction ¯ p mm : = p mm ( α ) / ∥ p mm ( α ) ∥ can be used to construct a strictly better objective than q . Note that p mm ( α ) exists because q is already a viable separating direction for tokens α . Specifically, consider the direction q ′ = q + ϵ ¯ p mm ∥ q + ϵ ¯ p mm ∥ . Observe that, q ′ lies within cone 2 ϵ ( q ), 5

We now argue that, there exists δ = δϵ > 0 such that for all R > R ϵ

To prove this, we study the margin Γ q ′ induced by q ′ and the maximum margin Γ δ induced within p ∈ N δ . Concretely, we will show that Γ q ′ > Γ δ and directly apply the first statement of Lemma 8 to conclude with (68).

Let Γ = 1 / ∥ p mm ( α ) ∥ be the margin induced by ¯ p mm . Note that Γ > Γ q by the optimality of p mm ( α ) and the fact that q , p mm ( α ). Consequently, we can lower and upper bound the margins via

where M = max i , t ,τ ∥ k it -k i τ ∥ .

Consequently, setting δ = ϵ 4 M ( Γ -Γ q ), we find that

Equipped with this inequality, we apply the first statement of Lemma 8 which concludes that 6 for some R ϵ = R ( ϵ 4 ( Γ -Γ q )) and all R > R ϵ , (68) holds. This in turn implies that, within C ϵ , the optimal solution is

either upper bounded by R ϵ in ℓ 2 norm (i.e. lim R →∞ ∥ ¯ p ( R ) ∥ < ∞ ) or

at least δ = δ ( ϵ ) > 0 away from q after ℓ 2-normalization i.e. ∥ ¯ p ( R ) ∥ ¯ p ( R ) ∥ -q ∥ ≥ δ .

In either scenario, we have proven that lim R →∞ ¯ p ( R ) R , q .

⋄ (A2) q is not neighbor-optimal. In this scenario, we will prove that ¯ p ( R ) is finite to obtain lim R →∞ ¯ p ( R ) / R = 0 , q . To start, assume that conic neighborhood ϵ 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 lim R →∞ ∥ ¯ p ( R ) ∥ → ∞ and ¯ p ( R ) ∈ C ϵ , since softmax will eventually perfectly select α (i.e. assigning probability 1 on token indices ( i , α i )), we would have

Note that, this is simply by selection of α and regardless of ¯ p ( R ) directionally converges to q . This means that, if there exists a finite p ∈ C ϵ such that L ( p ) < L ⋆ (i.e. outperforming the training loss of ∥ ¯ p ( R ) ∥ → ∞ ), then ∥ ¯ p ( R ) ∥ < ∞ . This would conclude the proof.

Thus, we will simply find such a p obeying L ( p ) < L ⋆ . To this aim, we first prove the following lemma.

Lemma 9 Given a fixed unit Euclidean norm vector p , if all directional neighbors of p consistently have higher scores for their associated selected tokens, i.e., γ i α i < γ i β for all ( i , α i ) and directional neighbor ( i , β ) , then there exists ¯ R such that for all R > ¯ R,

Proof. Define the maximum score di ff erence K + = sup i ∈ [ n ] , t , α i | γ i α i -γ it | . Also let M p be the set of directional neighbors achieving the minimum margin Γ p ; see (61). Define Γ it = k ⊤ i α i p -k ⊤ it p . Define δ p to be the margin di ff erence between the directional-neighbors and the second-most minimum-margin neighbors defined as

To proceed, setting κ = min( j ,β ) ∈M p γ j β -γ j α j > 0 and using (63), we can bound

Consequently, we have found that L ( R p ) < L ⋆ as soon as

This happens when R ≳ 1 /δ q (up to logarithmic terms) establishing the desired statement.

Based on this lemma, what we need is constructing a perturbation to modify q 's directional neighbors and make sure all of them have strictly better scores than their associated selected-tokens. Note that, once we construct a new candidate (say q 0 = q + perturbation), all su ffi ciently large scalings of q 0 will achieve L ( R · q 0 ) < L ⋆ . Thus, we can find a strictly better solution than L ⋆ for any norm lower bound R 0 - which is enforced within the definition of C ϵ .

Lemma 10 There are at most d directional neighbors i.e. |M q | ≤ d.

Proof. Directional neighbors are indices ( i , t ) obeying the inequality

Declare D to be the matrix with rows obtained by these key di ff erences k it -k i α i . D ∈ R M × d where M = |M q | . We then obtain Dq = -Γ q 1 M where 1 M is the all ones vector. If M > d , the equality Dq = -Γ q 1 M cannot be satisfied because 1 M is not in the range space of D by our assumption of general key embedding positions.

To proceed, |M q | ≤ d and let D be as defined in the lemma above. D is also full-rank by our assumption of general key positions. We use D to construct a perturbation as follows. Let M + q ⊂ M q be the set of directional neighbor with strictly higher scores than their associated selected-tokens. In other words, all ( j , β ) ∈ M + q obeys

Define the score di ff erence κ = min( j ,β ) ∈M + q γ j β -γ j α j > 0. We know κ > 0 because α is not neighboroptimal. Finally, define the indicator vector of 1 + with same dimension as cardinality |M q | . 1 + is 1 for the rows of D corresponding to M + q and is 0 otherwise. Finally, set the perturbation as

where we used the full-rankness of D during pseudo-inversion. To proceed, for a small ϵ 0 > 0, consider the candidate direction q 0 = q + ϵ 0 q ⊥ . We pick ϵ 0 = O ( ϵ ) su ffi ciently small to ensure q 0 ∈ C ϵ . To finalize, let us consider the margins of the tokens within q 0. Similar to Lemma 9, set δ q = min i ∈ [ n ] , t , α i , ( i , t ) < M q Γ it q -Γ q > 0. Let ¯ ϵ 0 = ∥ q 0 ∥ -1 = ∥ q + ϵ 0 q ⊥ ∥ -1. Using definition of q ⊥ , we have that

For ( i , t ) ∈ M + q , we achieve a margin of

For ( i , t ) ∈ M + q , we achieve a margin of Γ q 1 + ¯ ϵ 0 .

For ( i , t ) < M q , setting K = ∥ q ⊥ ∥ · sup i , t ,τ ∥ k i α i -k it ∥ , we achieve a margin of at most

In short, since ¯ ϵ 0 = O ( ϵ 0), setting ϵ 0 su ffi ciently small guarantees that M + q is the set of directional neighbors of q 0. Since M + q has strictly higher scores than their associated selected-tokens, applying Lemma 9 on q 0 shows that, L ( R · q 0 ) < L ⋆ for su ffi ciently large R implying ∥ ¯ p ( R ) ∥ < ∞ .

(B) q selects multiple tokens for some inputs i ∈ [ n ] : In this setting, we will again construct a perturbation to create a scenario where q 0 = q + ϵ q ⊥ selects a single token for each input i ∈ [ n ]. We will then employ margin analysis (first statement of Lemma 8) to conclude that q 0 outperforms a δ ≪ ϵ neighborhood of q .

Let I ⊂ [ n ] be the set of inputs for which q selects multiple tokens. Specifically, for each i ∈ I , there is T i ⊂ [ T ] such that |T i | ≥ 2 and for any i ∈ I and θ ∈ T i ,

From these multiply-selected token indices let us select the highest score one, namely, β i = arg max θ ∈T i γ i θ for i ∈ I . Now, define the unique optimal tokens for each input as α ∈ R n where α i : = β i for i ∈ I and α i = arg max t ∈ [ T ] k ⊤ it q for i < I . Define L ⋆ = 1 n P n i = 1 ℓ ( γ i α i ) as earlier.

Secondly, we construct a perturbation q ⊥ to show that q 0 = q + ϵ 0 q ⊥ can select tokens α asymptotically. To see this, define the matrix D where each (unique) row is given by k i α i -k i θ where θ ∈ T i , θ , α i , i ∈ I . Now note that, ( k i α i -k i θ ) ⊤ q = 0 for all θ ∈ T i , θ , α i , i ∈ I . Since keys are in general positions, this implies that D has at most d -1 rows and, thus, its rows are linearly independent. Consequently, choose q ⊥ = D † 1 where † denotes pseudo-inverse and 1 is the all ones vector. Also let Γ q be the margin of directional margin of q that is

With this choice and setting K = sup i , t ,τ ∥ k i α i -k it ∥ , we have that

For θ ∈ T i , θ , α i : ( k i α i -k i θ ) ⊤ q 0 = ( k i α i -k i θ ) ⊤ q ⊥ = ϵ 0 .

For all other ( i , t ) with t , α i : ( k i α i -k it ) ⊤ q 0 ≥ Γ q -K ∥ q ⊥ ∥ ϵ 0 .

Choosing ϵ 0 < Γ q / (1 + K ∥ q ⊥ ∥ ), together, these imply that,

α = ( α i ) n i = 1 is the selected-tokens of q 0 ,

q 0 achieves a directional margin of

q 0 is neighbor-optimal because directional neighbors are θ ∈ T i , θ , α i and γ i θ < γ i α i .

Note that, these conditions lay the groundwork for applying the first statement of Lemma 8 with p ← q 0. We next explore the optimal directions within N δ and show that q 0 strictly outperform them in terms of training loss.

To proceed, given small δ ≪ ϵ 0, let us study L R = min p ∈N δ , R ≤∥ p ∥ < ∞ L ( p ). Here, recall that p has a δ -small directional perturbation around q which can modify the token selections by breaking the ties between the multiply-selected token indices by q . However, thanks to the distinct token score assumption, for large R ≤ ∥ p ∥ , the optimal p ∈ N δ is guaranteed to (uniquely) select indices α i ∈ T i . Because all other p directions - which, asymptotically, either select other tokens θ ∈ T i , θ , α i or split the probabilities equally across a subset of T i - achieve a larger loss. For instance, q will split the probabilities equally across T i to achieve an asymptotic loss of

Thus, L q ⋆ > L ⋆ because ℓ GLYPH<16> 1 |T i | P θ ∈T i γ i θ GLYPH<17> > ℓ GLYPH<0> γ i α i GLYPH<1> = ℓ ( γ i α i ) where α i has the highest score i.e. γ i α i > 1 |T i | P θ ∈T i γ i θ . Set ˜ p ( R ) = arg min p ∈N δ , ∥ p ∥≤ R L ( p ). Consequently, there are two scenarios are:

lim R →∞ ∥ ˜ p ( R ) ∥ is finite. This already proves the statement of the theorem as ˜ p ( R ) / R → 0 within δ < ϵ neighborhood of q .

For su ffi ciently large R , the selected-tokens of ˜ p ( R ) are α = ( α i ) n i = 1 .

Proceeding with the second (remaining scenario), we study the directional margin of ˜ p ( R ). More broadly, for any p ∈ N δ and ¯ p = p / ∥ p ∥ with selected-tokens α , using the fact that ∥ ¯ p -q ∥ ≤ δ ≪ ϵ 0 , we can bound the directional margin as

For θ ∈ T i , θ , α i : ( k i α i -k i θ ) ⊤ ¯ p = ( k i α i -k i θ ) ⊤ ( ¯ p -q ) ≤ K δ .

For all other ( i , t ) with t , α i : ( k i α i -k it ) ⊤ ¯ p ≥ Γ q -K δ .

This means that, any such ¯ p ∈ N δ achieves a directional margin of at most

Applying Lemma 8 and setting δ = O ( ϵ 0), this implies that for

we have that L ( R · q 0 ) < min ∥ p ∥ = R , p ∈N δ L ( p ). Since this holds for all R , (68) holds (similar to Case (A1) ) and we conclude that whenever ∥ ¯ p ( R ) ∥ → ∞ , it doesn't directionally converge within N δ (i.e. δ > 0 neighborhood of q ) proving the advertised result.

B.6 Proof of Lemma 2

We prove a slightly general restatement where we require v ∈ range( W ⊤ ) - instead of full-rank W .

Lemma 11 Suppose for all i ∈ [ n ] and t , opt i , Y i = 1 and γ it < γ i opt i . Also suppose v ∈ range( W ⊤ ) . Then, p mm ⋆ exists - i.e. (ATT-SVM) is feasible for optimal indices α i ← opt i .

Proof. To establish the existence of p mm ⋆ , we only need to find a direction that demonstrates the feasibility of (ATT-SVM), i.e. we need to find p that satisfies the margin constraints. To begin, let's define the minimum score di ff erence:

We then set p = γ -1 ( W ⊤ ) † v where † denotes pseudo-inverse. By assumption W ⊤ p = γ -1 v . To conclude, observe that p is a feasible solution since k it = Wx it and for all i ∈ [ n ] and t , opt i , we have that

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 p R / R or v r / r fails to converge as R , r grows. Set Ξ = 1 / ∥ p mm ∥ , Γ = 1 / ∥ v mm ∥ , ˜ p mm = R Ξ p mm and ˜ v mm = r Γ v mm . The proof strategy is obtaining a contradiction by proving that (˜ v mm , ˜ p mm ) is a strictly better solution compared to ( v r , p R ) for large R , r . Without losing generality, we will set α i = 1 for all i ∈ [ n ] as the problem is invariant to tokens' permutation. Define q p i = 1 -s p i 1 to be the amount of non-optimality (cumulative probability of non-first tokens) where s p i = S ( K i p ) is the softmax probabilities.

Case 1: p R / R does not converge. Under this scenario there exists δ, γ = γ ( δ ) > 0 such that we can find arbitrarily large R with ∥ p R / R -˜ p mm / R ∥ ≥ δ and margin induced by p R / R is at most Ξ (1 -γ ) (from strong convexity of (ATT-SVM)). Following q p i definition above, set ˆ q max = sup i ∈ [ n ] q p R i to be

worst non-optimality in p R and q ⋆ max = sup i ∈ [ n ] q ˜ p mm i to be the same for ˜ p mm . Repeating the identical argument in Theorem 8 (specifically (84)), we can bound the non-optimality amount q ˜ p mm i of ˜ p mm as

Thus, q ⋆ max = max i ∈ [ n ] q ˜ p mm i ≤ T exp( -R Ξ ). Next without losing generality, assume first margin constraint is γ -violated by p R and min t , α 1 ( k 1 α 1 -k 1 t ) ⊤ p R ≤ Ξ R (1 -γ ). Denoting the amount of non-optimality of the first input as q p R 1 , we find

We similarly have q ⋆ max ≥ T -1 exp( -R Ξ ) to find that

In words, ˜ p mm contains exponentially less non-optimality compared to p R as R grows. The remainder of the proof di ff ers from Theorem 8 as we need to upper / lower bound the logistic loss of (˜ v mm , ˜ p mm ) and ( v r , p R ) respectively to conclude with the contradiction.

First, let us upper bound the logistic loss of (˜ v mm , ˜ p mm ). Set r i = X ⊤ i S ( K i ˜ p mm ). Observe that if ∥ r i -x i 1 ∥ ≤ ϵ i , we have that v mm satisfies the SVM constraints on r i with Yi · r ⊤ i v mm ≥ 1 -ϵ i / Γ . Consequently, setting ϵ max = sup i ∈ [ n ] ϵ i , v mm achieves a label-margin of Γ -ϵ max on the dataset ( Yi , r i ) i ∈ [ n ] . With this, we upper bound the logistic loss of (˜ v mm , ˜ p mm ) as follows. Let M = sup i ∈ [ n ] , t ,τ ∈ [ T ] ∥ x it -x i τ ∥ . Let us recall the fact (73) that worst-case perturbation is

Conversely, we obtain a lower bound for ( v r , p R ). Set r i = X ⊤ i S ( K i p R ). Using Assumption C, we find that solving (SVM) on ( Yi , r i ) i ∈ [ n ] achieves at most Γ -ν e -(1 -γ ) Ξ R / T margin. Consequently, we have

Observe that, this lower bound dominates the previous upper bound when R is large, namely, when (ignoring the multiplier 1 / 2 n for brevity)

Thus, we indeed obtain the desired contradiction since such large R is guaranteed to exist when p R / R ↛ p mm .

· Case 2: v r / r does not converge. This is the simpler scenario: There exists δ > 0 such that we can find arbitrarily large r obeying ∥ v r / r -v mm / ∥ v mm ∥∥ ≥ δ . If ∥ p R / R -Ξ p mm ∥ ↛ 0, then 'Case

This implies that

1' applies. Otherwise, we have ∥ p R / R -Ξ p mm ∥ → 0, thus we can assume ∥ p R / R -Ξ p mm ∥ ≤ ϵ for arbitrary choice of ϵ > 0.

On the other hand, due to the strong convexity of (SVM), for some γ : = γ ( δ ) > 0, v r achieves a margin of at most (1 -γ ) Γ r on the dataset ( Yi , x i 1 ) i ∈ [ n ]. Additionally, since ∥ p R / R -Ξ p mm ∥ ≤ ϵ , p R strictly separates all optimal tokens (for small enough ϵ > 0) and ˆ q max : = f ( ϵ ) → 0 as R → ∞ . Consequently, setting r i = X ⊤ i S ( K i p R ), for su ffi ciently large R > 0 setting M = sup i ∈ [ n ] , t ∈ [ T ] ∥ x it ∥ , we have that

This in turn implies that logistic loss is lower bounded by (following (75)),

Going back to (74), this exponentially dominates the upper bound of ( ˜ p mm , ˜ v mm ) whenever rMT exp( -Ξ R ) < r γ Γ / 2, (that is, whenever R , r are su ffi ciently large), again concluding the proof.

C.2 Proof of Theorem 6

We will prove this result in two steps. Our first claim restricts the optimization to the particular quadrant induced by min t , α i ( k i α i -k it ) ⊤ p R ≥ 0 under the theorem's condition S ( K i p R ) α i → 1.

Lemma 12 Suppose S ( K p ) → 1 . Then, there exists R such that for all R ≥ R , we have that,

Proof. Suppose the claim does not hold. Set s R i = S ( K i p R ). Fix R 0 such that s R i α i ≥ 0 . 9 for all R ≥ R 0 . On the other hand, there exists arbitrarily large R for which ( k i α i -k it ) ⊤ p R < 0 for some t , α i ∈ [ T ] and i ∈ [ n ]. At this ( R , i , t ) choices, we have that s R it ≥ s R i α i . Since s R it + s R i α i ≤ 1, we find s R i α i < 0 . 5 which contradicts with s R i α i ≥ 0 . 9.

Let Q be the set of p satisfying the quadrant constraint (77) - i.e. indices ( α i ) n i = 1 are selected. Let h R be the solution of regularization path of ( v , p ) subject to the constraint p ∈ Q . From Lemma 12, we know that, for some R 0 and all R ≥ R 0 , h R = p R . Thus, if the limit exists, we have that lim R →∞ h R / R = lim R →∞ p R / R .

To proceed, we will prove that lim R →∞ h R / R exists and is equal to p relax / ∥ p relax ∥ and simultaneously establish v r / r → v mm / ∥ v mm ∥ .

Lemma 13 lim R →∞ h R / R = p relax / ∥ p relax ∥ and lim r →∞ v r / r = v mm / ∥ v mm ∥ .

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 Ξ = 1 / ∥ p relax ∥ .

· Case 1: h R / R does not converge. Under this scenario there exists δ, γ = γ ( δ ) > 0 such that we can find arbitrarily large R with ∥ h R / R -Ξ p relax ∥ ≥ δ . This implies that margin induced by h R / R is at most Ξ (1 -γ ) over the support vectors S (from strong convexity of (10)). The reason is that, h R satisfies h ⊤ R ( k i α i -k it ) ≥ 0 for all t , α i by construction as h R ∈ Q . Thus, a constraint over the support vectors have to be violated (when normalized to the same ℓ 2 norm as ∥ p relax ∥ = 1 / Ξ ).

As usual, we will construct a solution strictly superior to h R and contradicts with its optimality.

Construction of competitor: Rather than using p relax direction, we will choose a 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.

Let p mm be the solution of (ATT-SVM) with α = ( α i ) n i = 1 (which was assumed to be separable). Observe that p mm ϵ = ϵ p mm + (1 -ϵ ) p relax satisfies the constraints of (78). Additionally, p mm ϵ would achieve a margin of 1 (1 -ϵ ) / Ξ+ ϵ/ ∆ = ∆Ξ ∆+ ϵ ( Ξ -∆ ) where ∆ = 1 / ∥ p mm ∥ . Using optimality of p ϵ -rlx , this implies that the reduced margin Ξ ϵ = 1 / ∥ p ϵ -rlx ∥ (by enforcing ϵ over non-support) over the support vectors is a Lipschitz function of ϵ . That is Ξ ϵ ≥ Ξ -ϵ M for some M ≥ 0. To proceed, choose an ϵ > 0 such that, it is strictly superior to margin induced by h R , that is,

To proceed, set ˜ p ϵ -rlx = R Ξ ϵ p ϵ -rlx . Let us recall the following notation from the proof of Theorem 5: s p i = S ( K i p ) and q p i = 1 -s i α i . Set ˆ q max = max i ∈S q h R i to be worst non-optimality of h R over support set . Similarly, define q ⋆ max = max i ∈S q ˜ p ϵ -rlx i to be the same for ˜ p ϵ -rlx . Repeating the identical arguments to (71), (72), (73), and using the fact that p ϵ -rlx achieves a margin Ξ (1 -γ 2 ) ≤ Ξ ϵ ≤ Ξ , we end up with the lines

In what follows, we will prove that ˜ p ϵ -rlx achieves a strictly smaller logistic loss contradicting with the optimality of p R (whenever ∥ h R / R -Ξ p relax ∥ ≥ δ ).

Upper bounding logistic loss. Let us now upper bound the logistic loss of (˜ v mm , ˜ p ϵ -rlx ) where ˜ v mm = r Γ v mm with v mm being the solution of (SVM) with r i ← x i α i and Γ = 1 / ∥ v mm ∥ . Set r i = X ⊤ i S ( K i ˜ p ϵ -rlx ). Set υ = min i ∈ ¯ S Yi · x ⊤ i α i v mm -1 to be the additional margin bu ff er that non-support vectors have access to. Also set M = sup i ∈ [ n ] , t ,τ ∈ [ T ] ∥ x it -x i τ ∥ . Observe that we can write

Non-supports achieve strong label-margin: Using above and (78) for all i ∈ ¯ S and t , α i , we have that s it ≤ e -ϵ Ξ ϵ R s i α i ≤ e -ϵ Ξ (1 -γ/ 2) R s i α i . Consequently, whenever R ≥ ¯ R 0 : = ( ϵ Ξ (1 -γ/ 2)) -1 log( TM Γ ),

This implies that, on i ∈ ¯ S

In words: Above a fixed ¯ R 0 that only depends on γ = γ ( δ ), features r i induced by all non-support indices i ∈ ¯ S achieve margin at least 1. What remains is analyzing the margin shrinkage over the support vectors as in Theorem 5.

Controlling support margin and combining bounds: Over S , suppose v mm satisfies the SVM constraints on r i with Yi · r ⊤ i v mm ≥ 1 -ϵ i / Γ . Consequently, setting ϵ max = sup i ∈ [ n ] ϵ i , v mm achieves a label-margin of Γ -ϵ max on the dataset ( Yi , r i ) i ∈ [ n ] . Next, we recall the fact (79b) that worst-case perturbation is ϵ max ≤ M exp( -Ξ (1 -0 . 5 γ ) R + log T ) = MT exp( -Ξ (1 -0 . 5 γ ) R ). With this and (80), we upper bound the logistic loss of (˜ v mm , ˜ p ϵ -rlx ) as follows.

Conversely, we obtain a lower bound for ( v r , h R ). Set r i = X ⊤ i S ( K i h R ). Recall the lower bound (79a) over the support vector set S . Combining this with our Assumption C over the support vectors

υ

of (SVM) implies that, solving (SVM) on ( Yi , ri ) i ∈ [ n ] achieves at most Γ -ν e -(1 -γ ) Ξ R / T margin. Consequently, we have

Observe that, this lower bound dominates the previous upper bound when R is large, namely, when (ignoring the multiplier 1 / 2 n for brevity)

Thus, we obtain the desired contradiction since ˜ p ϵ -rlx is a strictly better solution compared to p R = h R (once R is su ffi ciently large).

· Case 2: v r / r does not converge. This is the simpler scenario: There exists δ > 0 such that we can find arbitrarily large r obeying ∥ v r / r -v mm / ∥ v mm ∥∥ ≥ δ . First, note that, due to the strong convexity of (SVM), for some γ : = γ ( δ ) > 0, v r achieves a margin of at most ( Γ -γ ) r on the dataset ( Yi , x i 1 ) i ∈ [ n ] . By theorem's condition, we are provided that S ( K i p R ) α i → 1. This immediately implies that, for any choice of ϵ = γ/ 3 > 0, above some su ffi ciently large ( r 0 , R 0), we have that ∥ x p R i -r i ∥ ≤ ϵ . Following (81), this implies that, choosing ˜ v mm = r v mm / ∥ v mm ∥ achieves a logistic loss of at most e r γ/ 3 e -r Γ . Again using ∥ x p R i -r i ∥ ≤ ϵ , for su ffi ciently large ( r , R ) we have that

This in turn implies that logistic loss is lower bounded by (following (82)),

This dominates the above upper bound e r γ/ 3 e -r Γ of ˜ v mm whenever 1 2 n e γ r / 3 > 1 ⇐⇒ r > 3 γ log(2 n ), (that is, when r is su ffi ciently large), again concluding the proof.

D Regularization Path of Attention with Nonlinear Head

So far our 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 is showcasing this generalization. Specifically, we consider the prediction model f ( X ) = ψ ( X ⊤ S ( Kp )) where ψ ( · ) : R d → R generalizes the linear head v of our attention model. For instance, following exposition in Section 1.1, ψ ( · ) can represent a multilayer transformer with p being a tunable prompt at the input layer. Recall that ( X i , K i , Yi ) n i = 1 is the dataset of the input-key-label tuples. We consider the training risk

The challenge with nonlinear ψ ( · ) is that, we lack a clear score function (Def. 1) unlike the previous sections. The assumption below introduces a generic condition that splits the tokens of each X i into an optimal set O i and non-optimal set ¯ O i = [ T ] - O i . In words, non-optimal tokens are those that strictly increase the training risk L ( p ) if they are not fully suppressed by attention probabilities s p i .

Assumption D (Mixing non-optimal tokens hurt) There exists sets ( O i ) n i = 1 ⊂ [ T ] as follows. Let q p i = P t ∈ ¯ O i s p it be the sum of softmax similarities over the non-optimal set for p . Set q p max = max i ∈ [ n ] q p i . For any ∆ > 0 , there exists ρ < 0 such that:

This assumption is titled mixing hurts because the attention output X ⊤ i s p i is mixing the tokens of X i and our condition is that, to achieve optimal risk, this mixture should not contain any non-optimal tokens. In particular, we require that, a model p that contains exponentially less non-optimality (quantified via log( q max)) compared to p ′ is strictly preferable. As we discuss in the supplementary material, Theorem 1 is 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 i for each input.

Unlike (ATT-SVM), this problem is not necessarily convex when the optimal set O i is not a singleton. To see this, imagine n = d = 1 and T = 3: Set the two optimal tokens as k 1 = 1 and k 2 = -1 and the non-optimal token as k 3 = 0. The solution set of (ATT-SVM') is p mm ∈ {-1 , 1 } whereas their convex combination p = 0 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 8 Let G mm be the set of global minima of (ATT-SVM') . Suppose its margin Ξ : = 1 / ∥ p mm ∥ > 0 and Assumption D holds. Let dist ( · , · ) denote the ℓ 2 -distance between a vector and a set. Following (83) , define ¯ p ( R ) = arg min ∥ p ∥≤ R L ( p ) . We have that lim R →∞ dist GLYPH<16> ¯ p ( R ) Ξ R , G mm GLYPH<17> = 0 .

We note that Theorem 1 is a corollary of this result where O i 's and G mm are singleton. Based on this result, with multiple optimal tokens, Theorem 1 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 softmax-attention, (harmful) contribution of the non-optimal token with the minimum margin can dominate the contribution of all other tokens as R →∞ . This high-level approach is similar to earlier works on implicit bias of gradient descent with logistic loss [31, 22].

Pick p mm ∈ G mm and set p ⋆ R = R p mm ∥ p mm ∥ . This will be the baseline model that p R has to compete against. Also let ¯ p R = p R Ξ R . Now suppose dist ( ¯ p R , G mm ) ↛ 0 as R →∞ . Then, there exists δ > 0 such that, we can always find arbitrarily large R obeying dist ( ¯ p R , G mm ) ≥ δ .

Since ¯ p R is δ > 0 bounded away from G mm , and ∥ ¯ p R ∥ = ∥ p mm ∥ , ¯ p R strictly violates at least one of the inequality constraints in (ATT-SVM'). Otherwise, we would have ¯ p R ∈ G mm . Without losing generality, suppose ¯ p R violates the first margin constraint, that is, for some γ : = γ ( δ ) > 0, max α ∈O 1 min β ∈ ¯ O 1 ¯ p ⊤ R ( k 1 α -k 1 β ) ≤ 1 -γ . Now, we will argue that this will lead to a contradiction as R →∞ since we will show that L ( p ⋆ R ) < L ( p R ) for su ffi ciently large R .

First, let us control L ( p ⋆ R ). We study s ⋆ i = S ( K i p ⋆ R ) and let α i ∈ O i be the index α in (ATT-SVM') for which margin i = max α ∈O i min β ∈ ¯ O i ( k i α -k i β ) ⊤ p mm ≥ 1 is attained. Then, we bound the non-optimality amount q ⋆ i of p ⋆ R as

Thus, q ⋆ max = max i ∈ [ n ] q ⋆ i ≤ T exp( -Ξ R ). Secondly, we wish to control L ( p R ) by lower bounding the non-optimality in p R . Focusing on the first margin constraint, let α ∈ O 1 be the index in (ATT-SVM') for which margin 1 ≤ 1 -γ is attained. Denoting the amount of non-optimality of the first input as ˆ q 1 , we find 7

We similarly have q ⋆ max ≥ T -1 exp( -Ξ R ). In conclusion, for p R , p ⋆ R , denoting maximum nonoptimality by ˆ q max ≥ ˆ q 1 and q ⋆ max , we respectively obtained

The above inequalities satisfy Assumption D as follows where p ← p ⋆ R and p ′ ← p R : Set R 0 = 3 γ -1 Ξ -1 log T so that log T = γ Ξ R 0 3 . Secondly, set ρ 0 = -Ξ R 0 -log T . This way, ρ 0 ≥ log( q ⋆ max ) implies R ≥ R 0 and log T ≤ γ Ξ R 3 . Using the latter inequality, we bound the log T terms to obtain

log(ˆ q max) ≥ -(1 -2 γ/ 3)( Ξ R ), and

log( q ⋆ max ) ≤ -(1 -γ/ 3)( Ξ R ).

To proceed, we pick 1 +∆ = 1 -γ/ 3 1 -2 γ/ 3 implying ∆ : = γ 3 -2 γ . Finally, for this ∆ , there exists ρ ( ∆ ) which we need to ensure log(ˆ q max) ≤ ρ ( ∆ ). This can be guaranteed by picking su ffi ciently large R that ensures log( q ⋆ max ) ≤ -(1 -γ/ 3)( Ξ R ) ≤ ρ ( ∆ ) to satisfy all conditions of Assumption D. Since such large R exists by initial assumption dist ( ¯ p R , G mm ) ↛ 0, Assumption D in turn implies that L ( p ⋆ R ) < L ( p R ) contradicting with the optimality of p R in (83).

D.2 Application to 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 q p max = 0, then, Assumption D will hold.

Lemma 14 (Linear label mixing) Recall q p i = P t ∈ ¯ O t s p it from Assumption D. Suppose Yi ∈ {-1 , 1 } and

for some ( ν i ) n i = 1 > 0 . Here Zi = Zi ( p ) is the contribution of non-optimal tokens to prediction. For some C , ϵ > 0 and for all p ∈ R d , assume

Then, Assumption D holds for L ( p ) = 1 n P n i = 1 ℓ ( Yi · ψ ( X ⊤ i s p i )) when ℓ ( · ) is a strictly decreasing loss function with continuous derivative.

Proof. Recall the assumption Yi · ψ ( X ⊤ i s p i ) = ν i (1 -q p i ) + Zi with Zi obeying (85). Let us also write the loss function

Define q p max = sup i ≤ [ n ] q p i . Let M be the maximum absolute value of score over tokens. Let

Through Taylor's Theorem (integral remainder), we have that

Set L ⋆ = 1 n P n i = 1 ℓ ( ν i ). Set C + = B ( C + max i ∈ [ n ] ν i ) and C -= n -1 A ϵ min i ∈ [ n ] ν i . This also implies

Thus, to prove L ( p ′ ) > L ( p ), we simply need to establish the stronger statement C -q p ′ max > C + q p max .

Going back to the condition of Assumption D, any log( q p max ) ≤ (1 + ∆ ) log( q p ′ max ) obeys q p max ≤ ( q p ′ max ) 1 +∆ i.e. q p ′ max ≥ ( q p max ) (1 +∆ ) -1 . Following above, we wish to ensure q p ′ max > Θ q p max for such ( p , p ′ ) pairs where Θ = C + / C -> 1. This is guaranteed by

The above is satisfied by choosing a ρ ( ∆ ) : = -2(1 + ∆ -1 ) log( Θ ) in Assumption D. Thus, all p , p ′ with log( q p max ) ≤ ρ = ρ ( ∆ ) satisfies the condition of Assumption D 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 use 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 p mm . After obtaining p mm , we also verify that these indices satisfy our local-optimal definition. The examples we use in the paper are all trivial to verify (by construction).

In Figures 3(a) and 3(b), v mm (blue dashed) is solved using python package sklearn.svm via (SVM) based on the given label information, and red dashed line represents p relax direction instead, which is the solution of (10). Note that in both figures, v mm / ∥ v mm ∥ = [0 , 1]. Therefore, in Figure 3(a) all optimal tokens are support vectors and p relax = p mm . 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 p mm direction.

Failure of gradient descent's global convergence when n ≥ 2 (refer to Theorem 2) . Figure 8 provides a counter-example demonstrating that the n = 1 restriction is indeed necessary and tight to guarantee global convergence of the gradient descent iterates p ( t + 1) = p ( t ) -η ∇L ( p ( t )) on (ERM).

For this example, we use logistic loss in (ERM). We set n = T = d = 2, implying that there is only one nonoptimal token, thus Assumption B is satisfied. The red and blue lines represent GMM and LMM solutions, respectively. 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 mm ) represented by the blue line, the square teal token does not have any SVM-neighbors. The arrows indicate the two trajectories originating from di ff erent initializations. This demonstrates that the gradient descent iterates p ( t + 1) = p ( t ) -η ∇L ( p ( t )) on (ERM) with two di ff erent initializations converge to two di ff erent SVM solutions ( GMM and LMM ). Results validate the necessity of n = 1 in Theorem 2 to provide the gradient descent convergence to p mm ⋆ from any initialization.

6

4

2

0

2

4

6

2

1

0

1

2

3

4

5

Optimal tokens

Non-opt tokens

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 d ∈ { 2 , 5 , 10 , 100 , 300 , 500 } . These experiments use normalized gradient descent with learning rate 1 for 1000 iterations. Inputs x it and the linear head v are uniformly sampled from the unit sphere, while Yi is uniformly ± 1, and W is set to I .

The bar plot in Figure 9(a) 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 1 -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 non-saturated. For

2

5

10

100

300

500

d

0

20

40

60

80

100

Percentage

S

(

Kp

(

t

))

α <

1

p

(

t

)

→

p

p

(

t

)

→

p

⋆

(a) Perc. of di

ff

erent convergence scenarios for

p

(

t

)

0.2 0.0

0.2

0.4

0.6

0.8

1.0

1.2

Min score gap to neighbors (    )

0.0

0.1

0.2

0.3

0.4

Probability

d

=5

d

=10

d

=100

γ

(b) Prob. of

γ

:

=

min

i

∈

[

n

]

,

t

∈T

i

γ

i

α

i

-

γ

it

small d (e.g., d = 2), problem has small degrees of freedom to separate optimal tokens from the rest (i.e. no 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 ( t ) ∥ → ∞ ) and we observe that the selected tokens α almost always converges to an LMM direction (blue bar) - this is in line with Theorems 3 & 4. We also study the convergence to the globally-optimal GMM which is represented by the green bar: GMM is a 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-parameterized deep learning theory [57, 58, 59, 60] and motivates future research. The average correlation coe ffi cient between p ( t ) and its associated LMM / GMM direction is 0 . 997, suggesting that, whenever softmax saturates, gradient descent indeed directionally converges to a LMM solution p ∈ P mm , confirming Theorem 4.

Furthermore, we found that there exist problem instances, with saturated softmax and ∥ p ( t ) ∥ → ∞ , that do not converge to either LMM or GMM . We analyzed this phenomenon using the minimum score gap, γ : = min i ∈ [ n ] , t ∈T i γ i α i -γ it , where T i , i ∈ [ n ], represents the sets of SVM-neighbor tokens. Figure 9(b) provides the probability distribution of γ (with bins of width < 0 . 01) and demonstrates the rarity of such cases. Specifically, we found this happens less than 1% of the problems, that is, Prob( γ < 0) < 0 . 01. Figure 9(b) also reveals that, in these scenarios, even if γ < 0, it is typically close to zero i.e. even if there exists a SVM-neighbor with a higher score, it is only slightly so. This is 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-margin solution in average.

converges to an LMM i.e. p ( t ) / ∥ p ( t ) ∥ → p mm / ∥ p mm ∥ . We believe the rare occurrence of a negative score gap is due to small score di ff erences (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 p mm ( α ) (even if p mm ( α ) 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.

F.1 Related Work on Implicit Regularization

The introduction of Support Vector Machines (SVM), which utilize explicit regularization to choose maximum margin classifiers, represents one of the earliest relevant literature in this field [61]. The concept of maximizing the margin was later connected to 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 perceptron [63] originally introduced the concept of margins, the method itself does not possess an inherent bias as 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 ℓ 1-regularization, revealing that coordinate descent, when combined with the exponential loss, exhibits an inherent inclination towards ℓ 1-regularized 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, 29]. The works on classification typically utilize logistic loss or exponentially-tailed losses to establish connections to margin maximization. The results have also been extended to non-separable data by [30, 31, 21]. Additionally, several papers have explored the implicit bias of stochastic gradient descent [37, 38, 41, 42], as well as adaptive and momentum-based methods [43, 44, 45, 46].

While there are some similarities between our optimization approach for v and existing works, the optimization of p presents notable di ff erences. Firstly, our optimization problem is nonconvex and involves a 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 in 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-optimal tokens, we provide a tailored analysis that captures the unique characteristics of the attention model. Overall, our work o ff ers novel insights into the optimization of attention-based models and sheds light on the behavior of the attention mechanism during training.

F.2 Related Work on 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 a strong alternative to 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

Turing-completeness of Transformers [91, 92, 93, 51, 23, 94] with statistical guarantees [95, 96]. There is also a growing e ff ort towards a theoretical understanding of emergent abilities of language models - such as in-context learning [97, 98, 99] and chain-of-thought [100, 101, 102] - which are inherently related to the models ability to attend to the relevant information within the input sequence.

Focusing on the self-attention component, Edelman et al. [51] theoretically shows that a 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 to approximate softmax via a 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 to degenerate during the training process, with the result often being an output that is approximately a rank-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 a 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, o ff ering 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 a 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 a simplified attention structure where the attention matrix is not directly parameterized with respect to the input.

In contrast, our work o ff ers a comprehensive optimization-theoretic analysis of the attention model, establishing a formal connection to max-margin problems. While comparable works make assumptions on the dataset model, our results apply under minimal assumptions for general data and realistic conditions. Our analysis based on max-margin-equivalence allows us to gain a deeper understanding 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, T = 2 , 3 tokens and d = 2 , 3 dimensions (in contrast to [55, 56, 103, 17]). Notably, our work also presents the first 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 characterizing 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.

Peer reviews:
comments: 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.

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.
- Clear presentation with adequate explanation.

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.
- Lack the convergence of the gradient decent when jointly optimizing $(v,p)$.

QUESTIONS:
- Why Lemma 1 implies line 92-94?

LIMITATIONS:
See weaknesses.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 2
DATE: 2023-06-29

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nIn 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.\n\nQUESTIONS:\nMy 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.\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-04"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nsome 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.\n\nQUESTIONS:\nplease see the section above.\n\nLIMITATIONS:\nThe authors adequately discuss the limitations throughout the paper. I do not see any potential negative impact with this work.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\n1. 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).\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 8, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-07"}

comments: 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.

STRENGTHS:
1. This paper proposes a mathematical analysis for attention mechanims, which deepen our understanding of the operator.
2. 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. 
    However, I think it would be better to give conclusions and design guidances based on the authors' observations.
    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?

2. The authors discuss attention modules based on isolated simple operators. I think it is very helpful.
    However, I wonder if we extend to real large-scale attention based models, do the conclusions remain the same?

QUESTIONS:
Please address the questions in weakness part.

LIMITATIONS:
As indicated in weakness part.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 4
DATE: 2023-07-26

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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.\n\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThe work is highly theoretical and does not present any potential negative societal impacts to the best of my knowledge.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 1, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-27"}