SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Davoud Ataee Tarzanagh  
University of Pennsylvania  
tarzanaq@upenn.edu
Yingcong Li Xuechen Zhang  
University of California, Riverside  
{y1i692, xzhan394}@ucr.edu
Samet Oymak  
University of Michigan  
UC Riverside  
oymak@umich.edu
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, \text{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.
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(\pmb{X}) = \langle X\pmb{v}, \mathbb{S}(\pmb{X}\pmb{W}^{\top}\pmb{p}) \rangle$ . Here,  $\pmb{X}$  is the sequence of input tokens,  $\pmb{v}$  is the prediction head,  $\pmb{W}$  is the trainable key-query weights, and  $\mathbb{S}$
(a) Global convergence
Figure 1: The convergence behavior of the gradient descent on the attention weights  $\pmb{p}$  using the logistic loss in (ERM). The arrows  $(\longrightarrow)$  represent trajectories from different initializations. Here,  $(- - - )$  and  $(- - - )$  denote the globally- and locally-optimal max-margin directions (GMM, LMM).  $\gamma$  denotes the score of a token per Definition 1. Discussion is provided under Theorems 2 and 3.
(b) Global and local optimal directions
(c) Multiple inputs
denotes the softmax nonlinearity. For transformers,  $\pmb{p}$  corresponds to the [CLS] token or tunable prompt [16, 17, 18], whereas for RNN architectures [1],  $\pmb{p}$  corresponds to the hidden state. Given training data  $(Y_{i},X_{i})_{i = 1}^{n}$  with labels  $Y_{i}\in \{-1,1\}$  and inputs  $X_{i}\in \mathbb{R}^{T\times d}$ , we consider the empirical risk minimization with a decreasing loss function  $\ell (\cdot):\mathbb{R}\to \mathbb{R}$ ,
$$
\mathcal {L} (\boldsymbol {v}, \boldsymbol {p}, \boldsymbol {W}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(Y _ {i} \cdot f \left(\boldsymbol {X} _ {i}\right)\right), \text {w h e r e} f \left(\boldsymbol {X} _ {i}\right) = \boldsymbol {v} ^ {\top} \boldsymbol {X} _ {i} ^ {\top} \mathbb {S} \left(\boldsymbol {X} _ {i} \boldsymbol {W} ^ {\top} \boldsymbol {p}\right). \tag {1}
$$
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  $\pmb{p}$  and  $\pmb{W}$  admit a one-to-one mapping, thus we focus on optimizing  $\pmb{p}$  without losing generality. In Theorem 3, we prove that, under proper initialization:
Gradient descent on  $p$  converges in direction to a max-margin solution – namely (ATT-SVM) – that separates locally-optimal tokens from non-optimal ones.
We call these Locally-optimal Max-Margin (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 Globally-optimal Max-Margin (GMM) direction in Theorems 1 and 2, respectively. A central feature of our results is precisely quantifying optimality in terms of token scores  $\gamma_{t} = Y\cdot \nu^{\top}\boldsymbol{x}_{t}$  where  $\boldsymbol{x}_t$  is the  $t^{\mathrm{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_{i}^{att} = X_{i}^{\top}\mathbb{S}(X_{i}W^{\top}p)$  are separable based on their labels  $Y_{i}, v$  would implement a max-margin classifier. Building on this, we prove that  $p$  and  $v$  converge 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<sup>1</sup>, Section 5 discusses related literature, and Section 6 highlights limitations and future work.
Notations. For any integer  $N \geq 1$ , let  $[N] := \{1, \dots, N\}$ . We use lower-case and upper-case bold letters (e.g.  $\pmb{a}$  and  $\pmb{A}$ ) to represent vectors and matrices, respectively. The entries of  $\pmb{a}$  are denoted as  $\pmb{a}_i$ . We use  $\bar{\sigma}(\pmb{A})$  to denote the maximum singular value of  $\pmb{A}$ . We denote the minimum of two numbers  $a, b$  as  $a \wedge b$ , and the maximum as  $a \vee b$ . Big-O notation  $\mathcal{O}(\cdot)$  hides the universal constants. Throughout, we will use  $\mathcal{L}(\pmb{p})$  and  $\mathcal{L}(\pmb{v}, \pmb{p})$  to denote Objective (1) with fixed  $(\pmb{v}, \pmb{W})$  and  $\pmb{W}$ , respectively.
**Optimization.** Given an objective function  $\mathcal{L}:\mathbb{R}^d\to \mathbb{R}$  and an  $\ell_2$ -norm bound  $R$ , define the regularized solution as
The code for experiments can be found at https://github.com/ucr-optml/max_marginattention.
$$
\bar {p} (R) := \arg \min  _ {\| p \| \leq R} \mathcal {L} (p). \tag {2}
$$
Regularization path – the evolution of  $\bar{\pmb{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 \to \infty} \bar{\pmb{p}}(R)/R$ ) and gradient descent. For gradient descent, we assume the objective  $\mathcal{L}(\pmb{p})$  is smooth and describe the gradient descent process as
$$
\boldsymbol {p} (t + 1) = \boldsymbol {p} (t) - \eta (t) \nabla \mathcal {L} (\boldsymbol {p} (t)), \tag {3}
$$
where  $\eta (t)$  is the stepsize at time  $t$  and  $\nabla \mathcal{L}(\pmb {p}(t))$  is the gradient of  $\mathcal{L}$  at  $\pmb {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  $\mathbf{X} = [\pmb{x}_1,\dots ,\pmb{x}_T]^\top \in \mathbb{R}^{T\times d}$ , self-attention with key-query matrix  $\mathbf{W}\in \mathbb{R}^{d\times d}$ , and value matrix  $\mathbf{V}\in \mathbb{R}^{d\times v}$ , the self-attention model is defined as follows:
$$
f _ {s a} (\boldsymbol {X}) = \mathbb {S} \left(\boldsymbol {X} \boldsymbol {W} \boldsymbol {X} ^ {\top}\right) \boldsymbol {X} \boldsymbol {V}. \tag {4}
$$
Here,  $\mathbb{S}(\cdot)$  is the softmax nonlinearity that applies row-wise on the similarity matrix  $XW X^{\top}$ .
- 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 \in \mathbb{R}^d$  be the tunable token ([CLS] or prompt vector) and concatenate it to  $X$  to obtain  $X_p := [pX^\top]^\top \in \mathbb{R}^{(T+1)\times d}$ . Consider the cross-attention features obtained from  $X_p$  and  $X$  given by
$$
\left[ \begin{array}{c} f _ {\text {c l s}} ^ {\top} (X) \\ f _ {\text {s a}} (X) \end{array} \right] = \mathbb {S} (X _ {p} W X ^ {\top}) X V = \left[ \begin{array}{c} \mathbb {S} (p ^ {\top} W X ^ {\top}) \\ \mathbb {S} (X W X ^ {\top}) \end{array} \right] X V.
$$
The beauty of cross-attention is that it isolates the contribution of  $\pmb{p}$  under the upper term  $f_{cls}(\pmb {X}) =$ $\pmb{V}^{\top}\pmb{X}^{\top}\mathbb{S}(\pmb {X}\pmb{W}^{\top}\pmb {p})\in \mathbb{R}^{\nu}$ . In this work, we use the value weights for classification, thus we set  $\nu = 1$  and denote  $\nu = V\in \mathbb{R}^{d}$ . This brings us to our attention model of interest:
$$
f (\boldsymbol {X}) = \boldsymbol {v} ^ {\top} \boldsymbol {X} ^ {\top} \mathbb {S} (\boldsymbol {K} \boldsymbol {p}), \quad \text {w h e r e} \quad \boldsymbol {K} = \boldsymbol {X} \boldsymbol {W} ^ {\top}. \tag {5}
$$
Here,  $(\nu, 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 \gets \text{Identity}$ ) and  $W$  admit a one-to-one mapping.
Lemma 1 Fix  $\mathbf{u} \in \mathbb{R}^d \setminus \{\mathbf{0}\}$ . Let  $\psi : \mathbb{R}^d \to \mathbb{R}$  and  $\ell : \mathbb{R} \to \mathbb{R}$  be differentiable functions. On the same training data  $(Y_i, X_i)_{i=1}^n$ , define  $\tilde{\mathcal{L}}(\mathbf{p}) := 1/n \sum_{i=1}^n \ell(Y_i \cdot \psi(X_i^\top \mathbb{S}(X_i \mathbf{p})))$  and  $\mathcal{L}(\mathbf{W}) := 1/n \sum_{i=1}^n \ell(Y_i \cdot \psi(X_i^\top \mathbb{S}(X_i \mathbf{W}^\top \mathbf{u})))$ . Consider the gradient descent iterations on  $\mathbf{p}$  and  $\mathbf{W}$  with initial values  $\mathbf{p}(0)$  and  $\mathbf{W}(0) = \mathbf{u} \mathbf{p}(0)^\top / \| \mathbf{u} \|^2$  and stepsizes  $\eta$  and  $\eta / \| \mathbf{u} \|^2$ , respectively:
$$
\boldsymbol {p} (t + 1) = \boldsymbol {p} (t) - \eta \nabla \tilde {\mathcal {L}} (\boldsymbol {p} (t)),
$$
$$
\boldsymbol {W} (t + 1) = \boldsymbol {W} (t) - \frac {\eta}{\| \boldsymbol {u} \| ^ {2}} \nabla \mathcal {L} (\boldsymbol {W} (t)).
$$
We have that  $\pmb{W}(t) = \pmb{u}\pmb{p}(t)^{\top} / \| \pmb{u}\|^{2}$  for all  $t\geq 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,  $(Y_{i},X_{i})_{i = 1}^{n}$  denotes training dataset where  $Y_{i}\in \{-1,1\}$  and  $X_{i}\in \mathbb{R}^{T\times d}$ . We denote the key embeddings of  $X_{i}$  via  $\pmb {K}_i = \pmb {X}_i\pmb{W}^\top$  and explore the training risk
$$
\mathcal {L} (\boldsymbol {v}, \boldsymbol {p}) = \frac {1}{n} \sum_ {i = 1} ^ {n} \ell \left(Y _ {i} \cdot \boldsymbol {v} ^ {\top} \boldsymbol {X} _ {i} ^ {\top} \mathbb {S} (\boldsymbol {K} _ {i} \boldsymbol {p})\right). \tag {ERM}
$$
Importantly, our results apply to general tuples  $(Y_{i},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  $\pmb{x}_{it},\pmb{k}_{it}\in \mathbb{R}^d$ , respectively, for  $t\in [T]$ .
The highly nonlinear and nonconvex nature of the softmax operation makes the training problem in (ERM) a challenging nonconvex optimization problem for  $\pmb{p}$ , even with a fixed  $\nu$ . 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  $\pmb{p} \in \mathbb{R}^d$  given a fixed choice of  $\pmb{v} \in \mathbb{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  $\ell$ . We maintain the following assumption on the loss function throughout this section.
Assumption A (Well-behaved Loss) Over any bounded interval: (1)  $\ell :\mathbb{R}\to \mathbb{R}$  is strictly decreasing. (2)  $\ell^{\prime}$  is  $M_0$ -Lipschitz continuous and  $|\ell^{\prime}(u)|\leq M_{1}$ .
Assumption A includes many common loss functions, including the logistic loss  $\ell(u) = \log(1 + e^{-u})$ , exponential loss  $\ell(u) = e^{-u}$ , and correlation loss  $\ell(u) = -u$ . Assumption A implies that  $\mathcal{L}(\pmb{p})$  is  $L_{p}$ -smooth (see Lemma 6 in Supplementary), where
$$
L _ {p} := \frac {1}{n} \sum_ {i = 1} ^ {n} \left(M _ {0} \| \boldsymbol {v} \| ^ {2} \| \boldsymbol {W} \| ^ {2} \| \boldsymbol {X} _ {i} \| ^ {4} + 3 M _ {1} \| \boldsymbol {v} \| \| \boldsymbol {W} \| ^ {2} \| \boldsymbol {X} _ {i} \| ^ {3}\right). \tag {6}
$$
We now introduce a convex hard-margin SVM problem that separates one token of the input sequence from the rest, jointly solved over all inputs. We will show that this problem captures the optimization properties of softmax-attention. Fix indices  $\alpha = (\alpha_{i})_{i=1}^{n}$  and consider
$$
\boldsymbol {p} ^ {m m} (\alpha) = \arg \min  _ {\boldsymbol {p}} \| \boldsymbol {p} \| \quad \text {s u b j e c t t o} \quad \min  _ {t \neq \alpha_ {i}} \boldsymbol {p} ^ {\top} \left(\boldsymbol {k} _ {i \alpha_ {i}} - \boldsymbol {k} _ {i t}\right) \geq 1, \quad \text {f o r a l l} \quad 1 \leq i \leq n. (A T T - S V M)
$$
Note that existence of  $\pmb{p}^{mm}(\alpha)$  implies the separability of tokens  $\alpha$  from the others. Specifically, choosing direction  $\pmb{p}^{mm}(\alpha)$  will exactly select tokens  $(x_{i\alpha_i})_{i=1}^n$  at the attention output for each input sequence, that is,  $\lim_{R \to \infty} X_i^\top \mathbb{S}(R \cdot K_i \pmb{p}^{mm}(\alpha)) = x_{i\alpha_i}$ . We are now ready to introduce our main results that characterize the global and local convergence of the attention weights  $\pmb{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  $\pmb{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  $\ell(\cdot)$  have a clear definition.
Definition 1 (Token Scores, Optimality & GMM) The score of token  $\mathbf{x}_{it}$  of input  $X_{i}$  is defined as  $\gamma_{it} := Y_{i} \cdot v^{\top} x_{it}$ . The optimal tokens for input  $X_{i}$  are those tokens with highest scores given by
$$
o p t _ {i} \in \operatorname * {a r g m a x} _ {t \in [ T ]} \gamma_ {i t}.
$$
Globally-optimal max-margin (GMM) direction is defined as the solution of (ATT-SVM) with optimal indices  $(\mathbf{opt}_i)_{i=1}^n$  by  $\pmb{p}^{mm\star}$ .
It is worth noting that score definition simply uses the value embeddings  $\pmb{v}^{\top}\pmb{x}_{it}$  of the tokens. Note that multiple tokens within an input might attain the same score, thus  $\mathrm{opt}_i$  or  $p^{mm\star}$  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 \in [n]$  and  $t \neq \mathsf{opt}_i$ , the scores obey  $\gamma_{it} < \gamma_{iopt_i}$ . Then, the regularization path  $\bar{\pmb{p}}(R) = \arg \min_{\| \pmb{p} \| \leq R} \mathcal{L}(\pmb{p})$  converges to the GMM direction i.e.  $\lim_{R \to \infty} \bar{\pmb{p}}(R) / R = \pmb{p}^{mm\star} / \| \pmb{p}^{mm\star} \|$ .
Theorem 1 shows that as the regularization strength  $R$  increases towards the ridgeless problem  $\min_p\mathcal{L}(\pmb {p})$ , the optimal direction  $\bar{\pmb{p}} (R)$  aligns more closely with the max-margin solution  $\pmb{p}^{mm\star}$ . 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(\boldsymbol{X}) = \psi(\boldsymbol{X}^\top \mathbb{S}(\boldsymbol{X}\boldsymbol{W}^\top \boldsymbol{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\in [n]$  and  $t,\tau \neq \mathsf{opt}_i$  , the scores per Definition 1 obey  $\gamma_{it} = \gamma_{i\tau} <   \gamma_{i\mathrm{opt}_i}$
Theorem 2 (Global Convergence of Gradient Descent) Suppose Assumption A on the loss function  $\ell$  and Assumption B on the tokens' score hold. Then, the gradient descent iterates  $\pmb{p}(t + 1) = \pmb{p}(t) - \eta \nabla \mathcal{L}(\pmb{p}(t))$  on (ERM), with the stepsize  $\eta \leq 1 / L_p$  and any starting point  $\pmb{p}(0)$  satisfy  $\lim_{t\to \infty}\| \pmb{p}(t)\| = \infty$ . If  $n = 1$ , we also have  $\lim_{t\to \infty}\pmb{p}(t) / \| \pmb{p}(t)\| = \pmb{p}^{mm\star} / \| \pmb{p}^{mm\star}\|$ .
Theorem 2 shows that gradient descent will diverge in norm, and when  $n = 1$ , the normalized predictor  $\pmb{p}(t) / \| \pmb{p}(t) \|$  converges towards  $\pmb{p}^{mm\star}$ , 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  $\pmb{x} = [x_1, x_2, x_3]$ . We reserve the first two coordinates as key embeddings  $\pmb{k} = [x_1, x_2, 0]$  by setting  $W = \mathrm{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  $\pmb{v} = [0, 0, 1]$ . This way score becomes  $Y \cdot v^{\top} x = Y \cdot x_3$ , allowing us to assign any score (regardless of key embedding).
In Figure 1(a), the gray paths represent gradient descent trajectories from different initializations. The points  $(0,0)$  and  $(1,0)$  correspond to non-optimal tokens, while the point  $(-0.1,1)$  represents the optimal token. Notably, gradient descent iterates with various starting points converge towards the direction of the max-margin solution  $p^{mm\star}$  (depicted by -- -). Moreover, as the iteration count  $t$  increases, the inner product  $\langle p(t) / \| p(t)\| ,p^{mm\star} / \| p^{mm\star}\|\rangle$  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 \in [n]$  and  $t \neq \mathrm{opt}_i$ ,  $Y_i = 1$  and  $\gamma_{it} < \gamma_{i\mathrm{opt}_i}$ . Also assume  $\mathbf{W} \in \mathbb{R}^{d \times d}$  is full-rank. Then  $\mathbf{p}^{\mathrm{mm}}$  exists - i.e. (ATT-SVM) is feasible for optimal indices  $\alpha_i \gets \mathrm{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  $\pmb{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  $\alpha = (\alpha_{i})_{i=1}^{n}$  for which (ATT-SVM) is feasible to obtain  $p^{mm} = p^{mm}(\alpha)$ . Consider tokens  $\mathcal{T}_i \subset [T]$  such that  $(\boldsymbol{k}_{i\alpha_i} - \boldsymbol{k}_{it})^\top \boldsymbol{p}^{mm} = 1$  for all  $t \in \mathcal{T}_i$ . We refer to  $\mathcal{T}_i$  as SVM-neighbors of  $\boldsymbol{k}_{i\alpha_i}$ . Additionally, tokens with indices  $\alpha = (\alpha_i)_{i=1}^n$  are called locally-optimal if for all  $i \in [n]$  and  $t \in \mathcal{T}_i$  scores per Definition 1 obey  $\gamma_{i\alpha_i} > \gamma_{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  $\pmb{q}$  and a scalar  $\mu > 0$ , we define  $\mathrm{cone}_{\mu}(\pmb{q})$  as the set of vectors  $\pmb{p} \in \mathbb{R}^d$  such that the correlation coefficient between  $\pmb{p}$  and  $\pmb{q}$  is at least  $1 - \mu$ :