SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\operatorname {c o n e} _ {\mu} (\boldsymbol {q}) := \left\{\boldsymbol {p} \in \mathbb {R} ^ {d} \mid \left\langle \frac {\boldsymbol {p}}{\| \boldsymbol {p} \|}, \frac {\boldsymbol {q}}{\| \boldsymbol {q} \|} \right\rangle \geq 1 - \mu \right\}. \tag {7}
$$
Given  $R > 0$ , the intersection of  $\mathrm{cone}_{\mu}(\pmb{q})$  and the set  $\{\pmb{p} \in \mathbb{R}^d | \| \pmb{p} \| \geq R\}$  is denoted as  $C_{\mu,R}(\pmb{q})$ :
$$
\mathcal {C} _ {\mu , R} (\boldsymbol {q}) := \operatorname {c o n e} _ {\mu} (\boldsymbol {q}) \cap \left\{\boldsymbol {p} \in \mathbb {R} ^ {d} \mid \| \boldsymbol {p} \| \geq R \right\}. \tag {8}
$$
Figure 2: Gradient descent initialization  $\pmb{p}(0)$  inside the cone containing the locally-optimal solution  $\pmb{p}^{mm}$ .
Next, we demonstrate the existence of parameters  $\mu = \mu (\alpha) > 0$  and  $R > 0$  such that when  $R$  is sufficiently large, there are no stationary points within  $C_{\mu ,R}(\pmb{p}^{mm})$ . Further, the gradient descent initialized within  $C_{\mu ,R}(\pmb{p}^{mm})$  converges in direction to  $\pmb{p}^{mm} / \| \pmb{p}^{mm}\|$ ; refer to Figure 2 for a visualization.
Theorem 3 (Local Convergence of Gradient Descent) Suppose Assumption A on the loss function  $\ell$  holds and assume  $\alpha = (\alpha_{i})_{i=1}^{n}$  are indices of locally-optimal tokens per Definition 2. Then, there is a constant  $\mu = \mu(\alpha) \in (0,1)$  and  $R > 0$  such that  $C_{\mu,R}(\pmb{p}^{mm})$  does not contain any stationary points. Further, for any starting point  $\pmb{p}(0) \in C_{\mu,R}(\pmb{p}^{mm})$ , gradient descent iterates  $\pmb{p}(t+1) = \pmb{p}(t) - \eta \nabla \mathcal{L}(\pmb{p}(t))$  on (ERM) with stepsize  $\eta \leq 1/L_{p}$  satisfies  $\lim_{t \to \infty} \| \pmb{p}(t) \| = \infty$  and  $\lim_{t \to \infty} \pmb{p}(t) / \| \pmb{p}(t) \| = \pmb{p}^{mm} / \| \pmb{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 different 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 effectively separates the locally optimal tokens  $(1,0)$  from the non-optimal token  $(0,0)$ .
2.3 Regularization paths can only converge to locally-optimal max-margin directions
An important question arises regarding whether our definition of LMM (Definition 2) encompasses all possible convergence paths of the attention mechanism when  $\| \pmb{p} \| \to \infty$ . To address this, we introduce the set of LMM directions as follows:
$$
\mathcal {P} ^ {m m} := \left\{\frac {\boldsymbol {p} ^ {m m} (\boldsymbol {\alpha})}{\| \boldsymbol {p} ^ {m m} (\boldsymbol {\alpha}) \|} \mid \boldsymbol {\alpha} \text {i s l o c a l l y - o p t i m a l p e r D e f i n i t i o n 2} \right\}.
$$
The following theorem establishes the tightness of these directions: It demonstrates that for any candidate  $\pmb{q} \notin \mathcal{P}^{mm}$ , its local regularization path within an arbitrarily small neighborhood will provably not converge in the direction of  $\pmb{q}$ .
Theorem 4 Fix  $\pmb{q} \notin \mathcal{P}^{\mathrm{mm}}$  with unit  $\ell_2$  norm. Assume that token scores are distinct (namely  $\gamma_{it} \neq \gamma_{i\tau}$  for  $t \neq \tau$ ) and key embeddings  $\pmb{k}_{it}$  are in general position (see Theorem 7). Fix arbitrary  $\epsilon > 0, R_0 > 0$ . Define the local regularization path of  $\pmb{q}$  as its  $(\epsilon, R_0)$ -conic neighborhood:
$$
\bar {\boldsymbol {p}} (R) = \underset {\boldsymbol {p} \in C _ {\epsilon , R _ {0}} (\boldsymbol {q}), \| \boldsymbol {p} \| \leq R} {\arg \min } \mathcal {L} (\boldsymbol {p}), \quad \text {w h e r e} \quad C _ {\epsilon , R _ {0}} (\boldsymbol {q}) = \operatorname {c o n e} _ {\epsilon} (\boldsymbol {q}) \cap \left\{\boldsymbol {p} \in \mathbb {R} ^ {d} | \| \boldsymbol {p} \| \geq R _ {0} \right\}. \tag {9}
$$
Then, either  $\lim_{R\to \infty}\| \bar{\pmb{p}} (R)\| <  \infty$  or  $\lim_{R\to \infty}\bar{\pmb{p}} (R) / \| \bar{\pmb{p}} (R)\| \neq \pmb {q}.$  In both scenarios  $\lim_{R\to \infty}\bar{\pmb{p}} (R) / R\neq \pmb {q}.$
The result above nicely complements Theorem 3, which states that when gradient descent is initialized above a threshold  $(\| p(0)\| \geq 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\| \geq R_0$ ), the optimal solution path will not converge along  $q\notin \mathcal{P}^{mm}$ .
3 Joint Convergence of Head  $\nu$  and Attention Weights  $p$
In this section, we extend the preceding results to the general case of joint optimization of head  $\pmb{\nu}$  and attention weights  $\pmb{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  $\nu$ , logistic regression in  $\nu$  can exhibit its own implicit bias to a max-margin solution. Concretely, define the attention features  $x_{i}^{p} = X_{i}^{\top}\mathbb{S}(K_{i}p)$  and define the dataset  $\mathcal{D}^p = (Y_i,x_i^p)_{i = 1}^n$ . If this dataset  $\mathcal{D}^p$  is linearly separable, then fixing  $\pmb{p}$  and optimizing only  $\nu$  will converge in the direction of the standard max-margin classifier
$$
\boldsymbol {v} ^ {m m} = \underset {\boldsymbol {v} \in \mathbb {R} ^ {d}} {\arg \min } \| \boldsymbol {v} \| \quad \text {s u b j e c t t o} \quad Y _ {i} \cdot \boldsymbol {v} ^ {\top} \boldsymbol {r} _ {i} \geq 1, \quad \text {f o r a l l} \quad 1 \leq i \leq n, \tag {SVM}
$$
after setting inputs to the attention features  $\boldsymbol{r}_i \gets \boldsymbol{x}_i^p$  [22]. This motivates a clear question:
Under what conditions, optimizing  $\pmb{v}$ ,  $\pmb{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
(a) All inputs are support vectors
(b) (0.5,1.5) is not a support vector
(c) Probability evolutions in (a)
Figure 3: (a) and (b) Joint convergence of attention weights  $p$  ( $\longrightarrow$ ) and classifier head  $v$  ( $\longrightarrow$ ) to max-margin directions. (c) Averaged softmax probability evolution of optimal tokens and logistic probability evolution of output in (a).
tokens  $x_{i\alpha_i}, i \in [n]$  selected by  $p$ , when solving (SVM) with  $r_i \gets x_{i\alpha_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  $(\nu, p)$  asymptotically achieving zero training loss.
3.1 When all attention features are support vectors
In (SVM), define label margin to be  $1 / \| \pmb{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  $\alpha = (\alpha_{i})_{i=1}^{n}$  from each input data achieves the largest label margin. (2) The optimality of the  $\alpha$  choice is strict in the sense that mixing other tokens will shrink the label margin in (SVM).
Assumption C (Optimal Tokens) Let  $\Gamma > 0$  be the label margin when solving (SVM) with  $r_i \gets x_{i\alpha_i}$ . There exists  $\nu > 0$  such that for all  $p$ , solving (SVM) with  $r_i \gets x_i^p$  results in a label margin of at most  $\Gamma - \nu \cdot \max_{i \in [n]} (1 - s_{i\alpha_i})$  where  $s_i = \mathbb{S}(K_i p)$ .
Example: To gain intuition, let us fix  $\pmb{a} \in \mathbb{R}^d$  and consider the dataset obeying  $\pmb{x}_{i1} = Y_i \cdot \pmb{a}$  and  $\| \pmb{x}_{it} \| < \| \pmb{a} \|$  for all  $t \geq 2$  and all  $i \in [n]$ . For this dataset, we can choose  $\alpha_i = 1$ ,  $\pmb{v}^{mm} = \pmb{a} / \| \pmb{a} \|^2$ ,  $\Gamma = 1 / \| \pmb{v}^{mm} \| = \| \pmb{a} \|$  and  $\nu = \| \pmb{a} \| - \sup_{i \in [n], t \geq 2} \| \pmb{x}_{it} \|$ .
Theorem 5 Consider the ridge-constrained solutions  $(\pmb{v}_r, \pmb{p}_R)$  of (ERM) defined as
$$
(\boldsymbol{v}_{r},\boldsymbol{p}_{R}) = \operatorname *{arg  min}_{\| \boldsymbol {v}\| \leq r,\| \boldsymbol {p}\| \leq R}\mathcal{L}(\boldsymbol {v},\boldsymbol {p}).
$$
Suppose there are token indices  $\alpha = (\alpha_{i})_{i=1}^{n}$  for which  $\| \pmb{p}^{mm}(\alpha)\|$  exists (ATT-SVM is feasible) and Assumption C holds for some  $\Gamma, \nu > 0$ . Then,  $\lim_{R \to \infty} \pmb{p}_{R} / R = \pmb{p}^{mm} / \| \pmb{p}^{mm}\|$ , where  $\pmb{p}^{mm}$  is the solution of (ATT-SVM); and  $\lim_{r \to \infty} \pmb{v}_{r} / r = \pmb{v}^{mm} / \| \pmb{v}^{mm}\|$ , where  $\pmb{v}^{mm}$  is the solution of (SVM) with  $r_{i} = x_{i\alpha_{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  $\star$ ) are all support vectors of the (SVM) since  $\pmb{v}^{mm} = [0,1]$  is the optimal classifier direction (depicted by --). Because of this,  $\pmb{p}^{mm}$  will separate optimal  $\star$  tokens from tokens at the  $(0,0)$  coordinate via (ATT-SVM) and its direction is dictated by yellow and teal colored  $\star$ 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  $\pmb{p}$  diverges to  $\infty$  and ends up selecting one token per input. Suppose this  $\pmb{p}$  selects some coordinates  $\alpha = (\alpha_{i})_{i=1}^{n}$ . Let  $S \subset [n]$  be the set of indices where the associated token  $x_{i\alpha_{i}}$  is a support vector when solving (SVM). Set  $\bar{S} = [n] - S$ . Our intuition is as follows: Even if we slightly perturb this  $\pmb{p}$  choice and mix other tokens  $t \neq \alpha_{i}$  over the input set  $\bar{S} \subset [n]$ , since  $\bar{S}$  is not support vector for (SVM), we can preserve the label margin (by only preserving the support vectors  $S$ ). This means that  $\pmb{p}$  may not have to enforce max-margin constraint over inputs  $i \in \bar{S}$ , instead, it suffices to just select
(a) Evolution of softmax probability
(b) Evolution of attention weights
Figure 5: Trajectories of  $\pmb{p}$  with different loss functions and scores in Theorem 2.
Figure 4: Evolution of softmax probability and attention weights when training with normalized gradient descent or constant step size  $\eta$  respectively.
these tokens (asymptotically). This results in the following relaxed SVM problem:
$$
\boldsymbol {p} ^ {\text {r e l a x}} = \arg \min  _ {\boldsymbol {p}} \| \boldsymbol {p} \| \quad \text {s u c h t h a t} \quad \boldsymbol {p} ^ {\top} \left(\boldsymbol {k} _ {i \alpha_ {i}} - \boldsymbol {k} _ {i t}\right) \geq \left\{ \begin{array}{l l} 1 & \text {f o r a l l} \quad t \neq \alpha_ {i}, i \in S \\ 0 & \text {f o r a l l} \quad t \neq \alpha_ {i}, i \in \bar {S} \end{array} . \right. \tag {10}
$$
Here,  $\pmb{p}^{\top}(\pmb{k}_{i\alpha_i} - \pmb{k}_{it}) \geq 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  $\mathbb{S}(K_i\pmb{p}_R)_{\alpha_i} \to 1$ , i.e., the tokens  $(\alpha_i)_{i=1}^n$  are asymptotically selected. Let  $\pmb{v}^{mm}$  be the solution of (SVM) with  $\pmb{r}_i = \pmb{x}_{i\alpha_i}$  and  $\pmb{S}$  be its set of support vector indices. Suppose Assumption C holds over  $\pmb{S}$  i.e. having  $s_{i\alpha_i} < 1$  shrinks the margin when (SVM) is only solved over  $\pmb{S} \subset [n]$ . Then,  $\lim_{r \to \infty} \pmb{v}_r / r = \pmb{v}^{mm} / \| \pmb{v}^{mm} \|$  and  $\lim_{R \to \infty} \pmb{p}_R / R = \pmb{p}^{relax} / \| \pmb{p}^{relax} \|$ , where  $\pmb{p}^{relax}$  is the solution of (10) with  $(\alpha_i)_{i=1}^n$  choices.
To illustrate this numerically, consider Figure 3(b) which modifies Figure 3(a) by pushing the yellow  $\star$  to the northern position (0.5, 1.5). We still have  $\pmb{v}^{mm} = [0,1]$  however the yellow  $\star$  is no longer a support vector of (SVM). Thus,  $\pmb{p}$  solves the relaxed problem (10) which separates green and teal  $\star$ 's by enforcing the max-margin constraint on  $\pmb{p}$  (which is the red direction). Instead, yellow  $\star$  only needs to achieve positive correlation with  $\pmb{p}$  (unlike Figure 3(a) where it dictates the direction). We also display the direction of  $\pmb{p}^{mm}$  using a gray dashed line.
We further investigate the evolution of softmax and logistic output probabilities throughout the training process of Figure 3(a), and the results are illustrated in Figure 3(c). The averaged softmax probability of optimal tokens is represented by the red curve and is calculated as  $\frac{1}{n}\sum_{i=1}^{n}\max_{t\in[T]}\mathbb{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\sum_{i=1}^{n}1/(1 + e^{-Y_i \cdot f(X_i)})$ . This probability also reaches a value of 1, suggesting that the inputs are correctly classified.
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  $\mathbb{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  $\eta$  for training. The dataset model follows Figure 1(c). The softmax probability shown in Figure 4(a) is defined as  $\frac{1}{n}\sum_{i=1}^{n}\max_{t\in[T]}\mathbb{S}(K_i\pmb{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  $\|\pmb{p}\|$ .
The red curves in both figures represent the normalized gradient method, which updates the model parameters  $\pmb{p}$  using  $\pmb{p}(t + 1) = \pmb{p}(t) - \eta \nabla \mathcal{L}(\pmb{p}(t)) / \| \nabla \mathcal{L}(\pmb{p}(t))\|$  with  $\eta = 0.1$ . The blue curves correspond
(b) Epoch 0
(c) Epoch 100
(d) Epoch 200
(e) Epoch 300
(f) Epoch 400
Figure 7: Red curve is the sparsity level  $\widehat{\mathrm{nz}}(s)/T$  of the average attention map which takes values on [0,1]. A sparser vector implies that few key tokens receive significantly higher attention, while the majority of the tokens receive minimal attention. Blue curve is the Frobenius norm of attention weights  $\|W\|_F$  of the final layer. We display their evolutions over epochs.
Figure 6: Illustration of the progressive change in attention weights of the [CLS] token during training in the transformer model, using a specific input image shown in Figure 6(a).
to gradient descent with constant learning rate given by  $\pmb{p}(t + 1) = \pmb{p}(t) - \eta \nabla \mathcal{L}(\pmb{p}(t))$  with  $\eta = 1$ . Observe that the normalized gradient method achieves a softmax probability of 1 quicker as vanilla GD suffers from vanishing gradients. This is visible in Figure 4(b) where blue norm curve levels off.
Real experiments. To study softmax sparsity and the evolution of attention weights throughout training, we train a vision transformer (ViT-base) model [23] from scratch, utilizing the CIFAR-10 dataset [24] for 400 epochs with fixed learning rate  $3 \times 10^{-3}$ . ViT tokenizes an image into  $16 \times 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 \times 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  $\widehat{\mathrm{nz}}(s)$  where  $s$  is the softmax probability vector. The soft-sparsity is computed as the ratio of the  $\ell_1$ -norm to the squared  $\ell_2$ -norm, defined as  $\widehat{\mathrm{nz}}(s) = \| s \|_1 / \| s \|^2$ .  $\widehat{\mathrm{nz}}(s)$  takes values between 1 to  $T = 256$  and a smaller value indicates a sparser vector. Also note that  $\| s \|_1 = \sum_{t=1}^{T} s_t = 1$ . Together with sparsity, Figure 7 also displays the Frobenius norm of the combined key-query matrix  $\mathbf{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  $\| \mathbf{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\star}$ . However, it is worth noting that transient dynamics can exhibit bias towards certain input examples and their associated optimal tokens. We illustrate this idea in Fig 5(a), which displays the trajectories of the gradients for different scores and loss functions. We consider two optimal tokens ( $\star$ ) with scores  $\gamma_1 = 1$  and  $\gamma_2 = C$ , where  $C$  varies. For our analysis, we examine the correlation loss  $\ell(x) = -x$  and the logistic loss  $\ell(x) = \log(1 + e^{-x})$ .
In essence, as  $C$  increases, we can observe that the correlation loss  $\ell(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:  $\nabla \mathcal{L}_i(\pmb{p}) = \ell_i' \cdot \pmb{K}_i^\top \mathbb{S}'(\pmb{X}\pmb{p})\pmb{X}\pmb{v}$ , where  $\mathbb{S}'(\cdot)$  represents the derivative of the softmax function and  $\ell_i' := \ell' (Y_i \cdot \pmb{v}^\top \pmb{X}_i^\top \mathbb{S}(\pmb{X}_i\pmb{p}))$ . Assuming that  $\pmb{p}$  (approximately) selects the optimal tokens, this
simplifies to  $\ell_i' \approx \ell'(\gamma_i)$  and  $\|\nabla \mathcal{L}_i(\pmb{p})\| \propto |\ell'(\gamma_i)| \cdot \gamma_i$ . With the correlation loss,  $|\ell'| = 1$ , resulting in  $\|\nabla \mathcal{L}_i(\pmb{p})\| \propto \gamma_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.,  $|\ell'| = e^{-x} / (1 + e^{-x}) \approx e^{-x}$ . Consequently,  $\|\nabla \mathcal{L}_i(\pmb{p})\| \propto \gamma_i e^{-\gamma_i} \approx e^{-\gamma_i}$ , indicating that a smaller score leads to a larger gradient. These observations explain the empirical behavior we observe.
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  $\pmb{\nu}$  and existing works, the optimization of  $\pmb{p}$  stands out as significantly different. 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 differs from traditional models like MLPs and CNNs by leveraging global interactions for feature representations, showing exceptional empirical performance. However, the underlying mechanisms and learning processes of the attention layer remain unknown. Recent studies such as [51, 52, 53, 54, 23] have focused on specific aspects like representing sparse functions, convex-relaxations, and expressive power. In contrast to our nonconvex (ERM), [52] studies self-attention with linear activation instead of softmax, while [53] approximates softmax using a linear operation with unit simplex constraints. Their main objective is to derive convex reformulations for ERM-based training problem. [55, 56] have developed initial results 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  $\nu, 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.