Optimizing Attention with Mirror Descent: Generalized Max-Margin Token Selection

## Abstract

Abstract Attention mechanisms have revolutionized several domains of artificial intelligence, such as natural language processing and computer vision, by enabling models to selectively focus on relevant parts of the input data. While recent work has characterized the optimization dynamics of gradient descent (GD) in attention-based models and the structural properties of its preferred solutions, less is known about more general optimization algorithms such as mirror descent (MD). In this paper, we investigate the convergence properties and implicit biases of a family of MD algorithms tailored for softmax attention mechanisms, with the potential function chosen as the p 𝑝 p italic_p -th power of the ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT -norm. Specifically, we show that these algorithms converge in direction to a generalized hard-margin SVM with an ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT -norm objective when applied to a classification problem using a softmax attention model. Notably, our theoretical results reveal that the convergence rate is comparable to that of traditional GD in simpler models, despite the highly nonlinear and nonconvex nature of the present problem. Additionally, we delve into the joint optimization dynamics of the key-query matrix and the decoder, establishing conditions under which this complex joint optimization converges to their respective hard-margin SVM solutions. Lastly, our numerical experiments on real data demonstrate that MD algorithms improve generalization over standard GD and excel in optimal token selection.

## 1 Introduction

Attention mechanisms (bahdanau2014neural) have transformed natural language processing (NLP) and large language models (LLMs). Initially developed for encoder-decoder recurrent neural networks (RNNs), attention enables the decoder to focus on relevant input segments rather than relying solely on a fixed-length hidden state. This approach became fundamental in transformers (vaswani2017attention), where attention layers—computing softmax similarities among input tokens—are the architecture’s backbone. Transformers have driven rapid advancements in NLP with models like BERT (devlin2018bert) and ChatGPT (gpt4), and have become the preferred architecture for generative modeling (chen2021evaluating; ramesh2021zero), computer vision (dosovitskiy2021vit; radford2021learning), and reinforcement learning (driess2023palm; chen2021decision). This has led to increased exploration of the mathematical foundations of attention’s optimization.

To understand optimization dynamics of attention mechanisms, ataee2024max; tarzanagh2023transformers studied the implicit bias of gradient descent (GD) in a binary classification setting with a fixed linear decoder. This bias refers to the tendency of GD to learn specific weight characteristics when multiple valid solutions exist. For example, in linear logistic regression on separable data, GD favors solutions aligned with the max-margin class separator (soudry2018implicit; ji2018risk). Similarly, tarzanagh2023transformers; ataee2024max propose a model resembling a hard-margin Support Vector Machine (SVM)—specifically, ($\ell_{p}$-AttSVM) with $p=2$—which maximizes the margin between optimal and non-optimal input tokens based on their softmax logits. These studies show that as training progresses, the combined key-query weights $W(k)$ increasingly align with the locally optimal solution $W^{\alpha}_{\mathrm{mm}}$—the minimizer of ($\ell_{p}$-AttSVM) with $p=2$. Expanding on these insights, vasudeva2024implicit explores global directional convergence and the convergence rate of GD under specific conditions. sheen2024implicit further extends these findings by relaxing assumptions about the convergence of regularized paths for the $(W_{K},W_{Q})$ parameterization of the key-query matrix, showing that gradient flow minimizes the nuclear norm of the key-query weight $W=W_{K}W_{Q}^{\top}$.

Contributions. While the aforementioned works provide insights into the implicit bias and token selection properties of attention mechanisms, their analyses are limited to simplistic GD models. A broader understanding of general descent algorithms such as the mirror descent (MD) family and their token selection properties, is essential. We address this by examining a family of MD algorithms designed for softmax attention, where the potential function is the $p$-th power of the $\ell_{p}$-norm, termed $\ell_{p}$-AttGD. This generalizes both $\ell_{p}$-GD (azizan2018stochastic; sun2022mirror; sun2023unified) and attention GD (ataee2024max; tarzanagh2023transformers), enabling the exploration of key aspects of attention optimization via $\ell_{p}$-AttGD.

Implicit bias of $\ell_{p}$-AttGD for attention optimization.
Building on tarzanagh2023transformers; vasudeva2024implicit; sheen2024implicit, we consider a one-layer attention model for binary classification. Specifically, given a dataset $(X_{i},y_{i},z_{i})_{i=1}^{n}$ where $X_{i}\in\mathbb{R}^{T\times d}$ represents inputs with $T$ tokens, $y_{i}\in\{\pm 1\}$ is the label, and $z_{i}\in\mathbb{R}^{d}$ is the comparison token, we study a single-layer attention model $f(X_{i},z_{i}):=v^{\top}X^{\top}_{i}\sigma(X_{i}Wz_{i})$, where $\sigma(\cdot)$ is the softmax function, $W$ is the key-query matrix, and $v$ is a linear decoder. Our goal is to separate a locally optimal token $\alpha_{i}\in[T]$ of each input sequence $X_{i}$ from the rest via (ERM) with a smooth decreasing loss. We extend the SVM formulation in tarzanagh2023transformers to ($\ell_{p}$-AttSVM), defining a hard-margin SVM using the $\ell_{p}$-norm instead of the $\ell_{2}$-norm. The solution $W^{\alpha}_{\mathrm{mm}}$ separates the locally optimal tokens $(\alpha_{i})_{i=1}^{n}$ with generalized maximum margin (see Section 2). Theorem 3 provides sufficient conditions for $\ell_{p}$-AttGD to converge locally, in direction, to $W_{\mathrm{mm}}^{\alpha}$. Moreover, Theorem 2 shows that $\|W(k)\|_{p,p}$ diverges as $k\rightarrow\infty$. These results characterize the implicit bias towards ($\ell_{p}$-AttSVM) in separating locally optimal tokens, extending previous work to a broader class of algorithms. While Theorem 3 and ataee2024max; tarzanagh2023transformers offer insights into optimization dynamics for $p=2$, the finite-time convergence rate of $\ell_{p}$-AttGD for selecting locally optimal tokens remains unexplored.

Convergence rate of $\ell_{p}$-AttGD to the solution of ($\ell_{p}$-AttSVM). Theorem 4 establishes convergence rates for $\ell_{p}$-AttGD, showing that the iterates $W(k)$, for large $k$, satisfy that $D_{\psi}\left(W_{\mathrm{mm}}^{\alpha}/\|W_{\mathrm{mm}}^{\alpha}\|_{p,p},W(k)
/\|W(k)\|_{p,p}\right)$ decreases at an inverse poly-log rate, where $D_{\psi}(\cdot,\cdot)$ denotes the Bregman divergence (bregman); see Definition 1. Despite optimizing a highly nonlinear, nonconvex softmax function, we achieve a convergence rate similar to GD in linear binary classification (ji2018risk, Theorem 1.1). Compared to the recent polynomial rate $O(k^{-3/4})$ in (vasudeva2024implicit, Theorem 1) for optimizing attention, our rate is logarithmic and slower, but applicable to standard GD and MD for locally optimal token selection. Importantly, we do not require the near-orthogonality of tokens assumption used in vasudeva2024implicit.

Generalized Max-Margin Solutions and Joint Optimization of $(v,W)$. We study the joint problem under logistic loss using $\ell_{p}$-norm regularization path, where (ERM) is solved under $\ell_{p}$-norm constraints, examining the solution trajectory as these constraints relax. Since the problem is linear in $v$, if the attention features $\bar{X}_{i}=X_{i}^{\top}\sigma(X_{i}Wz_{i})$ are separable by their labels $y_{i}$, $v$ acts as a generalized max-margin classifier (azizan2021stochastic). Inspired by ataee2024max; tarzanagh2023transformers, we show that under suitable geometric conditions, $W$ and $v$ generated by $\ell_{p}$-norm regularization path converge to their respective max-margin solutions (Theorem 5 in the appendix).

Finally, we provide extensive numerical experiments on real and synthetic data, demonstrating that MD algorithms improve generalization over standard GD, excelling in optimal token selection and suppressing non-optimal tokens.

## 2 Preliminaries

##### Notations.

Let $N\geq 1$ and $[N]=\{1,2,\dots,N\}$. Vectors are denoted by lowercase letters (e.g., $a$), with components $a_{i}$, and matrices by uppercase letters (e.g., $A$). The minimum and maximum of scalars $a$ and $b$ are $a\wedge b$ and $a\vee b$, respectively. For a vector $v\in\mathbb{R}^{d}$, the $p$-norm is $\|v\|_{p}=(\sum_{i=1}^{d}|v_{i}|^{p})^{1/p}$. For a matrix $M\in\mathbb{R}^{d\times d}$, the $p,p$-norm is $\|M\|_{p,p}=(\sum_{i=1}^{d}\sum_{j=1}^{d}|M_{ij}|^{p})^{1/p}$. When $p=2$, these are the Euclidean norm for vectors and the Frobenius norm for matrices. For any two matrices $X,Y$ of the same dimensions, we define $\langle X,Y\rangle:=\text{trace}(X^{\top}Y)$. Throughout, for a differentiable function $f:\mathbb{R}^{d\times d}\rightarrow\mathbb{R}$, we define $D_{f}:\mathbb{R}^{d\times d}\times\mathbb{R}^{d\times d}\rightarrow\mathbb{R}$ as

$$ $\displaystyle D_{f}(W,V):=f(W)-f(V)-\langle\nabla f(V),W-V\rangle.$ (1) $$

Asymptotic notations $\mathcal{O}$ and $\Omega$ hide constant factors, and all logarithms are natural ($e$-base).

##### Single-head attention model.

Given input sequences $X,Z\in\mathbb{R}^{T\times d}$ with length $T$ and embedding dimension $d$, the output of a single-head (cross)-attention layer is computed as:

$$ $\displaystyle\sigma(XW_{Q}W_{K}^{\top}Z^{\top})XW_{V},$ $$

where $W_{Q},W_{K}\in\mathbb{R}^{d\times d_{1}}$, $W_{V}\in\mathbb{R}^{d\times d_{2}}$
are trainable key, query, value matrices, respectively; $\sigma(XW_{Q}W_{K}^{\top}Z^{\top})$ is the attention map; and $\sigma(\cdot):\mathbb{R}^{T\times T}\rightarrow\mathbb{R}^{T\times T}$ denotes the row-wise softmax function applied row-wise on $XW_{Q}W_{K}^{\top}Z^{\top}$. Similar to ataee2024max; tarzanagh2023transformers, we reparameterize the key-query product matrix as $W:=W_{Q}W_{K}^{\top}\in\mathbb{R}^{d\times d}$, and subsume the value weights $W_{V}$ within the prediction head $v\in\mathbb{R}^{d}$. Suppose the first token of $Z$, denoted by $z$, is used for prediction. Then, the attention model can be formulated as

$$ $\displaystyle f(X,z)=v^{\top}X^{\top}\sigma(XWz).$ (2) $$

Attention-based empirical risk minimization. We consider a one-layer attention model (2) for binary classification. Consider the dataset $(X_{i},y_{i},z_{i})_{i=1}^{n}$, where $X_{i}\in\mathbb{R}^{T\times d}$ is the input with $T$ tokens each of dimension $d$, $y_{i}\in\{\pm 1\}$ is the label, and $z_{i}\in\mathbb{R}^{d}$ is the token used for comparison. We use a smooth decreasing loss function $l:\mathbb{R}\rightarrow\mathbb{R}$ and study empirical risk minimization (ERM):

$$ $\displaystyle\min_{v\in\mathbb{R}^{d},W\in\mathbb{R}^{d\times d}}\quad\mathcal {L}(v,W):=\frac{1}{n}\sum_{i=1}^{n}l\left(y_{i}v^{\top}X_{i}^{\top}\sigma\left (X_{i}Wz_{i}\right)\right).$ (ERM) $$

Throughout, we will use $\mathcal{L}(W)$ to denote the objective of (ERM) with fixed $v$.

The highly nonlinear and nonconvex nature of the softmax operation makes the training problem described in (ERM) a challenging nonconvex optimization task for $W$, even with a fixed $v$. Next, we provide an assumption on the loss function necessary to demonstrate the convergence of MD for margin maximization within the attention mechanism.

###### Assumption A .

Within any closed interval, the loss function $l:\mathbb{R}\rightarrow\mathbb{R}$ is strictly decreasing and differentiable, and its derivative $l^{\prime}$ is bounded and Lipschitz continuous.

Assumption A aligns with the assumptions on loss functions in ataee2024max; tarzanagh2023transformers. Commonly used loss functions, such as $l(x)=e^{-x}$, $l(x)=-x$, and $l(x)=\log(1+e^{-x})$, satisfy this assumption.

Preliminaries on mirror descent. We review the MD algorithm (original-md) for solving attention-based (ERM). Mirror descent is defined using a potential function. We focus on differentiable and strictly convex potentials $\psi$ defined on the entire domain $\mathbb{R}^{d\times d}$. Note that in general, the potential function is a convex function of Legendre type (rockafellar2015convex, Section 26). We call $\nabla\psi$ the mirror map. The natural “distance” associated with the potential $\psi$ is given by the Bregman divergence (bregman1967relaxation).

###### Definition 1 (Bregman Divergence) .

For a strictly convex function $\psi:\mathbb{R}^{d\times d}\rightarrow\mathbb{R}$, the expression $D_{\psi}(\cdot,\cdot)$ defined in (1) is called the Bregman divergence.

An important example of a potential function is $\psi=\frac{1}{2}\|\cdot\|^{2}_{F}$. In this case, the Bregman divergence simplifies to $D_{\psi}(W,V)=\frac{1}{2}\|W-V\|^{2}_{F}$; For more details, see bauschke2017descent. MD with respect to the mirror map $\psi$ is a generalization of GD where the Bregman divergence is used as a measure of distance. Given a stepsize $\eta>0$, the MD algorithm is as follows:

$$ $\displaystyle W(k+1)\leftarrow\operatorname*{arg\,min}_{W\in\mathbb{R}^{d \times d}}\left\{{\eta}^{-1}D_{\psi}(W,W(k))+\left\langle\nabla\mathcal{L}(W(k )),{W}\right\rangle\right\}.$ ( MD ) $$

Equivalently, MD can be written as $\nabla\psi(W(k+1))=\nabla\psi(W(k))-\eta\nabla L(W(k))$; see bubeck2015convex; juditsky2011first. A useful fact about the Bregman divergence is that it is non-negative and $D_{\psi}(W,V)=0$ if and only if $W=V$.

Preliminaries on attention SVM. Following ataee2024max; tarzanagh2023transformers, we use the following definition of token scores.

###### Definition 2 (Token Score) .

For prediction head $v\in\mathbb{R}^{d}$, the score of token $X_{it}$ is $\gamma_{it}=y_{i}v^{\top}X_{it}$.

It is important to highlight that the score is determined solely based on the *value embeddings* $v^{\top}X_{it}$ of the tokens. The softmax function $\sigma(\cdot)$ minimizes (ERM) by selecting the token with the highest score (tarzanagh2023transformers, Lemma 2). Using (2), tarzanagh2023transformers defines globally optimal tokens $(\texttt{opt}_{i})_{i=1}^{n}$, with each $\texttt{opt}_{i}$ maximizing the score for $X_{i\texttt{opt}_{i}}$. For our MD analysis, we primarily consider locally optimal tokens, as they are more general than globally optimal ones. Locally optimal tokens (ataee2024max; tarzanagh2023transformers) are characterized by having scores that surpass those of nearby tokens. Intuitively, these are the tokens that locally minimize (ERM) upon selection and can be defined based on support tokens. Before presenting the mathematical notion of locally optimal tokens, we provide the formulation of the attention SVM problem. Given a set of (locally) optimal token indices $(\alpha_{i})_{i=1}^{n}$, tarzanagh2023transformers defines the following hard-margin attention SVM problem, which aims to separate, with maximal margin, (locally) optimal tokens from the rest of the tokens for every input sequence:

$$ $\begin{split}W^{\alpha}_{\mathrm{mm}}:=\operatorname*{arg\,min}_{W\in\mathbb{R }^{d\times d}}\|W\|_{F}~{}~{}\textnormal{s.t.}~{}~{}(X_{i\alpha_{i}}-X_{it})^{ \top}Wz_{i}\geq 1,~{}\textnormal{for all }t\in[T]-\{\alpha_{i}\},\,i\in[n]. \end{split}$ (3) $$

The constraint $(X_{i\alpha_{i}}-X_{it})^{\top}Wz_{i}\geq 1$ indicates that in the softmax probability vector $\sigma(X_{i}Wz_{i})$, the $\alpha_{i}$ component has a significantly higher probability compared to the rest, and so these problems solve for a sort of probability separator that has the lowest norm.

###### Definition 3 (Globally and Locally Optimal Tokens) .

Consider the dataset $(X_{i},y_{i},z_{i})_{i=1}^{n}$.

- 1.
The tokens with indices $\texttt{opt}=(\texttt{opt}_{i})_{i=1}^{n}$ are called globally optimal if they have the highest scores, given by $\texttt{opt}_{i}\in\arg\max_{t\in[T]}\gamma_{it}$.
- 2.
Fix token indices $(\alpha_{i})_{i=1}^{n}$ for
which (3) is feasible to obtain $W^{\alpha}_{\mathrm{mm}}$. Let the support tokens $\mathcal{T}_{i}$ for the $i^{th}$ data be the set of tokens $\tau$ such that $(X_{i\alpha_{i}}-X_{i\tau})^{\top}W^{\alpha}_{\mathrm{mm}}z_{i}=1$. The tokens with indices $(\alpha_{i})_{i=1}^{n}$ are called locally optimal if, for all $i\in[n]$ and $\tau\in\mathcal{T}_{i}$, the scores per Def. 2 obey $\gamma_{i\alpha_{i}}>\gamma_{i\tau}$.

It is worth noting that token scoring and optimal token identification can help us understand the importance of individual tokens and their impact on the overall objective. A token score measures how much a token contributes to a prediction or classification task, while an optimal token is defined as the token with the highest relevance in the corresponding input sequence (ataee2024max; tarzanagh2023transformers). For illustration, please refer to Figure 1.

###### Assumption A .

###### Definition 1 (Bregman Divergence) .

###### Definition 2 (Token Score) .

###### Definition 3 (Globally and Locally Optimal Tokens) .

## 3 Implicit Bias of Mirror Descent for Optimizing Attention

### 3.1 Optimizing Attention with Fixed Head v 𝑣 v italic_v

In this section, we assume the prediction head is fixed and focus on the directional convergence of MD and its token selection property through the training of the key-query matrix $W$. The analysis will later be expanded in Section 3.2 to include the joint optimization of both $v$ and $W$.

We investigate the theoretical properties of the main algorithm of interest, namely MD with $\psi(\cdot)=\frac{1}{p}\|\cdot\|_{p,p}^{p}$ for $p>1$ for training (ERM) with fixed $v$. We shall call this algorithm $\ell_{p}$-norm AttGD because it naturally generalizes attention training via GD to $\ell_{p}$ geometry, and for conciseness, we will refer to this algorithm by the shorthand $\ell_{p}$-AttGD. As noted by azizan2021stochastic, this choice of mirror potential is particularly of practical interest because the mirror map $\nabla\psi$ updates become separable in coordinates and thus can be implemented coordinate-wise independently of other coordinates.

$$ $\displaystyle\forall~{}~{}i,j\in[d],\quad\begin{cases}\left[W(k+1)\right]_{ij} \leftarrow\left|[W(k)]^{+}_{ij}\right|^{\frac{1}{p-1}}\cdot\textnormal{sign} \left([W(k)]^{+}_{ij}\right),\vspace{.3cm}\\ [W(k)]^{+}_{ij}:=|[W(k)]_{ij}|^{p-1}\textnormal{sign}([W(k)]_{ij})-\eta[\nabla L (W(k))]_{ij}.\end{cases}$ ( $\ell_{p}$ -AttGD ) $$

