SlowGuess/NeurIPS_2023_OCR huggingface dataset:
- 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  $\pmb{p}$  (and  $\pmb{\nu}$ ) at each iteration.
- Next, given the gradient solution  $\pmb{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  $\pmb{p}^{mm}$ . After obtaining  $\pmb{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),  $\pmb{\nu}^{mm}$  (blue dashed) is solved using python package sklearn.svm via (SVM) based on the given label information, and red dashed line represents  $p^{\text{relax}}$  direction instead, which is the solution of (10). Note that in both figures,  $\pmb{\nu}^{mm} / \| \pmb{\nu}^{mm}\| = [0,1]$ . Therefore, in Figure 3(a) all optimal tokens are support vectors and  $p^{\text{relax}} = p^{mm}$ . Whereas in Figure 3(b), yellow  $\star$  is not a support vector and only needs to satisfy positive correlation with  $\pmb{p}$ . Gray dashed line displays the  $p^{mm}$  direction.
Failure of gradient descent's global convergence when  $n \geq 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) - \eta \nabla \mathcal{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  $(\pmb{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 different initializations. This demonstrates that the gradient descent iterates  $\pmb{p}(t + 1) = \pmb{p}(t) - \eta \nabla \mathcal{L}(\pmb{p}(t))$  on (ERM) with two different initializations converge to two different SVM solutions (GMM and LMM). Results validate the necessity of  $n = 1$  in Theorem 2 to provide the gradient descent convergence to  $\pmb{p}^{mm\star}$  from any initialization.
Figure 8: The convergence behavior of the gradient descent on the attention weights  $\pmb{p}$  using the logistic loss in (ERM) with  $n = T = d = 2$ .
The convergence behavior of gradient descent under over-parameterization. To illustrate Theorems 3 & 4, we have investigated the convergence behavior of  $\pmb{p}(t)$  generated by gradient descent in Figure 9(a), using  $n = 4$ ,  $T = 6$ , and conducted 1,000 random trials for varying  $d \in \{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  $\nu$  are uniformly sampled from the unit sphere, while  $Y_{i}$  is uniformly  $\pm 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
(a) Perc. of different convergence scenarios for  $\pmb{p}(t)$
(b) Prob. of  $\underline{\gamma} := \min_{i \in [n], t \in T_i} \gamma_{i\alpha_i} - \gamma_{it}$
Figure 9: Convergence analysis of  $\pmb{p}(t)$  trained with random data using gradient descent. (a) shows three scenarios: (1) attention failing to select one token per input (i.e. softmax is not saturated); (2)  $\pmb{p}$  converging towards  $\pmb{p}^{mm}$ ; and (3)  $\pmb{p}^{mm}$  equating to  $\pmb{p}^{mm\star}$  with red, blue, and green bars, respectively. Considering saturated softmax instances where  $\pmb{p}(t)$  selects one token  $\alpha_{i}$  per-input, (b) presents histogram of the minimal score gap between  $\alpha_{i}$  and its corresponding neighbors  $\mathcal{T}_i$ .
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)\| \to \infty$ ) and we observe that the selected tokens  $\alpha$  almost always converge 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 coefficient 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 \in \mathcal{P}^{mm}$ , confirming Theorem 4.
Furthermore, we found that there exist problem instances, with saturated softmax and  $\| \pmb{p}(t) \| \to \infty$ , that do not converge to either LMM or GMM. We analyzed this phenomenon using the minimum score gap,  $\underline{\gamma} := \min_{i \in [n], t \in \mathcal{T}_i} \gamma_{i\alpha_i} - \gamma_{it}$ , where  $\mathcal{T}_i, i \in [n]$ , represents the sets of SVM-neighbor tokens. Figure 9(b) provides the probability distribution of  $\underline{\gamma}$  (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,  $\mathrm{Prob}(\underline{\gamma} < 0) < 0.01$ . Figure 9(b) also reveals that, in these scenarios, even if  $\underline{\gamma} < 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  $\underline{\gamma} < 0$ , gradient descent solution is still  $> 0.99$  correlated with the max-margin solution in average.
- In Figure 9, we again applied normalized gradient descent with a learning rate equal to 1.
Each trial involved randomly generated data and training for 1000 iterations as discussed in Theorem 4. The tokens selected by  $\pmb{p}$  were denoted as  $(\alpha_{i})_{i = 1}^{n}$ , where  $\alpha_{i} = \arg \max_{t\in [T]}\mathbb{S}(X_{i}\pmb {p})_{t}$ . The averaged softmax probabilities were calculated as  $\bar{s} := \frac{1}{n}\sum_{i = 1}^{n}\mathbb{S}(X_i\pmb {p})_{\alpha_i}$  (same as Figure 3(c)). The red bars in Figure 9(b) represent the values of  $\mathbb{P}(\bar{s}\leq 1 - 10^{-5})$  for each choice of  $d$ . Figure 10 displays the cumulative probability distribution of  $\underline{\gamma}$  from Figure 9(b), with the gray dashed line indicating  $\underline{\gamma} = 0$ . From this figure, we observe that the minimal score gap exhibit a sharp transition at zero (<1% of the instances have  $\underline{\gamma} < 0$ ), demonstrating that, in most random problem instances with  $\| \pmb {p}\| \to \infty$ $(\bar{s}\rightarrow 1)$ , problem directionally
Figure 10: Cumulative prob. of the gap  $\underline{\gamma} \coloneqq \min_{i \in [n], t \in \mathcal{T}_i} \gamma_{i\alpha_i} - \gamma_{it}$  in Figure 9(b).
converges to an LMM i.e.  $\pmb{p}(t) / \| \pmb{p}(t) \| \to \pmb{p}^{mm} / \| \pmb{p}^{mm} \|$ . We believe the rare occurrence of a negative score gap is due to small score differences (so that optimal token is not clearly distinguished) and finite number of gradient iterations we run. Interestingly, even in the negative score gap scenarios, gradient descent is aligned with  $\pmb{p}^{mm}(\alpha)$  (even if  $\pmb{p}^{mm}(\alpha)$  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  $\ell_1$ -regularization, revealing that coordinate descent, when combined with the exponential loss, exhibits an inherent inclination towards  $\ell_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  $\nu$  and existing works, the optimization of  $p$  presents notable differences. 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 offers 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 effort 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, offering flexibility and customization in designing attention-based models. In another closely related work, Oymak et al. [17] analyzed the same attention model as ours, denoted by (ERM). Specifically, they jointly optimize  $\nu$ ,  $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 offers 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.