In the following, we first identify the conditions that guarantee the convergence of $\ell_{p}$-AttGD. The intuition is that, for attention to exhibit implicit bias, the softmax nonlinearity should select the locally optimal token within each input sequence. tarzanagh2023transformers shows that under certain assumptions, training an attention model using GD causes its parameters’ direction to converge.

This direction can be found by solving a simpler optimization problem, such as (3), which selects the locally optimal token. Depending on the attention model’s parameterization, the attention SVM varies slightly. In this work, we generalize (3) using the $\ell_{p}$-norm as follows:

###### Definition 4 (Attention SVM with $\ell_{p}$ –norm Objective) .

For a dataset $\{(X_{i},y_{i},z_{i})\}_{i=1}^{n}$ with $y_{i}\in\{\pm 1\}$, $X_{i}\in\mathbb{R}^{T\times d}$, and token indices $(\alpha_{i})_{i=1}^{n}$, $\ell_{p}$-based attention SVM is defined as

$$ $\begin{split}W^{\alpha}_{\mathrm{mm}}&:=\operatorname*{arg\,min}_{W\in\mathbb{ R}^{d\times d}}\|W\|_{p,p}\\ ~{}~{}\textnormal{subj.~{}to}~{}&~{}(X_{i\alpha_{i}}-X_{it})^{\top}Wz_{i}\geq 1 ,~{}~{}\textnormal{for all}~{}~{}t\in[T]-\{\alpha_{i}\},~{}i\in[n].\end{split}$ ( $\ell_{p}$ -AttSVM) $$

Figure: Figure 1: Visualization of ($\ell_{p}$-AttSVM) for $p=3$.
Refer to caption: extracted/5918143/Images/svm_token_visualization.png

Problem ($\ell_{p}$-AttSVM) is strictly convex, so it has unique solutions when feasible. Furthermore, under mild overparameterization, $d\geq\max\{T-1,n\}$, the problem is almost always feasible (tarzanagh2023transformers, Theorem 1). Next, we assert that the solution to the ($\ell_{p}$-AttSVM) problems determines the direction that the attention model parameters approach as the training progresses.

###### Example 1 .

Consider the matrices $X_{1}=[5,~{}0;~{}0,~{}1]$ and $X_{2}=[-5,~{}0;~{}0,~{}-1]$ with $y_{1}=-y_{2}=1$. Let $X_{i1}$ be the optimal token and $X_{it}$ be the others. Solving Problem ($\ell_{p}$-AttSVM) with $p=3$ and setting $z_{i}=X_{i1}$, we obtain the solution $W_{mm}:=W^{\alpha}_{\mathrm{mm}}=[0.03846,~{}0;~{}-0.00769,~{}0]$. Figure 1 illustrates how the optimal tokens $X_{11}$ and $X_{21}$ are separated by the dashed decision boundaries. These boundaries are orthogonal to the vectors $W_{mm}z_{i}$ and indicate the hyperplanes that separate the sequences based on the optimal token in each case.

###### Theorem 1 ( $\ell_{p}$ –norm Regularization Path) .

Suppose Assumption A on the loss function holds. Consider the ridge-constrained solutions $W^{(R)}$ of (ERM) defined as

$$ $W^{(R)}:=\operatorname*{arg\,min}_{W\in\mathbb{R}^{d\times d}}~{}\mathcal{L}(W )\quad\textnormal{subj. to}\quad\|W\|_{p,p}\leq R.$ ( $\ell_{p}$ -AttRP ) $$

Then, $\lim_{R\to\infty}W^{(R)}/R=W^{\operatorname{opt}}_{\mathrm{mm}}/{\|W^{
\operatorname{opt}}_{\mathrm{mm}}\|_{p,p}}$, where $W^{\operatorname{opt}}_{\mathrm{mm}}$ is the solution of ($\ell_{p}$-AttSVM), with $\alpha_{i}$ replaced by $\operatorname{opt}_{i}$.

Theorem 1 shows that as the regularization strength $R$ increases, the optimal direction $W^{(R)}$ aligns more closely with the max-margin solution $W^{\alpha}_{\mathrm{mm}}$. This theorem, which allows for globally optimal tokens (see Definition 3), does not require any specific initialization for the $\ell_{p}$-AttRP algorithm and demonstrates that max-margin token separation is an essential feature of the attention mechanism.

Next, we present the convergence of MD applied to (ERM). Under certain initializations, the parameter’s $\ell_{p}$-norm increases to infinity during training, with its direction approaching the ($\ell_{p}$-AttSVM) solution. To describe these initializations, we introduce the concept of cone sets.

###### Definition 5 .

Given a square matrix $W\in\mathbb{R}^{d\times d}$, $\mu\in(0,1)$, and some $R>0$,

$$ $\displaystyle S_{p,\mu}(W)$ $\displaystyle:=\left\{W^{\prime}\in\mathbb{R}^{d\times d}\mid D_{\psi}\left( \frac{W}{\|W\|_{p,p}},\frac{W^{\prime}}{\|W^{\prime}\|_{p,p}}\right)\leq\mu \right\},$ (4a) $\displaystyle C_{p,\mu,R}(W)$ $\displaystyle:=S_{\mu}(W)\cap\left\{W^{\prime}\mid\|W\|_{p,p}\geq R\right\}.$ (4b) $$

These sets contain matrices with a similar direction to a reference matrix $W$, as captured by the inner product in $S_{\mu}(W)$. For $C_{p,\mu,R}(W)$, there is an additional constraint that the matrices must have a sufficiently high norm. We note that $S_{p,\mu}(W)$ and $C_{p,\mu,R}(W)$ reduce to their Euclidean variants as described in ataee2024max; tarzanagh2023transformers. With this definition, we present our first theorem about the norm of the parameter increasing during training.

###### Theorem 2 .

Suppose Assumption A holds. Let $(\alpha_{i})_{i=1}^{n}$ be locally optimal tokens as per Definition 3. Consider the sequence $W(k)$ generated by Algorithm $\ell_{p}$-AttGD. For a small enough stepsize $\eta$, if $W(0)\in C_{p,\mu,R}(W^{\alpha}_{\mathrm{mm}})$ for some dataset-dependent constants $\mu,R>0$, then we have $\lim_{k\rightarrow\infty}\|W(k)\|_{p,p}=\infty.$

###### Remark 1 .

The condition on the stepsize $\eta$ is that it must be sufficiently small so that $\psi(\cdot)-\eta\mathcal{L}(\cdot)$ remains convex for the matrices $W$ along the path traced by the iterates $W(k)$. Specifically, there exists an index $k$ and a real number $r\in[0,1]$ such that $W=rW(k)+(1-r)W(k+1)$. This restriction applies to all theorems in this paper that require a sufficiently small stepsize $\eta$.

This theorem implies that the parameters will increase and diverge to infinity, justifying the need to characterize the convergence of their direction.

###### Theorem 3 (Convergence of $\ell_{p}$ -AttGD ) .

Suppose Assumption A holds. Let $(\alpha_{i})_{i=1}^{n}$ be locally optimal tokens as per Definition 3. Consider the sequence $W(k)$ generated by Algorithm $\ell_{p}$-AttGD. For a small enough stepsize $\eta$, if $W(0)\in C_{p,\mu,R}(W^{\alpha}_{\mathrm{mm}})$ for some dataset-dependent constants $\mu>0,R>\textnormal{exp}(2)$, then

$$ $\lim_{k\rightarrow\infty}\frac{W(k)}{\|W(k)\|_{p,p}}=\frac{W^{\alpha}_{\mathrm {mm}}}{\|W^{\alpha}_{\mathrm{mm}}\|_{p,p}}.$ $$

These theorems show that as the parameters grow large enough and approach a locally optimal direction, they will keep moving toward that direction.

###### Theorem 4 (Convergence Rate of $\ell_{p}$ -AttGD ) .

Suppose Assumption A holds. Let $(\alpha_{i})_{i=1}^{n}$ be locally optimal tokens as per Definition 3. Consider the sequence $W(k)$ generated by Algorithm $\ell_{p}$-AttGD. For a small enough stepsize $\eta$, if $W(0)\in C_{p,\mu,R}(W^{\alpha}_{\mathrm{mm}})$ for some $\mu>0,R>\textnormal{exp}(2)$, then

$$ $D_{\psi}\left(\frac{W^{\alpha}_{\mathrm{mm}}}{\|W^{\alpha}_{\mathrm{mm}}\|_{p, p}},\frac{W(k)}{\|W(k)\|_{p,p}}\right)=\mathcal{O}\left(\begin{cases}\frac{ \log\log k}{\log k}&\text{if }p>2,\\ \frac{(\log\log k)^{2}}{\log k}&\text{if }p=2,\\ \frac{1}{(\log k)^{p-1}}&\textnormal{otherwise}.\end{cases}\right).$ (5) $$

Despite optimizing a highly nonlinear, nonconvex softmax function, we achieve a convergence rate similar to GD in linear binary classification (ji2018risk, Theorem 1.1) (up to a $\log\log k$ factor). The theorems we prove hinge on the parameter entering the set $W(k)\in C_{p,\mu,R}(W^{\alpha}_{\mathrm{mm}})$ with a high enough norm. Since we aim to show the parameter converges in direction to the cone center, $W_{\mathrm{mm}}^{\alpha}$, we need conditions ensuring the parameters remain in the cone. We formalize this in Lemma 17 and prove that for any $\mu>0$ and locally optimal tokens $(\alpha_{i})_{i=1}^{n}$ (Definition 3), there exist constants $R,\mu^{\prime}>0$ depending on the dataset and $\mu$, such that if $W(0)\in C_{p,\mu^{\prime},R}(W^{\alpha}_{\mathrm{mm}})$, then $W(k)\in C_{p,\mu,R}(W^{\alpha}_{\mathrm{mm}})$ for all $k$, meaning the iterates remain within a larger cone; see Figure 2.

Figure: Figure 2: Illustration of Lemma 17. $W(k)$, $\forall k>0$ are within the larger set.
Refer to caption: x1.png

For Theorem 2, we show in Lemma 11 that at any timestep $k\geq 0$, the norm of the $W$ parameter evolves in the following manner,

$$ $\displaystyle|W(k+1)\|_{p,p}^{p-1}$ $\displaystyle\geq\|W(k)\|_{p,p}^{p-1}$ $\displaystyle+\frac{\eta}{\|W(k)\|_{p,p}}\langle-\nabla\mathcal{L}(W(k)),W(k)\rangle.$ $$

With the above, to prove Theorem 2, it is enough to show that $\langle-\nabla\mathcal{L}(W(k)),W(k)\rangle$ is positive and large enough to keep the norm increasing to infinity. Specifically, in Lemma 9 we show that there exist dataset-dependent constants $R,\delta,\mu>0$ such that for all $W,V\in C_{p,\mu,R}(W^{\alpha}_{\mathrm{mm}})$ with $\|V\|_{p,p}=\|W^{\alpha}_{\mathrm{mm}}\|_{p,p}$,

$$ $-\langle\nabla\mathcal{L}(W),V\rangle=\Omega\left(\textnormal{exp}\left(-\frac {\|W\|_{p,p}}{\|W^{\alpha}_{\mathrm{mm}}\|_{p,p}}\left(1+\frac{1}{2}\delta \right)\right)\right)>0.$ $$

Theorem 3 is a direct consequence of Theorem 4, which extends the analysis that is done for Lemma 17 by providing a tighter bound on how thin the cone set $C_{p,\mu,R}$ may be for later iterates.

###### Definition 4 (Attention SVM with ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT –norm Objective) .

###### Example 1 .

###### Theorem 1 ( ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT –norm Regularization Path) .

###### Definition 5 .

###### Theorem 2 .

###### Remark 1 .

###### Theorem 3 (Convergence of ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT -AttGD ) .

###### Theorem 4 (Convergence Rate of ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT -AttGD ) .

### 3.2 Training Dynamics of Mirror Descent for Joint Optimization of W 𝑊 W italic_W and v 𝑣 v italic_v

This section explores the training dynamics of jointly optimizing the prediction head $v$ and attention weights $W$. Unlike Section 3.1, the main challenge here is the evolving token scores $\gamma$ influenced by the changing nature of $v$. This requires additional technical considerations beyond those in Section 3.1, which are also addressed in this section. Given stepsizes $\eta_{W},\eta_{v}>0$, we consider the following joint updates for $W$ and $v$ applied to (ERM), respectively: For all $i,j\in[d]$:

$$ $\displaystyle\begin{cases}\left[W(k+1)\right]_{ij}\leftarrow\left|[W(k)]^{+}_{ ij}\right|^{\frac{1}{p-1}}\cdot\textnormal{sign}\left([W(k)]^{+}_{ij}\right), \vspace{.3cm}\\ [W(k)]^{+}_{ij}:=|[W(k)]_{ij}|^{p-1}\textnormal{sign}([W(k)]_{ij})-\eta_{W}[ \nabla_{W}L(W(k),v(k))]_{ij},\vspace{.3cm}\\ [v(k+1)]_{i}\leftarrow\left|[v(k)]_{i}^{+}\right|^{\frac{1}{p-1}}\cdot \textnormal{sign}([v(k)]_{i}^{+}),\vspace{.3cm}\\ [v(k)]_{i}^{+}:=|[v(k)]_{i}|^{p-1}\textnormal{sign}([v(k)]_{i})-\eta_{v}[ \nabla_{v}L(W(k),v(k))]_{i}.\end{cases}$ ( $\ell_{p}$ -JointGD ) $$

We discuss the implicit bias and convergence for $v(k)$ below. From previous results (azizan2021stochastic), one can expect $v(k)$ to converge to the $\ell_{p}$-SVM solution, i.e., the max-margin classifier separating the set of samples $\{(X_{i\alpha_{i}},y_{i})\}_{i=1}^{n}$, where $X_{i\alpha_{i}}$ denote the (locally) optimal token for each $i\in[n]$. Consequently, we consider the following hard-margin SVM problem,

$$ $\displaystyle v_{\mathrm{mm}}=\arg\min_{v\in\mathbb{R}^{d}}\|v\|_{p}\quad\text {subj. to}\quad y_{i}X_{i\alpha_{i}}^{\top}v\geq 1\quad\textnormal{for all} \quad i\in[n].\quad$ ( $\ell_{p}$ -SVM) $$

Figure: Figure 3: Effect of token selection on margin size in ($\ell_{p}$-SVM) for Example 1. The first plot shows the largest class margin with optimal tokens $X_{11}$ and $X_{21}$. In subsequent plots, as different tokens are used, the class margin (light blue shaded area) decreases, reflecting suboptimal class separation.
Refer to caption: extracted/5918143/Images/svm_case1.png

In ($\ell_{p}$-SVM), define the *label margin* as $1/\|v_{\mathrm{mm}}\|_{p}$. The label margin quantifies the distance between the separating hyperplane and the nearest data point in the feature space. A larger label margin indicates better generalization performance of the classifier, as it suggests that the classifier has a greater separation between classes. From ($\ell_{p}$-SVM) and Definitions 2 and 3, an additional intuition by ataee2024max behind optimal tokens is that they maximize the label margin when selected; see Figure 3 in the appendix for a visualization. Selecting the locally optimal token indices $\alpha=(\alpha_{i})_{i=1}^{n}$ from each input data sequence achieves the largest label margin, meaning that including other tokens will reduce the label margin as defined in ($\ell_{p}$-SVM). In the Appendix G, we show that $W$ and $v$ generated by $\ell_{p}$-JointRP converge to their respective max-margin solutions under suitable geometric conditions (Theorem 5 in the appendix).

Figure: (a) $\ell_{1.75}$ Directional Bregman Div.
Refer to caption: extracted/5918143/Images/p1_75.jpg

Figure: (a) $\ell_{1.75}$ Optimization Path
Refer to caption: extracted/5918143/Images/p1_75_instance.jpg

## 4 Experimental Results

### 4.1 Synthetic Data Experiments

We describe the setup of the experiments for  $\ell_{p}$-AttGD and  $\ell_{p}$-JointGD and their results.

$\ell_{p}$-AttGD Experiment. To measure the directional distance between $W^{\mathrm{mm}}_{\alpha}$ (solution of ($\ell_{p}$-AttSVM)) and $W(k)$ (output of $\ell_{p}$-AttGD), we use a directional Bregman divergence, defined as $D_{\psi}(W/\|W\|_{p,p},V/\|V\|_{p,p})$ for $W,V\in\mathbb{R}^{d\times d}$. We compare the ($\ell_{p}$-AttSVM) solution with the $\ell_{q}$ optimization path for all $p,q\in\{1.75,2,3\}$ for synthetically generated data (described in detail in the Appendix). The experiment is repeated 100 times, and the average directional Bregman divergence is reported. A closer look at one sample trial is also provided.

Figure 4 shows the directional Bregman divergence between the ($\ell_{p}$-AttSVM) solution and the $\ell_{q}$ optimization path for each pair $p,q\in\{1.75,2,3\}$. In Figure 4(a), the divergence converges to 0 0 only for the ($\ell_{p}$-AttSVM) ($p=1.75$) solution, indicating that the $\ell_{1.75}$ path does not converge to the $p=2$ or $3$ solutions. The shrinking standard deviation shows consistent behavior. Similarly, Figures 4(b) and 4(c) show the divergence converging to 0 0 for the corresponding ($\ell_{p}$-AttSVM) solution, demonstrating that the $\ell_{p}$ optimization path converges to the ($\ell_{p}$-AttSVM) solution, with the direction of convergence changing with $p$. In addition to the directional Bregman divergence, we can also observe the convergence in direction for one of the trials directly by plotting how two of the entries of $W$ change during training simultaneously and plotting it on a Cartesian graph, then showing that the path it follows converges to the direction of the ($\ell_{p}$-AttSVM) solution. As we can see in Figure 5, each of the $\ell_{p}$ optimization paths follows the direction of the corresponding ($\ell_{p}$-AttSVM) solution.

$\ell_{p}$-JointGD Experiments. We use the data from Example  2 in the Appendix to train a model using  $\ell_{p}$-JointGD for $p=1.75,2,$ and $3$. The comparison between the iterates and the SVM solutions in Figure 6 shows that the iterates of $W$ and $v$ converge to the  $\ell_{p}$-AttSVM and  $\ell_{p}$-SVM directions, respectively, for each of $p=1.75,2,$ and $3$. These convergence are similar to Theorem 5, as in both this experiment and that theorem, we get that the iterates converge to the SVM problem solutions. In addition to these iterates, we record the evolution of the average softmax probability of the optimal token, along with the average logistic probability of the model, which we define to be $1/n\sum_{i=1}^{n}1/(1+e^{-\gamma_{i\alpha_{i}}})$.

Figure: (a) $\ell_{1.75}$ Optimization Paths
Refer to caption: extracted/5918143/Images/p1_75-joint-iterates.jpg

Figure: (a) $\ell_{1.75}$ Probabilties
Refer to caption: extracted/5918143/Images/p1_75-joint-prob.jpg

As shown in Figure 7, each average softmax probability converges to $1$, indicating that the attention mechanism produces a softmax vector converging to a one-hot vector during different  $\ell_{p}$-JointGD training. Moreover, the average logistic probability also converges to $1$, indicating the model’s prediction reaches $100\%$ accuracy.

### 4.2 Real Data Experiments

This section presents evidence of improved generalization and token selection from training an attention network with MD instead of GD, along with a hypothesis for this improvement.

**Table 1: Test accuracies of transformer classification models trained with $\ell_{1.1}$, $\ell_{2}$, and $\ell_{3}$-MD on the Stanford Large Movie Review Dataset. The model sizes refers to the number of layers in the transformer model and the number of attention heads per layer. $\ell_{1.1}$-MD provides superior generalization performance.**
| Algorithm | Model Size 3 | Model Size 4 | Model Size 6 |
| --- | --- | --- | --- |
| $\ell_{1.1}$-MD | 83.47 $\pm$ 0.09% | 83.36 $\pm$ 0.13% | 83.65 $\pm$ 0.13% |
| $\ell_{2}$-MD | $81.66\pm 0.09\%$ | $81.05\pm 0.17\%$ | $82.22\pm 0.13\%$ |
| $\ell_{3}$-MD | $82.57\pm 0.09\%$ | $82.40\pm 0.12\%$ | $81.97\pm 0.10\%$ |

We trained a transformer classification model on the Stanford Large Movie Review Dataset (maas-EtAl:2011:ACL-HLT2011) using MD with $\ell_{1.1}$, $\ell_{2}$, and $\ell_{3}$ potentials. The models are similar to the encoder module in vaswani2017attention, with the last layer being a linear classification layer on the feature representation of the first $[\texttt{CLS}]$ token. Table 1 summarizes the resulting test accuracy of several variants of that model when trained with the three algorithms, which shows that the $\ell_{1.1}$ potential MD outperforms the other MD algorithms, including the one with the $\ell_{2}$ potential, which is equivalent to the GD.

Figure: Figure 8: The attention map generated by the resulting models that were trained using $\ell_{1.1}$–MD and GD for five sample sentences. For three out of five of the sample sentences, the model trained using $\ell_{1.1}$–MD selects the optimal token better than the model trained using GD.
Refer to caption: extracted/5918143/Images/attention_map_cropped.jpg

To investigate this further, we look at how the attention layers of the model select the tokens from simple reviews that GPT-4o generated and investigate how much the attention layer focuses on a particular token that truly determines whether the whole review was a positive one or a negative one. We chose these pivotal tokens using GPT-4o as well. We do this procedure to the model trained using $\ell_{1.1}$–MD and the GD and tabulate the full results in Appendix H (we provide five of the results in Figure 8). We can see that the $\ell_{1.1}$–MD also outperforms the GD in token selection.

Figure: (a)
Refer to caption: extracted/5918143/Images/Parameter_K_1.1.jpg

Finally, we collect the training weights from the resulting models trained by $\ell_{1.1}$–MD and the GD and plot a histogram of their absolute values in Figure 9. Specifically, we take the histogram of the components of the key, query, and value matrices. The figures show that the resulting model that was trained using $\ell_{1.1}$–MD is sparser than the one trained using GD, which could hint at a potential explanation as to why $\ell_{1.1}$–MD can outperform the standard GD algorithm when it is used to train attention-based models.

## 5 Conclusion

We explored the optimization dynamics and generalization performance of a family of MD
algorithms for softmax attention mechanisms, focusing on $\ell_{p}$-AttGD, which generalizes GD by using the $p$-th power of the $\ell_{p}$-norm as the potential function. Our theoretical analysis and experiments show that $\ell_{p}$-AttGD converges to the solution of a generalized hard-margin SVM with an $\ell_{p}$-norm objective in classification tasks using a single-layer softmax attention model. This generalized SVM separates optimal from non-optimal tokens via linear constraints on token pairs. We also examined the joint problem under logistic loss with $\ell_{p}$-norm regularization and proved that $W$ and $v$ generated by $\ell_{p}$-norm regularization path converge to their respective generalized max-margin solutions. Finally, our numerical experiments on real data demonstrate that MD algorithms improve generalization over standard GD and excel in optimal token selection.

#### Acknowledgments

The authors acknowledge the MIT SuperCloud and Lincoln Laboratory Supercomputing Center for providing computing resources that have contributed to the results reported within this paper. This work was supported in part by MathWorks, the MIT-IBM Watson AI Lab, the MIT-Amazon Science Hub, and the MIT-Google Program for Computing Innovation.

## Appendix A Additional Related Work

Transformers Optimization.
Recently, the study of optimization dynamics of attention mechanisms has garnered significant attention (deora2023optimization; huang2023context; tian2023joma; fu2023transformers; li2024mechanics; ataee2024max; tarzanagh2023transformers; vasudeva2024implicit; sheen2024implicit; deng2023unmasking; makkuva2024attention; jeon2024information; zheng2023learn; collins2024context; chen2024provably; li2023theoretical; sheen2024implicit; ildiz2024self; vasudeva2024simplicity; bao2024self; chen2024training; huang2024non; wang2024transformers; zhao2024implicit). We discuss the works most closely related to this paper. Studies such as sahiner2022unraveling; ergen2022convexifying investigate the optimization of attention models through convex relaxations. jelassi2022vision demonstrate that Vision Transformers (ViTs) identify spatial patterns in binary classification via gradient methods. li2023theoretical provide sample complexity bounds and discuss attention sparsity in SGD for ViTs. oymak2023role and deora2023optimization explore optimization dynamics in prompt-attention and multi-head attention models, respectively. tian2023scan; tian2023joma study SGD dynamics and multi-layer transformer training. ataee2024max; tarzanagh2023transformers explored GD’s implicit bias for attention optimization.
vasudeva2024implicit discusses the global directional convergence and convergence rate of GD for attention optimization under specific data conditions. sheen2024implicit notes that gradient flow not only achieves minimal loss but also minimizes the nuclear norm of the key-query weight $W$. thrampoulidis2024implicit; li2024mechanics; zhao2024implicit also studied the optimization dynamics of attention mechanisms and provided the implicit bias of GD for next token prediction. Our work extends these findings and those of ataee2024max; tarzanagh2023transformers, focusing on the implicit bias of the general class of MD algorithms for attention training.

Implicit Bias of First Order Methods. In recent years, significant progress has been made in understanding the implicit bias of gradient descent on separable data, particularly highlighted by the works of soudry2018implicit; ji2018risk. For linear predictors, nacson2019convergence; ji2021characterizing; ji2021fast demonstrated that gradient descent methods rapidly converge to the max-margin predictor. Extending these insights to MLPs, ji2020directional; lyu2019gradient; chizat2020implicit have examined the implicit bias of GD and gradient flow using exponentially-tailed classification losses, and show convergence to the Karush-Kuhn-Tucker (KKT) points of the corresponding max-margin problem, both in finite ji2020directional; lyu2019gradient and infinite width scenarios chizat2020implicit. Further, the implicit bias of GD for training ReLU and Leaky-ReLU networks has been investigated, particularly on orthogonal data phuong2020inductive; frei2022implicit. Additionally, the implicit bias towards rank minimization in regression settings with square loss has been explored in vardi2021implicit; arora2019implicit; li2020towards.

Our work is closely related to the implicit bias of MD (gunasekar2018characterizing; azizan2018stochastic) for regression and classification, respectively. Specifically, sun2022mirror extended the findings of gunasekar2018characterizing; azizan2018stochastic to classification problems, and developed a class of algorithms exhibiting an implicit bias towards a generalized SVM with $\ell_{p}$ norms that effectively separates samples based on their labels; for a survey, we refer to vardi2023implicit.

## Appendix B Auxiliary lemmas

### B.1 Additional Notations

Consider the following constants for the proofs, depending on the dataset $(X_{i},Y_{i},z_{i})_{i=1}^{n}$, the parameter $v$, and the locally optimal token $(\alpha_{i})_{i=1}^{n}$:

$$ $\displaystyle\delta^{\prime}$ $\displaystyle:=\frac{1}{2}\min_{i\in[n]}\min_{\tau\in\bar{\mathcal{T}}_{i}} \left((X_{i\alpha_{i}}-X_{i\tau})^{\top}W^{\alpha}_{\mathrm{mm}}z_{i}-1\right)$ $\displaystyle\leq\frac{1}{2}\min_{i\in[n]}\min_{t\in\mathcal{T}_{i},\tau\in \bar{\mathcal{T}}_{i}}\left((X_{it}-X_{i\tau})^{\top}W^{\alpha}_{\mathrm{mm}}z _{i}\right);$ (6a) $\displaystyle\delta:=\min\{0.25,\delta^{\prime}\}.$ (6b) $$

When $\bar{\mathcal{T}}_{i}=\emptyset$ for all $i\in[n]$ (i.e. globally-optimal indices), we set $\delta^{\prime}=\infty$ as all non-neighbor related terms will disappear. Further, recalling Definition 4 and using $W^{\alpha}_{\mathrm{mm}}$—i.e., the minimizer of ($\ell_{p}$-AttSVM), we set

$$ $\displaystyle A^{\prime}$ $\displaystyle:=\|W_{\mathrm{mm}}^{\alpha}\|_{p,p}\max_{i\in[n],t\in[T]}\|X_{it }z_{i}^{\top}\|_{\frac{p}{p-1},\frac{p}{p-1}};$ $\displaystyle A$ $\displaystyle:=\max\{1,A^{\prime}\}.$ (7) $$

Recalling Definition 5, we provide the following initial radius $\mu=\mu_{0}$ which will be used later in Lemma 9:

$$ $\mu_{0}:=\left\{\begin{aligned} &\frac{1}{p}\left(\frac{\delta}{8A}\right)^{p} &\textnormal{if }p\geq 2,\\ &\frac{1}{p}\left(\frac{\delta(p-1)}{4Ad^{\frac{2}{p}-1}}\right)^{2}& \textnormal{otherwise}.\end{aligned}\right.$ (8) $$

Furthermore, define the following sums for $W$:

$$ $S_{i}(W):=\sum_{t\in\mathcal{T}_{i}}[\sigma(X_{i}Wz_{i})]_{t},\quad\textnormal {and}\quad Q_{i}(W):=\sum_{t\in\bar{\mathcal{T}}_{i}}[\sigma(X_{i}Wz_{i})]_{t}.$ $$

For the samples $i$ with non-empty supports $\mathcal{T}_{i}$, let

$$ $\displaystyle\gamma_{i}^{\text{gap}}:=\gamma_{i\alpha_{i}}-\max_{t\in\mathcal{ T}_{i}}\gamma_{it},\quad\textnormal{and}\quad\bar{\gamma}_{i}^{\text{gap}}:= \gamma_{i\alpha_{i}}-\min_{t\in\mathcal{T}_{i}}\gamma_{it}.$ (9) $$

Furthermore, we define the global score gap as

$$ $\displaystyle\Gamma$ $\displaystyle:=\sup_{i\in[n],t,\tau\in[T]}|\gamma_{it}-\gamma_{i\tau}|.$ (10) $$

### B.2 Lemma for Analyzing The ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT -Norm

In this section of the Appendix, we provide some analysis on comparing the $\ell_{p}$-norm, the $\ell_{p}$ Bregman divergence, and the $\ell_{2}$-norm of matrices. Since the $\ell_{2}$-norm of matrices are much easier to analyze and use, like in the inner product Cauchy-Schwarz inequality, having this comparison is valuable when analyzing the $\ell_{p}$-AttGD.

###### Lemma 1 .

For any $d\times d$ matrix $W$, let $w$ denote its vectorization. Then,

$$ $\displaystyle\|w\|_{p}\in\left[d^{\frac{2}{p}-1}\|w\|_{2},\|w\|_{2}\right]$ $$

for $p\geq 2$, and for $1<p\leq 2$, $\|w\|_{p}$ is in a similar interval, with the two ends switched.

###### Proof.

Let $w_{1},w_{2},...,w_{d^{2}}$ be the entries of $w$. Therefore, for $p\geq 2$,

$$ $\displaystyle\|w\|_{p}$ $\displaystyle=\sqrt[p]{\sum_{i=1}^{d^{2}}|w_{i}|^{p}}$ $\displaystyle=\sqrt[p]{\sum_{i=1}^{d^{2}}(|w_{i}|^{2})^{p/2}},$ $$

and because $\frac{p}{2}\geq 1$, we would have

$$ $\displaystyle\sqrt[p]{\sum_{i=1}^{d^{2}}(|w_{i}|^{2})^{p/2}}$ $\displaystyle\leq\sqrt[p]{\left(\sum_{i=1}^{d^{2}}|w_{i}|^{2}\right)^{p/2}}$ $\displaystyle=\sqrt[p]{\|w\|_{2}^{p}}=\|w\|_{2}.$ $$

Therefore, $\|w\|_{p}\leq\|w\|_{2}$ whenever $p\geq 2$. A similar argument will get us $\|w\|_{p}\geq\|w\|_{2}$ whenever $1<p\leq 2$, so one end of the interval is solved for each case, now for the other end.

Using the power-mean inequality, we can get that whenever $p\geq 2$,

$$ $\sqrt[p]{\frac{1}{d^{2}}\sum_{i=1}^{d^{2}}|w_{i}|^{p}}\geq\sqrt{\frac{1}{d^{2} }\sum_{i=1}^{d^{2}}|w_{i}|^{2}},$ $$

$$ $d^{-\frac{2}{p}}\|w\|_{p}\geq d^{-1}\|w\|_{2},$ $$

$$ $\|w\|_{p}\geq d^{\frac{2}{p}-1}\|w\|_{2}.$ $$

Similarly, for $1<p\leq 2$,

$$ $\|w\|_{p}\leq d^{\frac{2}{p}-1}\|w\|_{2}.$ $$

∎

###### Lemma 2 .

Let $W_{1},W_{2}\in\mathbb{R}^{d\times d}$ be two matrices such that $\|W_{1}\|_{p,p}=\|W_{2}\|_{p,p}=1$. Then, the following inequalities hold:

- L1.
For $p\geq 2$,
$D_{\psi}(W_{1},W_{2})\geq\frac{1}{p\times 2^{p}}\|W_{1}-W_{2}\|_{p,p}^{p},$
- L2.
For $p\in(1,2)$,
$D_{\psi}(W_{1},W_{2})\geq\frac{(p-1)^{2}}{p}\|W_{1}-W_{2}\|_{2,2}^{2}.$

Here, $D_{\psi}(\cdot,\cdot)$ denotes the Bregman divergence given in Definition 1.

###### Proof.

Let $W_{1}=(x_{ij})_{i,j\in[d]}$ and $W_{2}=(y_{ij})_{i,j\in[d]}$, then from Definition 1, we have

$$ $\displaystyle D_{\psi}(W_{1},W_{2})$ $\displaystyle=\frac{1}{p}\sum_{i,j\in[d]}|x_{ij}|^{p}-\frac{1}{p}\sum_{i,j\in[ d]}|y_{ij}|^{p}-\sum_{i,j\in[d]}|y_{ij}|^{p-1}(x_{ij}-y_{ij})\operatorname{ sign}(y_{ij})$ $\displaystyle=\sum_{i,j\in[d]}\left(\frac{1}{p}|x_{ij}|^{p}+\frac{p-1}{p}|y_{ ij}|^{p}-|y_{ij}|^{p-1}|x_{ij}|\operatorname{sign}(x_{ij}y_{ij})\right).$ $$

Therefore, it is enough to prove that whenever $x,y\in[-1,1]$, the expression

$$ $\displaystyle\frac{1}{p}|x|^{p}+\frac{p-1}{p}|y|^{p}-|x\|y|^{p-1}\operatorname {sign}(xy)$ (11) $$

is at least $\frac{1}{p2^{p}}|x-y|^{p}$ if $p\geq 2$, or is at least $\frac{(p-1)^{2}}{p}|x-y|^{2}$ if $p\in(1,2)$. We split the argument into two cases, the first is when the signs of $x$ and $y$ are the same, and the second for when they are not.

Case 1: $\operatorname{sign}(xy)=1$, so both $x$ and $y$ have the same sign, WLOG both are non-negative. Let us fix the value $\Delta\in[-1,1]$ and find the minimum value of (11) when we constraint $x$ and $y$ to be positive and $x-y=\Delta$. Therefore, that expression can be written as

$$ $\frac{(y+\Delta)^{p}+(p-1)y^{p}}{p}-(y+\Delta)y^{p-1},$ $$

the first derivative with respect to $y$ is

$$ $\displaystyle(y+\Delta)^{p-1}+(p-1)y^{p-1}$ $\displaystyle-y^{p-1}-(p-1)(y+\Delta)y^{p-2}$ $\displaystyle=(y+\Delta)^{p-1}-y^{p-1}-(p-1)\Delta y^{p-2}.$ $$

Since the function $t\mapsto t^{p-1}$ is convex for $p\geq 2$, and concave for $p\in(1,2)$, then that derivative is always non-negative when $p\geq 2$ and always negative when $p\in(1,2)$.

Sub-Case 1.1: $p\geq 2$. In this subcase, (11) reaches its minimum when $(x,y)=(\Delta,0)$ or $(0,-\Delta)$, depending on the sign of $\Delta$, plugging them in gets us the minimum, which is $\frac{1}{p}|\Delta|^{p}$ when $\Delta\geq 0$ or $\frac{p-1}{p}|\Delta|^{p}$ otherwise.

Sub-Case 1.2: $p\in(1,2)$. In this subcase, (11) reaches its minimum when $(x,y)=(1,1-\Delta)$ if $\Delta$ is non-negative or $(1+\Delta,1)$ otherwise. When $\Delta$ is non-negative, the desired minimum is

$$ $\displaystyle\frac{1+(p-1)(1-\Delta)^{p}}{p}-(1-\Delta)^{p-1}$ $\displaystyle=\frac{1}{p}(1-(1-\Delta)^{p-1}-(p-1)\Delta(1-\Delta)^{p-1})$ $\displaystyle\geq\frac{1}{p}((p-1)\Delta-(p-1)\Delta(1-\Delta)^{p-1})$ $\displaystyle=\frac{(p-1)\Delta}{p}(1-(1-\Delta)^{p-1})\geq\frac{(p-1)^{2}}{p} \Delta^{2}.$ $$

Combining the results from the subcases, we get that the expression in (11) is lower-bounded by $\frac{1}{p}|x-y|^{p}$ when $p\geq 2$, or $\frac{(p-1)^{2}}{p}|x-y|^{2}$ otherwise, which sufficiently satisfies the desired bounds for case 1.

Case 2: $\operatorname{sign}(xy)=-1$, so $x$ and $y$ has opposite sign. The expression in (11) can be simplified to

$$ $\frac{1}{p}|x|^{p}+\frac{p-1}{p}|y|^{p}+|x||y|^{p-1},$ $$

and we want to prove that it is at least $\frac{1}{p2^{p}}(|x|+|y|)^{p}$ when $p\geq 2$, or is at least $\frac{(p-1)^{2}}{p}(|x|+|y|)^{2}$ when $p\in(1,2)$. In the case that $p\geq 2$, one of $|x|$ or $|y|$ is at least $\frac{|x|+|y|}{2}$, so the above is at least $\frac{1}{p}\left(\frac{|x|+|y|}{2}\right)^{p}=\frac{1}{p2^{p}}(|x|+|y|)^{p}$. Otherwise,

$$ $\displaystyle\frac{1}{p}|x|^{p}+\frac{p-1}{p}|y|^{p}+|x|y|^{p-1}$ $\displaystyle=\frac{|x|(|x|^{p-1}+|y|^{p-1})+(p-1)|y|^{p-1}(|x|+|y|)}{p}$ $\displaystyle\geq\frac{(|x|+|y|)(|x|+(p-1)|y|^{p-1})}{p}$ $\displaystyle\geq\frac{(|x|+|y|)((p-1)|x|+(p-1)|y|)}{p}$ $\displaystyle=\frac{p-1}{p}(|x|+|y|)^{2}\geq\frac{(p-1)^{2}}{p}(|x|+|y|)^{2}.$ $$

Therefore, we have proven the bound for this case.
∎

###### Lemma 3 .

For any $x\geq y\geq 0$, we we have

$$ $\frac{p-1}{p}x^{p}-\frac{p-1}{p}y^{p}\geq y(x^{p-1}-y^{p-1}).$ $$

###### Proof.

$$ $\frac{d}{dx}\left(\frac{p-1}{p}x^{p}-\frac{p-1}{p}y^{p}\right)=(p-1)x^{p-1},$ $$

$$ $\frac{d}{dx}y(x^{p-1}-y^{p-1})=(p-1)x^{p-2}y\leq(p-1)x^{p-1},$ $$

so as we increase $x$, the left side grows faster than the right side, so we simply need to prove that the inequality holds at $x=y$, which is trivially true.
∎

###### Lemma 4 .

For any $x\geq y\geq 0$, we we have that if $q\geq 1$

$$ $x^{q}-y^{q}\leq qx^{q-1}(x-y),$ $$

and if $0<q<1$,

$$ $x^{q}-y^{q}\leq qy^{q-1}(x-y)$ $$

###### Proof.

$$ $\frac{d}{dx}(x^{q}-y^{q})=qx^{q-1},$ $$

$$ $\frac{d}{dx}qx^{q-1}(x-y)=q(q-1)x^{q-2}(x-y)+qx^{q-1},\text{ and}$ $$

$$ $\frac{d}{dx}qy^{q-1}(x-y)=qy^{q-1}.$ $$

When $q\geq 1$,

$$ $\frac{d}{dx}(x^{q}-y^{q})\geq\frac{d}{dx}qx^{q-1}(x-y),$ $$

so because we have

$$ $x^{q}-y^{q}=qx^{q-1}(x-y)=0$ $$

when $x=y$, then

$$ $x^{q}-y^{q}\geq qx^{q-1}(x-y)$ $$

when $x\geq y\geq 0$ if $q\geq 1$. We can use a similar argument for the $0<q<1$ case.
∎

###### Lemma 1 .

###### Proof.

###### Lemma 2 .

###### Proof.

###### Lemma 3 .

###### Proof.

###### Lemma 4 .

###### Proof.

### B.3 Lemma for Analyzing ERM Objective and Its Gradient

In this section of the Appendix, we analyze the objective function. We especially want to know about its gradient and the inner product of this gradient with the matrices of the cone set, as was mentioned before in the main body of the paper. The first one bounds the loss objective,

###### Lemma 5 .

Under Assumption A, $\mathcal{L}(W)$ is bounded from above by $\mathcal{L}_{max}$ and below by $\mathcal{L}_{min}$ for some dataset-dependent constants $\mathcal{L}_{max}$ and $\mathcal{L}_{min}$ that are finite.

###### Proof.

It is enough to show the same thing for each of the loss contributions of each sample, $l_{i}(y_{i}v^{\top}X_{i}^{\top}\sigma(X_{i}Wz_{i}))$. By Assumption A, we simply need to show that $y_{i}v^{\top}X_{i}^{\top}\sigma(X_{i}Wz_{i})$ is bounded by dataset-dependent bounds. However, $W$ only affects the softmax, so the above expression is bounded above by $\max_{t\in[T]}\gamma_{it}$ and bounded below by $\min_{t\in[T]}\gamma_{it}$, which are dataset dependent.
∎

###### Lemma 6 .

If we denote $h_{i}:=X_{i}Wz_{i}$ and $l_{i}^{\prime}:=l^{\prime}(\gamma_{i}^{\top}\sigma(h_{i}))$, then

$$ $\nabla\mathcal{L}(W)=\frac{1}{n}\sum_{i=1}^{n}l_{i}^{\prime}X_{i}^{\top}( \operatorname{diag}(\sigma(h_{i}))-\sigma(h_{i})\sigma(h_{i})^{\top})\gamma_{i }z_{i}^{\top},$ $$

where $\mathcal{L}(W)$ denotes the objective of (ERM) with fixed $v$.

###### Proof.

We first calculate the derivatives of each term in the sum of $\mathcal{L}(W)$. The derivative of the $i$-th term for the $W_{j_{1}j_{2}}$ component is

$$ $\displaystyle\frac{\partial}{\partial W_{j_{1}j_{2}}}l(y_{i}v^{\top}X_{i}^{ \top}\sigma(X_{i}Wz_{i}))$ $\displaystyle=l_{i}^{\prime}\gamma_{i}^{\top}\frac{\partial}{\partial W_{j_{1} j_{2}}}\sigma(X_{i}Wz_{i})$ $\displaystyle=l_{i}^{\prime}\gamma_{i}^{\top}\nabla\sigma(h_{i})X_{i,:,j_{1}}^ {\top}z_{ij_{2}}$ $\displaystyle=l_{i}^{\prime}X_{i,:,j_{1}}\nabla\sigma(h_{i})^{\top}\gamma_{i}z _{ij_{2}}.$ $$

Therefore, the derivative for the $j_{2}$-th row of $W$ is

$$ $l_{i}^{\prime}X_{i}^{\top}\nabla\sigma(h_{i})^{\top}\gamma_{i}z_{ij_{2}}.$ $$

Next, the full gradient for the $i$-th term equals

$$ $l_{i}^{\prime}X_{i}^{\top}\nabla\sigma(h_{i})^{\top}\gamma_{i}z_{i}^{\top}.$ $$

To finish the proof, we calculate the derivative of $\sigma(h_{i})$. The derivative of the $j_{1}$-th component of $\sigma(h_{i})$ with respect to $h_{ij_{2}}$ is

$$ $\displaystyle\frac{\partial}{\partial h_{ij_{2}}}\left(\frac{e^{h_{ij_{1}}}}{ \sum_{l=1}^{T}e^{h_{il}}}\right)$ $\displaystyle=\frac{e^{h_{ij_{1}}}1_{j_{1}=j_{2}}}{\sum_{l=1}^{T}e^{h_{il}}}- \frac{e^{h_{ij_{1}}}e^{h_{ij_{2}}}}{\left(\sum_{l=1}^{T}e^{h_{il}}\right)^{2}}$ $\displaystyle=\sigma(h_{i})_{j_{1}}1_{j_{1}=j_{2}}-\sigma(h_{i})_{j_{1}}\sigma (h_{i})_{j_{2}}.$ $$

Thus, the derivative of $\sigma(h_{i})$ is a matrix in $\mathbb{R}^{T\times T}$ defined as

$$ $\operatorname{diag}(\sigma(h_{i}))-\sigma(h_{i})\sigma(h_{i})^{\top}.$ $$

Therefore, the full gradient is

$$ $\frac{1}{n}\sum_{i=1}^{n}l_{i}^{\prime}X_{i}^{\top}(\operatorname{diag}(\sigma (h_{i}))-\sigma(h_{i})\sigma(h_{i})^{\top})\gamma_{i}z_{i}^{\top}.$ $$

∎

###### Lemma 7 .

Under Assumption A, $\|\nabla\mathcal{L}(W)\|_{p,p}$ is bounded by a dataset-dependent constant $L$.

###### Proof.

Using the expression in Lemma 6, since $l^{\prime}$ is bounded and the entries in $\sigma(h_{i})$ is always between 0 0 and $1$, then the entries of $\nabla\mathcal{L}(W)$ is bounded by a dataset-dependent bounded, which directly implies this lemma statement.
∎

In the following lemma, we analyze the behaviors of the ($\ell_{p}$-AttSVM) constraint $(X_{it}-X_{i\tau})^{\top}Wz_{i}$ for all $W\in S_{p,\mu_{0}}(W_{\mathrm{mm}}^{\alpha})$ satisfying $\|W\|_{p,p}=\|W_{\mathrm{mm}}^{\alpha}\|_{p,p}$, the result of which is a generalization of (tarzanagh2023transformers, Equation 64) for a general $\ell_{p}$ norm.

###### Lemma 8 .

Let $\alpha=(\alpha_{i})_{i=1}^{n}$ be locally optimal tokens as per Definition 3, and let $W^{\alpha}_{\mathrm{mm}}$ be the ($\ell_{p}$-AttSVM) solution. Let $(\mathcal{T}_{i})_{i=1}^{n}$ be the index set of all support tokens per Definition 3. Let $\bar{\mathcal{T}}_{i}=[T]-\mathcal{T}_{i}-\{\alpha_{i}\}$. For any $W\in S_{p,\mu_{0}}(W_{\mathrm{mm}}^{\alpha})$ with $\mu_{0}$ defined in (8) and $\|W\|_{p,p}=\|W_{\mathrm{mm}}^{\alpha}\|_{p,p}$, we have

$$ $\displaystyle(X_{it}-X_{i\tau})^{\top}Wz_{i}$ $\displaystyle\geq\frac{3}{2}\delta>0,$ (12a) $\displaystyle(X_{i\alpha_{i}}-X_{i\tau})^{\top}Wz_{i}$ $\displaystyle\geq 1+\frac{3}{2}\delta,$ (12b) $\displaystyle 1+\frac{1}{2}\delta\geq(X_{i\alpha_{i}}-X_{it})^{\top}Wz_{i}$ $\displaystyle\geq 1-\frac{1}{2}\delta,$ (12c) $$

for all $t\in\mathcal{T}_{i}$ and $\tau\in\bar{\mathcal{T}}_{i}$

###### Proof.

Let

$$ $\bar{W}:=\frac{W}{\|W\|_{p,p}}\quad\textnormal{and}\quad\bar{W}_{\mathrm{mm}}^ {\alpha}:=\frac{W_{\mathrm{mm}}^{\alpha}}{\|W_{\mathrm{mm}}^{\alpha}\|_{p,p}}.$ $$

Using Lemma 2 and the definition of $S_{p,\mu_{0}}(W_{\mathrm{mm}}^{\alpha})$ in (4a), when $p\geq 2$,

$$ $\displaystyle\|\bar{W}-\bar{W}_{\mathrm{mm}}^{\alpha}\|_{p,p}^{p}$ $\displaystyle\leq 2^{p}pD_{\psi}(\bar{W}_{\mathrm{mm}}^{\alpha},\bar{W})$ $\displaystyle\leq 2^{p}p\mu_{0}$ $\displaystyle=\left(\frac{\delta}{4A}\right)^{p},$ $$

which implies that

$$ $\|\bar{W}-\bar{W}_{\mathrm{mm}}^{\alpha}\|_{p,p}\leq\frac{\delta}{4A}.$ $$

When $p\in(1,2)$, we can also use Lemmas 1 and 2 to obtain

$$ $\displaystyle\|\bar{W}-\bar{W}_{\mathrm{mm}}^{\alpha}\|_{p,p}$ $\displaystyle\leq d^{\frac{2}{p}-1}\|\bar{W}-\bar{W}_{\mathrm{mm}}^{\alpha}\|_ {2,2}$ $\displaystyle\leq d^{\frac{2}{p}-1}\frac{\sqrt{p}}{p-1}\sqrt{D_{\psi}(\bar{W}_ {\mathrm{mm}}^{\alpha},\bar{W})}$ $\displaystyle\leq d^{\frac{2}{p}-1}\frac{\sqrt{p}}{p-1}\sqrt{\mu_{0}}=\frac{ \delta}{4A},$ $$

where the last inequality uses the definition of $S_{p,\mu_{0}}(W_{\mathrm{mm}}^{\alpha})$ in (4a).

Therefore, either way, we have

$$ $\|W-W_{\mathrm{mm}}^{\alpha}\|_{p,p}\leq\frac{\delta}{4A}\|W_{\mathrm{mm}}^{ \alpha}\|_{p,p}.$ $$

We will proceed to show a bound on $(X_{it_{1}}-X_{it_{2}})^{\top}(W-W_{\mathrm{mm}}^{\alpha})z_{i}$ for any $i\in[n]$ and any token indices $t_{1},t_{2}\in[T]$. To do that, let us focus on the term $X_{it_{1}}^{\top}(W-W_{\mathrm{mm}}^{\alpha})z_{i}$ first,

$$ $\displaystyle\left|X_{it_{1}}^{\top}(W-W_{\mathrm{mm}}^{\alpha})z_{i}\right|$ $\displaystyle=\left|\langle W-W_{\mathrm{mm}}^{\alpha},X_{it_{1}}z_{i}^{\top} \rangle\right|$ $\displaystyle\leq\|W-W_{\mathrm{mm}}^{\alpha}\|_{p,p}\cdot\|X_{it_{1}}z_{i}^{ \top}\|_{\frac{p}{p-1},\frac{p}{p-1}}$ $\displaystyle\leq\frac{\delta}{4A}\|W_{\mathrm{mm}}^{\alpha}\|_{p,p}\cdot\|X_{ it_{1}}z_{i}^{\top}\|_{\frac{p}{p-1},\frac{p}{p-1}}$ $\displaystyle\leq\frac{\delta}{4A}\cdot A$ $\displaystyle=\frac{\delta}{4}.$ $$

The first inequality above uses Hölder’s Inequality. We now have

$$ $\left|(X_{it_{1}}-X_{it_{2}})^{\top}(W-W_{\mathrm{mm}}^{\alpha})z_{i}\right| \leq\frac{1}{2}\delta.$ $$

To obtain the first inequality of the lemma in (12a), for all $t\in\mathcal{T}_{i}$ and $\tau\in\bar{\mathcal{T}}_{i}$, we have

$$ $\displaystyle(X_{it}-X_{i\tau})^{\top}Wz_{i}$ $\displaystyle\geq(X_{it}-X_{i\tau})^{\top}W_{\mathrm{mm}}^{\alpha}z_{i}+(X_{it }-X_{i\tau})^{\top}(W-W_{\mathrm{mm}}^{\alpha})z_{i}$ $\displaystyle\geq 2\delta^{\prime}-\frac{1}{2}\delta\geq\frac{3}{2}\delta.$ $$

To get the second inequality in (12b), for all $\tau\in\bar{\mathcal{T}}_{i}$, we have

$$ $\displaystyle(X_{i\alpha_{i}}-X_{i\tau})^{\top}Wz_{i}$ $\displaystyle\geq(X_{i\alpha_{i}}-X_{i\tau})^{\top}W_{\mathrm{mm}}^{\alpha}z_{ i}+(X_{i\alpha_{i}}-X_{i\tau})^{\top}(W-W_{\mathrm{mm}}^{\alpha})z_{i}$ $\displaystyle\geq 1+2\delta^{\prime}-\frac{1}{2}\delta\geq 1+\frac{3}{2}\delta.$ $$

Finally, to get the last inequality in (12c), for all $t\in\mathcal{T}_{i}$, we have

$$ $\displaystyle\left|(X_{i\alpha_{i}}-X_{it})^{\top}Wz_{i}-1\right|$ $\displaystyle=\left|(X_{i\alpha_{i}}-X_{it})^{\top}W_{\mathrm{mm}}^{\alpha}z_{ i}+(X_{i\alpha_{i}}-X_{it})^{\top}(W-W_{\mathrm{mm}}^{\alpha})z_{i}-1\right|$ $\displaystyle=|(X_{i\alpha_{i}}-X_{it})^{\top}(W-W_{\mathrm{mm}}^{\alpha})z_{i }|\leq\frac{1}{2}\delta,$ $$

which implies that

$$ $1+\frac{1}{2}\delta\geq(X_{i\alpha_{i}}-X_{it})^{\top}Wz_{i}\geq 1-\frac{1}{2}\delta.$ $$

∎

The following two lemmas aim at bounding the correlation between the gradient and the attention matrix parameter, each of which is a generalization of (tarzanagh2023transformers, Lemmas 13 and 14) for the generalized $\ell_{p}$ norm.

###### Lemma 9 .

Suppose Assumption A holds. Let $\alpha=(\alpha_{i})_{i=1}^{n}$ be locally optimal tokens as per Definition 3, and let $W^{\alpha}_{\mathrm{mm}}$ be the solution to ($\ell_{p}$-AttSVM). There exists a dataset-dependent constant
$R_{\delta}=\mathcal{O}(1/\delta)$ such that for all $W,V\in C_{p,\mu_{0},R_{\delta}}(W^{\alpha}_{\mathrm{mm}})$ with $\|V\|_{p,p}=\|W^{\alpha}_{\mathrm{mm}}\|_{p,p}$, $\delta$ and $\mu_{0}$ defined in (6) and (8), respectively,

$$ $\displaystyle-\langle\nabla\mathcal{L}(W),V\rangle=\Omega\left(\textnormal{exp }\left(-\frac{\|W\|_{p,p}}{\|W^{\alpha}_{\mathrm{mm}}\|_{p,p}}\left(1+\frac{1} {2}\delta\right)\right)\right)>0.$ $$

###### Proof.

Let

$$ $\displaystyle h_{i}:=X_{i}Wz_{i},~{}~{}\tilde{h}_{i}:=X_{i}Vz_{i},~{}~{}l_{i}^ {\prime}:=l^{\prime}(\gamma_{i}^{\top}\sigma(h_{i})),~{}~{}\textnormal{and}~{} ~{}s_{i}=\sigma(h_{i}).$ $$

Therefore,

$$ $\displaystyle\langle\nabla\mathcal{L}(W),V\rangle$ $\displaystyle=\frac{1}{n}\sum_{i=1}^{n}l_{i}^{\prime}\langle X_{i}^{\top}( \operatorname{diag}(s_{i})-s_{i}s_{i}^{\top})\gamma_{i}z_{i}^{\top},V\rangle$ $\displaystyle=\frac{1}{n}\sum_{i=1}^{n}l_{i}^{\prime}\langle(\operatorname{ diag}(s_{i})-s_{i}s_{i}^{\top})\gamma_{i},X_{i}Vz_{i}\rangle$ $\displaystyle=\frac{1}{n}\sum_{i=1}^{n}l_{i}^{\prime}\langle(\operatorname{ diag}(s_{i})-s_{i}s_{i}^{\top})\gamma_{i},\tilde{h}_{i}\rangle$ $\displaystyle=\frac{1}{n}\sum_{i=1}^{n}l_{i}^{\prime}\tilde{h}_{i}^{\top}( \operatorname{diag}(s_{i})-s_{i}s_{i}^{\top})\gamma_{i},$ $$

$$ $\displaystyle-\langle\nabla\mathcal{L}(W),V\rangle=\frac{1}{n}\sum_{i=1}^{n}(- l_{i}^{\prime})\tilde{h}_{i}^{\top}(\operatorname{diag}(s_{i})-s_{i}s_{i}^{ \top})\gamma_{i}.$ (13) $$

The value $\gamma_{i}^{\top}\sigma(h_{i})$ for any $i\in[n]$ must be bounded, and the bound is only dataset-dependent, so by Assumption A, $l_{i}^{\prime}$ is bounded for any $i\in[n]$ by some bound that is dataset-dependent. Furthermore, because $l$ is decreasing, $-l^{\prime}$ is always non-negative, so an easier approach is to lower-bound the following for each $i\in[n]$,

$$ $\tilde{h}_{i}^{\top}s_{i}s_{i}^{\top}\gamma_{i}-\tilde{h}_{i}^{\top} \operatorname{diag}(s_{i})\gamma_{i}.$ $$

Next, we can get for all $i\in[n]$ and $t\in[T]$ that

$$ $\displaystyle\tilde{h}_{it}=X_{it}^{\top}Vz_{i}$ $\displaystyle=\langle X_{it}z_{i}^{\top},V\rangle$ $\displaystyle\leq\|V\|_{p,p}\|X_{it}z_{i}^{\top}\|_{\frac{p}{p-1}}$ $\displaystyle\leq A,$ $$

where $A$ is defined in (B.1).

Therefore, if we drop the $i$ notation and let $\alpha_{i}=1$, and use (tarzanagh2023transformers, Lemma 7),

$$ $\left|\tilde{h}^{\top}ss^{\top}\gamma-\tilde{h}^{\top}\operatorname{diag}(s) \gamma-\sum_{t=2}^{T}(\tilde{h}_{1}-\tilde{h}_{t})s_{t}(\gamma_{1}-\gamma_{t}) \right|\leq 2\Gamma A(1-s_{1})^{2}.$ $$

Let us attempt to remove the non-support tokens from the sum above by bounding the sum of the term for the non-supports,

$$ $\left|\sum_{t\in\bar{\mathcal{T}}}(\tilde{h}_{1}-\tilde{h}_{t})s_{t}(\gamma_{1 }-\gamma_{t})\right|\leq 2\max_{t\in[T]}\{|\tilde{h}_{t}|\}Q(W)\Gamma\leq 2AQ( W)\Gamma.$ $$

Therefore,

$$ $\left|\tilde{h}^{\top}ss^{\top}\gamma-\tilde{h}^{\top}\operatorname{diag}(s) \gamma-\sum_{t\in\mathcal{T}}(\tilde{h}_{1}-\tilde{h}_{t})s_{t}(\gamma_{1}- \gamma_{t})\right|\leq 2\Gamma A((1-s_{1})^{2}+Q(W)),$ $$

which implies that

$$ $\tilde{h}^{\top}ss^{\top}\gamma-\tilde{h}^{\top}\operatorname{diag}(s)\gamma \geq\sum_{t\in\mathcal{T}}(\tilde{h}_{1}-\tilde{h}_{t})s_{t}(\gamma_{1}-\gamma _{t})-2\Gamma A((1-s_{1})^{2}+Q(W)).$ $$

Using Lemma 8, we have

$$ $\displaystyle\tilde{h}^{\top}ss^{\top}\gamma-\tilde{h}^{\top}\operatorname{ diag}(s)\gamma\geq\left(1-\frac{1}{2}\delta\right)\sum_{t\in\mathcal{T}}s_{t}( \gamma_{1}-\gamma_{t})-2\Gamma A((1-s_{1})^{2}+Q(W)).$ (14) $$

To proceed, we can upper-bound $1-s_{1}$ and $Q(W)$. For bounding $1-s_{1}$, let $\tau>1$ be some index that maximizes $X_{\tau}^{\top}Wz$, so

$$ $\displaystyle 1-s_{1}=\frac{\sum_{t=2}^{T}e^{X_{t}^{\top}Wz}}{\sum_{t=1}^{T}e^ {X_{t}^{\top}Wz}}$ $\displaystyle\leq\frac{(T-1)e^{X_{\tau}^{\top}Wz}}{(T-1)e^{X_{\tau}^{\top}Wz}+ e^{X_{1}^{\top}Wz}}$ $\displaystyle\leq\frac{T}{T+e^{(X_{1}-X_{\tau})^{\top}Wz}}$ $\displaystyle\leq\frac{T}{T+e^{\frac{\|W\|_{p,p}}{\|W_{\mathrm{mm}}^{\alpha}\| _{p,p}}(1-\frac{1}{2}\delta)}}$ $\displaystyle\leq\frac{T}{e^{\frac{\|W\|_{p,p}}{\|W_{\mathrm{mm}}^{\alpha}\|_{ p,p}}(1-\frac{1}{2}\delta)}},$ $$

with the last inequality using the third inequality Lemma 8.

For ease of notation, denote

$$ $\displaystyle R^{\prime}:=\frac{\|W\|_{p,p}}{\|W_{\mathrm{mm}}^{\alpha}\|_{p,p }}.$ (15) $$

To upper bound $Q(W)$, we use a method similar to that for $1-s_{1}$, but we utilize the second inequality of Lemma 8 instead of the first. This gives:

$$ $Q(W)\leq\frac{T}{T+e^{(1+\frac{3}{2}\delta)R^{\prime}}}\leq\frac{T}{e^{(1+ \frac{3}{2}\delta)R^{\prime}}}.$ $$

Therefore, we have

$$ $\displaystyle 2\Gamma A((1-s_{1})^{2}+Q(W))$ $\displaystyle\leq 2\Gamma A\left(\frac{T^{2}}{e^{(2-\delta)R^{\prime}}}+\frac{ T}{e^{(1+\frac{3}{2}\delta)R^{\prime}}}\right)$ $\displaystyle\leq\frac{2\Gamma AT(T+1)}{e^{(1+\frac{3}{2}\delta)R^{\prime}}}.$ (16) $$

Now it is time to lower-bound the sum on the right side of Equation (14). When the set of supports is empty, that sum is zero. However, if it is not empty,

$$ $\sum_{t\in\mathcal{T}}s_{t}(\gamma_{1}-\gamma_{t})\geq S(W)\gamma^{\textnormal {gap}}.$ $$

If we let $\tau\in\mathcal{T}$ be the support index that minimizes $X_{\tau}^{\top}Wz$, then

$$ $\displaystyle S(W)=\frac{\sum_{t\in\mathcal{T}}e^{X_{t}^{\top}Wz}}{\sum_{t=1}^ {T}e^{X_{t}^{\top}Wz}}\geq\frac{e^{X_{\tau}^{\top}Wz}}{Te^{X_{1}^{\top}Wz}}$ $\displaystyle=\frac{1}{Te^{(X_{1}-X_{\tau})^{\top}Wz}}$ $\displaystyle\geq\frac{1}{Te^{(1+\frac{1}{2}\delta)R^{\prime}}},$ $$

with the last inequality coming from the third inequality of Lemma 8.

Therefore,

$$ $\sum_{t\in\mathcal{T}}s_{t}(\gamma_{1}-\gamma_{t})\geq\frac{\gamma^{ \textnormal{gap}}}{Te^{(1+\frac{1}{2}\delta)R^{\prime}}}>0.$ $$

Using Equation (14), we get that if the support index set is empty,

$$ $\tilde{h}^{\top}ss^{\top}\gamma-\tilde{h}^{\top}\operatorname{diag}(s)\gamma \geq-\frac{2\Gamma AT(T+1)}{e^{(1+\frac{3}{2}\delta)R^{\prime}}},$ $$

otherwise,

$$ $\tilde{h}^{\top}ss^{\top}\gamma-\tilde{h}^{\top}\operatorname{diag}(s)\gamma \geq\frac{\gamma^{\textnormal{gap}}}{Te^{(1+\frac{1}{2}\delta)R^{\prime}}} \left(1-\frac{1}{2}\delta\right)-\frac{2\Gamma AT(T+1)}{e^{(1+\frac{3}{2} \delta)R^{\prime}}}.$ $$

Plugging everything back into Equation (13), and considering that some samples will have non-empty support index sets, we have:

$$ $\displaystyle-\langle\mathcal{L}(W),V\rangle$ $\displaystyle\geq-\frac{\min_{i\in\mathcal{T}_{i}}\{\gamma^{\text{gap}}_{i}\}} {nTe^{(1+\frac{1}{2}\delta)R^{\prime}}}\left(1-\frac{1}{2}\delta\right)\max_{i =1}^{n}\{l_{i}^{\prime}\}$ $\displaystyle\quad+\frac{2\Gamma AT(T+1)}{e^{(1+\frac{3}{2}\delta)R^{\prime}}} \sum_{i=1}^{n}l_{i}^{\prime}=\Omega\left(e^{-(1+\frac{1}{2}\delta)R^{\prime}} \right).$ (17) $$

Let

$$ $\displaystyle\bar{L}:=\frac{\sum_{i=1}^{n}l_{i}^{\prime}}{\max_{i=1}^{n}\{l_{i }^{\prime}\}}.$ (18) $$

Note that using Assumption A, $\bar{L}$ is positive. Hence, using (18) and (B.3), the term $-\langle\mathcal{L}(W),V\rangle$ is positive when

$$ $R^{\prime}\geq\frac{1}{\delta}\log\left(\frac{2\Gamma\bar{L}AT^{2}(T+1)n}{\min _{i\in\mathcal{T}_{i}}\{\gamma^{\text{gap}}_{i}\}\left(1-\frac{1}{2}\delta \right)}\right),$ $$

or equivalently, from (15), we have

$$ $\|W\|_{p,p}\geq\frac{\|W^{\alpha}_{\mathrm{mm}}\|_{p,p}}{\delta}\log\left( \frac{2\Gamma\bar{L}AT^{2}(T+1)}{\min_{i\in\mathcal{T}_{i}}\{\gamma^{\text{gap }}_{i}\}\left(1-\frac{1}{2}\delta\right)}\right).$ $$

∎

Finally, we introduce the following lemma to help understand the correlation between the gradient of the objective and the parameter.

###### Lemma 10 .

Suppose Assumption A holds. Let $\alpha=(\alpha_{i})_{i=1}^{n}$ be locally optimal tokens as per Definition 3, let $W^{\alpha}_{\mathrm{mm}}$ be the ($\ell_{p}$-AttSVM) solution, and let $R_{\delta}$ be the constant from Lemma 9. For any choice of $\pi\in(0,1)$, there exists $R_{\pi}$ that depends on $\pi$ defined as

$$ $R_{\pi}:=\max\left\{R_{\delta},\mathcal{O}\left(\frac{1}{\pi\delta}\log\frac{ \delta}{\pi}\right)\right\},$ $$

such that for all $W\in C_{p,\mu_{0},R_{\pi}}(W_{\mathrm{mm}}^{\alpha})$,

$$ $\left\langle\nabla\mathcal{L}(W),\frac{W}{\|W\|_{p,p}}\right\rangle\geq(1+\pi) \left\langle\nabla\mathcal{L}(W),\frac{W_{\mathrm{mm}}^{\alpha}}{\|W_{\mathrm{ mm}}^{\alpha}\|_{p,p}}\right\rangle.$ $$

###### Proof.

Let

$$ $\begin{split}h_{i}&:=X_{i}Wz_{i},\quad\tilde{h}_{i}:=X_{i}W_{\mathrm{mm}}^{ \alpha}z_{i},\quad l_{i}^{\prime}:=l^{\prime}(\gamma_{i}^{\top}\sigma(h_{i})), \\ s_{i}&:=\sigma(h_{i}),\quad\bar{W}:=\frac{\|W_{\mathrm{mm}}^{\alpha}\|_{p,p}W} {\|W\|_{p,p}},\quad\text{and}\quad\bar{h}_{i}:=X_{i}\bar{W}z_{i}.\end{split}$ (19) $$

By decomposing $\mathcal{L}(W)$ into its sum and using Lemma 6, the main inequality is equivalent to the following,

$$ $\displaystyle\sum_{i=1}^{n}(-l_{i}^{\prime})\langle X_{i}^{\top}(\operatorname {diag}(s_{i})$ $\displaystyle-s_{i}s_{i}^{\top})\gamma_{i}z_{i}^{\top},\bar{W}\rangle$ $\displaystyle\leq(1+\pi)\sum_{i=1}^{n}(-l_{i}^{\prime})\langle X_{i}^{\top}( \operatorname{diag}(s_{i})-s_{i}s_{i}^{\top})\gamma_{i}z_{i}^{\top},W_{\mathrm {mm}}^{\alpha}\rangle,$ $$

which implies that

$$ $\displaystyle\sum_{i=1}^{n}(-l_{i}^{\prime})\langle(\operatorname{diag}(s_{i})$ $\displaystyle-s_{i}s_{i}^{\top})\gamma_{i},X_{i}\bar{W}z_{i}\rangle$ $\displaystyle\leq(1+\pi)\sum_{i=1}^{n}(-l_{i}^{\prime})\langle(\operatorname{ diag}(s_{i})-s_{i}s_{i}^{\top})\gamma_{i},X_{i}W_{\mathrm{mm}}^{\alpha}z_{i}\rangle.$ $$

Using (19), we get

$$ $\displaystyle\sum_{i=1}^{n}(-l_{i}^{\prime})\langle(\operatorname{diag}(s_{i}) -s_{i}s_{i}^{\top})\gamma_{i},\bar{h}_{i}\rangle$ $\displaystyle\leq(1+\pi)\sum_{i=1}^{n}(-l_{i}^{\prime})\langle(\operatorname{ diag}(s_{i})-s_{i}s_{i}^{\top})\gamma_{i},\tilde{h}_{i}\rangle,$ $$

which gives

$$ $\displaystyle\sum_{i=1}^{n}(-l_{i}^{\prime})\bar{h}_{i}^{\top}(\operatorname{ diag}(s_{i})-s_{i}s_{i}^{\top})\gamma_{i}$ $\displaystyle\leq(1+\pi)\sum_{i=1}^{n}(-l_{i}^{\prime})\tilde{h}_{i}^{\top}( \operatorname{diag}(s_{i})-s_{i}s_{i}^{\top})\gamma_{i}.$ $$

Hence,

$$ $\displaystyle\sum_{i=1}^{n}(-l_{i}^{\prime})$ $\displaystyle\left[(1+\pi)\left(\tilde{h}_{i}^{\top}\operatorname{diag}(s_{i}) \gamma_{i}-\tilde{h}_{i}^{\top}s_{i}s_{i}^{\top}\gamma_{i}\right)-\left(\bar{h }_{i}^{\top}\operatorname{diag}(s_{i})\gamma_{i}-\bar{h}_{i}^{\top}s_{i}s_{i}^ {\top}\gamma_{i}\right)\right]\geq 0.$ $$

Using a similar technique as the one we used to prove Lemma 9,

$$ $\displaystyle\Big{|}\tilde{h}_{i}^{\top}\operatorname{diag}(s_{i})\gamma_{i}- \tilde{h}_{i}^{\top}s_{i}s_{i}^{\top}\gamma_{i}$ $\displaystyle-\sum_{t\in\mathcal{T}_{i}}(\tilde{h}_{i\alpha_{i}}-\tilde{h}_{it })s_{it}(\gamma_{i\alpha_{i}}-\gamma_{it})\Big{|}$ $\displaystyle\leq 2\Gamma A((1-s_{i\alpha_{i}})^{2}+Q_{i}(W)).$ $$

Similarly,

$$ $\displaystyle\Big{|}\bar{h}_{i}^{\top}\operatorname{diag}(s_{i})\gamma_{i}- \bar{h}_{i}^{\top}s_{i}s_{i}^{\top}\gamma_{i}$ $\displaystyle-\sum_{t\in\mathcal{T}_{i}}(\bar{h}_{i\alpha_{i}}-\bar{h}_{it})s_ {it}(\gamma_{i\alpha_{i}}-\gamma_{it})\Big{|}$ $\displaystyle\leq 2\Gamma A((1-s_{i\alpha_{i}})^{2}+Q_{i}(W)).$ $$

Therefore, it is enough to prove that

$$ $\displaystyle\begin{split}\sum_{i=1}^{n}(-l_{i}^{\prime})&\left((1+\pi)\left( \sum_{t\in\mathcal{T}_{i}}(\tilde{h}_{i\alpha_{i}}-\tilde{h}_{it})s_{it}( \gamma_{i\alpha_{i}}-\gamma_{it})-2\Gamma A((1-s_{i\alpha_{i}})^{2}+Q_{i}(W)) \right)\right.\\ &\left.-\left(\sum_{t\in\mathcal{T}_{i}}(\bar{h}_{i\alpha_{i}}-\bar{h}_{it})s_ {it}(\gamma_{i\alpha_{i}}-\gamma_{it})+2\Gamma A((1-s_{i\alpha_{i}})^{2}+Q_{i} (W))\right)\right),\end{split}$ (20) $$

Using the fact that $\pi<1$ and using Equation (B.3), we get another lower-bound

$$ $\displaystyle\sum_{i=1}^{n}\sum_{t\in\mathcal{T}_{i}}(-l_{i}^{\prime})(1+\pi-( \bar{h}_{i\alpha_{i}}-\bar{h}_{it}))s_{it}(\gamma_{i\alpha_{i}}-\gamma_{it})+ \frac{6\Gamma AT(T+1)}{e^{(1+\frac{3}{2}\delta)R^{\prime}}}\sum_{i=1}^{n}l_{i} ^{\prime},$ (21) $$

with $R^{\prime}$ again being $\frac{\|W\|_{p,p}}{\|W_{\mathrm{mm}}^{\alpha}\|_{p,p}}$. Next, we analyze the softmax probability $s_{it}$, and lower and upper-bound them in terms of $R^{\prime}$ and $\bar{h}_{i\alpha_{i}}-\bar{h}_{it}$. For the lower-bound,

$$ $\displaystyle s_{it}=\frac{e^{\bar{h}_{it}R^{\prime}}}{\sum_{\tau\in[T]}e^{ \bar{h}_{i\tau}R^{\prime}}}$ $\displaystyle\geq\frac{e^{\bar{h}_{it}R^{\prime}}}{Te^{\bar{h}_{i\alpha_{i}}R^ {\prime}}}$ $\displaystyle=\frac{1}{T}e^{-(\bar{h}_{i\alpha_{i}}-\bar{h}_{it})R^{\prime}}.$ $$

For the upper-bound,

$$ $\displaystyle s_{it}=\frac{e^{\bar{h}_{it}R^{\prime}}}{\sum_{\tau\in[T]}e^{ \bar{h}_{i\tau}R^{\prime}}}$ $\displaystyle\leq\frac{e^{\bar{h}_{it}R^{\prime}}}{e^{\bar{h}_{i\alpha_{i}}R^{ \prime}}}$ $\displaystyle=e^{-(\bar{h}_{i\alpha_{i}}-\bar{h}_{it})R^{\prime}}.$ $$

In both bounds, the main inequality derivation stems from the fact that $\bar{h}_{i\alpha_{i}}>\bar{h}_{i\tau}$ for all $\tau\in[T]$, which we obtain from Lemma 8. Now, we analyze the left double-summation in Equation (21). To analyze the sum, let $\mathcal{I}$ be the subset of $[n]\times[T]$ that contains all $(i,t)$ such that $t\in\mathcal{T}_{i}$. Furthermore, let

$$ $\displaystyle\mathcal{I}_{1}$ $\displaystyle:=\left\{(i,t)\in\mathcal{I}\mid\bar{h}_{i\alpha_{i}}-\bar{h}_{it }\leq 1\right\},$ $\displaystyle\mathcal{I}_{2}$ $\displaystyle:=\left\{(i,t)\in\mathcal{I}\mid 1<\bar{h}_{i\alpha_{i}}-\bar{h}_ {it}\leq 1+\pi\right\},$ $\displaystyle\mathcal{I}_{3}$ $\displaystyle:=\left\{(i,t)\in\mathcal{I}\mid\bar{h}_{i\alpha_{i}}-\bar{h}_{it }>1+\pi\right\}.$ $$

Therefore, we can split the sum above into the sum over $\mathcal{I}_{1}$, $\mathcal{I}_{2}$, and $\mathcal{I}_{3}$. The set $\mathcal{I}_{1}$ in particular must be non-empty because $\|\bar{W}\|_{p,p}=\|W_{\mathrm{mm}}^{\alpha}\|_{p,p}$, meaning that one of the constraints in the $\ell_{p}$-AttSVM problem must either be fulfilled exactly or violated.

The sum over $\mathcal{I}_{1}$ must be positive and is at least

$$ $-\frac{\pi}{T}\min_{i\in\mathcal{I}_{1}}\{\gamma_{i}^{gap}\}e^{-R^{\prime}} \max_{i=1}^{n}\{l_{i}^{\prime}\}.$ $$

The sum over $\mathcal{I}_{2}$ must be non-negative, and the sum over $\mathcal{I}_{3}$ is negative can be bounded from below using Lemma 8

$$ $\frac{1}{2}\delta\max_{i\in\mathcal{I}_{3}}\{\bar{\gamma}^{gap}_{i}\}Te^{-(1+ \pi)R^{\prime}}\sum_{i=1}^{n}l_{i}^{\prime}.$ $$

Putting things together into Equation (21), we get that we want the following to be non-negative

$$ $\displaystyle-\frac{\pi}{T}\min_{i\in\mathcal{I}_{1}}\{\gamma_{i}^{gap}\}e^{-R ^{\prime}}\max_{i=1}^{n}\{l_{i}^{\prime}\}$ $\displaystyle+\frac{1}{2}\delta\max_{i\in\mathcal{I}_{3}}\{\bar{\gamma}^{gap}_ {i}\}Te^{-(1+\pi)R^{\prime}}\sum_{i=1}^{n}l_{i}^{\prime}$ $\displaystyle+6\Gamma AT(T+1)e^{-(1+\frac{3}{2}\delta)R^{\prime}}\sum_{i=1}^{n }l_{i}^{\prime}.$ $$

This can be achieved when

$$ $R^{\prime}\geq\frac{1}{\min\{\pi,\frac{3}{2}\delta\}}\log\left(\frac{\frac{1}{ 2}\delta\max_{i\in\mathcal{I}_{3}}\{\bar{\gamma}^{gap}_{i}\}T^{2}+6\Gamma AT^{ 2}(T+1)}{\pi\min_{i\in\mathcal{I}_{1}}\{\gamma_{i}^{gap}\}\max_{i=1}^{n}\{l_{i }^{\prime}\}}\sum_{i=1}^{n}l_{i}^{\prime}\right),$ $$

or equivalently,

$$ $\|W\|_{p,p}\geq\frac{\|W^{\alpha}_{\mathrm{mm}}\|_{p,p}}{\min\{\pi,\frac{3}{2} \delta\}}\log\left(\frac{\frac{1}{2}\delta\max_{i\in\mathcal{I}_{3}}\{\bar{ \gamma}^{gap}_{i}\}T^{2}+6\Gamma AT^{2}(T+1)}{\pi\min_{i\in\mathcal{I}_{1}}\{ \gamma_{i}^{gap}\}\max_{i=1}^{n}\{l_{i}^{\prime}\}}\sum_{i=1}^{n}l_{i}^{\prime }\right),$ $$

which means that such dataset dependent $R_{\pi}$ exists.
∎

###### Lemma 5 .

###### Proof.

###### Lemma 6 .

###### Proof.

###### Lemma 7 .

###### Proof.

###### Lemma 8 .

###### Proof.

###### Lemma 9 .

###### Proof.

###### Lemma 10 .

###### Proof.

### B.4 Lemma for Analyzing ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT -AttGD

We introduce the lemmas for analyzing $\ell_{p}$-AttGD. The first we prove is Lemma 11, which describes the lower bound of the $W$ parameter at every iterate.

###### Lemma 11 .

Suppose Assumption A holds. For the sequence $\{W(k)\}_{k\geq 0}$ generated by $\ell_{p}$-AttGD, we have

$$ $\|W(k+1)\|_{p,p}^{p-1}\geq\|W(k)\|_{p,p}^{p-1}+\frac{\eta}{\|W(k)\|_{p,p}} \langle-\nabla\mathcal{L}(W(k)),W(k)\rangle.$ $$

###### Proof.

With $\psi(W)=\frac{1}{p}\|W\|_{p,p}$, the derivative $\nabla\psi(\cdot)$ is computed as follows:

$$ $\nabla\psi(W)=(\operatorname{sign}(W_{ij})|W_{ij}|^{p-1})_{1\leq i,j\leq d}.$ $$

Thus, we have

$$ $\langle\nabla\psi(W),W\rangle=\sum_{i,j}\operatorname{sign}(W_{ij})|W_{ij}|^{p -1}W_{ij}=\|W\|_{p,p}^{p}.$ $$

Using this fact, we take the inner product of both sides of (23) with $W(k)$:

$$ $\langle\nabla\psi(W(k+1)),W(k)\rangle=\langle\nabla\psi(W(k)),W(k)\rangle+\eta \langle-\nabla\mathcal{L}(W(k)),W(k)\rangle,$ $$

$$ $\displaystyle\langle\nabla\psi(W(k+1)),W(k)\rangle=\|W(k)\|_{p,p}^{p}+\eta \langle-\nabla\mathcal{L}(W(k)),W(k)\rangle.$ (22) $$

The left side of the above equation is upper-bounded by

$$ $\sum_{i,j}\operatorname{sign}(W_{ij}(k+1))|W_{ij}(k+1)|^{p-1}W_{ij}(k)\leq\sum _{i,j}|W_{ij}(k+1)|^{p-1}|W_{ij}(k)|.$ $$

Using Hölder’s inequality:

$$ $\displaystyle\sum_{i,j}|W_{ij}(k+1)|^{p-1}|W_{ij}(k)|$ $\displaystyle\leq\Big{(}\sum_{i,j}(|W_{ij}(k+1)|^{p-1})^{\frac{p}{p-1}}\Big{)} ^{\frac{p-1}{p}}\Big{(}\sum_{i,j}|W_{ij}(k)|^{p}\Big{)}^{\frac{1}{p}}$ $\displaystyle=\|W(k+1)\|_{p,p}^{p-1}\|W(k)\|_{p,p}.$ $$

Combining this result with (22), we get:

$$ $\|W(k+1)\|_{p,p}^{p-1}\geq\|W(k)\|_{p,p}^{p-1}+\frac{\eta}{\|W(k)\|_{p,p}} \langle-\nabla\mathcal{L}(W(k)),W(k)\rangle.$ $$

∎

Next, we show several tools for analyzing the algorithm further and for analyzing the Bregman divergence. The following three specifically are from sun2022mirror and azizan2018stochastic, and so the proofs are omitted.

###### Lemma 12 .

For any $W\in\mathbb{R}^{d\times d}$, the following identities hold for MD:

$$ $\displaystyle D_{\psi}(W,W(k))$ $\displaystyle=D_{\psi}(W,W(k+1))+D_{\psi-\eta\mathcal{L}}(W(k+1),W(k))$ $\displaystyle\quad-\eta\langle\nabla\mathcal{L}(W(k)),W-W(k)\rangle-\eta \mathcal{L}(W(k))+\eta\mathcal{L}(W(k+1)).$ (23) $$

###### Lemma 13 .

Suppose Assumptions A hold and $\eta$ is small enough. For the sequence $\{W(k)\}_{k\geq 0}$ generated by $\ell_{p}$-AttGD, we have

$$ $\displaystyle\frac{p-1}{p}\|W(k+1)\|_{p,p}^{p}-\frac{p-1}{p}\|W(k)\|_{p,p}^{p}$ $\displaystyle+\eta\mathcal{L}(W(k+1))-\eta\mathcal{L}(W(k))$ $\displaystyle\leq\langle-\eta\nabla\mathcal{L}(W(k)),W(k)\rangle.$ (24) $$

###### Lemma 14 .

Suppose Assumptions A hold. Consider the sequence $W(k)$ generated by Algorithm $\ell_{p}$-AttGD. Given that the step size $\eta$ is sufficiently small, then the ERM objective $\mathcal{L}(W(k))$ is decreasing in $k$.

This following is a well-known lemma, so the proof is omitted.

###### Lemma 15 (Bregman Divergences Cosine Law) .

For any $w,w^{\prime},w^{\prime\prime}$ that are all vectors or matrices with the same dimensionalities, we have

$$ $D_{\psi}(w,w^{\prime})=D_{\psi}(w,w^{\prime\prime})+D_{\psi}(w^{\prime\prime}, w^{\prime})-\langle\nabla\psi(w^{\prime})-\nabla\psi(w^{\prime\prime}),w-w^{ \prime\prime}\rangle.$ $$

The following is adapted from sun2022mirror for the case of our attention model. Our proof is quite similar, except that we use our version of the gradient correlation lemma.

###### Lemma 16 .

Suppose Assumptions A hold. Consider the sequence $W(k)$ generated by Algorithm $\ell_{p}$-AttGD. For any $\pi\in(0,1)$, if $W(k)\in C_{p,\mu_{0},R_{\pi}}(W^{\alpha}_{\mathrm{mm}})$, with $R_{\pi}$ being the constant from Lemma 10, then for a small enough step size $\eta$,

$$ $\displaystyle\begin{split}\langle\nabla\psi(W(k+1))-\nabla\psi(W(k)),\bar{W}_{ \mathrm{mm}}^{\alpha}\rangle&\geq\frac{1}{1+\pi}(\|W(k+1)\|_{p,p}^{p-1}-\|W(k) \|_{p,p}^{p-1})\\ &+\frac{\eta}{\|W(k)\|_{p,p}}(\mathcal{L}(W(k+1))-\mathcal{L}(W(k))).\end{split}$ (25) $$

###### Proof.

Let $\bar{W}^{\alpha}_{\mathrm{mm}}=\frac{W^{\alpha}_{\mathrm{mm}}}{\|W^{\alpha}_{
\mathrm{mm}}\|_{p,p}}$. Using the $\ell_{p}$-AttGD algorithm equation,

$$ $\langle\nabla\psi(W(k+1))-\nabla\psi(W(k)),\bar{W}_{\mathrm{mm}}^{\alpha} \rangle=\langle-\eta\nabla\mathcal{L}(W(k)),\bar{W}^{\alpha}_{\mathrm{mm}}\rangle.$ $$

Then, using Lemma 10, we get that

$$ $\langle-\eta\nabla\mathcal{L}(W(k)),\bar{W}^{\alpha}_{\mathrm{mm}}\rangle\geq \frac{1}{(1+\pi)\|W(k)\|_{p,p}}\langle-\eta\nabla\mathcal{L}(W(k)),W(k)\rangle,$ $$

and using Lemma 13, we get that this is lower-bounded by

$$ $\frac{p-1}{p(1+\pi)\|W(k)\|_{p,p}}(\|W(k+1)\|_{p,p}^{p}-\|W(k)\|_{p,p}^{p})+ \frac{\eta}{(1+\pi)\|W(k)\|_{p,p}}(\mathcal{L}(W(k+1))-\mathcal{L}(W(k))).$ $$

By Lemma 9, $\langle-\eta\nabla\mathcal{L}(W(k)),W(k)\rangle>0$, so by Lemma 11, $\|W(k+1)\|_{p,p}\geq\|W(k)\|_{p,p}$. Therefore, we can use Lemma 3 to get that the above is lower-bounded by

$$ $\frac{1}{1+\pi}(\|W(k+1)\|_{p,p}^{p-1}-\|W(k)\|_{p,p}^{p-1})+\frac{\eta}{(1+ \pi)\|W(k)\|_{p,p}}(\mathcal{L}(W(k+1))-\mathcal{L}(W(k))).$ $$

From Lemma 14, we get that we can lower-bound the above further using the right hand side of (25).
∎

With all these lemmas in hand, we provide the following Lemma 17.

###### Lemma 17 .

Suppose Assumptions A holds and that the step size $\eta$ is sufficiently small.
For any $\mu\in(0,\mu_{0}]$ and any locally optimal tokens $(\alpha_{i})_{i=1}^{n}$ as per Definition 3, there exists constants $R_{\mu}$ and $\mu^{\prime}\in(0,\mu]$ that depends on the dataset and $\mu$ such that if $C_{1}$ is the wider cone $C_{p,\mu,R_{\mu}}(W^{\alpha}_{\mathrm{mm}})$ and $C_{2}$ is the thinner cone $C_{p,\mu^{\prime},R_{\mu}}(W^{\alpha}_{\mathrm{mm}})$, then if $W(0)\in C_{2}$, then $W(k)\in C_{1}$ for all positive indices $k$.

###### Proof.

Let $\pi$ be some positive real number that we determine later, and let $R_{\pi}$ be as described in Lemma 10.

For the proof, we use induction with the assumption that $W(k)\in C_{p,\mu,R_{\pi}}(W_{\mathrm{mm}}^{\alpha})$ for all $k=0,\ldots,K-1$. We aim to find the correct $\mu^{\prime}$ and $R_{\mu}$ such that $W(K)\in C_{p,\mu,R_{\pi}}(W_{\mathrm{mm}}^{\alpha})$.

Denote $\bar{W}(k):=\frac{W(k)}{\|W(k)\|_{p,p}}$, so

$$ $\displaystyle D_{\psi}(\bar{W}_{\mathrm{mm}}^{\alpha},\bar{W}(k))$ $\displaystyle=\frac{1}{p}\|\bar{W}_{\mathrm{mm}}^{\alpha}\|_{p,p}-\frac{1}{p} \|\bar{W}(k)\|_{p,p}-\langle\nabla\psi(\bar{W}(k)),\bar{W}_{\mathrm{mm}}^{ \alpha}-\bar{W}(k)\rangle$ $\displaystyle=1-\langle\nabla\psi(\bar{W}(k)),\bar{W}_{\mathrm{mm}}^{\alpha}\rangle.$ $$

So now, let us analyze the term $\langle\nabla\psi(\bar{W}(K)),\bar{W}_{\mathrm{mm}}^{\alpha}\rangle$ using the inductive hypothesis on $k=0,1,...,K-1$. Lemma 16 tells us that

$$ $\displaystyle\begin{split}\langle\nabla\psi(W(k+1))-\nabla\psi(W(k)),\bar{W}_{ \mathrm{mm}}^{\alpha}\rangle&\geq\frac{\|W(k+1)\|_{p,p}^{p-1}-\|W(k)\|_{p,p}^{ p-1}}{(1+\pi)}\\ &+\frac{\eta}{\|W(k)\|_{p,p}}(\mathcal{L}(W(k+1))-\mathcal{L}(W(k))).\end{split}$ (26) $$

Since this is true for all $k=0,1,...,K-1$, and since $\|W(k)\|_{p,p}$ is increasing in $k$, we can sum all the above inequalities and get the following,

$$ $\displaystyle\langle\nabla\psi(W(K))-\nabla\psi(W(0)),\bar{W}_{\mathrm{mm}}^{ \alpha}\rangle$ $\displaystyle\geq\frac{\|W(K)\|_{p,p}^{p-1}-\|W(0)\|_{p,p}^{p-1}}{(1+\pi)}$ $\displaystyle+\frac{\eta}{\|W(0)\|_{p,p}}(\mathcal{L}(W(K))-\mathcal{L}(W(0))).$ $$

Rearranging this, we get

$$ $\displaystyle\|W(K)\|_{p,p}^{p-1}-\langle\nabla\psi(W(K)),\bar{W}_{\mathrm{mm} }^{\alpha}\rangle$ $\displaystyle\leq\|W(0)\|_{p,p}^{p-1}-\langle\nabla\psi(W(0)),\bar{W}_{\mathrm {mm}}^{\alpha}\rangle$ $\displaystyle+\frac{\pi}{1+\pi}(\|W(K)\|_{p,p}^{p-1}-\|W(0)\|_{p,p}^{p-1})$ $\displaystyle+\frac{\eta}{\|W(0)\|_{p,p}}(\mathcal{L}(W(0))-\mathcal{L}(W(K))).$ $$

Dividing by $\|W(K)\|_{p,p}^{p-1}$, we get

$$ $\displaystyle\begin{split}D_{\psi}(\bar{W}_{\mathrm{mm}}^{\alpha},\bar{W}(K))& \leq\frac{\|W(0)\|_{p,p}^{p-1}}{\|W(K)\|_{p,p}^{p-1}}D_{\psi}(\bar{W}_{\mathrm {mm}}^{\alpha},\bar{W}(0))+\frac{\pi}{1+\pi}\left(1-\frac{\|W(0)\|_{p,p}^{p-1} }{\|W(K)\|_{p,p}^{p-1}}\right)\\ &+\frac{\eta}{\|W(K)\|_{p,p}^{p-1}\|W(0)\|_{p,p}}(\mathcal{L}(W(0))-\mathcal{L }(W(K)))\\ &\leq\mu^{\prime}+\pi+\frac{\eta(\mathcal{L}(W(0))-\mathcal{L}(W(K)))}{R_{\mu} ^{p}}.\end{split}$ (27) $$

Therefore, we can simply choose $\mu^{\prime}=\frac{1}{3}\mu$, $\pi$ be any real number below $\frac{1}{3}\mu$, and have $R_{\mu}$ big enough so that $\frac{\eta(\mathcal{L}(W(0))-\mathcal{L}(W(K)))}{R_{\mu}^{p}}\leq\frac{1}{3}\mu$ and $R_{\mu}\geq R_{\pi}$, such $R_{\mu}$ exists because $\mathcal{L}$ is bounded.
∎

###### Lemma 11 .

###### Proof.

###### Lemma 12 .

###### Lemma 13 .

###### Lemma 14 .

###### Lemma 15 (Bregman Divergences Cosine Law) .

###### Lemma 16 .

###### Proof.

###### Lemma 17 .

###### Proof.

### B.5 Lemma for Analyzing Rate of Convergence

###### Lemma 18 .

Suppose Assumptions A holds.
Let $R_{\delta}$ be from Lemma 9, let $c$ be from Lemma 16, let $\mu^{\prime}$ and $R_{\mu}$ be from Lemma 17 when $\mu=\mu_{0}$, and let $R:=\max\{R_{\mu},R_{\delta},e^{1/c}\}$. For any $W\in\mathbb{R}^{d\times d}$, denote $\bar{W}:=W/\|W\|_{p,p}$. If the initialization $W(0)$ is in $C_{p,\mu^{\prime},R}(W_{\mathrm{mm}}^{\alpha})$, then for a sufficiently small step size $\eta$, the sequence $\{W(k)\}_{k\geq 0}$ generated by $\ell_{p}$-AttGD satisfies

$$ $D_{\psi}(\bar{W}^{\alpha}_{\mathrm{mm}},\bar{W}(k))=\left\{\begin{aligned} & \mathcal{O}\left(\frac{\log\|W(k)\|_{p,p}}{\|W(k)\|_{p,p}}\right)&\textnormal{ if }p>2,\\ &\mathcal{O}\left(\frac{(\log\|W(k)\|_{p,p})^{2}}{\|W(k)\|_{p,p}}\right)& \textnormal{if }p=2,\\ &\mathcal{O}\left(\frac{1}{\|W(k)\|_{p,p}^{p-1}}\right)&\textnormal{otherwise} .\end{aligned}\right.$ (28) $$

###### Proof.

Using Lemma 10, setting $c$ as the dataset dependent constant hidden by the $\mathcal{O}$ notation for $R_{\pi}$, we can get that by setting $\pi=\min\{\frac{c\log\|W(k)\|_{p,p}}{\delta\|W(k)\|_{p,p}},1\}$, we can use the result of Lemma 16 on $k$, so rearranging that result, we get

$$ $\displaystyle\|W(k+1)\|_{p,p}^{p-1}-\langle\nabla\psi(W(k+1)),\bar{W}_{\mathrm {mm}}^{\alpha}\rangle$ $\displaystyle\leq\|W(k)\|_{p,p}^{p-1}-\langle\nabla\psi(W(k)),\bar{W}_{\mathrm {mm}}^{\alpha}\rangle$ $\displaystyle+\frac{\pi}{1+\pi}(\|W(k+1)\|_{p,p}^{p-1}-\|W(k)\|_{p,p}^{p-1})$ $\displaystyle+\frac{\eta}{\|W(k)\|_{p,p}}(\mathcal{L}(W(k))-\mathcal{L}(W(k+1) )).$ $$

From Lemma 9 and Lemma 11, $\|W(k)\|_{p,p}$ is increasing, so focusing on the second line, we can use Lemma 4 and get

$$ $\displaystyle\frac{\pi}{1+\pi}(\|W(k+1)\|_{p,p}^{p-1}-\|W(k)\|_{p,p}^{p-1})$ $\displaystyle\leq\pi(\|W(k+1)\|_{p,p}^{p-1}-\|W(k)\|_{p,p}^{p-1})$ $\displaystyle\leq\frac{cp}{\delta\|W(k)\|_{p,p}}\max\{\|W(k)\|_{p,p}^{p-2},\|W (k+1)\|_{p,p}^{p-2}\}$ $\displaystyle\times\log\|W(k)\|_{p,p}$ $\displaystyle\times(\|W(k+1)\|_{p,p}-\|W(k)\|_{p,p}).$ $$

From Lemma 7, we know that for all index $k$,

$$ $\displaystyle\|W(k+1)\|_{p,p}\leq\|W(k)\|_{p,p}+\eta L,$ (29) $$

so we can use integral approximation when bounding the sums of $\Delta(k)$’s. Let

$$ $\displaystyle\Delta(k)$ $\displaystyle=\frac{cp}{\delta\|W(k)\|_{p,p}}\max\{\|W(k)\|_{p,p}^{p-2},\|W(k+ 1)\|_{p,p}^{p-2}\}\log\|W(k)\|_{p,p}$ $\displaystyle\times(\|W(k+1)\|_{p,p}-\|W(k)\|_{p,p}),$ $$

so we can get that

$$ $\displaystyle\begin{split}\|W(K)\|_{p,p}^{p-1}-\langle\nabla\psi(W(K)),\bar{W} ^{\alpha}_{\mathrm{mm}}\rangle&\leq\|W(0)\|_{p,p}^{p-1}-\langle\nabla\psi(W(0) ),\bar{W}^{\alpha}_{\mathrm{mm}}\rangle\\ &+\sum_{k=0}^{k-1}\Delta(k)+\frac{\eta}{c}(\mathcal{L}(W(0))-\mathcal{L}(W(K)) ),\end{split}$ $$

$$ $\displaystyle\begin{split}\|W(K)\|_{p,p}^{p-1}D_{\psi}(\bar{W}^{\alpha}_{ \mathrm{mm}},\bar{W}(K))&\leq\|W(0)\|_{p,p}^{p-1}D_{\psi}(\bar{W}^{\alpha}_{ \mathrm{mm}},\bar{W}(0))\\ &+\sum_{k=0}^{k-1}\Delta(k)+\frac{\eta}{c}(\mathcal{L}(W(0))-\mathcal{L}(W(K)) ).\end{split}$ (30) $$

When $p>2$, we have

$$ $\displaystyle\Delta(k)=\frac{cp}{\delta\|W(k)\|_{p,p}}(\|W(k)\|_{p,p}+\eta L)^ {p-2}\log\|W(k)\|_{p,p}(\|W(k+1)\|_{p,p}-\|W(k)\|_{p,p}).$ $$

We can see that

$$ $\frac{d}{dx}(x+\eta L)^{p-2}(\log x-\log c)>\frac{p-2}{x}(x+\eta L)^{p-2}\log x$ $$

for all $x>0$, so from Equation (29), we can get that

$$ $\sum_{k=0}^{K-1}\Delta(k)=O(\|W(K)\|^{p-2}\log\|W(K)\|_{p,p}).$ $$

When $p=2$, we have

$$ $\displaystyle\Delta(k)=\frac{cp}{\|W(k)\|_{p,p}}\log\|W(k)\|_{p,p}(\|W(k+1)\|_ {p,p}-\|W(k)\|_{p,p}).$ $$

We can see that

$$ $\frac{d}{dx}(\log x)^{2}>\frac{2}{x}(\log x)$ $$

for all $x\geq c$, so from Equation (29), we can get that

$$ $\sum_{k=0}^{K-1}\Delta(k)=O((\log\|W(K)\|_{p,p})^{2}).$ $$

When $p<2$, we have

$$ $\displaystyle\Delta(k)=cp\|W(k)\|_{p,p}^{p-3}\log\|W(k)\|_{p,p}(\|W(k+1)\|_{p, p}-\|W(k)\|_{p,p}).$ $$

From Equation (29), we can get that

$$ $\sum_{k=0}^{K-1}\Delta(k)=O(1).$ $$

Combining the above cases with Equation (30), we get that

$$ $\|W(K)\|_{p,p}^{p-1}D_{\psi}(\bar{W}^{\alpha}_{\mathrm{mm}},\bar{W}(K))=\begin {cases}O(\|W(K)\|_{p,p}^{p-2}\log\|W(K)\|_{p,p})&\text{if }p>2,\\ O((\log\|W(K)\|_{p,p})^{2})&\text{ if }p=2,\\ O(1)&\text{otherwise}\end{cases},$ $$

Dividing both sides by $\|W(K)\|_{p,p}^{p-1}$ gives (28).
∎

###### Lemma 19 .

Suppose Assumptions A holds. Let $\mu^{\prime}$ be that from Lemma 17 if $\mu=\mu_{0}$, and let $R$ the maximum of the $R_{\mu}$ from 17 and $R_{\delta}$ 9. Let $\{W(k)\}_{k\geq 0}$ be the sequence generated by $\ell_{p}$-AttGD. If the initialization $W(0)$ is in $C_{p,\mu^{\prime},R}(W_{\mathrm{mm}}^{\alpha})$, then with a small enough step size $\eta$, we have the following for each $k\geq 0$,

$$ $\displaystyle\|W(k)\|_{p,p}=\Omega(\log k).$ $$

###### Proof.

For each $k\geq 0$, Lemma 11 gives

$$ $\|W(k+1)\|_{p,p}^{p-1}\geq\|W(k)\|_{p,p}^{p-1}+\frac{\eta}{\|W(k)\|_{p,p}} \langle-\nabla\mathcal{L}(W(k)),W(k)\rangle.$ $$

Lemma 17 gives us that $W(k)\in C_{p,\mu,R}(W_{\mathrm{mm}}^{\alpha})$ for each $k\geq 0$, so by Lemma 9,

$$ $\frac{\eta}{\|W(k)\|_{p,p}}\langle-\nabla\mathcal{L}(W(k)),W(k)\rangle=\Omega \left(e^{-\frac{\|W(k)\|_{p,p}}{\|W_{\mathrm{mm}}^{\alpha}\|_{p,p}}(1+\frac{1} {2}\delta)}\right),$ $$

so there exists dataset dependent constants $R_{1},R_{2}>0$ such that

$$ $\frac{\eta}{\|W(k)\|_{p,p}}\langle-\nabla\mathcal{L}(W(k)),W(k)\rangle\geq R_{ 1}e^{-R_{2}\|W(k)\|_{p,p}},$ $$

so for each $k\geq 0$,

$$ $\|W(k+1)\|_{p,p}^{p-1}\geq\|W(k)\|_{p,p}^{p-1}+R_{1}e^{-R_{2}\|W(k)\|_{p,p}}.$ $$

Set $k_{0}=0$, and let $k_{i+1}$ be the lowest indices such that $\|W(k_{i+1})\|_{p,p}\geq\|W(k_{i})\|_{p,p}+1$ for all index $i\geq 0$. Therefore,

$$ $k_{i+1}-k_{i}\leq\frac{(\|W(k_{i})\|_{p,p}+1)^{p-1}-\|W(k_{i})\|_{p,p}^{p-1}}{ R_{1}e^{-R_{2}(\|W(k_{i})\|_{p,p}+1)}}=e^{O(\|W(k_{i})\|_{p,p})}.$ $$

Therefore,

$$ $\|W(k)\|_{p,p}=\Omega(\log k).$ $$

∎

###### Lemma 18 .

###### Proof.

###### Lemma 19 .

###### Proof.

## Appendix C Proof of Theorem 1

###### Proof.

The proof is similar to the proof of (ataee2024max, Theorem 1). Specifically, we need to show that $f(X)=v^{\top}X^{\top}\sigma(XW)$ satisfies the assumptions of (ataee2024max, Lemma 14), where the nonlinear head is replaced by the linear term $v$. This holds independently of the choice of algorithm or the attention SVM solution. Thus, we omit the details and refer to the proof of (ataee2024max, Theorem 1).
∎

###### Proof.

## Appendix D Proof of Theorem 2

###### Proof.

It is enough to show the existence of such constants $\mu,R>0$ such that if $W(0)$ is in $C_{p,\mu,R}(W^{\alpha}_{\mathrm{mm}})$, then the norm diverges to infinity. As discussed in Lemma 11, for any timestep $k$,

$$ $\displaystyle\|W(k+1)\|_{p}^{p-1}\geq\|W(k)\|_{p}^{p-1}-\frac{\eta}{\|W(k)\|_{ p}}\langle\nabla\mathcal{L}(W(k)),W(k)\rangle.$ (31) $$

Let $R_{1}$ be the $R$ from Lemma 9, set $\mu$ and $R_{2}$ to be the $\mu^{\prime}$ and $R$ for $\mu=\mu_{0}$ of Lemma 17, and set $R:=\max\{R_{1},R_{2}\}$. From Lemma 17, we know that $W(k)\in C_{p,\mu_{0},R}(W^{\mathrm{mm}}_{\alpha})$ for any timestep $k$, so from Lemma 9,

$$ $\langle\nabla\mathcal{L}(W(k)),W(k)\rangle<0,$ $$

for all timesteps $k$.

Therefore, the $l_{p}$-norm is always increasing, but this does not immediately imply that the $l_{p}$-norm will approach infinity; it could converge to a finite value. However, if $\|W(k)\|_{p}$ converges to a finite value, then again by Lemma 9, we get a lower bound for $-\frac{\eta}{\|W(k)\|_{p}}\langle\nabla\mathcal{L}(W(k)),W(k)\rangle$ at any timestep $k$. Therefore, by Equation (31),

$$ $\lim_{k\rightarrow\infty}\|W(k)\|_{p}^{p-1}=\infty,$ $$

a contradiction, so $\|W(k)\|_{p}$ converges to infinity.
∎

###### Proof.

## Appendix E Proof of Theorem 3

###### Proof.

This is a direct consequence of Theorem 4.
∎

###### Proof.

## Appendix F Proof of Theorem 4

###### Proof.

Let $R$ be the one from Lemma 18. Given $W(0)\in C_{p,\mu,R}(W^{\alpha}_{\mathrm{mm}})$, by Lemma 18, we have

$$ $D_{\psi}(\bar{W}^{\alpha}_{\mathrm{mm}},\bar{W}(k))=\left\{\begin{aligned} & \mathcal{O}\left(\frac{\log\|W(k)\|_{p,p}}{\|W(k)\|_{p,p}}\right)&\text{if }p> 2,\\ &\mathcal{O}\left(\frac{(\log\|W(k)\|_{p,p})^{2}}{\|W(k)\|_{p,p}}\right)&\text {if }p=2,\\ &\mathcal{O}\left(\frac{1}{\|W(k)\|_{p,p}^{p-1}}\right)&\text{otherwise}.\end{ aligned}\right.$ $$

From Lemma 19, we know that

$$ $\|W(k)\|_{p,p}=\Omega(\log k).$ $$

The derivative $\frac{d}{dx}\left(\frac{\log x}{x}\right)=\frac{1-\log x}{x^{2}}$ is negative when $x>e$, so $\frac{\log x}{x}$ is decreasing when $x>e$. Similarly, $\frac{(\log x)^{2}}{x}$ is decreasing when $x>e^{2}$.

Thus when $p>2$, for a large enough $k$,

$$ $\displaystyle D_{\psi}(\bar{W}^{\alpha}_{\mathrm{mm}},\bar{W}(k))=O\left(\frac {\log\log k}{\log k}\right).$ (32a) Similarly, when $p=2$ , for a large enough $k$ , $\displaystyle D_{\psi}(\bar{W}^{\alpha}_{\mathrm{mm}},\bar{W}(k))=O\left(\frac {(\log\log k)^{2}}{\log k}\right).$ (32b) Finally, when $1<p<2$ , $\displaystyle D_{\psi}(\bar{W}^{\alpha}_{\mathrm{mm}},\bar{W}(k))=O\left(\frac {1}{(\log k)^{p-1}}\right).$ (32c) $$

∎

###### Proof.

## Appendix G On the Convergence of the ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT Regularization Path for Joint W 𝑊 W italic_W and v 𝑣 v italic_v

In this section, we extend the results of Theorem 1 to the case of joint optimization of head $v$ and attention weights $W$ using a logistic loss function.

###### Assumption B .

Let $\Gamma,\Gamma^{\prime}>0$ denote the label margins when solving ($\ell_{p}$-SVM) with $X_{i\alpha_{i}}$ and its replacement with $X^{\top}_{i}\sigma(X_{i}Wz_{i})$, for all $i\in[n]$, respectively. There exists $\nu>0$ such that for all $i\in[n]$ and $W\in\mathbb{R}^{d\times d}$,

$$ $\displaystyle\Gamma-\Gamma^{\prime}\geq\nu\cdot(1-s_{i\alpha_{i}}),\quad \textnormal{where}\quad s_{i\alpha_{i}}=[\sigma(X_{i}Wz_{i})]_{\alpha_{i}}.$ $$

Assumption B is similar to ataee2024max and highlights that selecting optimal tokens is key to maximizing the classifier’s label margin. When attention features, a weighted combination of all tokens, are used, the label margin shrinks based on how much attention is given to the optimal tokens. The term $\nu\cdot(1-s_{i\alpha_{i}})$ quantifies this minimum shrinkage. If the attention mechanism fails to focus on these tokens (i.e., low $s_{i\alpha_{i}}$), the margin decreases, reducing generalization. This assumption implies that optimal performance is achieved when attention converges on the most important tokens, aligning with the max-margin attention SVM solution.

Similar to how we provided the characterization of convergence for the regularization path of $\ell_{p}$-AttGD, we offer a similar characterization here for $\ell_{p}$-JointGD.

###### Theorem 5 (Joint $\ell_{p}$ –norm Regularization Path) .

Consider (ERM) with a logistic loss $l(x)=\log(1+e^{-x})$, and define

$$ $(v^{(r)},W^{(R)}):=\operatorname*{arg\,min}_{(v,W)}~{}\mathcal{L}(v,W)~{}~{}~{ }~{}\textnormal{subj. to}\quad\|W\|_{p,p}\leq R~{}~{}\textnormal{and}~{}~{}\|v \|_{p}\leq r.$ ( $\ell_{p}$ -JointRP ) $$

Suppose there are token indices $\alpha=(\alpha_{i})_{i=1}^{m}$ for which $W^{\alpha}_{\mathrm{mm}}$ of ($\ell_{p}$-AttSVM) exists and Assumption B holds for some $\Gamma,\nu>0$. Then,

$$ $\displaystyle\lim_{(r,R)\to(\infty,\infty)}\left(\frac{v^{(r)}}{r},\frac{W^{(R )}}{R}\right)=\left(\frac{v_{\mathrm{mm}}}{\|v_{\mathrm{mm}}\|_{p}},\frac{W^{ \alpha}_{\mathrm{mm}}}{\|W^{\alpha}_{\mathrm{mm}}\|_{p,p}}\right).$ (33) $$

Here, $v_{\mathrm{mm}}$ and $W^{\alpha}_{\mathrm{mm}}$ are the solution of ($\ell_{p}$-SVM) and ($\ell_{p}$-AttSVM), respectively.

###### Proof.

The proof follows a similar approach to (ataee2024max, Theorem 5), adapted to the $\ell_{p}$-norm. We provide the revised version for the generalized attention SVM, tracking the required changes. Without loss of generality, we set $\alpha_{i}=1$ for all $i\in[n]$, and we use $W_{\mathrm{mm}}$ instead of $W_{\mathrm{mm}}^{\alpha}$. Suppose the claim is incorrect, meaning either $W^{(R)}/R$ or $v^{(r)}/r$ fails to converge as $R$ and $r$ grow. Define

$$ $\displaystyle\Xi$ $\displaystyle=\frac{1}{\|\bar{W}_{\mathrm{mm}}\|_{p,p}},\qquad\Gamma=\frac{1}{ \|v_{\mathrm{mm}}\|_{p}},$ $\displaystyle\bar{W}_{\mathrm{mm}}$ $\displaystyle:=R\Xi W_{\mathrm{mm}},\qquad\bar{v}_{\mathrm{mm}}:=r\Gamma v_{ \mathrm{mm}}$ (34) $$

Our strategy is to show that $(\bar{v}_{\mathrm{mm}},\bar{W}_{\mathrm{mm}})$ is a strictly better solution compared to $(v^{(r)},W^{(R)})$ for large $R$ and $r$, leading to a contradiction.

$\bullet$ Case 1: $W^{(R)}/R$ does not converge to $\bar{W}_{\mathrm{mm}}/R$.
In this case, there exists $\delta,\gamma=\gamma(\delta)>0$ such that we can find arbitrarily large $R$ with

$$ $\|W^{(R)}/R-\bar{W}_{\mathrm{mm}}/R\|\geq\delta$ $$

and the margin induced by $W^{(R)}/R$ is at most $\Xi(1-\gamma)$.

We bound the amount of non-optimality $q^{*}_{i}$ of $\bar{W}_{\mathrm{mm}}$:

$$ $\displaystyle q^{*}_{i}:=\frac{\sum_{t\neq\alpha_{i}}\exp(X_{it}^{\top}\bar{W} _{\mathrm{mm}}z_{i})}{\sum_{t\in[T]}\exp(X_{it}^{\top}\bar{W}_{\mathrm{mm}}z_{ i})}$ $\displaystyle\leq\frac{\sum_{t\neq\alpha_{i}}\exp(X_{it}^{\top}\bar{W}_{ \mathrm{mm}}z_{i})}{\exp(X_{i\alpha_{i}}^{\top}\bar{W}_{\mathrm{mm}}z_{i})}$ $\displaystyle\leq T\exp(-\Xi R).$ $$

Thus,

$$ $q^{*}_{\max}:=\max_{i\in[n]}q^{*}_{i}\leq T\exp(-\Xi R).$ (35a) Next, assume without loss of generality that the first margin constraint is $\gamma$ -violated by $W^{(R)}$ , meaning $\min_{t\neq\alpha_{1}}~{}(X_{1\alpha_{1}}-X_{1t})^{\top}W^{(R)}z_{1}\leq\Xi R( 1-\gamma).$ Denoting the amount of non-optimality of the first input of $W^{(R)}$ as $\hat{q}_{1}$ , we find $\displaystyle\hat{q}_{1}:=\frac{\sum_{t\neq\alpha_{1}}\exp(X_{1t}^{\top}W^{(R) }z_{1})}{\sum_{t\in[T]}\exp(X_{1t}^{\top}W^{(R)}z_{1})}$ $\displaystyle\geq\frac{1}{T}\frac{\sum_{t\neq\alpha_{1}}\exp(X_{1t}^{\top}W^{( R)}z_{1})}{\exp(X_{1\alpha_{1}}^{\top}W^{(R)}z_{1})}$ $\displaystyle\geq T^{-1}\exp(-(1-\gamma)R\Xi).$ This implies that $\hat{q}_{\max}:=\max_{i\in[n]}q^{*}_{i}\geq T^{-1}\exp(-\Xi R(1-\gamma)).$ (35b) We similarly have $q^{*}_{\max}\geq T^{-1}\exp(-\Xi R).$ (35c) $$

Thus, (35) gives the following relationship between the upper and lower bounds on non-optimality:

$$ $\displaystyle-(1-\gamma)\Xi R-\log T$ $\displaystyle\leq\log(\hat{q}_{\max}),$ $\displaystyle-\Xi R-\log T$ $\displaystyle\leq\log(q^{*}_{\max})\leq-\Xi R+\log T.$ (36) $$

In other words, $\bar{W}_{\mathrm{mm}}$ has exponentially less non-optimality compared to $W^{(R)}$ as $R$ grows. To proceed, we need to upper and lower bound the logistic loss of $(\bar{v}_{\mathrm{mm}},\bar{W}_{\mathrm{mm}})$ and $(v^{(r)},W^{(R)})$ respectively, to establish a contradiction.

$\bullet$ Sub-Case 1.1: Upper bound for $\mathcal{L}(\bar{v}_{\mathrm{mm}},\bar{W}_{\mathrm{mm}})$.
We now bound the logistic loss for the limiting solution. Set $\bar{r}_{i}=X_{i}^{\top}\sigma(X_{i}\bar{W}_{\mathrm{mm}}z_{i})$. If $\|\bar{r}_{i}-X_{i1}\|_{p}\leq\epsilon_{i}$, then $v_{\mathrm{mm}}$ satisfies the SVM constraints on $\bar{r}_{i}$ with $Y_{i}\cdot\bar{r}_{i}^{\top}v_{\mathrm{mm}}\geq 1-\epsilon_{i}/\Gamma$. Setting $\epsilon_{\max}=\sup_{i\in[n]}\epsilon_{i}$, $v_{\mathrm{mm}}$ achieves a label-margin of $\Gamma-\epsilon_{\max}$ on the dataset $(Y_{i},\bar{r}_{i})_{i\in[n]}$. Let $M=\sup_{i\in[n],t,\tau\in[T]}\|X_{it}-X_{i\tau}\|_{p}$. Recalling (G), the worst-case perturbation is

$$ $\epsilon_{\max}\leq M\exp(-\Xi R+\log T)=MT\exp(-\Xi R).$ $$

This implies the upper bound on the logistic loss:

$$ $\displaystyle\mathcal{L}(\bar{v}_{\mathrm{mm}},\bar{W}_{\mathrm{mm}})$ $\displaystyle\leq\max_{i\in[n]}\log(1+\exp(-Y_{i}\bar{r}_{i}^{\top}\bar{v}_{ \mathrm{mm}}))$ $\displaystyle\leq\max_{i\in[n]}\exp(-Y_{i}\bar{r}_{i}^{\top}\bar{v}_{\mathrm{ mm}})$ $\displaystyle\leq\exp(-r\Gamma+r\epsilon_{\max})$ $\displaystyle\leq e^{rMT\exp(-\Xi R)}e^{-r\Gamma}.$ (37) $$

$\bullet$ Sub-Case 1.2: Lower bound for $\mathcal{L}(v^{(r)},W^{(R)})$.
We now bound the logistic loss for the finite solution. Set $\bar{r}_{i}=X_{i}^{\top}\sigma(X_{i}W^{(R)}z_{i})$. Using Assumption B, solving ($\ell_{p}$-SVM) on $(y_{i},\bar{r}_{i})_{i\in[n]}$ achieves at most $\Gamma-\nu e^{-(1-\gamma)\Xi R}/T$ margin. Consequently, we have:

$$ $\displaystyle\mathcal{L}(v^{(r)},W^{(R)})$ $\displaystyle\geq\frac{1}{n}\max_{i\in[n]}\log(1+\exp(-Y_{i}\bar{r}_{i}^{\top} v^{(r)}))$ $\displaystyle\geq\left(\frac{1}{2n}\max_{i\in[n]}\exp(-Y_{i}\bar{r}_{i}^{\top} v^{(r)})\right)\wedge\log 2$ $\displaystyle\geq\left(\frac{1}{2n}\exp(-r(\Gamma-\nu e^{-(1-\gamma)\Xi R}/T)) \right)\wedge\log 2$ $\displaystyle\geq\left(\frac{1}{2n}e^{r(\nu/T)\exp(-(1-\gamma)\Xi R)}e^{-r \Gamma}\right)\wedge\log 2.$ $$

Observe that this lower bound dominates the upper bound from (37) when $R$ is large, specifically when (ignoring the multiplier $1/2n$ for simplicity):

$$ $(\nu/T)e^{-(1-\gamma)\Xi R}\geq MTe^{-\Xi R}\implies R\geq\frac{1}{\gamma\Xi} \log\left(\frac{MT^{2}}{\nu}\right).$ $$

Thus, we obtain the desired contradiction since such a large $R$ is guaranteed to exist when $W^{(R)}/R\not\rightarrow\bar{W}_{\mathrm{mm}}$. Therefore, $W^{(R)}/R$ must converge to $\bar{W}_{\mathrm{mm}}/R$.

$\bullet$ Case 2: Suppose $v^{(r)}/r$ does not converge. In this case, there exists $\delta>0$ such that we can find arbitrarily large $r$ obeying $\text{dist}(v^{(r)}/r,\bar{v}_{\mathrm{mm}}/r)\geq\delta$. If $\text{dist}(W^{(R)}/R,\Xi W_{\mathrm{mm}})\not\rightarrow 0$, then "Case 1" applies. Otherwise, we have $\text{dist}(W^{(R)}/R,\Xi W_{\mathrm{mm}})\rightarrow 0$, thus we can assume $\text{dist}(W^{(R)}/R,\Xi W_{\mathrm{mm}})\leq\epsilon$ for an arbitrary choice of $\epsilon>0$.

On the other hand, due to the strong convexity of ($\ell_{p}$-AttSVM), for some $\gamma:=\gamma(\delta)>0$, $v^{(r)}$ achieves a margin of at most $(1-\gamma)\Gamma r$ on the dataset $(Y_{i},X_{i1})_{i\in[n]}$, where $X_{i1}$ denotes the optimal token for each $i\in[n]$.
Additionally, since $\text{dist}(W^{(R)}/R,\Xi W_{\mathrm{mm}})\leq\epsilon$, $W^{(R)}$ strictly separates all optimal tokens (for small enough $\epsilon>0$) and $\hat{q}_{\max}:=f(\epsilon)\rightarrow 0$ as $R\rightarrow\infty$. Note that $f(\epsilon)$ quantifies the non-optimality of $W^{(R)}$ compared to $W_{\mathrm{mm}}$; as $\epsilon\to 0$, meaning $W^{(R)}/R$ converges to $\Xi W_{\mathrm{mm}}/R$, $f(\epsilon)\to 0$. Consequently, setting $r_{i}=X_{i}^{\top}\sigma(X_{i}W^{(R)}z_{i})$, for sufficiently large $R>0$ and setting $M=\sup_{i\in[n],t\in[T]}\|X_{it}\|$, we have that

$$ $\displaystyle\min_{i\in[n]}~{}~{}y_{i}(v^{(r)})^{\top}r_{i}$ $\displaystyle\leq\min_{i\in[n]}~{}~{}y_{i}(v^{(r)})^{\top}X_{i1}+\sup_{i\in[n] }|(v^{(r)})^{\top}(X_{it}-X_{i1})|$ $\displaystyle\leq(1-\gamma)\Gamma r+Mf(\epsilon)r$ $\displaystyle\leq(1-\gamma/2)\Gamma r.$ (38) $$

This in turn implies that logistic loss is lower bounded by

$$ $\mathcal{L}(v^{(r)},W^{(R)})\geq\left(\frac{1}{2n}e^{\gamma\Gamma r/2}e^{- \Gamma r}\right)\wedge\log 2.$ $$

Now, using (37), this exponentially dominates the upper bound of $(\bar{W}_{\mathrm{mm}},\bar{v}_{\mathrm{mm}})$ whenever $rMT\exp(-\Xi R)<r\gamma\Gamma/2$, completing the proof.
∎

###### Assumption B .

###### Theorem 5 (Joint ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT –norm Regularization Path) .

###### Proof.

## Appendix H Implementation Details

The experiments were run on an Intel i7 core and a single V100 GPU using the pytorch and huggingface libraries. They should be runnable on any generic laptop. The GitHub repository can be found at [https://github.com/aaronk2002/AttentionMD](https://github.com/aaronk2002/AttentionMD).

### H.1 ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT -AttGD Experiment

The dataset $(X_{i},Y_{i},z_{i})_{i=1}^{n}$ is generated randomly: $X_{i}$ and $z_{i}$ are sampled from the unit sphere, and $Y_{i}$ is uniformly sampled from $\{\pm 1\}$. Additionally, $v$ is randomly selected from the unit sphere. We use $n=6$ samples, $T=8$ tokens per sample, and $d=10$ dimensions per token, fulfilling the overparameterized condition for the $\ell_{p}$-AttSVM problem to be almost always feasible.

The model is trained with parameters initialized near the origin, using Algorithms $\ell_{p}$-AttGD with $p=1.75,2$, and $3$, and a learning rate of $0.1$. Training lasted for $1,500$ epochs for $p=1.75$, $2,000$ epochs for $p=2$, and $20,000$ epochs for $p=3$. Gradients are normalized to accelerate convergence without altering results significantly. We refer to the parameter histories as the $\ell_{1.75},\ell_{2},$ and $\ell_{3}$ optimization paths. We compute the chosen tokens $(\alpha_{i})_{i=1}^{n}$ for the ($\ell_{p}$-AttSVM) problem by selecting the token with the highest softmax probability for each sample. This process is repeated for $p=1.75$, $2$, and $3$.

### H.2 ℓ p subscript ℓ 𝑝 \ell_{p} roman_ℓ start_POSTSUBSCRIPT italic_p end_POSTSUBSCRIPT -JointGD Experiment

We use the following example dataset for the experiment on joint optimization.

###### Example 2 .

Let $n=2$, $T=3$, $d=2$. Let $y_{1}=1$, $y_{2}=-1$. Let:

$$ $\displaystyle X_{1}=\begin{pmatrix}X_{11}\\ X_{12}\\ X_{13}\end{pmatrix}=\begin{pmatrix}-5.4&2.4\\ 2.8&4.2\\ 2.6&-0.2\end{pmatrix},\quad\textnormal{and}\quad X_{2}=\begin{pmatrix}X_{21}\\ X_{22}\\ X_{23}\end{pmatrix}=\begin{pmatrix}0.8&-4.4\\ -2.2&-0.8\\ 1.8&0.2\end{pmatrix}.$ (39) $$

Let $z_{1}=X_{11}$, $z_{2}=X_{21}$.

We use learning rates $0.1$ and we trained the model for $1,500$ epochs for when $p=1.75$, $2,000$ epochs for $p=2$, and $20,000$ epochs for $p=3$. As it was done in the previous experiment, we obtain the parameter history for each $p$, and compute the optimal token for the  ($\ell_{p}$-AttSVM) and  $\ell_{p}$-SVM problems.

###### Example 2 .

### H.3 Addendum to the Attention Map Results

## Peer Reviews

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: 1. What does the comparison token $z_i$ mean in practice (Line 55)? For example, for a real-world application problem, how to set the comparison token $z_i$ given the input $X$ and label $y$?
soundness: 3
strengths: 1. Theoretical Contributions: The paper provides a solid theoretical foundation for understanding the convergence properties and implicit bias of MD in attention models. The extension to $\ell_p$-norm objectives adds flexibility in modeling and opens up new avenues for optimizing attention mechanisms.

2. Generalization of Attention Optimization: The approach generalizes previous work on attention models by using MD with a broad class of potential functions, allowing a deeper exploration of the optimization landscape beyond vanilla gradient descent.
summary: The paper introduces a novel approach to optimizing attention mechanisms using Mirror Descent (MD), specifically focusing on a generalized max-margin token selection strategy for softmax attention models. The authors propose a family of MD algorithms, termed $\ell_{p}$-AttGD, which generalize traditional gradient descent by using the $p$-th power of the $\ell_p$-norm as the potential function. The main contributions include proving the convergence of these algorithms to generalized max-margin Support Vector Machine (SVM) solutions for optimizing the attention mechanism, both for fixed and jointly optimized parameters (key-query matrix and decoder).
weaknesses: 1. As far as I know, mirror descent is not a popular optimization algorithm for training deep learning models. I agree that a simplified model (e.g., the one layer model considered in this paper and previous work) could provide valuable insights on understanding transformers, but it is not clear what is the role/implication the $\ell_p$-norm for deep learning. If possible, it would helpful if the authors could highlight a few practical mirror descent-based optimizers in the revision.

2. For Line431, for the $\ell_{1,1}$/$\ell_3$ optimizer, the authors provide code for implementing the $\ell_{1,1}$/$\ell_3$ optimizer for optimizing the transformer. In term of efficiency, would the $\ell_{1,1}$/$\ell_3$ optimizer be as efficient as popular optimizers like AdamW [LH2019]? Could the authors provide comparison with AdamW in Table 1?

[LH2019] Decoupled Weight Decay Regularization. Ilya Loshchilov, Frank Hutter.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 4
questions: - Can the authors clarify the questions above regarding feasibility? 
- Can the authors compare the training curves with that of adam?
- The model sizes do not seem to be specified. The architectural details should be put in the appendix


**Minor remarks and typos**

- L68: ERM refers to an equation much later in the manuscript
- L71: W(k) has not bee introduced thus far
- L113: the softmax acts here on a matrix, but later in (2) it acts on vectors.
- L139: (Blair, 1985)
- L165: (Tarzanagh et al. (2024; 2023)).  
- L165: the sense of nearby is unclear here.
- L169: it would be good to specify to which set the indices belong. Generally, it would be clearer to indicate in which space lives each newly introduced variable.
- L176: “the αi component has a significantly higher probability” : this statement should also involve the norm of W, right? Because if all values in the pre-softmax vector are high, then the probability vector will be almost constant.
- Eq 4b: typo W -> W’
- L280: $exp(2)$ could be $\exp(2)$
- In eq.5, it would be good to recall that $D_\psi$ depends on $p$.
- All figures should be in pdf format
- “In the Appendix G, we show that W and v generated by ℓp-JointRP converge to their respective…” it would strengthen that section to put it in the main text.
- Fig4: a log scale would make things clearer, especially highlighting the practical convergence rates.
soundness: 2
strengths: - The paper is well written 
- It is quite easy to understand
- The fact that MD achieves sparser weights and leads to better generalization is interesting.
summary: This paper proposes to study the dynamics of a simplified one-layer attention network. The implicit bias of gradient descent for such a problem is already known, this paper proposes to study what happens when the training algorithm is mirror descent with p-norms. The paper extends previous results known for gradient descent to this setting, and then demonstrate that using mirror descent to train a transformer on a sentiment analysis dataset yields better generalization.
weaknesses: - The motivation is quite unclear. “A broader understanding of general descent algorithms, including the mirror descent (MD) family and their token selection properties, is essential.” why? In practice, nobody trains attention layers with mirror descent. The observation that it works better than gradient descent is not very strong in my opinion, because in practice transformers cannot be trained with gradient descent either. In the experiments, it would be worthwhile to compare the proposed methods to the widely used adam algorithm. Also, providing train losses values for the real data experiment would be enlightening.
- There are many unclear things about global and optimal tokens (def.3). First, the fact that there is a feasible solution is very important, precise conditions for it to happen must be put in the paper. The authors simply state that “under mild overparameterization d ≥ max{T − 1, n}, the problem is almost always feasible (Tarzanagh et al., 2023, Theorem 1).” These conditions must be put in the paper. Also, it is unclear whether the statement is “for all $\alpha_i$, the problem is feasible” or “the problem is feasible for at least one set of $\alpha_i$.
- I cannot understand why Theorem 1 does not require those conditions: clearly, $W^{opt}_{mm}$ must exist for it to hold, and there is nothing in the assumptions guaranteeing it.
- It is never properly stated in the paper if “globally optimal tokens” are also “locally optimal”. It seems not to be the case, since globally optimal tokens depend on $v$ while locally optimal tokens do not. This also causes issues in thm 1: globally optimal tokens do not seem to be always feasible, so W^{opt}_{mm} is not always defined, while the lp-AttRp always has a solution. For instance, it could very well have a global solution of fixed norm.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: See weaknesses.
soundness: 2
strengths: The paper attempts to provide a theoretical analysis of mirror descent for attention training, extending prior work focused on gradient descent. It derives convergence results to a generalized hard-margin SVM and establishes convergence rates. The use of $l_p$ norms offers a degree of generality in the theoretical analysis.
summary: This paper investigates the optimization dynamics of mirror descent (MD) for training attention mechanisms, specifically focusing on $l_p$-AttGD. The authors claim that $l_p$-AttGD converges directionally to a generalized hard-margin SVM with an $l_p$ norm objective when applied to binary classification with a single-layer attention model. Some experiments on synthetic and real data are presented.
weaknesses: 1. The core idea of connecting attention optimization to SVM-like objectives is not new and has been explored in prior work, notably in "A Primal-Dual Framework for Transformers and Neural Networks" by Nguyen et al. and related papers. These prior works establish the fundamental link between attention and SVMs, including the optimization perspective. While this paper extends the analysis to mirror descent, the incremental contribution feels minimal and lacks motivation. Other core ideas of analyzing the implicit bias of GD/MD algorithms for softmax attention is already present in the cited works (Tarzanagh et al., 2023, 2024; Vasudeva et al., 2024a; Sheen et al., 2024) without fundamentally changing the nature of the problem or leading to significantly different insights.

2. The synthetic experiments are too simplistic and lack the complexity needed to represent realistic attention training scenarios. The real-data experiments, while showing some improvements in generalization, are insufficient to support the claims. The demonstration of improved token selection is based on a handful of examples (Figure 8) without any rigorous quantitative evaluation. Comparisons with GD do not include other commonly used optimization algorithms like Adam, making it impossible to judge the relative merits of MD for attention training.

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 3
questions: Please refer to the weakness.
soundness: 3
strengths: 1. The paper has a strong theoretical foundation, it provides rigorous mathematical analysis and proofs for the convergence properties of mirror descent in attention optimization, extending previous work on gradient descent to a more general framework.
2. This paper provides a novel algorithmic insight, the introduction of $\ell_p$-AttGD generalizes both $\ell_p$-GD and attention GD, offering new perspectives on attention optimization and token selection.
3. This work provides a complete theoretical treatment of the optimization dynamics by examining both fixed-head and joint optimization scenarios.
summary: This paper examines how mirror descent (MD) algorithms approach convergence and exhibit bias when optimizing attention mechanisms in softmax attention, with potential functions using $\ell_p$-norms as a basis for analysis. Key theoretical findings include demonstrating convergence towards hard margin Support Vector Machine (SVM) solutions and achieving convergence rates similar to gradient descent, in more straightforward models. The research builds upon findings related to descent by expanding the scope to encompass a wider range of optimization algorithms and shedding light on properties related to token selection.
weaknesses: 1. The empirical evaluation is limited, the paper includes experiments, and they are primarily focused on synthetic data and a single real-world dataset (Stanford Large Movie Review Dataset). More diverse real-world applications would strengthen the practical implications.
2. Theoretical results are highly dependent on assumption, Theorem 2, Theorem 3, and Theorem 4, rely on specific assumptions about initialization and step sizes, which may limit their practical applicability.
3. The paper does not thoroughly address the computational overhead of implementing mirror descent compared to standard gradient descent.