$\displaystyle\leq\frac{1}{n}\sum_{i=1}^{n}\left\|\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\bm{h}_{i})\right)\cdot{\bm{z}}_{i}\bm{\gamma}_{i}^{\top}\mathbb{S}^{\prime}(\bm{h}_{i}){\bm{X}}_{i}-\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\dot{\bm{h}}_{i})\right)\cdot{\bm{z}}_{i}\bm{\gamma}_{i}^{\top}\mathbb{S}^{\prime}(\dot{\bm{h}}_{i}){\bm{X}}_{i}\right\|_{F}$
$\displaystyle\leq\frac{1}{n}\sum_{i=1}^{n}\left\|{\bm{z}}_{i}\bm{\gamma}_{i}^{\top}\mathbb{S}^{\prime}(\dot{\bm{h}}_{i}){\bm{X}}_{i}\right\|_{F}~{}\left|\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\bm{h}_{i})\right)-\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\dot{\bm{h}}_{i})\right)\right|$
$\displaystyle+\frac{1}{n}\sum_{i=1}^{n}\left|\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\bm{h}_{i})\right)\right|~{}\left\|{\bm{z}}_{i}\bm{\gamma}_{i}^{\top}\mathbb{S}^{\prime}(\bm{h}_{i}){\bm{X}}_{i}-{\bm{z}}_{i}\bm{\gamma}_{i}^{\top}\mathbb{S}^{\prime}(\dot{\bm{h}}_{i}){\bm{X}}_{i}\right\|_{F}$
$\displaystyle\leq\frac{1}{n}\sum_{i=1}^{n}M_{0}~{}\|\bm{\gamma}_{i}\|^{2}~{}\|{\bm{z}}_{i}\|~{}\|{\bm{X}}_{i}\|~{}\left\|\mathbb{S}(\bm{h}_{i})-\mathbb{S}(\dot{\bm{h}}_{i})\right\|$
$\displaystyle+\frac{1}{n}\sum_{i=1}^{n}M_{1}~{}\|\bm{\gamma}_{i}\|~{}\|{\bm{z}}_{i}\|~{}\|{\bm{X}}_{i}\|~{}\left\|\mathbb{S}^{\prime}(\bm{h}_{i})-\mathbb{S}^{\prime}(\dot{\bm{h}}_{i})\right\|_{F},$
(17)

where the second inequality follows from the fact that $|ab-
cd|\leq|d||a-c|+|a||b-d|$ and the third inequality uses Assumption A and (15).

Substituting (16a) and (A.3) into (A.3), we get

$\displaystyle\left\|\nabla\mathcal{L}(\bm{W})-\nabla\mathcal{L}(\dot{\bm{W}})\right\|_{F}$
$\displaystyle\leq\frac{1}{n}\sum_{i=1}^{n}\left(M_{0}~{}\|\bm{\gamma}_{i}\|^{2}\|{\bm{z}}_{i}\|^{2}\|{\bm{X}}_{i}\|^{2}+3M_{1}\|\bm{\gamma}_{i}\|~{}\|{\bm{z}}_{i}\|^{2}~{}\|{\bm{X}}_{i}\|^{2}\right)\|\bm{W}-\dot{\bm{W}}\|_{F}$
$\displaystyle\leq\frac{1}{n}\sum_{i=1}^{n}\left(M_{0}~{}\|\bm{v}\|^{2}\|{\bm{z}}_{i}\|^{2}\|{\bm{X}}_{i}\|^{4}+3M_{1}\|\bm{v}\|~{}\|{\bm{z}}_{i}\|^{2}~{}\|{\bm{X}}_{i}\|^{3}\right)\|\bm{W}-\dot{\bm{W}}\|_{F}$
$\displaystyle\leq L_{\bm{W}}~{}\|\bm{W}-\dot{\bm{W}}\|_{F},$

where $L_{\bm{W}}$ is defined in (3).

Let $\bm{g}_{i}={\bm{X}}_{i}{\bm{K}}{\bm{Q}}^{\top}{\bm{z}}_{i}$. We have

$\displaystyle\nabla_{{\bm{K}}}{\cal{L}}({\bm{K}},{\bm{Q}})=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\bm{g}_{i})\right)\cdot{\bm{z}}_{i}\bm{\gamma}_{i}^{\top}\mathbb{S}^{\prime}(\bm{g}_{i}){\bm{X}}_{i}{\bm{Q}},$
(18a)
$\displaystyle\nabla_{{\bm{Q}}}{\cal{L}}({\bm{K}},{\bm{Q}})=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\bm{g}_{i})\right)\cdot{\bm{X}}_{i}^{\top}\mathbb{S}^{\prime}(\bm{g}_{i})\bm{\gamma}_{i}{\bm{z}}_{i}^{\top}{\bm{K}}.$
(18b)

By the similar argument as in (A.3), for any ${\bm{Q}}$ and
$\dot{\bm{Q}}\in\mathbb{R}^{d\times m}$, we have

$\displaystyle\left\|\nabla_{{\bm{Q}}}\mathcal{L}({\bm{K}},{\bm{Q}})-\nabla_{{\bm{Q}}}\mathcal{L}({\bm{K}},\dot{{\bm{Q}}})\right\|_{F}$
$\displaystyle\leq\frac{\|{\bm{K}}\|}{n}\sum_{i=1}^{n}\left\|\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\bm{h}_{i})\right)\cdot{\bm{z}}_{i}\bm{\gamma}_{i}^{\top}\mathbb{S}^{\prime}(\bm{h}_{i}){\bm{X}}_{i}-\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\dot{\bm{h}}_{i})\right)\cdot{\bm{z}}_{i}\bm{\gamma}_{i}^{\top}\mathbb{S}^{\prime}(\dot{\bm{h}}_{i}){\bm{X}}_{i}\right\|_{F}$
$\displaystyle\leq L_{\bm{W}}\|{\bm{K}}\|~{}\|{\bm{Q}}-\dot{{\bm{Q}}}\|_{F}.$
(19)

Similarly, for any ${\bm{K}},\dot{\bm{K}}\in\mathbb{R}^{d\times m}$, we get

$\left\|\nabla_{{\bm{K}}}\mathcal{L}({\bm{K}},{\bm{Q}})-\nabla_{{\bm{K}}}\mathcal{L}(\dot{{\bm{K}}},{\bm{Q}})\right\|_{F}\leq
L_{\bm{W}}\|{\bm{Q}}\|~{}\|{\bm{K}}-\dot{{\bm{K}}}\|_{F}.$

###  A.4 A Useful Lemma for Gradient Descent Analysis

###### Lemma 7

For any ${\bm{X}}\in\mathbb{R}^{T\times d}$,
$\bm{W},\bm{V}\in\mathbb{R}^{d\times d}$ and
${\bm{z}},\bm{v}\in\mathbb{R}^{d}$, let ${\bm{a}}={\bm{X}}\bm{V}{\bm{z}}$,
$\bm{s}=\mathbb{S}({\bm{X}}\bm{W}{\bm{z}})$, and $\bm{\gamma}={\bm{X}}\bm{v}$.
Set

$\Gamma=\sup_{t,\tau\in[T]}|\bm{\gamma}_{t}-\bm{\gamma}_{\tau}|~{}~{}~{}\textnormal{and}~{}~{}~{}A=\sup_{t\in[T]}\|{{\bm{a}}_{t}}\|.$

We have that

$\left|{\bm{a}}^{\top}\textnormal{diag}(\bm{s})\bm{\gamma}-{\bm{a}}^{\top}\bm{s}\bm{s}^{\top}\bm{\gamma}-\sum_{t\geq
2}^{T}({\bm{a}}_{1}-{\bm{a}}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\right|\leq
2\Gamma A(1-\bm{s}_{1})^{2}.$

Proof. The proof is similar to [TLZO23, Lemma 4], but for the sake of
completeness, we provide it here. Set
$\bar{\gamma}=\sum_{t=1}^{T}\bm{\gamma}_{t}\bm{s}_{t}$. We have

$\displaystyle\bm{\gamma}_{1}-\bar{\gamma}=\sum_{t\geq
2}^{T}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\bm{s}_{t},~{}~{}\textnormal{and}~{}~{}|\bm{\gamma}_{1}-\bar{\gamma}|\leq\Gamma(1-\bm{s}_{1}).$

Then,

$\displaystyle{\bm{a}}^{\top}\text{diag}(\bm{s})\bm{\gamma}-{\bm{a}}^{\top}\bm{s}\bm{s}^{\top}\bm{\gamma}$
$\displaystyle=\sum_{t=1}^{T}{\bm{a}}_{t}\bm{\gamma}_{t}\bm{s}_{t}-\sum_{t=1}^{T}{\bm{a}}_{t}\bm{s}_{t}\sum_{t=1}^{T}\bm{\gamma}_{t}\bm{s}_{t}$
$\displaystyle={\bm{a}}_{1}\bm{s}_{1}(\bm{\gamma}_{1}-\bar{\gamma})-\sum_{t\geq
2}^{T}{\bm{a}}_{t}\bm{s}_{t}(\bar{\gamma}-\bm{\gamma}_{t}).$ (20)

Since

$\left|\sum_{t\geq
2}^{T}{\bm{a}}_{t}\bm{s}_{t}(\bar{\gamma}-\bm{\gamma}_{t})-\sum_{t\geq
2}^{T}{\bm{a}}_{t}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\right|\leq
A\Gamma(1-\bm{s}_{1})^{2},$

we obtain555For simplicity, we use $\pm$ on the right hand side to denote the
upper and lower bounds.

$\displaystyle{\bm{a}}^{\top}\text{diag}(\bm{s})\bm{\gamma}-{\bm{a}}^{\top}\bm{s}\bm{s}^{\top}\bm{\gamma}$
$\displaystyle={\bm{a}}_{1}\bm{s}_{1}(\bm{\gamma}_{1}-\bar{\gamma})-\sum_{t\geq
2}^{T}{\bm{a}}_{t}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\pm
A\Gamma(1-\bm{s}_{1})^{2}$ $\displaystyle={\bm{a}}_{1}\bm{s}_{1}\sum_{t\geq
2}^{T}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\bm{s}_{t}-\sum_{t\geq
2}^{T}{\bm{a}}_{t}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\pm
A\Gamma(1-\bm{s}_{1})^{2}$ $\displaystyle=\sum_{t\geq
2}^{T}({\bm{a}}_{1}\bm{s}_{1}-{\bm{a}}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\pm
A\Gamma(1-\bm{s}_{1})^{2}$ $\displaystyle=\sum_{t\geq
2}^{T}({\bm{a}}_{1}-{\bm{a}}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\pm
2A\Gamma(1-\bm{s}_{1})^{2}.$

Here, $\pm$ on the right handside uses the fact that

$\left|\sum_{t\geq
2}^{T}({\bm{a}}_{1}\bm{s}_{1}-{\bm{a}}_{1})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\right|\leq(1-\bm{s}_{1})A\Gamma\sum_{t\geq
2}^{T}\bm{s}_{t}=(1-\bm{s}_{1})^{2}A\Gamma.$

##  Appendix B Global Convergence of Gradient Descent

###  B.1 Divergence of $\|W(k)\|_{F}$

The next lemma establishes the descent property of gradient descent for
$\mathcal{L}(\bm{W})$ under Assumption A.

######  Lemma 8 (Descent Lemma)

Under Assumption A, if $\eta\leq 2/L_{\bm{W}}$, then for any initialization
$\bm{W}(0)$, Algorithm W-GD satisfies:

$\displaystyle\mathcal{L}(\bm{W}(k+1))-\mathcal{L}(\bm{W}(k))\leq-\frac{\eta}{2}\|{\nabla\mathcal{L}(\bm{W}(k))}\|_{F}^{2},$
(21)

for all $k\geq 0$. Additionally, it holds that
$\sum_{k=0}^{\infty}\|{\nabla\mathcal{L}\left(\bm{W}(k)\right)}\|_{F}^{2}<\infty$,
and
$\lim_{k\rightarrow\infty}\|{\nabla\mathcal{L}\left(\bm{W}\left(k\right)\right)}\|_{F}^{2}=0$.

Proof. The proof is similar to [TLZO23, Lemma 5].

The lemma below reveals that the correlation between the training loss’s
gradient at any arbitrary matrix $\bm{W}$ and the attention SVM solution
$\bm{W}^{\textsl{mm}}$ is negative. Consequently, for any finite $\bm{W}$,
$\left<\nabla{\cal{L}}(\bm{W}),\bm{W}^{\textsl{mm}}\right>$ cannot be equal to
zero.

###### Lemma 9

Let $\bm{W}^{\textsl{mm}}$ be the SVM solution of (Att-SVM). Suppose
Assumptions A and B hold. Then, for all $\bm{W}\in\mathbb{R}^{d\times d}$, the
training loss (W-ERM) obeys
$\left<\nabla{\cal{L}}(\bm{W}),\bm{W}^{\textsl{mm}}\right>\leq-c<0$, for some
constant $c>0$ depending on the data, the head $\bm{v}$, and a loss derivative
bound.

Proof. Let

$\bm{\bar{h}}_{i}={\bm{X}}_{i}\bm{W}^{\textsl{mm}}{\bm{z}}_{i},~{}~{}~{}\bm{\gamma}_{i}=Y_{i}\cdot{\bm{X}}_{i}\bm{v},~{}~{}~{}\textnormal{and}~{}~{}~{}\bm{h}_{i}={\bm{X}}_{i}\bm{W}{\bm{z}}_{i}.$
(22)

Let us recall the gradient evaluated at $\bm{W}$ which is given by

$\displaystyle\nabla{\cal{L}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\bm{h}_{i})\right)\cdot{\bm{X}}_{i}^{\top}\mathbb{S}^{\prime}(\bm{h}_{i})\bm{\gamma}_{i}{\bm{z}}_{i}^{\top},$
(23)

which implies that

$\begin{split}\left<\nabla{\cal{L}}(\bm{W}),\bm{W}^{\textsl{mm}}\right>&=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\bm{h}_{i})\right)\cdot\left\langle{\bm{X}}_{i}^{\top}\mathbb{S}^{\prime}(\bm{h}_{i})\bm{\gamma}_{i}{\bm{z}}_{i}^{\top},\bm{W}^{\textsl{mm}}\right\rangle\\\
&=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}_{i}\cdot{{\operatorname{trace}}}\left((\bm{W}^{\textsl{mm}})^{\top}{\bm{X}}_{i}^{\top}\mathbb{S}^{\prime}(\bm{h}_{i})\bm{\gamma}_{i}{\bm{z}}_{i}^{\top}\right)\\\
&=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}_{i}\cdot\bm{\bar{h}}_{i}^{\top}\mathbb{S}^{\prime}(\bm{h}_{i})\bm{\gamma}_{i}\\\
&=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}_{i}\cdot\left(\bm{\bar{h}}^{\top}_{i}\text{diag}(\bm{s}_{i})\bm{\gamma}_{i}-\bm{\bar{h}}^{\top}_{i}\bm{s}_{i}\bm{s}^{\top}_{i}\bm{\gamma}_{i}\right).\end{split}$
(24)

Here, let
$\ell^{\prime}_{i}:=\ell^{\prime}(\bm{\gamma}_{i}^{\top}\mathbb{S}(\bm{h}_{i}))$,
$\bm{s}_{i}=\mathbb{S}(\bm{h}_{i})$ and the third equality uses
${{\operatorname{trace}}}\left(\bm{b}{\bm{a}}^{\top}\right)={\bm{a}}^{\top}\bm{b}$.

In order to move forward, we will establish the following result, with a focus
on the equal score condition (Assumption B.2): Let $\gamma=\bm{\gamma}_{t\geq
2}$ be a constant, and let $\bm{\gamma}_{1}$ and $\bar{\bm{h}}_{1}$ represent
the largest indices of vectors $\bm{\gamma}$ and $\bm{\bar{h}}$ respectively.
For any vector $\bm{s}$ that satisfies $\sum_{t\in[T]}\bm{s}_{t}=1$ and
$\bm{s}_{t}>0$, we aim to prove that
$\bm{\bar{h}}^{\top}\text{diag}(\bm{s})\bm{\gamma}-\bm{\bar{h}}^{\top}\bm{s}\bm{s}^{\top}\bm{\gamma}>0$.
To demonstrate this, we proceed by writing the following:

$\begin{split}\bm{\bar{h}}^{\top}\text{diag}(\bm{s})\bm{\gamma}-\bm{\bar{h}}^{\top}\bm{s}\bm{s}^{\top}\bm{\gamma}&=\sum_{t=1}^{T}\bar{\bm{h}}_{t}\bm{\gamma}_{t}\bm{s}_{t}-\sum_{t=1}^{T}\bar{\bm{h}}_{t}\bm{s}_{t}\sum_{t=1}^{T}\bm{\gamma}_{t}\bm{s}_{t}\\\
&=\left(\bar{\bm{h}}_{1}\bm{\gamma}_{1}\bm{s}_{1}+\gamma\sum_{t\geq
2}^{T}\bar{\bm{h}}_{t}\bm{s}_{t}\right)-\Big{(}\bm{\gamma}_{1}\bm{s}_{1}+\gamma(1-\bm{s}_{1})\Big{)}\left(\bar{\bm{h}}_{1}\bm{s}_{1}+\sum_{t\geq
2}^{T}\bar{\bm{h}}_{t}\bm{s}_{t}\right)\\\
&=\bar{\bm{h}}_{1}(\bm{\gamma}_{1}-\gamma)\bm{s}_{1}(1-\bm{s}_{1})-(\bm{\gamma}_{1}-\gamma)\bm{s}_{1}\sum_{t\geq
2}^{T}\bar{\bm{h}}_{t}\bm{s}_{t}\\\
&=(\bm{\gamma}_{1}-\gamma)(1-\bm{s}_{1})\bm{s}_{1}\left[\bar{\bm{h}}_{1}-\frac{\sum_{t\geq
2}^{T}\bar{\bm{h}}_{t}\bm{s}_{t}}{\sum_{t\geq 2}^{T}\bm{s}_{t}}\right]\\\
&\geq(\bm{\gamma}_{1}-\gamma)(1-\bm{s}_{1})\bm{s}_{1}(\bar{\bm{h}}_{1}-\max_{t\geq
2}\bar{\bm{h}}_{t}).\end{split}$ (25)

To proceed, define

${\gamma}_{\textsl{gap}}^{i}=\bm{\gamma}_{i\texttt{opt}_{i}}-\max_{t\neq\texttt{opt}_{i}}\bm{\gamma}_{it}~{}~{}\textnormal{and}~{}~{}{\bar{h}}_{\textsl{gap}}^{i}=\bar{\bm{h}}_{i\texttt{opt}_{i}}-\max_{t\neq\texttt{opt}_{i}}\bar{\bm{h}}_{it}.$

With these, we obtain

$\bm{\bar{h}}^{\top}_{i}\text{diag}(\bm{s}_{i})\bm{\gamma}_{i}-\bm{\bar{h}}^{\top}_{i}\bm{s}_{i}\bm{s}^{\top}_{i}\bm{\gamma}_{i}\geq{\gamma}_{\textsl{gap}}^{i}{\bar{h}}_{\textsl{gap}}^{i}(1-\bm{s}_{i\texttt{opt}_{i}})\bm{s}_{i\texttt{opt}_{i}}.$
(26)

Note that

$\begin{split}&{\bar{h}}_{\textsl{gap}}^{i}=\min_{t\neq\texttt{opt}_{i}}~{}(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it})^{\top}\bm{W}^{\textsl{mm}}{\bm{z}}_{i}\geq
1,\\\
&{\gamma}_{\textsl{gap}}^{i}=\min_{t\neq\texttt{opt}_{i}}~{}\bm{\gamma}_{i\texttt{opt}_{i}}-\bm{\gamma}_{it}>0,\\\
&\bm{s}_{i\texttt{opt}_{i}}(1-\bm{s}_{i\texttt{opt}_{i}})>0.\end{split}$

Hence,

$c_{0}:=\min_{i\in[n]}\left\\{\left(\min_{t\neq\texttt{opt}_{i}}~{}(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it})^{\top}\bm{W}^{\textsl{mm}}{\bm{z}}_{i}\right)\cdot\left(\min_{t\neq\texttt{opt}_{i}}~{}\bm{\gamma}_{i\texttt{opt}_{i}}-\bm{\gamma}_{it}\right)\cdot\bm{s}_{i\texttt{opt}_{i}}(1-\bm{s}_{i\texttt{opt}_{i}})\right\\}>0.$
(27)

It follows from (26) and (27) that

$\min_{i\in[n]}\left\\{\bm{\bar{h}}^{\top}_{i}\text{diag}(\bm{s}_{i})\bm{\gamma}_{i}-\bm{\bar{h}}^{\top}_{i}\bm{s}_{i}\bm{s}^{\top}_{i}\bm{\gamma}_{i}\right\\}\geq
c_{0}>0.$ (28)

Further, by our assumption $\ell^{\prime}_{i}<0$. Since by Assumption A,
$\ell^{\prime}$ is continuous and the domain is bounded, the maximum is
attained and negative, and thus

$-c_{1}=\max_{x}\ell^{\prime}(x),\qquad\textnormal{for some}\quad c_{1}>0.$
(29)

Hence, using (28) and (29) in (24), we obtain

$\begin{split}\left<\nabla{\cal{L}}(\bm{W}),\bm{W}^{\textsl{mm}}\right>\leq-c<0,\quad\textnormal{where}\quad
c=c_{1}\cdot c_{0}.\end{split}$ (30)

In the scenario that Assumption B.1 holds (all tokens are support),
$\bm{\bar{h}}_{t}=\bm{x}_{it}^{\top}\bm{W}^{\textsl{mm}}{\bm{z}}_{i}$ is
constant for all $t\geq 2$. Hence, following similar steps as in (25)
completes the proof.

####  B.1.1 Proof of Theorem 3

It follows from Lemma 8 that under Assumption A, $\eta\leq 2/L_{\bm{W}}$, and
for any initialization $\bm{W}(0)$, the gradient descent sequence
$\bm{W}(k+1)=\bm{W}(k)-\eta\nabla\mathcal{L}(\bm{W}(k))$ satisfies
$\lim_{k\rightarrow\infty}\|{\nabla\mathcal{L}\left(\bm{W}\left(k\right)\right)}\|_{F}^{2}=0$.

Further, it follows from Lemma 9 that
$\left<\nabla{\cal{L}}(\bm{W}),\bm{W}^{\textsl{mm}}\right><0$ for all
$\bm{W}\in\mathbb{R}^{d\times d}$. Hence, for any finite $\bm{W}$,
$\left<\nabla{\cal{L}}(\bm{W}),\bm{W}^{\textsl{mm}}\right>$ cannot be equal to
zero. Therefore, there are no finite critical points $\bm{W}$, for which
$\nabla\mathcal{L}(\bm{W})=0$ which contradicts Lemma 8. This implies that
$\left\|\bm{W}\left(k\right)\right\|\rightarrow\infty$. $\hfill\blacksquare$

###  B.2  Global Convergence: Equal Score

The following lemma illustrates that when non-optimal tokens within an input
share the same scores, the negative gradient of the loss function at $\bm{W}$
becomes more correlated with the max-margin solution ($\bm{W}^{\textsl{mm}}$)
than with $\bm{W}$ itself.

###### Lemma 10

Let $\bm{W}^{\textsl{mm}}$ be the SVM solution of (Att-SVM). Suppose
Assummption B.2 on the tokens’ score hold and $\ell(\cdot)$ is strictly
decreasing and differentiable. For any choice of $\pi>0$, there exists
$R:=R_{\pi}$ such that, for any $\bm{W}$ with $\|{\bm{W}}\|_{F}\geq R$, we
have

$\left<\nabla{\cal{L}}(\bm{W}),\frac{\bm{W}}{\|{\bm{W}}\|_{F}}\right>\geq(1+\pi)\left<\nabla{\cal{L}}(\bm{W}),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right>.$

Proof. Let
${\bm{\bar{W}}}=\|{\bm{W}^{\textsl{mm}}}\|_{F}\bm{W}/\|{\bm{W}}\|_{F}$,
$M=\sup_{i,t}\|{\bm{x}_{it}{\bm{z}}_{i}^{\top}}\|$,
$\Theta=1/\|{\bm{W}^{\textsl{mm}}}\|_{F}$,
$\bm{s}_{i}=\mathbb{S}({\bm{X}}_{i}\bm{W}{\bm{z}}_{i})$,
$\bm{h}_{i}={\bm{X}}_{i}{\bm{\bar{W}}}{\bm{z}}_{i}$,
$\bm{\bar{h}}_{i}={\bm{X}}_{i}\bm{W}^{\textsl{mm}}{\bm{z}}_{i}$, and
$\bm{\gamma}_{i}=\bm{\gamma}_{i,t\geq 2}$. Without losing generality assume
$\alpha_{i}=\texttt{opt}_{i}=1$ for all $i\in[n]$. Repeating the proof of
Lemma 9 yields

$\displaystyle\left<\nabla{\cal{L}}(\bm{W}),\bm{W}^{\textsl{mm}}\right>$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}_{i}\cdot(\bm{\gamma}_{i1}-\bm{\gamma}_{i})(1-\bm{s}_{i1})\bm{s}_{i1}\left[\bm{\bar{h}}_{i1}-\frac{\sum_{t\geq
2}^{T}\bm{\bar{h}}_{it}\bm{s}_{it}}{\sum_{t\geq 2}^{T}\bm{s}_{it}}\right],$
$\displaystyle\left<\nabla{\cal{L}}(\bm{W}),{\bm{\bar{W}}}\right>$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}_{i}\cdot(\bm{\gamma}_{i1}-\bm{\gamma}_{i})(1-\bm{s}_{i1})\bm{s}_{i1}\left[\bm{h}_{i1}-\frac{\sum_{t\geq
2}^{T}\bm{h}_{it}\bm{s}_{it}}{\sum_{t\geq 2}^{T}\bm{s}_{it}}\right].$

Focusing on a single example $i\in[n]$ with $\bm{s},\bm{h},\bm{\bar{h}}$
vectors (dropping subscript $i$), given $\pi$, for sufficiently large $R$, we
wish to show that

$\displaystyle\left[\bm{h}_{1}-\frac{\sum_{t\geq
2}^{T}\bm{h}_{t}\bm{s}_{t}}{\sum_{t\geq
2}^{T}\bm{s}_{t}}\right]\leq(1+\pi)\cdot\left[\bm{\bar{h}}_{1}-\frac{\sum_{t\geq
2}^{T}\bm{\bar{h}}_{t}\bm{s}_{t}}{\sum_{t\geq 2}^{T}\bm{s}_{t}}\right].$ (31)

We consider two scenarios.

Scenario 1:
$\|{{\bm{\bar{W}}}-\bm{W}^{\textsl{mm}}}\|_{F}\leq\varepsilon:=\pi/(2M)$. In
this scenario, for any token, we find that

$|\bm{h}_{t}-\bm{\bar{h}}_{t}|=|\bm{x}_{t}^{\top}({\bm{\bar{W}}}-\bm{W}^{\textsl{mm}}){\bm{z}}_{t}|\leq
M\|{{\bm{\bar{W}}}-\bm{W}^{\textsl{mm}}}\|_{F}\leq M\varepsilon.$

Consequently, we obtain

$\bm{\bar{h}}_{1}-\frac{\sum_{t\geq
2}^{T}\bm{\bar{h}}_{t}\bm{s}_{t}}{\sum_{t\geq
2}^{T}\bm{s}_{t}}\geq\bm{h}_{1}-\frac{\sum_{t\geq
2}^{T}\bm{h}_{t}\bm{s}_{t}}{\sum_{t\geq
2}^{T}\bm{s}_{t}}-2M\varepsilon=\bm{h}_{1}-\frac{\sum_{t\geq
2}^{T}\bm{h}_{t}\bm{s}_{t}}{\sum_{t\geq 2}^{T}\bm{s}_{t}}-\pi.$

Also noticing $\bm{\bar{h}}_{1}-\frac{\sum_{t\geq
2}^{T}\bm{\bar{h}}_{t}\bm{s}_{t}}{\sum_{t\geq 2}^{T}\bm{s}_{t}}\geq 1$ (thanks
to $\bm{W}^{\textsl{mm}}$ satisfying $\geq 1$ margin), this implies (31).

Scenario 2:
$\|{{\bm{\bar{W}}}-\bm{W}^{\textsl{mm}}}\|_{F}\geq\varepsilon:=\pi/(2M)$. In
this scenario, for some $\delta=\delta(\varepsilon)$ and $\tau\geq 2$, we have
that

$\displaystyle\bm{h}_{1}-\bm{h}_{\tau}\leq 1-2\delta.$

Recall that $\bm{s}=\mathbb{S}(\bar{R}\bm{h})$ where
$\bar{R}=\|{\bm{W}}\|_{F}/\|{\bm{W}^{\textsl{mm}}}\|_{F}$. To proceed, split
the tokens into two groups: Let ${\cal{N}}$ be the group of tokens obeying
$(\bm{x}_{1}-\bm{x}_{t})^{\top}{\bm{\bar{W}}}{\bm{z}}\geq 1-\delta$ for
$t\in{\cal{N}}$ and $[T]-{\cal{N}}$ be the rest. Observe that

$\frac{\sum_{t\in{\cal{N}}}\bm{s}_{t}}{\sum_{t\geq
2}^{T}\bm{s}_{t}}\leq\frac{\sum_{t\in{\cal{N}}}\bm{s}_{t}}{\bm{s}_{\tau}}\leq
T\frac{e^{\delta\bar{R}}}{e^{2\delta\bar{R}}}=Te^{-\bar{R}\delta}.$

Set $\bar{M}=M/\Theta$ and note that
$\|{\bm{h}_{t}}\|\leq\|{\bm{W}^{\textsl{mm}}}\|_{F}\cdot\|{\bm{x}_{t}{\bm{z}}^{\top}}\|\leq\bar{M}$.
Using $(\bm{x}_{1}-\bm{x}_{t})^{\top}{\bm{\bar{W}}}{\bm{z}}<1-\delta$ over
$t\in[T]-{\cal{N}}$ and plugging in the above bound, we obtain

$\displaystyle\frac{\sum_{t\geq
2}^{T}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}}{\sum_{t\geq 2}^{T}\bm{s}_{t}}$
$\displaystyle=\frac{\sum_{t\in[T]-{\cal{N}}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}}{\sum_{t\geq
2}^{T}\bm{s}_{t}}+\frac{\sum_{t\in{\cal{N}}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}}{\sum_{t\geq
2}^{T}\bm{s}_{t}}$ $\displaystyle\leq(1-\delta)+2\bar{M}Te^{-\bar{R}\delta}.$

Using the fact that $\bm{\bar{h}}_{1}-\frac{\sum_{t\geq
2}^{T}\bm{\bar{h}}_{t}\bm{s}_{t}}{\sum_{t\geq 2}^{T}\bm{s}_{t}}\geq 1$, the
above implies (31) with $\pi^{\prime}=2\bar{M}Te^{-\bar{R}\delta}-\delta$. To
proceed, choose

$R_{\pi}=\delta^{-1}\Theta^{-1}\log(2\bar{M}T/\pi)\quad\textnormal{to
ensure}\quad\pi^{\prime}\leq\pi.$ (32)

####  B.2.1 Proof of Theorem 4

Given any $\epsilon\in(0,1)$, let $\pi=\epsilon/(1-\epsilon)$. It follows from
Theorem 3 that $\lim_{k\to\infty}\|{\bm{W}(k)}\|_{F}=\infty$. Hence, we can
choose $k_{0}$ such that for any $k\geq k_{0}$, it holds that
$\|{\bm{W}(k)}\|_{F}>R_{\epsilon}\vee 1/2$ for some parameter $R_{\epsilon}$.
Now for any $k\geq k_{0}$, it follows from Lemma 10 that

$\displaystyle\left\langle-\nabla\mathcal{L}(\bm{W}(k)),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle\geq(1-\epsilon)\left\langle-\nabla\mathcal{L}(\bm{W}(k)),\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}}\right\rangle.$

Multiplying both sides by the stepsize $\eta$ and using the gradient descent
update, we get

$\begin{split}\left\langle\bm{W}(k+1)-\bm{W}(k),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle&\geq(1-\epsilon)\left\langle\bm{W}(k+1)-\bm{W}(k),\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}}\right\rangle\\\
&=\frac{(1-\epsilon)}{2\|{\bm{W}(k)}\|_{F}}\left(\|{\bm{W}(k+1)}\|_{F}^{2}-\|{\bm{W}(k)}\|_{F}^{2}-\|{\bm{W}(k+1)-\bm{W}(k)}\|_{F}^{2}\right)\\\
&\geq(1-\epsilon)\left(\frac{1}{2\|{\bm{W}(k)}\|_{F}}\left(\|{\bm{W}(k+1)}\|_{F}^{2}-\|{\bm{W}(k)}\|_{F}^{2}\right)-\|{\bm{W}(k+1)-\bm{W}(k)}\|_{F}^{2}\right)\\\
&\geq(1-\epsilon)\left(\|{\bm{W}(k+1)}\|_{F}-\|{\bm{W}(k)}\|_{F}-\|{\bm{W}(k+1)-\bm{W}(k)}\|_{F}^{2}\right)\\\
&\geq(1-\epsilon)\Big{(}\|{\bm{W}(k+1)}\|_{F}-\|{\bm{W}(k)}\|_{F}-2\eta\left(\mathcal{L}(\bm{W}(k))-\mathcal{L}(\bm{W}(k+1))\right)\Big{)}.\end{split}$
(33)

Here, the second inequality is obtained from $\|{\bm{W}(k)}\|_{F}\geq 1/2$;
the third inequality follows since for any $a,b>0$, we have
$(a^{2}-b^{2})/(2b)-(a-b)\geq 0$; and the last inequality uses Lemma 8.

Summing the above inequality over $k\geq k_{0}$ gives

$\displaystyle\left\langle\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle\geq
1-\epsilon+\frac{C(\epsilon,\eta)}{\|{\bm{W}(k)}\|_{F}},$

for some finite constant $C(\epsilon,\eta)$ defined as

$C(\epsilon,\eta):=\left\langle\bm{W}(k_{0}),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle-(1-\epsilon)\|{\bm{W}(k_{0})}\|_{F}-2\eta(1-\epsilon)(\mathcal{L}(\bm{W}(k_{0}))-\mathcal{L}_{\star}),$
(34)

where $\mathcal{L}_{\star}\leq\mathcal{L}\left(\bm{W}\left(k\right)\right)$
for all $k\geq 0$.

Since $\left\|\bm{W}\left(k\right)\right\|\rightarrow\infty$, we get

$\displaystyle\liminf_{k\to\infty}\left\langle\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle\geq
1-\epsilon.$ (35)

Given that $\epsilon$ is arbitrary, we can consider the limit as $\epsilon$
approaches zero. Thus,
$\bm{W}(k)/\|{\bm{W}(k)}\|_{F}\to\bm{W}^{\textsl{mm}}/\|{\bm{W}^{\textsl{mm}}}\|_{F}$.
$\hfill\blacksquare$

###  B.3 Global Convergence: Non-Equal Scores

####  B.3.1 Proof of Theorem 5

We provide the proof in six steps as shown below.
Step 1: Defining the original and equally-scored problems.  Given the original
dataset $({\bm{X}}_{i},{\bm{z}}_{i},Y_{i})$ with scores $\bm{\gamma}_{it}$,
define an approximate dataset
$(\tilde{{\bm{X}}}_{i},\tilde{{\bm{z}}}_{i},Y_{i})$ as follows. Without loss
of generality assume $\texttt{opt}_{i}=1,\forall i\in[n]$ and let
${\bf{P}}^{\perp}_{\bm{v}}$ denote the projection onto the subspace orthogonal
to head $\bm{v}$. For all $i\in[n]$, set $\tilde{\bm{x}}_{i1}=\bm{x}_{i1}$ and
$\tilde{\bm{x}}_{it}=\big{(}\frac{\bm{v}^{\top}\bm{x}_{i2}}{\|\bm{v}\|_{2}^{2}}\big{)}\,\bm{v}+{{\bf{P}}^{\perp}_{\bm{v}}}\bm{x}_{it}$.
Note this construction guarantees for the new dataset that
$\tilde{\texttt{opt}}_{i}=1$ and equal scores
$\tilde{\bm{\gamma}}_{it}=\bm{\gamma}_{i2}$ for all $t>1$. Additionally, for
Cross-Attention set $\tilde{{\bm{z}}}_{i}={\bm{z}}_{i}$ and for Self-Attention
set $\tilde{{\bm{z}}}_{i}=\tilde{\bm{x}}_{i,\tau}$ for query corresponding to
$\tau$-th token. Following this setting, we define ${\cal{L}}(\bm{W})$ and
$\tilde{{\cal{L}}}(\tilde{\bm{W}})$ as the ERM objectives of the original and
equally-scored problems, respectively, as follows:

$\displaystyle{\cal{L}}(\bm{W})$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\ell\left(Y_{i}\cdot\bm{v}^{\top}{\bm{X}}_{i}^{\top}\mathbb{S}({\bm{X}}_{i}\bm{W}{\bm{z}}_{i})\right),$
(36a) $\displaystyle\tilde{{\cal{L}}}(\tilde{\bm{W}})$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\ell\left(Y_{i}\cdot\bm{v}^{\top}\tilde{{\bm{X}}}_{i}^{\top}\mathbb{S}(\tilde{{\bm{X}}}_{i}\tilde{\bm{W}}\tilde{{\bm{z}}}_{i})\right).$
(36b)

Let $\bm{W}^{\textsl{mm}}$ and $\tilde{\bm{W}}^{\textsl{mm}}$ denote the
solutions of (Att-SVM) for the original and equally-scored SVM problems,
respectively:

$\displaystyle\bm{W}^{\textsl{mm}}$
$\displaystyle=\arg\min_{\bm{W}}\|{\bm{W}}\|_{F}\quad\text{subj.
to}\quad(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it})^{\top}\bm{W}{\bm{z}}_{i}\geq
1\quad\text{for all}\quad t\neq\texttt{opt}_{i},\quad i\in[n],$ (37a)
$\displaystyle\tilde{\bm{W}}^{\textsl{mm}}$
$\displaystyle=\arg\min_{\tilde{\bm{W}}}\|{\tilde{\bm{W}}}\|_{F}\quad\text{subj.
to}\quad(\bm{x}_{i\texttt{opt}_{i}}-\tilde{\bm{x}}_{it})^{\top}\tilde{\bm{W}}\tilde{{\bm{z}}}_{i}\geq
1\quad\text{for all}\quad t\neq\texttt{opt}_{i},\quad i\in[n].$ (37b)

Note that there exists $\Delta_{0}(\delta)$ that depends solely on the
original problem and $\delta$, and can be made arbitrarily small by decreasing
$\delta$, such that

$\|\tilde{\bm{W}}^{\textsl{mm}}-\bm{W}^{\textsl{mm}}\|_{F}/\|\tilde{\bm{W}}^{\textsl{mm}}\|_{F}\leq\Delta_{0}(\delta).$
(38)

To proceed, let $\bm{\gamma}_{i}=Y_{i}\cdot{\bm{X}}_{i}\bm{v}$,
$\tilde{\bm{\gamma}}_{i}=Y_{i}\cdot\tilde{\bm{X}}_{i}\bm{v}$,
$\bm{h}_{i}={\bm{X}}_{i}\bm{W}{\bm{z}}_{i}$ and
$\tilde{\bm{h}}_{i}=\tilde{{\bm{X}}}_{i}\tilde{\bm{W}}\tilde{{\bm{z}}}_{i}$.
Following Lemma 4 and under Assumption A, we have

$\displaystyle\left\|\nabla\mathcal{L}(\bm{W})-\nabla\tilde{\mathcal{L}}(\tilde{\bm{W}})\right\|_{F}$
$\displaystyle\leq\frac{1}{n}\sum_{i=1}^{n}\left\|\ell^{\prime}\left(\bm{\gamma}_{i}^{\top}\mathbb{S}(\bm{h}_{i})\right)\cdot{\bm{z}}_{i}\bm{\gamma}_{i}^{\top}\mathbb{S}^{\prime}(\bm{h}_{i}){\bm{X}}_{i}-\ell^{\prime}\left(\tilde{\bm{\gamma}}_{i}^{\top}\mathbb{S}(\tilde{\bm{h}}_{i})\right)\cdot\tilde{{\bm{z}}}_{i}\tilde{\bm{\gamma}}_{i}^{\top}\mathbb{S}^{\prime}(\tilde{\bm{h}}_{i})\tilde{{\bm{X}}}_{i}\right\|_{F}$
$\displaystyle\leq\frac{1}{n}\sum_{i=1}^{n}M_{0}~{}\left\|\tilde{{\bm{z}}}_{i}\tilde{\bm{\gamma}}_{i}^{\top}\mathbb{S}^{\prime}(\tilde{\bm{h}}_{i})\tilde{{\bm{X}}}_{i}\right\|_{F}\left\|\bm{\gamma}_{i}^{\top}\mathbb{S}(\bm{h}_{i})-\tilde{\bm{\gamma}}_{i}^{\top}\mathbb{S}(\tilde{\bm{h}}_{i})\right\|$
(39)
$\displaystyle+\frac{1}{n}\sum_{i=1}^{n}M_{1}~{}\left\|{\bm{z}}_{i}\bm{\gamma}_{i}^{\top}\mathbb{S}^{\prime}(\bm{h}_{i}){\bm{X}}_{i}-\tilde{{\bm{z}}}_{i}\tilde{\bm{\gamma}}_{i}^{\top}\mathbb{S}^{\prime}(\tilde{\bm{h}}_{i})\tilde{{\bm{X}}}_{i}\right\|_{F}.$

Step 2: Monitoring the fluctuations during iterations until $\bm{W}(k)$ enters
the local cone around $\bm{W}^{\textsl{mm}}$. Fix
$\bm{W}(0)=\tilde{\bm{W}}(0)$. Algorithm W-GD applied to ${\cal{L}}(\bm{W})$
and $\tilde{{\cal{L}}}(\tilde{\bm{W}})$ defined in (36) implies that

$\displaystyle\tilde{\bm{W}}(k+1)$
$\displaystyle=\tilde{\bm{W}}(k)-\eta\nabla\tilde{{\cal{L}}}(\tilde{\bm{W}}(k)),$
(40a) $\displaystyle\bm{W}(k+1)$
$\displaystyle=\bm{W}(k)-\eta\nabla{\cal{L}}(\bm{W}(k)).$ (40b)

It follows from Theorem 7 that there exist parameters
$\mu=\mu(\texttt{opt})\in(0,1)$ and $R=R_{\mu}>0$ such that
$\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}})$ does not contain any stationary
points, where

$\displaystyle\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}})=\left\\{\bm{W}\in\mathbb{R}^{d\times
d}~{}:~{}\left\langle\frac{\bm{W}}{\|{\bm{W}}\|_{F}},\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle\geq
1-\mu,~{}\|{\bm{W}}\|_{F}\geq R\right\\}.$ (41)

We will prove the following claim.

###### Claim 1

There exists $k\geq 1$ such that
$\bm{W}(k)\in\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}})$ , where $\bm{W}(k)$ is
defined in (40b).

To proceed, choose $\tilde{\mu}\in(0,1)$ such that $\tilde{\mu}=\mu/2$. For
the equally-scored problem, Theorem 4 implies the existence of
$k=k_{\tilde{\mu}}$ such that if we run gradient descent in (40a) with any
$\eta\leq 1/L_{W}$ for $k$ iterations, then $\|{\tilde{\bm{W}}(k)}\|_{F}\geq
R_{\tilde{\mu}}$ for some $R_{\tilde{\mu}}>0$, and

$\left\langle\frac{\tilde{\bm{W}}(k)}{\|{\tilde{\bm{W}}(k)}\|_{F}},\frac{\tilde{\bm{W}}^{\textsl{mm}}}{\|{\tilde{\bm{W}}^{\textsl{mm}}}\|_{F}}\right\rangle\geq
1-\tilde{\mu}.$ (42)

Let
$\Delta_{1,\tau}(\delta):=\left\|\nabla\mathcal{L}(\bm{W}(\tau))-\nabla\tilde{\mathcal{L}}(\tilde{\bm{W}}(\tau))\right\|_{F}$.
Since $\bm{W}(0)=\tilde{\bm{W}}(0)$, it follows from (B.3.1) that there exists
$\Delta_{1}(\delta)$ that depends on $\tilde{\mu}=\mu/2$ and $\delta$ and can
be made arbitrarily small by decreasing $\delta$ such that

$\begin{split}\frac{\|{\tilde{\bm{W}}(k)-\bm{W}(k)}\|_{F}}{\|{\tilde{\bm{W}}(k)}\|_{F}}&\leq\frac{\eta}{R_{\tilde{\mu}}}\sum_{\tau=0}^{k-1}\left\|\nabla\mathcal{L}(\bm{W}(\tau))-\nabla\tilde{\mathcal{L}}(\tilde{\bm{W}}(\tau))\right\|_{F}\leq\frac{\eta}{R_{\tilde{\mu}}}\sum_{\tau=0}^{k-1}\Delta_{1,\tau}(\delta)\leq\eta\Delta_{1}(\delta).\end{split}$
(43)

Hence,

$\begin{split}\left\langle\frac{\bm{W}(k)}{\|{\tilde{\bm{W}}(k)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}}{\|{\tilde{\bm{W}}^{\textsl{mm}}}\|_{F}}\right\rangle&=\left\langle\frac{\tilde{\bm{W}}(k)}{\|{\tilde{\bm{W}}(k)}\|_{F}}+\frac{\bm{W}(k)-\tilde{\bm{W}}(k)}{\|{\tilde{\bm{W}}(k)}\|_{F}},\frac{\tilde{\bm{W}}^{\textsl{mm}}}{\|{\tilde{\bm{W}}^{\textsl{mm}}}\|_{F}}+\frac{\bm{W}^{\textsl{mm}}-\tilde{\bm{W}}^{\textsl{mm}}}{\|{\tilde{\bm{W}}^{\textsl{mm}}}\|_{F}}\right\rangle\\\
&=\left\langle\frac{\tilde{\bm{W}}(k)}{\|{\tilde{\bm{W}}(k)}\|_{F}},\frac{\tilde{\bm{W}}^{\textsl{mm}}}{\|{\tilde{\bm{W}}^{\textsl{mm}}}\|_{F}}\right\rangle+\left\langle\frac{\bm{W}(k)-\tilde{\bm{W}}(k)}{\|{\tilde{\bm{W}}(k)}\|_{F}},\frac{\tilde{\bm{W}}^{\textsl{mm}}}{\|{\tilde{\bm{W}}^{\textsl{mm}}}\|_{F}}\right\rangle\\\
&+\left\langle\frac{\tilde{\bm{W}}(k)}{\|{\tilde{\bm{W}}(k)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}-\tilde{\bm{W}}^{\textsl{mm}}}{\|{\tilde{\bm{W}}^{\textsl{mm}}}\|_{F}}\right\rangle+\left\langle\frac{\bm{W}(k)-\tilde{\bm{W}}(k)}{\|{\tilde{\bm{W}}(k)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}-\tilde{\bm{W}}^{\textsl{mm}}}{\|{\tilde{\bm{W}}^{\textsl{mm}}}\|_{F}}\right\rangle\\\
&\geq
1-\tilde{\mu}+\left(-\eta\Delta_{1}(\delta)-\Delta_{0}(\delta)-\eta\Delta_{1}(\delta)\Delta_{0}(\delta)\right)\\\
&\geq 1-2\tilde{\mu}+\tilde{\mu}-\Delta(\delta)\\\ &\geq
1-\mu+\tilde{\mu}-\Delta(\delta),\end{split}$ (44)

where

$\Delta(\delta):=\eta\Delta_{1}(\delta)+\Delta_{0}(\delta)+\eta\Delta_{1}(\delta)\Delta_{0}(\delta).$
(45)

Further,

$\begin{split}\frac{\|{\bm{W}(k)}\|_{F}}{\|{\tilde{\bm{W}}(k)}\|_{F}}&\leq
1+\frac{\|{\tilde{\bm{W}}(k)-\bm{W}(k)}\|_{F}}{\|{\tilde{\bm{W}}(k)}\|_{F}}\leq
1+\eta\Delta_{1}(\delta),\\\
\frac{\|{\bm{W}^{\textsl{mm}}}\|_{F}}{\|{\tilde{\bm{W}}^{mm}}\|_{F}}&\leq
1+\frac{\|{\bm{W}^{\textsl{mm}}-\tilde{\bm{W}}^{\textsl{mm}}}\|_{F}}{\|{\tilde{\bm{W}}^{\textsl{mm}}}\|_{F}}\leq
1+\Delta_{0}(\delta).\end{split}$ (46)

Now, it follows from (46) and (44) for $k=k_{\tilde{\mu}}$,

$\begin{split}\left\langle\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle&\geq\frac{1}{1+\eta\Delta_{1}(\delta)}\cdot\frac{1}{1+\Delta_{0}(\delta)}\left(1-\mu+{\tilde{\mu}}-\Delta(\delta)\right)\\\
&=\frac{1}{1+\Delta(\delta)}\Big{(}1-\mu+\tilde{\mu}-\Delta(\delta)\Big{)}\\\
&\geq
1-\mu+\frac{1}{1+\Delta(\delta)}\Big{(}\tilde{\mu}-(2-\mu)\Delta(\delta)\Big{)}\\\
&\geq 1-\mu.\end{split}$ (47)

Here, the last inequality is obtained by choosing some sufficiently small
$\delta>0$ such that $\tilde{\mu}-(2-\mu)\Delta(\delta)\geq 0$.

Hence, we have proven Claim 1 and we have shown that for $k=k_{\tilde{\mu}}$,
$\bm{W}(k)\in\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}})$, where $\bm{W}(k)$ is
defined in (40).
Steps 2-6: The remaining steps of the proof can be completed by utilizing
Theorem 7 for the original problem. In this approach, we use $\bm{W}(k)$ for
$k=k_{\tilde{\mu}}$ as the initialization within the cone.
$\hfill\blacksquare$

###  B.4 Global Convergence: Zero Initialization

To establish a foundation for providing the convergence of GD to the globally
optimal solution $\bm{W}^{\textsl{mm}}$, we present the following definitions.
Given $\mu\geq 0$, consider the following subset of the sphere and its
associated cone:

$\displaystyle\mathcal{S}_{\mu}:=\left\\{\bm{W}\in\mathbb{R}^{d\times
d}~{}\Big{|}~{}\left<(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it}){\bm{z}}_{i}^{\top},\frac{\bm{W}}{\|{\bm{W}}\|_{F}}\right>\geq\frac{\mu}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\quad\textnormal{for
all}\quad t\neq\texttt{opt}_{i},\quad i\in[n]\right\\},$ (48a)
$\displaystyle\bar{\mathcal{C}}_{\mu,R}:=\left\\{\bm{W}\in\mathcal{S}_{\mu}~{}\Big{|}~{}\|{\bm{W}}\|_{F}\geq
R\right\\}.$ (48b)

The following lemma shows that as $\pi$ approaches zero, the negative gradient
of the loss function at
$\bm{W}\in\bar{\mathcal{C}}_{\mu,R}(\bm{W}^{\textsl{mm}})$ becomes more
correlated with the max-margin solution ($\bm{W}^{\textsl{mm}}$) than with
$\bm{W}$ itself.

###### Lemma 11

Suppose Assumption A holds and let $\texttt{opt}=(\texttt{opt}_{i})_{i=1}^{n}$
be the unique optimal tokens with $\bm{W}^{\textsl{mm}}$ denoting the SVM
solution. Fix any $\mu>0$ (per Lemma 12). For any choice of $\pi>0$, there
exists $R:=R_{\pi,\mu}$ such that, for any
$\bm{W}\in\bar{\mathcal{C}}_{\mu,R}(\bm{W}^{\textsl{mm}})$, we have

$\left<\nabla{\cal{L}}(\bm{W}),\frac{\bm{W}}{\|{\bm{W}}\|_{F}}\right>\geq(1+\pi)\left<\nabla{\cal{L}}(\bm{W}),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right>.$

Here, $\bar{\mathcal{C}}_{\mu,R}(\bm{W}^{\textsl{mm}})$ is the cone defined at

Proof. Let
${\bm{\bar{W}}}=\|{\bm{W}}\|_{F}\bm{W}/\|{\bm{W}^{\textsl{mm}}}\|_{F}$,
$\bm{h}_{i}={\bm{X}}_{i}{\bm{\bar{W}}}{\bm{z}}_{i}$,
$\bm{\bar{h}}_{i}={\bm{X}}_{i}\bm{W}^{\textsl{mm}}{\bm{z}}_{i}$, and
$\bm{s}_{i}=\mathbb{S}({\bm{X}}_{i}\bm{W}{\bm{z}}_{i})$. To establish the
result, we will prove that, for sufficiently large $R=R_{\pi}$, for any
$\bm{W}\in\bar{\mathcal{C}}_{\mu,R}(\bm{W}^{\textsl{mm}})$ and any $i\in[n]$,

$\displaystyle\left<\bm{h}_{i},\mathbb{S}^{\prime}({\bm{X}}_{i}\bm{W}{\bm{z}}_{i})\bm{\gamma}_{i}\right>\leq(1+\pi)\left<\bm{\bar{h}}_{i},\mathbb{S}^{\prime}({\bm{X}}_{i}\bm{W}{\bm{z}}_{i})\bm{\gamma}_{i}\right>.$
(49)

Once (49) holds for all $i$, the same conclusion will hold for the gradient
correlations via (55). Moving forward, we shall again focus on a single point
$i\in[n]$ and drop all subscripts $i$ i.e. we will use vectors
$\bm{h},\bm{\bar{h}},\bm{s}$. Also assume $\texttt{opt}=\texttt{opt}_{i}=1$
without losing generality as above. Following (57), for all
$\bm{V}\in\mathcal{S}_{\mu}$ with
$\|{\bm{V}}\|_{F}=\|{\bm{W}^{\textsl{mm}}}\|_{F}$ and
$\tilde{\bm{h}}={\bm{X}}\bm{V}{\bm{z}}$, we have found

$\displaystyle\big{|}\tilde{\bm{h}}^{\top}\text{diag}(\bm{s})\bm{\gamma}-\tilde{\bm{h}}^{\top}\bm{s}\bm{s}^{\top}\bm{\gamma}-\sum_{t\neq\texttt{opt}}(\tilde{\bm{h}}_{1}-\tilde{\bm{h}}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\big{|}\leq
2\Gamma A(1-\bm{s}_{1})^{2}.$ (50)

Recalling $\bm{\bar{h}}_{1}-\bm{\bar{h}}_{t}\geq 1$, we note that
$\sum_{t\neq\texttt{opt}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\leq\sum_{t\neq\texttt{opt}}(\bm{\bar{h}}_{1}-\bm{\bar{h}}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$.
Now plugging in $\bm{h},\bm{\bar{h}}$ in the bound above and assuming $\pi\leq
1$ (w.l.o.g.), (49) is implied by the following stronger inequality

$\displaystyle 6\Gamma
A(1-\bm{s}_{1})^{2}+\sum_{t\neq\texttt{opt}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$
$\displaystyle\leq(1+\pi)\sum_{t\neq\texttt{opt}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$
$\displaystyle\leq(1+\pi)\sum_{t\neq\texttt{opt}}(\bm{\bar{h}}_{1}-\bm{\bar{h}}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t}).$

First, we claim that
$0.5\pi\sum_{t\neq\texttt{opt}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\geq
6\Gamma A(1-\bm{s}_{1})^{2}$. The proof of this claim directly follows the
earlier argument, namely following (58), (59), (56) we have that
$1-\bm{s}_{1}\leq Te^{-R\mu\Theta}$ and
$\bm{\gamma}_{1}-\bm{\gamma}_{t}\geq\gamma^{\textsl{gap}}_{\min}$. This leads
to the choice (for $C=12$)

$\displaystyle R_{\pi}=\frac{1}{\mu\Theta}\log\big{(}\frac{C\cdot T\Gamma
A}{\pi\gamma^{\textsl{gap}}_{\min}}\big{)}.$ (51)

Following this control over the perturbation term $6\Gamma
A(1-\bm{s}_{1})^{2}$, to conclude with the result, what remains is
establishing the comparison

$\displaystyle\sum_{t\neq\texttt{opt}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\leq(1+0.5\pi)\sum_{t\neq\texttt{opt}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t}).$
(52)

To proceed, we split the problem into two scenarios.

Scenario 1:
$\|{{\bm{\bar{W}}}-\bm{W}^{\textsl{mm}}}\|_{F}\leq\varepsilon=\frac{\pi}{4A\Theta}$
for some $\varepsilon>0$. In this scenario, for any token, we find that

$|\bm{h}_{t}-\bm{\bar{h}}_{t}|\leq A\Theta\varepsilon=\pi/4.$

Consequently, we obtain

$\bm{h}_{1}-\bm{h}_{t}\leq\bm{\bar{h}}_{1}-\bm{\bar{h}}_{t}+2A\Theta\varepsilon=1+0.5\pi.$

Since $\bm{\gamma}_{1}-\bm{\gamma}_{t}\geq 0$, for all $t\neq 1$, we find
$(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\leq(1+0.5\pi)\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$.
This implies the desired result (52).

Scenario 2:
$\|{{\bm{\bar{W}}}-\bm{W}^{\textsl{mm}}}\|_{F}\geq\varepsilon=\frac{\pi}{4A\Theta}$.
Since ${\bm{\bar{W}}}$ is not the max-margin solution, in this scenario, for
some $\nu=\nu(\varepsilon)>0$ and $\tau\neq 1$, we have that
$\bm{h}_{1}-\bm{h}_{\tau}\leq 1-2\nu$. Here $\tau=\arg\max_{\tau\neq
1}\bm{x}_{\tau}{\bm{\bar{W}}}{\bm{z}}$ denotes the nearest point to
$\bm{h}_{1}$ (along the ${\bm{\bar{W}}}$ direction). Recall that
$\bm{s}=\mathbb{S}(C^{\textsl{mm}}\bm{h})$, where
$C^{\textsl{mm}}=\|{\bm{W}}\|_{F}/\|{\bm{W}^{\textsl{mm}}}\|_{F}$. To proceed,
split the tokens into two groups: Let ${\cal{N}}$ be the group of tokens
obeying $(\bm{x}_{1}-\bm{x}_{\tau})\bm{W}{\bm{z}}\leq 1-\nu$ and
$[T]-\\{1\\}-{\cal{N}}$ be the rest of the non-optimal tokens. Observe that

$\frac{\sum_{t\not\in\\{\texttt{opt}\\}\cup{\cal{N}}}\bm{s}_{t}}{\sum_{t\neq\texttt{opt}}\bm{s}_{t}}\leq\frac{\sum_{t\not\in\\{\texttt{opt}\\}\cup{\cal{N}}}\bm{s}_{t}}{\sum_{t=\tau}\bm{s}_{t}}\leq
T\frac{e^{\nu C^{\textsl{mm}}}}{e^{2\nu
C^{\textsl{mm}}}}=Te^{-C^{\textsl{mm}}\nu}.$

Thus, using $|\bm{h}_{1}-\bm{h}_{t}|\leq 2A$ and recalling the definition of
$\gamma^{\textsl{gap}}$, observe that

$\sum_{t\not\in\\{\texttt{opt}\\}\cup{\cal{N}}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\leq\frac{2\Gamma
ATe^{-C^{\textsl{mm}}\nu}}{\gamma^{\textsl{gap}}}\sum_{t\neq\texttt{opt}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t}).$

Plugging this into (52), we obtain

$\displaystyle\sum_{t\neq\texttt{opt}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$
$\displaystyle=\sum_{t\in{\cal{N}}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})+\sum_{t\not\in\\{\texttt{opt}\\}\cup{\cal{N}}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$
$\displaystyle\leq\sum_{t\in{\cal{N}}}(1-\nu)\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})+\sum_{t\not\in\\{\texttt{opt}\\}\cup{\cal{N}}}2A\Gamma
Te^{-C^{\textsl{mm}}\nu}$ $\displaystyle\leq\left(1-\nu+\frac{2\Gamma
ATe^{-C^{\textsl{mm}}\nu}}{\gamma^{\textsl{gap}}}\right)\sum_{t\neq\texttt{opt}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$
$\displaystyle\leq\left(1+\frac{2\Gamma
ATe^{-C^{\textsl{mm}}\nu}}{\gamma^{\textsl{gap}}}\right)\sum_{t\neq\texttt{opt}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t}).$

Consequently, the proof boils down to ensuring the perturbation term
$\frac{2\Gamma ATe^{-C^{\textsl{mm}}\nu}}{\gamma^{\textsl{gap}}}\leq 0.5\pi$.
Since $C^{\textsl{mm}}\geq R\Theta$, this is guaranteed for all inputs
$i\in[n]$ by recalling
$\gamma^{\textsl{gap}}_{\min}=\min_{i\in[n]}\gamma^{\textsl{gap}}_{i}$ and
choosing

$R\geq R_{\pi}=\frac{1}{\nu\Theta}\log(\frac{4\Gamma
AT}{\gamma^{\textsl{gap}}_{\min}\pi}),$

where $\nu=\nu(\frac{\pi}{4A\Theta})$ depends only on $\pi$ and global problem
variables.

Combining this with the prior $R_{\pi}$ choice of (51) (by taking their
maximum), we conclude with the statement.

###### Lemma 12

Suppose Assumption A holds and let $\texttt{opt}=(\texttt{opt}_{i})_{i=1}^{n}$
be the unique globally-optimal indices with $\bm{W}^{\textsl{mm}}$ denoting
the SVM solution. Define the margin $\Theta=1/\|{\bm{W}^{\textsl{mm}}}\|_{F}$.
For any $\mu>0$, there exists sufficiently large $R=R_{\mu}={\cal{O}}(1/\mu)$
(see (59)) such that:

  1. L1.

There is no stationary point within $\bar{\mathcal{C}}_{\mu,R}$, where
$\bar{\mathcal{C}}_{\mu,R}$ is defined in (48).

  2. L2.

Let $\bm{s}_{i}=\mathbb{S}({\bm{X}}_{i}\bm{W}{\bm{z}}_{i})$. For all
$\bm{V}\in\mathcal{S}_{\mu},\bm{W}\in\bar{\mathcal{C}}_{\mu,R}$, there exist
$C,c>0$ such that

$\displaystyle
C\cdot\max_{i\in[n]}\left(1-\bm{s}_{i\texttt{opt}_{i}}\right)\geq-\left<\nabla{\cal{L}}(\bm{W}),\bm{V}\right>\geq
c\cdot\mu\cdot\min_{i\in[n]}\left(1-\bm{s}_{i\texttt{opt}_{i}}\right).$

Proof. Let us introduce the norm upper bound

$\begin{split}&A=\max_{i\in[n],t,\tau\in[T]}\frac{(\|{\bm{x}_{it}}\|\vee\|{\bm{x}_{it}-\bm{x}_{i\tau}}\|)\cdot\|{{\bm{z}}_{i}}\|}{\Theta}.\end{split}$
(53)

The following inequalities hold for all
$\bm{V}\in\mathcal{S}_{\mu},~{}\|{\bm{V}}\|_{F}=\|{\bm{W}^{\textsl{mm}}}\|_{F}$
and all $i\in[n],t\neq\texttt{opt}_{i}$:

$\begin{split}A\geq(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it})^{\top}\bm{V}{\bm{z}}_{i}&\geq\mu.\end{split}$
(54)

To proceed, we write the gradient correlation following (13) and (24)

$\displaystyle\left<\nabla{\cal{L}}(\bm{W}),\bm{V}\right>$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}_{i}\cdot\left<\bm{h}_{i},\mathbb{S}^{\prime}(\tilde{\bm{h}}_{i})\bm{\gamma}_{i}\right>,$
(55)

where we denoted
$\ell^{\prime}_{i}=\ell^{\prime}(Y_{i}\cdot\bm{v}^{\top}{\bm{X}}_{i}^{\top}\mathbb{S}(\tilde{\bm{h}}_{i}))$,
$\bm{h}_{i}={\bm{X}}_{i}\bm{V}{\bm{z}}_{i}$,
$\tilde{\bm{h}}_{i}={\bm{X}}_{i}\bm{W}{\bm{z}}_{i}$,
$\bm{s}_{i}=\mathbb{S}(\tilde{\bm{h}}_{i})$. It follows from (53) that
$A\geq\max_{i\in[n],t\in[T]}\|{\bm{h}_{it}}\|$. Using (54), we can bound the
softmax probabilities $\bm{s}_{i}=\mathbb{S}(\tilde{\bm{h}}_{i})$ as follows,
for all $i\in[n]$:

$\displaystyle S_{i}:=\sum_{\tau\neq\texttt{opt}_{i}}\bm{s}_{i\tau}\leq
Te^{-R\mu\Theta}\bm{s}_{i\texttt{opt}_{i}}\leq Te^{-R\mu\Theta}.$ (56)

Recall scores $\bm{\gamma}_{it}=Y_{i}\cdot\bm{v}^{\top}\bm{x}_{it}$. Define
the score gaps:

$\gamma^{\textsl{gap}}_{i}=\bm{\gamma}_{i\texttt{opt}_{i}}-\max_{t\neq\texttt{opt}_{i}}\bm{\gamma}_{it},~{}~{}~{}\bar{\gamma}^{\textsl{gap}}_{i}=\bm{\gamma}_{i\texttt{opt}_{i}}-\min_{t\neq\texttt{opt}_{i}}\bm{\gamma}_{it},~{}~{}~{}\textnormal{and}~{}~{}~{}\Gamma=\sup_{i\in[n],t,\tau\in[T]}|\bm{\gamma}_{it}-\bm{\gamma}_{i\tau}|.$

Let us focus on a fixed datapoint $i\in[n]$, assume (without losing
generality) $\texttt{opt}_{i}=1$, and drop subscripts $i$. Directly applying
Lemma 7, we obtain

$\big{|}\bm{h}^{\top}\text{diag}(\bm{s})\bm{\gamma}-\bm{h}^{\top}\bm{s}\bm{s}^{\top}\bm{\gamma}-\sum_{t\geq
2}^{T}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\big{|}\leq
2\Gamma A(1-\bm{s}_{1})^{2}.$

To proceed, let us upper/lower bound the gradient correlation. Since
$A\geq\bm{h}_{1}-\bm{h}_{t}\geq\mu>0$ from (54), setting
$S:=\sum_{t\neq\texttt{opt}_{i}}\bm{s}_{t}=1-\bm{s}_{1}$, we find

$A\cdot
S\cdot\bar{\gamma}^{\textsl{gap}}\geq\sum_{t\neq\texttt{opt}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\geq\mu\cdot
S\cdot\gamma^{\textsl{gap}}.$ (57)

Next we show that $S=1-\bm{s}_{1}$ dominates $(1-\bm{s}_{1})^{2}=S^{2}$ for
large $R$. Specifically, we wish for

$\displaystyle\mu S\gamma^{\textsl{gap}}/2\geq 2\Gamma A(1-\bm{s}_{1})^{2}\iff
S\geq\frac{4}{\mu}\frac{\Gamma A}{\gamma^{\textsl{gap}}}S^{2}\iff
S\leq\frac{\mu\gamma^{\textsl{gap}}}{4\Gamma A}.$ (58)

Using (56), what we wish is ensured for all $i\in[n]$, by guaranteeing
$Te^{-R\mu\Theta}\leq\frac{\mu\gamma^{\textsl{gap}}}{4\Gamma A}$. That is, by
choosing

$\displaystyle R:=R_{\mu}=\frac{1}{\mu\Theta}\log\big{(}\frac{4T\Gamma
A}{\mu\gamma^{\textsl{gap}}_{\min}}\big{)},$ (59)

where $\gamma^{\textsl{gap}}_{\min}=\sup_{i\in[n]}\gamma^{\textsl{gap}}_{i}$
is the global scalar corresponding to the worst case score gap over all
inputs.

With the above choice of $R$, we guaranteed

$2A(1-\bm{s}_{1})\cdot\bar{\gamma}^{\textsl{gap}}\geq 2A\cdot
S\cdot\bar{\gamma}^{\textsl{gap}}\geq\sum_{t\neq\texttt{opt}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\geq\frac{\mu\cdot
S\cdot\gamma^{\textsl{gap}}}{2}\geq\frac{\mu(1-\bm{s}_{1})\gamma^{\textsl{gap}}}{2}.$

via (58) and (57). Since this holds over all inputs, going back to the
gradient correlation (55) and averaging above over all inputs $i\in[n]$ and
plugging back the indices $i$, we obtain the advertised bound by setting
$q_{i}=1-\bm{s}_{i\texttt{opt}_{i}}$ (where we have used $\texttt{opt}_{i}=1$
above, without losing generality)

$\displaystyle\frac{2A}{n}\sum_{i\in[n]}\ell^{\prime}_{i}\cdot
q_{i}\cdot\bar{\gamma}^{\textsl{gap}}_{i}\geq\left<\nabla{\cal{L}}(\bm{W}),\bm{V}\right>\geq\frac{\mu}{2n}\sum_{i\in[n]}\ell^{\prime}_{i}\cdot
q_{i}\cdot\gamma^{\textsl{gap}}_{i}.$ (60)

####  B.4.1 Proof of Theorem 6

We provide the proof in four steps:
Step 1: There are no stationary points within
$\bar{\mathcal{C}}_{\mu,R_{\mu}}(\bm{W}^{\textsl{mm}})$ for sufficiently large
$R_{\mu}$. This step directly follows from Lemma 12: for all
$\bm{V},\bm{W}\in\mathcal{S}_{\mu}(\bm{W}^{\textsl{mm}})$ with
$\|{\bm{V}}\|_{F}=\|{\bm{W}^{\textsl{mm}}}\|_{F}$ and $\|{\bm{W}}\|_{F}\geq
R_{\mu}$, it is guaranteed that there exists a value of $R_{\mu}$ such that
$\left\langle\bm{V},-\nabla{\cal{L}}(\bm{W})\right\rangle$ is strictly
positive.
Step 2: Consider an arbitrary value of $\epsilon\in(0,\mu/2)$ and let
$1/(1+\pi)=1-\epsilon$. It follows from Lemma 11 that, there exists
$R_{\epsilon}$ such that all
$\bm{W}\in\bar{\mathcal{C}}_{\mu,R_{\epsilon}}(\bm{W}^{\textsl{mm}})$ satisfy

$\displaystyle\left\langle-\nabla\mathcal{L}(\bm{W}),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle\geq(1-\epsilon)\left\langle-\nabla\mathcal{L}(\bm{W}),\frac{\bm{W}}{\|{\bm{W}}\|_{F}}\right\rangle.$
(61)

Step 3: Updates remain inside the cone
$\bar{\mathcal{C}}_{\mu,R_{\mu}}(\bm{W}^{\textsl{mm}})$. By leveraging the
results from Step 1 and Step 2, we demonstrate that the gradient iterates
remain within this cone. Note that by our assumption and for sufficiently
large $\eta_{0}$, we get
$\bm{W}(1)\in\bar{\mathcal{C}}_{\mu,R}(\bm{W}^{\textsl{mm}})$. Here, $R$ is
chosen such that

$\displaystyle R\geq R_{\mu}\vee R_{\epsilon}\vee 1/2.$ (62)

We provide the proof by induction. By the above argument,
$\bm{W}(1)\in\bar{\mathcal{C}}_{\mu,R}(\bm{W}^{\textsl{mm}})$. Let
$\bm{W}(k)\in\bar{\mathcal{C}}_{\mu,R}(\bm{W}^{\textsl{mm}})$. For all $k\geq
1$, we have that

$\displaystyle\min_{t\neq\texttt{opt}_{i},~{}i\in[n]}~{}~{}\left\langle(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it}){\bm{z}}_{i}^{\top},-\nabla\mathcal{L}(\bm{W}(k))\right\rangle\geq\bar{\rho},$
(63a)
$\displaystyle\frac{1}{1-\epsilon}\left\langle{\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}},-\nabla\mathcal{L}(\bm{W}(k))\right\rangle\geq\rho.$
(63b)

for some positive constants $\bar{\rho}$ and $\rho$.

To proceed, for all $t\neq\texttt{opt}_{i},~{}i\in[n]$, we obtain

$\begin{split}\left\langle(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it}){\bm{z}}_{i}^{\top},\frac{\bm{W}(k+1)}{\|{\bm{W}(k)}\|_{F}}\right\rangle&=\left\langle(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it}){\bm{z}}_{i}^{\top},\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}}-\frac{\eta}{\|{\bm{W}(k)}\|_{F}}\nabla\mathcal{L}(\bm{W}(k))\right\rangle\\\
&\geq\mu\Theta-\frac{\eta}{\|{\bm{W}(k)}\|_{F}}\left\langle(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it}){\bm{z}}_{i}^{\top},\nabla\mathcal{L}(\bm{W}(k))\right\rangle\\\
&\geq\mu\Theta+\frac{\eta\bar{\rho}}{\|{\bm{W}(k)}\|_{F}}.\end{split}$ (64)

Here, the first inequality follows from the induction assumption
$\bm{W}(k)\in\bar{\mathcal{C}}_{\mu,R}(\bm{W}^{\textsl{mm}})$, and the second
inequality uses (63a).

From Lemma 12, we have
$\left\langle\nabla{\cal{L}}(\bm{W}(k)),\bm{W}(k)\right\rangle<0$ which
implies that $\|{\bm{W}(k+1)}\|_{F}\geq\|{\bm{W}(k)}\|_{F}$. Hence,

$\displaystyle\|{\bm{W}(k+1)}\|_{F}$
$\displaystyle\leq\frac{1}{{2\|{\bm{W}(k)}\|_{F}}}\left(\|{\bm{W}(k+1)}\|_{F}^{2}+\|{\bm{W}(k)}\|_{F}^{2}\right)$
$\displaystyle=\frac{1}{2\|{\bm{W}(k)}\|_{F}}\left(2\|{W(k)}\|_{F}^{2}-2\eta\left\langle\nabla{\cal{L}}(\bm{W}(k)),\bm{W}(k)\right\rangle+\eta^{2}\|{\nabla{\cal{L}}(\bm{W}(k))}\|_{F}^{2}\right)$
$\displaystyle\leq\|{\bm{W}(k)}\|_{F}-\frac{\eta}{\|{\bm{W}(k)}\|_{F}}\left\langle\nabla{\cal{L}}(\bm{W}(k)),\bm{W}(k)\right\rangle+\eta^{2}\|\nabla{\cal{L}}(\bm{W}(k))\|^{2},$

where the last inequality follows from $\|{\bm{W}(k)}\|_{F}\geq R\geq 1/2$.

Thus,

$\begin{split}\frac{\|{\bm{W}(k+1)}\|_{F}}{\|{\bm{W}(k)}\|_{F}}&\leq
1-\eta\left\langle\nabla{\cal{L}}(\bm{W}(k)),\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}}\right\rangle+\eta^{2}\frac{\|\nabla\mathcal{L}(\bm{W}(k))\|^{2}}{\|{\bm{W}(k)}\|_{F}}\\\
&\leq
1-\frac{\eta}{1-\epsilon}\left\langle\nabla\mathcal{L}(\bm{W}(k)),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle+\eta^{2}\frac{\|\nabla\mathcal{L}(\bm{W}(k))\|^{2}}{\|{\bm{W}(k)}\|_{F}}\\\
&\leq
1+\frac{\eta\rho}{\|{\bm{W}(k)}\|_{F}}+\frac{\eta^{2}\|\nabla\mathcal{L}(\bm{W}(k))\|^{2}}{\|{\bm{W}(k)}\|_{F}}=:C_{1}(\rho,\eta).\end{split}$
(65)

Now, it follows from (64) and (65) that

$\begin{split}\min_{t\neq\texttt{opt}_{i},~{}i\in[n]}~{}~{}\left\langle(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it}){\bm{z}}_{i}^{\top},\frac{\bm{W}(k+1)}{\|{\bm{W}(k+1)}\|_{F}}\right\rangle&\geq\frac{1}{C_{1}(\bar{\rho},\eta)}\left(\mu\Theta+\frac{\eta\bar{\rho}}{\|{\bm{W}(k)}\|_{F}}\right)\\\
&\geq\mu\Theta+\frac{\eta}{C_{1}(\bar{\rho},\eta)}\left(\frac{\bar{\rho}-\rho\mu\Theta}{\|{\bm{W}(k)}\|_{F}}-\eta\mu\Theta\frac{\|\nabla\mathcal{L}(\bm{W}(k))\|^{2}}{\|{\bm{W}(k)}\|_{F}}\right)\\\
&\geq\mu\Theta.\end{split}$ (66)

Here, we leverage the fact that the parameter $\mu$ can be chosen arbitrarily,
and then we select the step size $\eta$ to satisfy the following condition:

$\eta\leq\frac{\bar{\rho}-\rho\mu\Theta}{\mu\Theta}\frac{1}{\|\nabla\mathcal{L}(\bm{W}(k))\|^{2}}.$

Hence, (66) gives $\bm{W}(k+1)\in\bar{\mathcal{C}}_{\mu,R}$.
Step 4: The correlation of $\bm{W}(k)$ and $\bm{W}^{\textsl{mm}}$ increases
over $k$. It follows from Theorem 3 that
$\left\|\bm{W}\left(k\right)\right\|\rightarrow\infty$. Hence, we can choose
$k_{0}$ such that for any $k\geq k_{0}$, it holds that $\|{\bm{W}(k)}\|_{F}>R$
for some $R\geq R_{\mu}\vee R_{\epsilon}\vee 1/2$. Now, following similar
steps in (33) and (34), for some constant $C(\epsilon,\eta)$, we obtain

$\displaystyle\left\langle\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle\geq
1-\epsilon+\frac{C(\epsilon,\eta)}{\|{\bm{W}(k)}\|_{F}}.$

Consequently,

$\displaystyle\liminf_{t\to\infty}\left\langle\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle\geq
1-\epsilon.$

Since $\epsilon\in(0,\mu/2)$ is arbitrary, we get
$\bm{W}(k)/\|{\bm{W}(k)}\|_{F}\to\bm{W}^{\textsl{mm}}/\|{\bm{W}^{\textsl{mm}}}\|_{F}$.
$\hfill\blacksquare$

##  Appendix C Local Convergence of Gradient Descent

###### Lemma 13

Suppose Assumption A on the loss function $\ell$ holds, and let
${\bm{\alpha}}=(\alpha_{i})_{i=1}^{n}$ be locally optimal tokens according to
Definition 2. Let $\bm{W}^{\textsl{mm}}=\bm{W}^{\textsl{mm}}({\bm{\alpha}})$
denote the SVM solution obtained via (Att-SVM) by applying the Frobenius norm
and replacing $(\texttt{opt}_{i})_{i=1}^{n}$ with
$\bm{\alpha}=(\alpha_{i})_{i=1}^{n}$. To provide a basis for discussing local
convergence gradient descent, we establish a cone centered around
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$ using the following construction. For
parameters $\mu\in(0,1)$ and $R>0$, let

$\displaystyle\texttt{cone}_{\mu}(\bm{W}^{\textsl{mm}})$
$\displaystyle=\left\\{\bm{W}\in\mathbb{R}^{d\times
d}~{}:~{}\left\langle\frac{\bm{W}}{\|{\bm{W}}\|_{F}},\frac{{\bm{W}^{\textsl{mm}}}}{\|{{\bm{W}^{\textsl{mm}}}}\|_{F}}\right\rangle\geq
1-\mu\right\\},$ $\displaystyle\mathcal{C}_{\mu,R}({\bm{W}^{\textsl{mm}}})$
$\displaystyle=\texttt{cone}_{\mu}(\bm{W}^{\textsl{mm}})\cap\left\\{\bm{W}\in\mathbb{R}^{d\times
d}~{}:~{}\|{\bm{W}}\|_{F}\geq R\right\\}.$

There exists a scalar $\mu=\mu({\bm{\alpha}})>0$ such that for sufficiently
large $R=R_{\mu}$:

  1. L1.

There is no stationary point within
$\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}})$.

  2. L2.

Let $\bm{s}_{i}=\mathbb{S}({\bm{X}}_{i}\bm{W}{\bm{z}}_{i})$. For all
$\bm{V},\bm{W}\in\texttt{cone}_{\mu}(\bm{W}^{\textsl{mm}})$ with
$\|{\bm{V}}\|_{F}=\|{\bm{W}^{\textsl{mm}}}\|_{F}$, there exist $C,c>0$ such
that

$\displaystyle
C\cdot\max_{i\in[n]}\left(1-\bm{s}_{i\alpha_{i}}\right)\geq-\left<\nabla{\cal{L}}(\bm{W}),\bm{V}\right>\geq
c\cdot\min_{i\in[n]}\left(1-\bm{s}_{i\alpha_{i}}\right).$

Proof. Let $(\mathcal{T}_{i})_{i=1}^{n}$ be the set of all support indices per
Definition 2. Let $\bar{\mathcal{T}}_{i}=[T]-\mathcal{T}_{i}-\\{\alpha_{i}\\}$
be the non-support indices. Let

$\begin{split}&\Theta=1/\|{\bm{W}^{\textsl{mm}}}\|_{F},\\\
&\delta=\frac{1}{2}\min_{i\in[n]}\min_{t\in\mathcal{T}_{i},\tau\in\bar{\mathcal{T}}_{i}}(\bm{x}_{it}-\bm{x}_{i\tau})^{\top}\bm{W}^{\textsl{mm}}{\bm{z}}_{i},\\\
&A=\max_{i\in[n],t\in[T]}\frac{\|{\bm{x}_{it}{\bm{z}}_{i}^{\top}}\|}{\Theta},\\\
&\mu=\mu(\delta)=\frac{1}{8}\left(\frac{\min(0.5,\delta)}{A}\right)^{2}.\end{split}$
(67)

Since $\bm{W}^{\textsl{mm}}$ is the max-margin model ensuring
$(\bm{x}_{i\alpha_{i}}-\bm{x}_{it})^{\top}\bm{W}^{\textsl{mm}}{\bm{z}}_{i}\geq
1$, the following inequalities hold for all
$\bm{V}\in\texttt{cone}_{\mu}(\bm{W}^{\textsl{mm}}),~{}\|{\bm{V}}\|_{F}=\|{\bm{W}^{\textsl{mm}}}\|_{F}$
and all $i\in[n],t\in\mathcal{T}_{i},\tau\in\bar{\mathcal{T}}_{i}$:

$\begin{split}(\bm{x}_{it}-\bm{x}_{i\tau})^{\top}\bm{V}{\bm{z}}_{i}&\geq\delta>0,\\\
(\bm{x}_{i\alpha_{i}}-\bm{x}_{i\tau})^{\top}\bm{V}{\bm{z}}_{i}&\geq
1+\delta,\\\
\frac{3}{2}\geq(\bm{x}_{i\alpha_{i}}-\bm{x}_{it})^{\top}\bm{V}{\bm{z}}_{i}&\geq\frac{1}{2}.\end{split}$
(68)

Here, we used
$\|{\bm{V}-\bm{W}^{\textsl{mm}}}\|_{F}^{2}/\|{\bm{W}^{\textsl{mm}}}\|_{F}^{2}\leq
2\mu$ which implies
$\|{\bm{V}-\bm{W}^{\textsl{mm}}}\|_{F}\leq\sqrt{2\mu}/\Theta$.

To proceed, we write the gradient correlation following (13) and (25)

$\displaystyle\left<\nabla{\cal{L}}(\bm{W}),\bm{V}\right>$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}_{i}\cdot\bm{h}_{i}^{\top}\mathbb{S}^{\prime}(\tilde{\bm{h}}_{i})\bm{\gamma}_{i},$
(69)

where we denoted
$\ell^{\prime}_{i}=\ell^{\prime}(Y_{i}\cdot\bm{v}^{\top}{\bm{X}}_{i}^{\top}\mathbb{S}(\tilde{\bm{h}}_{i}))$,
$\bm{h}_{i}={\bm{X}}_{i}\bm{V}{\bm{z}}_{i}$,
$\tilde{\bm{h}}_{i}={\bm{X}}_{i}\bm{W}{\bm{z}}_{i}$,
$\bm{s}_{i}=\mathbb{S}(\tilde{\bm{h}}_{i})$.

Using (68), for all $t\in\mathcal{T}_{i},\tau\in\bar{\mathcal{T}}_{i}$, for
all $\bm{W}\in\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}})$, we have that

$\displaystyle\tilde{\bm{h}}_{i\alpha_{i}}-\tilde{\bm{h}}_{i\tau}\geq
R\Theta(1+\delta),$
$\displaystyle\tilde{\bm{h}}_{it}-\tilde{\bm{h}}_{i\tau}\geq R\Theta\delta,$
$\displaystyle\tilde{\bm{h}}_{i\alpha_{i}}-\tilde{\bm{h}}_{it}\geq R\Theta/2.$

Consequently, we can bound the softmax probabilities
$\bm{s}_{i}=\mathbb{S}(\tilde{\bm{h}}_{i})$ over non-support indices as
follows: For all $i\in[n]$ and any $t_{i}\in\mathcal{T}_{i}$

$\displaystyle S_{i}:=\sum_{\tau\in\mathcal{T}_{i}}\bm{s}_{i\tau}\leq
Te^{-R\Theta/2}\bm{s}_{i\alpha_{i}}\leq Te^{-R\Theta/2},$ (70a) $\displaystyle
Q_{i}:=\sum_{\tau\in\bar{\mathcal{T}}_{i}}\bm{s}_{i\tau}\leq
Te^{-R\Theta\delta}\bm{s}_{it_{i}}\leq Te^{-R\Theta\delta}S_{i}.$ (70b)

Recall scores $\bm{\gamma}_{it}=Y_{i}\cdot\bm{v}^{\top}\bm{x}_{it}$. Define
the score gaps over support indices:

$\gamma^{\textsl{gap}}_{i}=\bm{\gamma}_{i\alpha_{i}}-\max_{t\in\mathcal{T}_{i}}\bm{\gamma}_{it}~{}~{}~{}\textnormal{and}~{}~{}~{}\bar{\gamma}^{\textsl{gap}}_{i}=\bm{\gamma}_{i\alpha_{i}}-\min_{t\in\mathcal{T}_{i}}\bm{\gamma}_{it}.$

It follows from (67) that

$\displaystyle
A=\max_{i\in[n],t\in[T]}\frac{\|{\bm{x}_{it}{\bm{z}}_{i}^{\top}}\|}{\Theta}\geq\max_{i\in[n],t\in[T]}\|{\bm{h}_{it}}\|.$

Define the ${\bm{\alpha}}$-dependent global scalar
$\Gamma=\sup_{i\in[n],t,\tau\in[T]}|\bm{\gamma}_{it}-\bm{\gamma}_{i\tau}|$.

Let us focus on a fixed datapoint $i\in[n]$, assume (without losing
generality) $\alpha_{i}=1$, and drop subscripts $i$. Directly applying Lemma
7, we obtain

$\big{|}\bm{h}^{\top}\text{diag}(\bm{s})\bm{\gamma}-\bm{h}^{\top}\bm{s}\bm{s}^{\top}\bm{\gamma}-\sum_{t\geq
2}^{T}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\big{|}\leq
2\Gamma A(1-\bm{s}_{1})^{2}.$

To proceed, let us decouple the non-support indices within $\sum_{t\geq
2}^{T}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$ via

$\big{|}\sum_{t\in\bar{\mathcal{T}}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\big{|}\leq
2Q\Gamma A.$

Aggregating these, we found

$\displaystyle\big{|}\bm{h}^{\top}\text{diag}(\bm{s})\bm{\gamma}-\bm{h}^{\top}\bm{s}\bm{s}^{\top}\bm{\gamma}-\sum_{t\in\mathcal{T}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\big{|}\leq
2\Gamma A((1-\bm{s}_{1})^{2}+Q).$ (71)

To proceed, let us upper/lower bound the gradient correlation. Since
$1.5\geq\bm{h}_{1}-\bm{h}_{t}\geq 0.5$, we find

$1.5\cdot
S\cdot\bar{\gamma}^{\textsl{gap}}\geq\sum_{t\in\mathcal{T}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\geq
0.5\cdot S\cdot\gamma^{\textsl{gap}},$

where recall the definition of $S,Q$ (having dropped subscripts) in (70a).
Next we claim that $S$ dominates $((1-\bm{s}_{1})^{2}+Q)$ for large $R$.
Specifically, we wish for

$\displaystyle\frac{S\cdot\gamma^{\textsl{gap}}}{4}\geq 4\Gamma
A\max((1-\bm{s}_{1})^{2},Q)\iff S\geq 16\frac{\Gamma
A}{\gamma^{\textsl{gap}}}\max((1-\bm{s}_{1})^{2},Q).$ (72)

Now choose $R\geq\delta^{-1}\log(T)/\Theta$ to ensure $Q\leq S$ since $Q\leq
Te^{-R\Theta\delta}S$ from (70a). Consequently

$(1-\bm{s}_{1})^{2}=(Q+S)^{2}\leq 4S^{2}\leq 4STe^{-R\Theta/2}.$

Combining these, what we wish is ensured by guaranteeing

$\displaystyle S\geq 16\frac{\Gamma
A}{\gamma^{\textsl{gap}}}\max(4STe^{-R\Theta/2},Te^{-R\Theta\delta}S).$ (73)

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

$\displaystyle R=\frac{\max(2,\delta^{-1})}{\Theta}\log\big{(}\frac{64T\Gamma
A}{\gamma^{\textsl{gap}}_{\min}}\big{)},$ (74)

where $\gamma^{\textsl{gap}}_{\min}=\min_{i\in[n]}\gamma^{\textsl{gap}}_{i}$
is the global scalar which is the worst case score gap over all inputs. With
the above choice of $R$, we guaranteed

$2(1-\bm{s}_{1})\cdot\bar{\gamma}^{\textsl{gap}}\geq 2\cdot
S\cdot\bar{\gamma}^{\textsl{gap}}\geq\sum_{t\in\mathcal{T}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\geq\frac{S\cdot\gamma^{\textsl{gap}}}{4}\geq\frac{(1-\bm{s}_{1})\gamma^{\textsl{gap}}}{8}.$

via (72) and (71). Since this holds over all inputs, going back to the
gradient correlation (69) and averaging above over all inputs $i\in[n]$ and
plugging back the indices $i$, we obtain the advertised bound by setting
$q_{i}=1-\bm{s}_{i\alpha_{i}}$ (where we set $\alpha_{i}=1$ above without
losing generality)

$\displaystyle\frac{2}{n}\sum_{i\in[n]}\ell^{\prime}_{i}\cdot
q_{i}\cdot\bar{\gamma}^{\textsl{gap}}_{i}\geq\left<\nabla{\cal{L}}(\bm{W}),\bm{V}\right>\geq\frac{1}{8n}\sum_{i\in[n]}\ell^{\prime}_{i}\cdot
q_{i}\cdot\gamma^{\textsl{gap}}_{i}.$ (75)

###### Lemma 14

Suppose Assumption A on the loss function $\ell$ holds, and let
${\bm{\alpha}}=(\alpha_{i})_{i=1}^{n}$ be locally optimal tokens according to
Definition 2. Let $\bm{W}^{\textsl{mm}}=\bm{W}^{\textsl{mm}}({\bm{\alpha}})$
denote the SVM solution obtained via (Att-SVM) by replacing
$(\texttt{opt}_{i})_{i=1}^{n}$ with $\bm{\alpha}=(\alpha_{i})_{i=1}^{n}$. For
any choice of $\pi>0$, there exists $R:=R_{\pi}$ such that, for any
$\bm{W}\in\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}}(\alpha))$, we have

$\left<\nabla{\cal{L}}(\bm{W}),\frac{\bm{W}}{\|{\bm{W}}\|_{F}}\right>\geq(1+\pi)\left<\nabla{\cal{L}}(\bm{W}),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right>.$

Proof. Let
${\bm{\bar{W}}}=\|{\bm{W}^{\textsl{mm}}}\|_{F}\bm{W}/\|{\bm{W}}\|_{F}$,
$\bm{h}_{i}={\bm{X}}_{i}{\bm{\bar{W}}}{\bm{z}}_{i}$, and
$\bm{\bar{h}}_{i}={\bm{X}}_{i}\bm{W}^{\textsl{mm}}{\bm{z}}_{i}$. To establish
the result, we will prove that, for sufficiently large $R=R_{\pi}$, for any
$\bm{W}\in\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}})$ and for any $i\in[n]$,

$\displaystyle\left<\bm{h}_{i},\mathbb{S}^{\prime}({\bm{X}}_{i}\bm{W}{\bm{z}}_{i})\bm{\gamma}_{i}\right>\leq(1+\pi)\left<\bm{\bar{h}}_{i},\mathbb{S}^{\prime}({\bm{X}}_{i}\bm{W}{\bm{z}}_{i})\bm{\gamma}_{i}\right>.$
(76)

Once (76) holds for all $i$, the same conclusion will hold for the gradient
correlations via (69). Moving forward, we shall again focus on a single point
$i\in[n]$ and drop all subscripts $i$. Also, assume $\alpha=\alpha_{i}=1$
without losing generality (same as above).

Following (71), for all $\bm{V}\in\texttt{cone}_{\mu}$ with
$\|{\bm{V}}\|_{F}=\|{\bm{W}^{\textsl{mm}}}\|_{F}$ and
$\tilde{\bm{h}}={\bm{X}}\bm{V}{\bm{z}}$, and
$\bm{s}=\mathbb{S}(\tilde{\bm{h}})$, we have found

$\displaystyle\big{|}\tilde{\bm{h}}^{\top}\text{diag}(\bm{s})\bm{\gamma}-\tilde{\bm{h}}^{\top}\bm{s}\bm{s}^{\top}\bm{\gamma}-\sum_{t\in\mathcal{T}}(\tilde{\bm{h}}_{1}-\tilde{\bm{h}}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\big{|}\leq
2\Gamma A((1-\bm{s}_{1})^{2}+Q).$ (77)

Plugging in $\bm{h},\bm{\bar{h}}$ in the bound above and assuming $\pi\leq 1$
(w.l.o.g.), (76) is implied by the following stronger inequality

$\displaystyle 6\Gamma
A((1-\bm{s}_{1})^{2}+Q)+\sum_{t\in\mathcal{T}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$
$\displaystyle\leq(1+\pi)\sum_{t\in\mathcal{T}}(\bm{\bar{h}}_{1}-\bm{\bar{h}}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$
$\displaystyle=(1+\pi)\sum_{t\in\mathcal{T}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t}).$

First, we claim that
$0.5\pi\sum_{t\in\mathcal{T}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\geq
6\Gamma A((1-\bm{s}_{1})^{2}+Q)$. The proof of this claim directly follows the
earlier argument, namely, following (72), (74) and (73) which leads to the
choice

$R_{\pi}=\frac{\max(2,\delta^{-1})}{\Theta}\log\big{(}\frac{C\cdot T\Gamma
A}{\pi\gamma^{\textsl{gap}}_{\min}}\big{)}$

for some constant $C>0$.

Following this control over the perturbation term $6\Gamma
A((1-\bm{s}_{1})^{2}+Q)$, to conclude with the result, what remains is proving
the comparison

$\displaystyle\sum_{t\in\mathcal{T}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\leq(1+0.5\pi)\sum_{t\in\mathcal{T}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t}).$
(78)

To proceed, we split the problem into two scenarios.

Scenario 1:
$\|{{\bm{\bar{W}}}-\bm{W}^{\textsl{mm}}}\|_{F}\leq\varepsilon=\frac{\pi}{4A\Theta}$
for some $\varepsilon>0$. In this scenario, for any token, we find that

$|\bm{h}_{t}-\bm{\bar{h}}_{t}|\leq A\Theta\varepsilon=\pi/4.$

Consequently, we obtain

$\bm{h}_{1}-\bm{h}_{t}\leq\bm{\bar{h}}_{1}-\bm{\bar{h}}_{t}+2A\Theta\varepsilon=1+0.5\pi.$

Similarly, $\bm{h}_{1}-\bm{h}_{t}\geq 1-0.5\pi\geq 0.5$. Since all terms
$\bm{h}_{1}-\bm{h}_{t},\bm{s}_{t},\bm{\gamma}_{1}-\bm{\gamma}_{t}$ in (78) are
nonnegative and
$(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\leq(1+0.5\pi)\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$,
the above implies the desired result (78).

Scenario 2:
$\|{{\bm{\bar{W}}}-\bm{W}^{\textsl{mm}}}\|_{F}\geq\varepsilon=\frac{\pi}{4A\Theta}$.
Since ${\bm{\bar{W}}}$ is not (locally) max-margin, in this scenario, for some
$\nu=\nu(\varepsilon)>0$ and $\tau\in\mathcal{T}$, we have that
$\bm{h}_{1}-\bm{h}_{\tau}\leq 1-2\nu$. Here
$\tau=\arg\max_{\tau\in\mathcal{T}}\bm{x}_{\tau}{\bm{\bar{W}}}{\bm{z}}$
denotes the nearest point to $\bm{h}_{1}$ (along the ${\bm{\bar{W}}}$
direction). Note that a non-support index $\tau\in\bar{\mathcal{T}}$ cannot be
closest because $\bm{W}\in\mathcal{C}_{\mu}$ and (68) holds. Recall that
$\bm{s}=\mathbb{S}(\bar{R}\bm{h})$ where $\bar{R}=\|{\bm{W}}\|_{F}\Theta\geq
R\Theta$. To proceed, split the tokens into two groups: Let ${\cal{N}}$ be the
group of tokens obeying $(\bm{x}_{1}-\bm{x}_{\tau})\bm{W}{\bm{z}}\leq 1-\nu$
and $\mathcal{T}-{\cal{N}}$ be the rest of the support indices. Observe that

$\frac{\sum_{t\in\mathcal{T}-{\cal{N}}}\bm{s}_{t}}{\sum_{t\in\mathcal{T}}\bm{s}_{t}}\leq\frac{\sum_{t\in\mathcal{T}-{\cal{N}}}\bm{s}_{t}}{\sum_{t=\tau}\bm{s}_{t}}\leq
T\frac{e^{\nu\bar{R}}}{e^{2\nu\bar{R}}}=Te^{-\bar{R}\nu}.$

Thus, using $|\bm{h}_{1}-\bm{h}_{t}|\leq 2A$ and recalling the definition of
$\gamma^{\textsl{gap}}$, observe that

$\sum_{t\in\mathcal{T}-{\cal{N}}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})\leq\frac{2\Gamma
ATe^{-\bar{R}\nu}}{\gamma^{\textsl{gap}}}\sum_{t\in\mathcal{T}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t}).$

Plugging this into (78), we obtain

$\displaystyle\sum_{t\in\mathcal{T}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$
$\displaystyle=\sum_{t\in{\cal{N}}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})+\sum_{t\in\mathcal{T}-{\cal{N}}}(\bm{h}_{1}-\bm{h}_{t})\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$
$\displaystyle\leq\sum_{t\in{\cal{N}}}(1-\nu)\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})+\sum_{t\in\mathcal{T}-{\cal{N}}}2A\Gamma
Te^{-\bar{R}\nu}$ $\displaystyle\leq\left(1-\nu+\frac{2\Gamma
ATe^{-\bar{R}\nu}}{\gamma^{\textsl{gap}}}\right)\sum_{t\in\mathcal{T}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t})$
$\displaystyle\leq\left(1+\frac{2\Gamma
ATe^{-\bar{R}\nu}}{\gamma^{\textsl{gap}}}\right)\sum_{t\in\mathcal{T}}\bm{s}_{t}(\bm{\gamma}_{1}-\bm{\gamma}_{t}).$

Consequently, the proof boils down to ensuring the perturbation term
$\frac{2\Gamma ATe^{-\bar{R}\nu}}{\gamma^{\textsl{gap}}}\leq 0.5\pi$.
Recalling $\bar{R}\geq R\Theta$, this is guaranteed for all inputs $i\in[n]$
by recalling
$\gamma^{\textsl{gap}}_{\min}=\min_{i\in[n]}\gamma^{\textsl{gap}}_{i}$ and
choosing

$R\geq R_{\pi}=\frac{1}{\nu\Theta}\log(\frac{4\Gamma
AT}{\gamma^{\textsl{gap}}_{\min}\pi}),$

where $\nu=\nu(\frac{\pi}{4A\Theta})$ depends only on $\pi$ and global problem
variables.

Combining this with the prior $R_{\pi}$ choice (by taking maximum), we
conclude with the statement.

###  C.1 Proof of Theorem 7

We provide the proof in four steps:
Step 1: There are no stationary points within
$\mathcal{C}_{\mu,R_{\mu}}(\bm{W}^{\textsl{mm}})$ for sufficiently large
$R_{\mu}$. We begin by proving that there are no stationary points within
$\mathcal{C}_{\mu,R_{\mu}}(\bm{W}^{\textsl{mm}})$ for a specific radius
$R_{\mu}$. Let $(\mathcal{T}_{i})_{i=1}^{n}$ denote the sets of support
indices as defined in Definition 2. We define
$\bar{\mathcal{T}}_{i}=[T]-\mathcal{T}_{i}-\\{\alpha_{i}\\}$ as the tokens
that are non-support indices. Additionally, let $\mu$ be defined as in (67).
Then, for all $\bm{V},\bm{W}\in\texttt{cone}_{\mu}(\bm{W}^{\textsl{mm}})$ with
$\|{\bm{V}}\|=\|{\bm{W}^{\textsl{mm}}}\|$, it follows from Lemma 13 that there
exists $R_{\mu}$ such that $-\bm{V}^{\top}\nabla{\cal{L}}(\bm{W})$ is strictly
positive for $\|{\bm{W}}\|\geq R_{\mu}$.
Step 2: Consider an arbitrary value of $\epsilon\in(0,\mu/2)$ and let
$1/(1+\pi)=1-\epsilon$. It follows from Lemma 13 that, there exists
$R_{\epsilon}$ such that all
$\bm{W}\in\mathcal{C}_{\mu,R_{\epsilon}}(\bm{W}^{\textsl{mm}})$ satisfy

$\displaystyle\left\langle-\nabla\mathcal{L}(\bm{W}),\frac{\bm{W}^{\textsl{mm}}}{\|\bm{W}^{\textsl{mm}}\|}\right\rangle\geq(1-\epsilon)\left\langle-\nabla\mathcal{L}(\bm{W}),\frac{\bm{W}}{\|\bm{W}\|}\right\rangle.$
(79)

Step 3: Updates remain inside the cone. By leveraging the results from Step 1
and Step 2, we demonstrate that the gradient iterates, with appropriate step
size, starting from $\bm{W}(0)\in\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}})$,
remain within this cone. Here, $R$ is chosen such that

$\displaystyle R\geq R_{\mu}\vee R_{\epsilon}\vee 1/2.$ (80)

We proceed by induction. Suppose that the claim holds up to iteration $k\geq
0$. This implies that $\bm{W}(k)\in\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}})$.
Hence, there exists scalar $\mu=\mu({\bm{\alpha}})\in(0,1)$ and $R$ such that
${\texttt{corr\\_coef}}(\bm{W}(k),\bm{W}^{\textsl{mm}})\geq 1-\mu$ and
$\|{\bm{W}(k)}\|\geq R$. For all $k\geq 1$, there exists a constant $\rho$
such that

$\displaystyle\rho\leq-\frac{1}{1-\epsilon}\left\langle\nabla\mathcal{L}(\bm{W}(k)),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle.$
(81)

Note that $\rho>0$ due to Step 1 and Lemma 13. This together with the gradient
descent update rule gives

$\begin{split}\left\langle\frac{\bm{W}(k+1)}{\|{\bm{W}(k)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle&=\left\langle\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}}-\frac{\eta}{\|{\bm{W}(k)}\|_{F}}\nabla\mathcal{L}(\bm{W}(k)),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle\\\
&\geq
1-\mu-\frac{\eta}{\|{\bm{W}(k)}\|_{F}}\left\langle\nabla\mathcal{L}(\bm{W}(k)),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle\\\
&\geq 1-\mu+\frac{\eta\rho(1-\epsilon)}{\|{\bm{W}(k)}\|_{F}}.\end{split}$
(82a) Note that from Lemma 13, we have
$\left\langle\nabla{\cal{L}}(\bm{W}(k)),\bm{W}(k)\right\rangle<0$ which
implies that $\|{\bm{W}(k+1)}\|_{F}\geq\|{\bm{W}(k)}\|_{F}$. Hence,
$\displaystyle\|{\bm{W}(k+1)}\|_{F}$
$\displaystyle\leq\frac{1}{{2\|{\bm{W}(k)}\|_{F}}}\left(\|{\bm{W}(k+1)}\|_{F}^{2}+\|{\bm{W}(k)}\|_{F}^{2}\right)$
$\displaystyle=\frac{1}{2\|{\bm{W}(k)}\|_{F}}\left(2\|{W(k)}\|_{F}^{2}-2\eta\left\langle\nabla{\cal{L}}(\bm{W}(k)),\bm{W}(k)\right\rangle+\eta^{2}\|{\nabla{\cal{L}}(\bm{W}(k))}\|_{F}^{2}\right)$
$\displaystyle\leq\|{\bm{W}(k)}\|_{F}-\frac{\eta}{\|{\bm{W}(k)}\|_{F}}\left\langle\nabla{\cal{L}}(\bm{W}(k)),\bm{W}(k)\right\rangle+\eta^{2}\|\nabla{\cal{L}}(\bm{W}(k))\|^{2},$
which gives
$\begin{split}\frac{\|{\bm{W}(k+1)}\|_{F}}{\|{\bm{W}(k)}\|_{F}}&\leq
1-\eta\left\langle\nabla{\cal{L}}(\bm{W}(k)),\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}}\right\rangle+\eta^{2}\frac{\|\nabla\mathcal{L}(\bm{W}(k))\|^{2}}{\|{\bm{W}(k)}\|_{F}}\\\
&\leq
1-\frac{\eta}{1-\epsilon}\left\langle\nabla\mathcal{L}(\bm{W}(k)),\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle+\eta^{2}\frac{\|\nabla\mathcal{L}(\bm{W}(k))\|^{2}}{\|{\bm{W}(k)}\|_{F}}\\\
&\leq
1+\frac{\eta\rho}{\|{\bm{W}(k)}\|_{F}}+\frac{\eta^{2}\|\nabla\mathcal{L}(\bm{W}(k))\|^{2}}{\|{\bm{W}(k)}\|_{F}}=:C_{1}(\rho,\eta).\end{split}$
(82b)

Here, the second inequality follows from (79) and (81).

Now, it follows from (82a) and (82b) that

$\begin{split}\left\langle\frac{\bm{W}(k+1)}{\|{\bm{W}(k+1)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle&\geq\frac{1}{C_{1}(\rho,\eta)}\left(1-\mu+\frac{\eta\rho(1-\epsilon)}{\|{\bm{W}(k)}\|_{F}}\right)\\\
&\geq
1-\mu+\frac{\eta}{C_{1}(\rho,\eta)}\left(\frac{\rho(\mu-\epsilon)}{\|{\bm{W}(k)}\|_{F}}-\eta(1-\mu)\frac{\|\nabla\mathcal{L}(\bm{W}(k))\|^{2}}{\|{\bm{W}(k)}\|_{F}}\right)\\\
&\geq 1-\mu,\end{split}$ (83)

where the last inequality uses

$\eta\leq\zeta(\mu):=\frac{(\mu-\epsilon)\rho}{1-\mu}\frac{1}{c^{2}}\leq\frac{(\mu-\epsilon)\rho}{1-\mu}\frac{1}{\|\nabla\mathcal{L}(\bm{W}(k))\|^{2}}$

for some constant $c$ satisfying $\|\nabla\mathcal{L}(\bm{W}(k))\|\leq c$ due
to Lemma 13.

Hence,
$\bm{W}(k+1)\in\texttt{cone}_{\mu}\bm{W}^{\textsl{mm}})\bigcap\\{\bm{W}~{}|~{}\|{\bm{W}}\|_{F}\geq
R_{\mu}\\}$.
Step 4: The correlation of $\bm{W}(k)$ and $\bm{W}^{\textsl{mm}}$ increases
over $k$. It follows Theorem 3
$\left\|\bm{W}\left(k\right)\right\|\rightarrow\infty$. Hence, we can choose
$k_{0}$ such that for any $k\geq k_{0}$, it holds that $\|{\bm{W}(k)}\|_{F}>R$
for some $R\geq R_{\mu}\vee R_{\epsilon}\vee 1/2$. Now, following similar
steps in (33) and (34), we obtain

$\displaystyle\left\langle\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle\geq
1-\epsilon+\frac{C(\epsilon,\eta)}{\|{\bm{W}(k)}\|_{F}},$

for some finite constant $C(\epsilon,\eta)$.

Consequently,

$\displaystyle\liminf_{k\to\infty}\left\langle\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}},\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}\right\rangle\geq
1-\epsilon.$

Since $\epsilon\in(0,\mu/2)$ is arbitrary, we get
$\bm{W}(k)/\|{\bm{W}(k)}\|_{F}\to\bm{W}^{\textsl{mm}}/\|{\bm{W}^{\textsl{mm}}}\|_{F}$.
$\hfill\blacksquare$

###  C.2 Proof of Lemma 6

Suppose $\tau$ described by (8) exists and set $\lambda$ accordingly. Let
$\mathcal{S}\subset[T]$ denote the top $\tau$ indices of $\bm{\gamma}$ with
largest scores. Denote ${\bm{X}}^{1}\in\mathbb{R}^{\tau\times d}$ to be the
sequence corresponding to $\mathcal{S}$ and
${\bm{X}}^{2}\in\mathbb{R}^{(T-\tau)\times d}$ to be the sequence
corresponding to $[T]-\mathcal{S}$. Similarly, denote the subvectors
$\bm{\gamma}_{1},\bm{s}^{(1)}\in\mathbb{R}^{\tau}$ and
$\bm{\gamma}_{2},\bm{s}^{(2)}\in\mathbb{R}^{T-\tau}$ and define the
probability over $\mathcal{S}$ as $S_{1}=\sum_{i\in\mathcal{S}}s_{i}$. The
orthogonality and unit norm assumption on the tokens imply

$1\geq\|{{\bm{X}}^{\top}\bm{s}}\|^{2}=\sum_{i=1}^{T}s_{i}^{2}\geq
S_{1}^{2}/\tau+(1-S_{1})^{2}/(T-\tau).$

Also note that
$\bm{v}^{\top}{\bm{X}}\bm{s}=\bm{\gamma}_{1}^{\top}\bm{s}^{(1)}+\bm{\gamma}_{2}^{\top}\bm{s}^{(2)}$.
With these, we can write the objective
${\cal{L}}(\bm{s}):=\ell(h({\bm{X}}\bm{s}))$ as follows

${\cal{L}}(\bm{s})=\lambda\sum_{i=1}^{T}s_{i}^{2}-\bm{\gamma}_{1}^{\top}\bm{s}^{(1)}-\bm{\gamma}_{2}^{\top}\bm{s}^{(2)}.$

Note that, for fixed $\bm{\gamma}$ and over all permutations of entries of
$\bm{s}$, $\bm{\gamma}^{\top}\bm{s}$ is maximized when $\bm{s}$ and
$\bm{\gamma}$ are aligned namely, when the entries of $\bm{s}$ are sorted
according to the entries of $\bm{\gamma}$. Otherwise, we could swap two
unsorted entries of $\bm{s}$ (i.e. with unaligned $\bm{\gamma}$ entries) to a
sorted position to obtain a strictly better optimal (where we also used the
fact that $\bm{s}$ has nonnegative entries). Thus, we can assume the entries
of $\bm{s}^{\star}$ are sorted according to $\bm{\gamma}$. Specifically, the
largest $\tau$ entries of $\bm{s}^{\star}$ lie on the set $\mathcal{S}$.

$\bullet$ We first show that $\bm{s}:=\bm{s}^{\star}$ cannot have more than
$\tau$ entries. To prove this, we compare $\bm{s}$ against the baseline
$\bar{\bm{s}}$ where $\bar{\bm{s}}^{1}=\bm{s}^{(1)}/S_{1}$ and
$\bar{\bm{s}}^{2}=0$ so that $\bar{\bm{s}}$ is $\tau$-sparse. In this
scenario, $\bar{\bm{s}}$ yields the objective

${\cal{L}}(\bar{\bm{s}})=\frac{\lambda}{S_{1}^{2}}\sum_{i\in\mathcal{S}}s_{i}^{2}-\frac{1}{S_{1}}\bm{\gamma}_{1}^{\top}\bm{s}^{(1)}.$

We claim that ${\cal{L}}(\bar{\bm{s}})<{\cal{L}}(\bm{s})$. To see this, we
first observe that
$\bm{\gamma}_{1}^{\top}\bm{s}^{(1)}/S_{1}\geq\bm{\gamma}_{2}^{\top}\bm{s}^{(2)}/(1-S_{1})+\bm{\gamma}^{\textsl{gap}}_{\tau}$.
This implies

$(1/S_{1}-1)\bm{\gamma}_{1}^{\top}\bm{s}^{(1)}-\bm{\gamma}_{2}^{\top}\bm{s}^{(2)}\geq(1-S_{1})\bm{\gamma}^{\textsl{gap}}_{\tau}.$

Recalling $\sum_{i\in\mathcal{S}}s_{i}^{2}\leq S_{1}^{2}/\tau$, we can now
utilize the following chain of implications

$\displaystyle{\cal{L}}(\bar{\bm{s}})<{\cal{L}}(\bm{s})$ $\displaystyle\iff$
$\displaystyle\frac{\lambda}{S_{1}^{2}}\sum_{i\in\mathcal{S}}s_{i}^{2}-\frac{1}{S_{1}}\bm{\gamma}_{1}^{\top}\bm{s}^{(1)}<\lambda\sum_{i=1}^{T}s_{i}^{2}-\bm{\gamma}_{1}^{\top}\bm{s}^{(1)}-\bm{\gamma}_{2}^{\top}\bm{s}^{(2)}$
$\displaystyle\impliedby$
$\displaystyle\lambda(1/S_{1}^{2}-1)\sum_{i\in\mathcal{S}}s_{i}^{2}<(1/S_{1}-1)\bm{\gamma}_{1}^{\top}\bm{s}^{(1)}-\bm{\gamma}_{2}^{\top}\bm{s}^{(2)}$
$\displaystyle\impliedby$
$\displaystyle\lambda(1/S_{1}^{2}-1)\sum_{i\in\mathcal{S}}s_{i}^{2}<(1-S_{1})\bm{\gamma}^{\textsl{gap}}_{\tau}$
$\displaystyle\impliedby$
$\displaystyle\lambda(1-S_{1}^{2})/\tau<(1-S_{1})\bm{\gamma}^{\textsl{gap}}_{\tau}$
$\displaystyle\impliedby$
$\displaystyle\lambda(1+S_{1})/\tau<\bm{\gamma}^{\textsl{gap}}_{\tau}$
$\displaystyle\impliedby$ $\displaystyle
2\lambda/\tau<\bm{\gamma}^{\textsl{gap}}_{\tau}$ $\displaystyle\impliedby$
$\displaystyle\lambda<\tau\bm{\gamma}^{\textsl{gap}}_{\tau}/2.$

$\bullet$ We next prove that there are at least two nonzeros in the optimal
solution. Denote the largest and second largest entry of $\bm{\gamma}$ by
$\bar{\gamma}_{1}$ and $\bar{\gamma}_{2}$ respectively. For
$\bm{s}^{\texttt{one}}\in{\bm{\Delta}}_{T}$ containing a single nonzero (i.e.
one-hot vector), the best achievable risk is given by

${\cal{L}}(\bm{s}^{\texttt{one}})=\lambda-\bar{\gamma}_{1}.$

On the other hand consider the $2$-sparse reference solution
$\bm{s}^{\texttt{ref}}$ which assigns equal likelihood over the top two
entries. This achieves

${\cal{L}}(\bm{s}^{\texttt{ref}})=\frac{\lambda}{2}-\bm{\gamma}^{\top}\bm{s}^{\texttt{ref}}\leq\frac{\lambda}{2}-\frac{\bar{\gamma}_{1}+\bar{\gamma}_{2}}{2}.$

The latter is superior as soon as

$\frac{\lambda}{2}-\frac{\bar{\gamma}_{1}+\bar{\gamma}_{2}}{2}<\lambda-\bar{\gamma}_{1}\iff\lambda>\bm{\gamma}^{\textsl{gap}}_{1}.$

Thus, we conclude with the statement by selecting
$\tau\bm{\gamma}^{\textsl{gap}}_{\tau}/2>\lambda>\gamma^{\textsl{gap}}_{1}$.
$\hfill\blacksquare$

##  Appendix D Convergence of Regularization Path for Sequence-to-Sequence
Setting

In this section, we provide proofs for the regularization path analysis. We
first provide a more general formulation of the optimization problem that
allows for regressing multiple token outputs. To distinguish from (W-ERM)&(KQ-
ERM), let us call this more general version Sequence Empirical Risk
Minimization (SERM).

Problem definition: Rather than a single input sequence ${\bm{X}}$, let us
allow for two input sequences ${\bm{X}}\in\mathbb{R}^{T\times d}$ and
$\bm{Z}\in\mathbb{R}^{K\times d}$ with $(\bm{x}_{t})_{t=1}^{T}$ and
$({\bm{z}}_{k})_{k=1}^{K}$. The cross-attention admits ${\bm{X}},\bm{Z}$ and
outputs $K$ tokens. We will also allow for $K$ separate prediction heads
$(h_{k})_{k=1}^{K}$ for individual cross attention outputs which strictly
generalizes the setting where we used single prediction head $h(\cdot)$.
Denote the training labels associated to each token as
${\bm{Y}}=(Y_{ik})_{i=1}^{K}$. Given $n$ samples
$({\bm{Y}}_{i},{\bm{X}}_{i},\bm{Z}_{i})_{i=1}^{n}$, for a decreasing loss
function $\ell(\cdot)$, minimize the empirical risk by the prediction of first
attention output either univariate ($\bm{W}\in\mathbb{R}^{d\times d}$) or
bivariate (${\bm{K}},{\bm{Q}}\in\mathbb{R}^{d\times m}$) fashion:

$\displaystyle{\cal{L}}(\bm{W})$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\sum_{k=1}^{K}\ell(Y_{ik}\cdot
h_{k}({\bm{X}}_{i}^{\top}\mathbb{S}({\bm{X}}_{i}\bm{W}{\bm{z}}_{ik}))),$
(SERM-W) $\displaystyle{\cal{L}}({\bm{K}},{\bm{Q}})$
$\displaystyle=\frac{1}{n}\sum_{i=1}^{n}\ell(Y_{ik}\cdot
h_{k}({\bm{X}}_{i}^{\top}\mathbb{S}({\bm{X}}_{i}{\bm{K}}{\bm{Q}}^{\top}{\bm{z}}_{ik}))).$
(SERM-KQ)

In order to recover the single-output self-attention model, we can simply set
$K=1$ and ${\bm{z}}_{i1}\leftarrow\bm{x}_{i1}$ and $h_{k}\leftarrow h$. To
proceed, we introduce the more general version of the (Att-SVM) problem, which
we refer to as _Sequential Cross-Attention SVM_ to preserve consistent
phrasing. Suppose ${\bm{K}},{\bm{Q}}\in\mathbb{R}^{d\times m}$ with $m\leq d$
and let $\mathcal{R}_{m}$ denote the set of rank-$m$ matrices in
$\mathbb{R}^{d\times d}$. Given indices
${\bm{\alpha}}=(\alpha_{ik})_{{ik}=(1,1)}^{(n,K)}$, consider the SVM with
${\diamond}$-norm constraint

$\bm{W}^{\textsl{mm}}_{{\diamond},{\bm{\alpha}}}\in\arg\min_{\bm{W}\in\mathcal{R}_{m}}\|\bm{W}\|_{\diamond}\quad\text{s.t.}\quad\min_{t\neq\alpha_{ik}}(\bm{x}_{i\alpha_{ik}}-\bm{x}_{it})^{\top}\bm{W}{\bm{z}}_{ik}\geq
1\quad\forall\quad i\in[n],k\in[K].$ (SAtt-SVM)

When solution is non-unique, we denote the solution set by
${\cal{W}}^{\textsl{mm}}({\bm{\alpha}})$. In what follows, we denote
${\bm{F}}_{ikt}:=\bm{x}_{it}{\bm{z}}_{{ik}}^{\top}$ and given ${\bm{\alpha}}$,
we denote $\bm{x}^{\bm{\alpha}}_{ik}=\bm{x}_{i\alpha_{ik}}$ and
${\bm{F}}^{\bm{\alpha}}_{ik}:=\bm{x}_{i\alpha_{ik}}{\bm{z}}_{{ik}}^{\top}$.
With this notation, we can equivalently write

$\bm{W}^{\textsl{mm}}_{{\diamond},{\bm{\alpha}}}\in\arg\min_{\bm{W}\in\mathcal{R}_{m}}\|\bm{W}\|_{\diamond}\quad\text{s.t.}\quad\min_{t\neq\alpha_{ik}}\left<{\bm{F}}^{\bm{\alpha}}_{ik}-{\bm{F}}_{ikt},\bm{W}\right>\geq
1\quad\forall~{}i\in[n],k\in[K].$ (SAtt-SVM’)

######  Definition 4 (Support indices and Locally-Optimal Indices)

Fix token indices ${\bm{\alpha}}=(\alpha_{ik})_{ik=(1,1)}^{(n,k)}$ for which
(SAtt-SVM) is feasible to obtain
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}:=\bm{W}^{\textsl{mm}}_{{\diamond},{\bm{\alpha}}}$.
Define token scores as

$\bm{\gamma}_{ikt}=Y_{ik}\cdot
h_{k}(\bm{x}_{it}),\quad\bm{\gamma}^{\bm{\alpha}}_{ik}:=\bm{\gamma}_{ik\alpha_{ik}}=Y_{ik}\cdot
h_{k}(\bm{x}^{\bm{\alpha}}_{ik}).$

Consider tokens $\mathcal{T}_{ik}\subset[T]$ such that
$\left<{\bm{F}}^{\bm{\alpha}}_{ik}-{\bm{F}}_{ikt},\bm{W}^{\textsl{mm}}_{\bm{\alpha}}\right>=1$
for all $t\in\mathcal{T}_{ik}$. $\mathcal{T}_{ik}$ is allowed to be an empty
set. We refer to $\mathcal{T}_{ik}$ as support indices of
${\bm{F}}^{\bm{\alpha}}_{ik}=\bm{x}^{\bm{\alpha}}_{ik}{\bm{z}}_{ik}^{\top}$
and define its complement
$\bar{\mathcal{T}}_{ik}=[T]-\mathcal{T}_{ik}-\\{\alpha_{ik}\\}$. Additionally,
token indices ${\bm{\alpha}}=(\alpha_{ik})_{ik=(1,1)}^{(n,k)}$ are called
locally-optimal if for all $i\in[n],k\in[K]$ and $t\in\mathcal{T}_{ik}$, token
scores obey $\bm{\gamma}^{\bm{\alpha}}_{ik}>\bm{\gamma}_{ikt}$. Associated
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$ is called a locally-optimal direction.
Finally, let $\texttt{opt}_{ik}\in\arg\max_{t\in[T]}\bm{\gamma}_{ikt}$ be the
optimal indices and define the associated $\bm{W}^{\textsl{mm}}(\texttt{opt})$
to be a globally-optimal direction.

######  Lemma 15 (Mapping regularization path of $({\bm{K}},{\bm{Q}})$ to
$\bm{W}$)

Let ${\bm{K}},{\bm{Q}}\in\mathbb{R}^{d\times m}$ and consider regularization
path solutions of (SERM-W) and (SERM-KQ)

$\displaystyle{\bm{\bar{W}}}_{R}\in\underset{\bm{W}\in\mathcal{R}_{m}:\|{\bm{W}}\|_{\star}\leq
R}{\arg\min}{\cal{L}}(\bm{W})$ (84)
$\displaystyle{\bm{\bar{K}}}_{R},{\bm{\bar{Q}}}_{R}\in\underset{\|{{\bm{K}}}\|_{F}^{2}+\|{{\bm{Q}}}\|_{F}^{2}\leq
2R}{\arg\min}{\cal{L}}({\bm{K}},{\bm{Q}}).$ (85)

For all $R\geq 0$, there is a one-to-one map between the set of solutions
${\bm{\bar{W}}}_{R}$ of (84) and ${\bm{\bar{K}}}_{R}{\bm{\bar{Q}}}_{R}^{\top}$
of (85).

Proof. To prove the mapping, first fix a ${\bm{\bar{W}}}_{R}$ solution with
rank $m$, set ${\cal{L}}_{F}={\cal{L}}({\bm{\bar{W}}}_{R})$ and show the
existence of ${\bm{K}},{\bm{Q}}$ with
${\bm{K}}{\bm{Q}}^{\top}={\bm{\bar{W}}}_{R}$ feasible for (85) and
${\cal{L}}({\bm{K}},{\bm{Q}})\leq{\cal{L}}_{F}$. Use the singular value
decomposition ${\bm{\bar{W}}}_{R}={\bm{U}}{\bm{{\Sigma}}}{\bm{V}}^{\top}$ with
${\bm{{\Sigma}}}\in\mathbb{R}^{m\times m}$ being diagonal matrix of singular
values. Set ${\bm{K}}={\bm{U}}\sqrt{{\bm{{\Sigma}}}}$ and
${\bm{Q}}={\bm{V}}\sqrt{{\bm{{\Sigma}}}}$. Observe that
${\bm{K}}{\bm{Q}}^{\top}=\bm{W}$ and

$\|{{\bm{K}}}\|_{F}^{2}=\|{{\bm{Q}}}\|_{F}^{2}=\sum_{i=1}^{m}\sqrt{{\bm{{\Sigma}}}_{ii}}^{2}=\|{{\bm{\bar{W}}}_{R}}\|_{\star}\leq
R.$

Thus, ${\bm{K}},{\bm{Q}}$ achieves
${\cal{L}}({\bm{K}},{\bm{Q}})={\cal{L}}_{F}$. Conversely, given
${\bm{\bar{K}}}_{R},{\bm{\bar{Q}}}_{R}$ with
${\cal{L}}_{\star}={\cal{L}}({\bm{\bar{K}}}_{R},{\bm{\bar{Q}}}_{R})$,
$\bm{W}={\bm{\bar{K}}}_{R}{\bm{\bar{Q}}}_{R}^{\top}$ obeys
${\cal{L}}(\bm{W})={\cal{L}}_{\star}$ and, using the standard nuclear norm
inequality, we have

$\|{\bm{W}}\|_{\star}=\|{{\bm{\bar{K}}}_{R}{\bm{\bar{Q}}}_{R}^{\top}}\|_{\star}\leq\frac{1}{2}(\|{{\bm{\bar{K}}}_{R}}\|_{F}^{2}+\|{{\bm{\bar{Q}}}_{R}}\|_{F}^{2})=R.$

This shows $\bm{W}$ is feasible for (84). Combining the two findings above, we
find that optimal costs are equal (${\cal{L}}_{\star}={\cal{L}}_{F}$) and for
any $({\bm{\bar{K}}}_{R},{\bm{\bar{Q}}}_{R})$ solution there exists a
${\bm{\bar{W}}}_{R}$ solution and vice versa.

###  D.1 Regularization Path Analysis

To proceed with our analysis, let us define the set of optimal solutions
$\bm{W}^{\textsl{mm}}_{{\bm{\alpha}}}:=\bm{W}^{\textsl{mm}}_{{\diamond},{\bm{\alpha}}}$
to (SAtt-SVM). Let us denote this set by
${\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}:={\cal{W}}^{\textsl{mm}}_{\diamond}({\bm{\alpha}})$.
Note that, if the ${\diamond}$-norm is not strongly-convex,
${\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}$ may not be singleton.

###  D.2 Local Regularization Path and Proof of Theorem 9

We first recall _local regularization path_ which solves the
${\diamond}$-norm-constrained problem over a conic slice $\mathcal{C}$, namely
${\bm{\bar{W}}}(R)=\min_{\|{\bm{W}}\|_{\diamond}\leq
R,\bm{W}\in\mathcal{C}}{\cal{L}}(\bm{W})$. We will show that proper local RP
directionally converge to locally-optimal directions. Setting the cone to be
the rank-$m$ manifold $\mathcal{R}_{m}$ (and more specifically the set of all
matrices $\mathbb{R}^{d\times d}$), this will also establish the convergence
of global RP to the globally-optimal direction.

Our cone definition $\texttt{cone}_{\varepsilon}({\bm{\alpha}})$ is induced by
a token selection ${\bm{\alpha}}=(\alpha_{ik})_{ik=(1,1)}^{(n,k)}$ and has a
simple interpretation: It prioritizes tokens with lower score than
${\bm{\alpha}}$ over tokens with high-score than ${\bm{\alpha}}$. This way
lower score tokens create a barrier for ${\bm{\alpha}}$ and prevents
optimization to move towards higher score tokens.

######  Definition 5 (Low&High Score Tokens and Separating Cone)

Given $\alpha\in[T]$, input sequence ${\bm{X}}$ with label $Y$,
$h(\cdot):\mathbb{R}^{d}\rightarrow\mathbb{R}$, and score
$\bm{\gamma}_{t}=Y\cdot h(\bm{x}_{t})$ for all $t\in[T]$, define the low and
high score tokens as

$\texttt{low}^{\alpha}({\bm{X}})=\left\\{t\in[T]{~{}\big{|}~{}}\bm{\gamma}_{t}<\bm{\gamma}_{\alpha}\\},\quad\texttt{high}^{\alpha}({\bm{X}})=\\{t\in[T]-\\{\alpha\\}{~{}\big{|}~{}}\bm{\gamma}_{t}\geq\bm{\gamma}_{\alpha}\right\\}.$

For input ${\bm{X}}_{ik}$ and index $\alpha_{ik}$, we use the shorthand
notations $\texttt{low}_{ik}^{\alpha},\texttt{high}_{ik}^{\alpha}$. Finally
define $\texttt{cone}_{\varepsilon}({\bm{\alpha}})$ as

$\displaystyle\texttt{cone}_{\varepsilon}({\bm{\alpha}})=\left\\{\bm{W}\in\mathcal{R}_{m}{~{}\big{|}~{}}\min_{i\in[n]}\max_{t\in\texttt{low}_{ik}^{\alpha}}\min_{\tau\in\texttt{high}_{ik}^{\alpha}}\left<{\bm{F}}_{ikt}-{\bm{F}}_{ik\tau},\bm{W}\right>\geq\varepsilon\|{\bm{W}}\|_{F}\right\\}.$
(86)

###### Lemma 16

Consider the cone definition of (86) and suppose an SVM solution
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$ exists. If indices ${\bm{\alpha}}$ are
locally-optimal,
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}\in\texttt{cone}_{\varepsilon}({\bm{\alpha}})$
for all sufficiently small $\varepsilon>0$. Otherwise,
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}\not\in\texttt{cone}_{\varepsilon}({\bm{\alpha}})$
for all $\varepsilon>0$. Additionally, suppose optimal indices
$\texttt{opt}_{ik}\in\arg\max_{t\in[T]}\bm{\gamma}_{ikt}$ are unique. Then,
$\texttt{cone}_{\varepsilon}(\texttt{opt})=\mathcal{R}_{m}$.

Proof. Suppose ${\bm{\alpha}}$ is locally optimal. Observe that, thanks to
local optimality, $\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$ obeys

$\min_{t\in\mathcal{T}_{ik}}\left<{\bm{F}}_{ikt},\bm{W}^{\textsl{mm}}_{\bm{\alpha}}\right>>\max_{\tau\not\in\mathcal{T}_{ik}\cup\\{\alpha_{ik}\\}}\left<{\bm{F}}_{ik\tau},\bm{W}^{\textsl{mm}}_{\bm{\alpha}}\right>,$

for all $i\in[n]$. Next, observe that
$\mathcal{T}_{ik}\subseteq\texttt{low}_{ik}^{\alpha}$ and
$\texttt{high}_{ik}^{\alpha}\subseteq\bar{\mathcal{T}}_{ik}=[T]-\mathcal{T}_{ik}-\\{\alpha_{ik}\\}$.
Thus, the inequality (86) holds for small enough $\varepsilon>0$.

Conversely, suppose ${\bm{\alpha}}$ is not locally-optimal. Fix support index
$t\in\mathcal{T}_{ik}$ with $t\in\texttt{high}_{ik}^{\alpha}$. Since
$t\in\mathcal{T}_{ik}$, observe that

$\left<{\bm{F}}_{ikt},\bm{W}^{\textsl{mm}}_{\bm{\alpha}}\right>\geq\max_{\tau\neq\alpha_{ik}}\left<{\bm{F}}_{ik\tau},\bm{W}^{\textsl{mm}}_{\bm{\alpha}}\right>.$

In other words, for this $i\in[n]$, we found

$\max_{\tau\in\texttt{low}_{ik}^{\alpha}}\left<{\bm{F}}_{ik\tau}-{\bm{F}}_{ikt},\bm{W}^{\textsl{mm}}_{\bm{\alpha}}\right>\leq
0,$

violating (86) definition for any $\varepsilon>0$. To show the final claim,
observe that, setting ${\bm{\alpha}}:=\texttt{opt}$, we have that
$\texttt{high}_{ik}^{\alpha}=\emptyset$ for all $i\in[n]$ as
$\texttt{opt}_{ik}$ are unique optimal indices. Thus, there is no constraint
enforced on the cone definition in (86) making it equal to the rank-$m$
manifold $\mathcal{R}_{m}$.

Our main assumption regarding prediction head is a monotonicity condition
which is a strict generalization of linearity: We ask for $h$ to preserve the
order of token scores under convex combinations.

######  Assumption E ($h$ preserves the top score)

The functions $(h_{k})_{k=1}^{K}$ are $L_{h}$-Lipschitz in Euclidean distance.
Given ${\bm{\alpha}}=(\alpha_{ik})_{ik=(1,1)}^{(n,k)}$, there exists a scalar
$c:=c_{\bm{\alpha}}>0$ such that for all $i\in[n],k\in[K]$ the following
holds: Consider any convex combination

$\bm{x}(\bm{s})=\sum_{t\in\texttt{low}_{ik}^{\alpha}\cup\\{\alpha_{ik}\\}}s_{t}\cdot\bm{x}_{it}\quad\text{where}\quad\sum_{t\in\texttt{low}_{ik}^{\alpha}\cup\\{\alpha_{ik}\\}}s_{t}=1,~{}s_{t}\geq
0.$

We have that $Y_{ik}\cdot
h_{k}(\bm{x}(\bm{s}))\leq\bm{\gamma}^{\bm{\alpha}}_{ik}-c(1-s_{\alpha_{ik}})$
where $\bm{\gamma}^{\bm{\alpha}}_{ik}=Y_{ik}\cdot
h_{k}(\bm{x}^{\bm{\alpha}}_{{ik}})$ is the score of $\alpha_{ik}$.

This condition states that convex combinations of tokens with scores lower
than $\alpha_{ik}$ cannot achieve a score higher than $\alpha_{ik}$. Here,
$1-s_{\alpha_{ik}}$ term denotes the total share of non-optimal tokens. We
require this condition to hold over the training dataset rather than the full
domain $\mathbb{R}^{d}$. Crucially, it is a strict generalization of the
linearity assumption: Any linear $h_{k}$ satisfies Assumption E by setting
$c_{\bm{\alpha}}>0$ to be the difference between the score of $\alpha_{ik}$
and the largest score within $\texttt{low}_{ik}^{\alpha}$ i.e.

$\displaystyle
c_{\bm{\alpha}}:=\min_{i\in[n],k\in[K]}\\{\bm{\gamma}^{\bm{\alpha}}_{ik}-\max_{t\in\texttt{low}_{ik}^{\alpha}}\bm{\gamma}_{ikt}\\}>0.$
(87)

This can be seen by writing
$h_{k}(\bm{x}(\bm{s}))=\sum_{t\in\texttt{low}_{ik}^{\alpha}\cup\\{\alpha_{ik}\\}}s_{t}\cdot\bm{\gamma}_{ikt}=\bm{\gamma}^{\bm{\alpha}}_{ik}+\sum_{t\in\texttt{low}_{ik}^{\alpha}}s_{t}\cdot(\bm{\gamma}_{ikt}-\bm{\gamma}^{\bm{\alpha}}_{ik})\leq\bm{\gamma}^{\bm{\alpha}}_{ik}-c_{\bm{\alpha}}(1-s_{\alpha_{ik}})$.
To provide a nonlinear example, consider the setting all labels are $Y_{ik}=1$
and $h$ is an arbitrary convex function. Thanks to convexity, we can write
$h(\bm{x}(\bm{s}))\leq\sum_{t=1}^{T}s_{t}\cdot
h(\bm{x}_{t})=\sum_{t\in\texttt{low}_{ik}^{\alpha}\cup\\{\alpha_{ik}\\}}s_{t}\cdot\bm{\gamma}_{ikt}$.
Thus, we can use the same choice of $c_{\bm{\alpha}}$ in (87).

We remark that, in Section 6, we derive formulae describing general inductive
bias of attention without enforcing any assumption on $h$. These formulae
allow for arbitrary output tokens generated by the transformer model trained
by gradient descent. This includes the setting where gradient descent selects
and composes multiple tokens from each sequence rather than a single token
$\alpha_{ik}$.

The following result is our main theorem regarding the convergence of
regularization path to locally-optimal directions when restricted over the
cone $\texttt{cone}_{\varepsilon}({\bm{\alpha}})$.

######  Theorem 10 (Convergence of Local Regularization Path)

Suppose (SAtt-SVM) is feasible and
${\bm{\alpha}}=(\alpha_{ik})_{ik=(1,1)}^{(n,k)}$ are locally-optimal token
indices. Suppose Assumptions A&E hold. Recall
$\texttt{cone}_{\varepsilon}({\bm{\alpha}})$ of (86) and consider the norm-
constrained cone

$\mathcal{C}_{\varepsilon,R_{0}}^{\diamond}:=\texttt{cone}_{\varepsilon}({\bm{\alpha}})\bigcap\\{\bm{W}{~{}\big{|}~{}}\|{\bm{W}}\|_{\diamond}\geq
R_{0}\\}.$

Define the conic regularization path
${\bm{\bar{W}}}(R)=\min_{\mathcal{C}_{\varepsilon,R_{0}}^{\diamond},\|{\bm{W}}\|_{\diamond}\leq
R}{\cal{L}}(\bm{W})$. Let ${\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}$ be its set
of minima and $\Xi_{\diamond}>0$ be the associated margin i.e.
$\Xi_{\diamond}=1/\|{{\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}}\|_{\diamond}$.
For any sufficiently small $\varepsilon>0$ and sufficiently large
$R_{0}={\cal{O}}(1/\varepsilon)>0$,
$\lim_{R\rightarrow\infty}\texttt{dist}\left(\frac{{\bm{\bar{W}}}(R)}{R\Xi_{\diamond}},{\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}\right)=0$.
Additionally, suppose optimal indices
$\texttt{opt}=(\texttt{opt}_{ik})_{ik=(1,1)}^{(n,k)}$ are unique and set
${\bm{\alpha}}\leftarrow\texttt{opt}$. Then, the same RP convergence guarantee
holds with $\mathcal{C}_{\varepsilon,R_{0}}^{\diamond}=\mathcal{R}_{m}$.

Proof. We will prove that ${\bm{\bar{W}}}(R)$ is the optimal direction and
also $\|{{\bm{\bar{W}}}(R)}\|_{\diamond}\rightarrow\infty$. Define the
absolute constant

$c_{\diamond}=\min_{\|{\bm{W}}\|_{\diamond}=1}\|{\bm{W}}\|_{F}.$

This guarantees that for any $\bm{W}$ we have $\|{\bm{W}}\|_{F}\geq
c_{\diamond}\|{\bm{W}}\|_{\diamond}$. Also denote
$\varepsilon_{\diamond}=c_{\diamond}\varepsilon$. Let us first determine the
$\varepsilon$ parameter: Fix
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}\in{\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}$.
For general ${\bm{\alpha}}$, we can choose any $\varepsilon>0$ that is
sufficiently small to guarantee
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}\in\texttt{cone}_{\varepsilon}({\bm{\alpha}})$
based on Lemma 16. For ${\bm{\alpha}}=\texttt{opt}$, our analysis will
entirely avoid using $\varepsilon$, specifically, observe that
$\texttt{cone}_{\varepsilon}({\bm{\alpha}})=\mathcal{R}_{m}$ based on Lemma
16.

Step 1: Let us first prove that ${\bm{\bar{W}}}(R)$ achieves the optimal risk
as $R\rightarrow\infty$ – rather than problem having finite optima. Define
norm-normalized
$\bar{\bm{W}}^{\textsl{mm}}=\Xi_{\diamond}\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$.
Note that $\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$ separates tokens
${\bm{\alpha}}$ from rest of the tokens for each $i,k\in[n]\times[K]$. Thus,
we have that

$\displaystyle\lim_{R\rightarrow\infty}{\cal{L}}({\bm{\bar{W}}}(R))\leq\lim_{R\rightarrow\infty}{\cal{L}}(R\cdot\bar{\bm{W}}^{\textsl{mm}}):={\cal{L}}_{\star}=\frac{1}{n}\sum_{i=1}^{n}\sum_{k=1}^{K}\ell(\bm{\gamma}^{\bm{\alpha}}_{ik}).$
(88)

On the other hand, for any choice of
$\bm{W}\in\texttt{cone}_{\varepsilon}({\bm{\alpha}})$, set
$\bm{x}^{\bm{W}}_{ik}=\sum_{t=1}^{T}\mathbb{S}({\bm{X}}_{i}\bm{W}{\bm{z}}_{ik})_{t}\bm{x}_{t}$.
Set softmax probabilities
$\bm{s}^{({ik})}=\mathbb{S}({\bm{X}}_{i}\bm{W}{\bm{z}}_{ik})$. Recalling
$\texttt{low}_{ik}^{\alpha},\texttt{high}_{ik}^{\alpha}$ definitions, we can
decompose the attention features as

$\displaystyle\bm{x}^{\bm{W}}_{ik}=\bm{s}^{({ik})}_{\alpha_{ik}}\bm{x}^{\bm{\alpha}}_{ik}+\sum_{t\in\texttt{low}_{ik}^{\alpha}}\bm{s}^{({ik})}_{t}\bm{x}_{it}+\sum_{\tau\in\texttt{high}_{ik}^{\alpha}}\bm{s}^{({ik})}_{\tau}\bm{x}_{{i\tau}}.$
(89)

When ${\bm{\alpha}}=\texttt{opt}$, note that we simply have
$\texttt{high}_{ik}^{\alpha}=\emptyset$. This will be important for setting
$R_{0}=0$ and $\mathcal{C}_{\varepsilon,R_{0}}^{\diamond}=\mathcal{R}_{m}$ in
the proof for opt indices.

Set
$\gamma^{\textsl{gap}}_{ikt}=\bm{\gamma}_{ikt}-\bm{\gamma}^{\bm{\alpha}}_{ik}=Y_{ik}\cdot(h_{k}(\bm{x}_{it})-h_{k}(\bm{x}^{\bm{\alpha}}_{ik}))$.
Building on $L_{h}$-Lipschitzness of the prediction head $h_{k}(\cdot)$, we
define

$\displaystyle
B:=\max_{i\in[n],k\in[K]}\max_{t,\tau\in[T]}L_{h}\cdot\|{\bm{x}_{it}-\bm{x}_{i\tau}}\|\geq|\gamma^{\textsl{gap}}_{ikt}|.$
(90)

Define $P^{ik}:=\sum_{t\in\texttt{low}_{ik}^{\alpha}}\bm{s}^{({ik})}_{t}$,
$Q^{ik}:=\sum_{t\in\texttt{high}_{ik}^{\alpha}}\bm{s}^{({ik})}_{t}$, and
$\bm{\gamma}^{\bm{W}}_{ik}=Y_{ik}\cdot h_{k}(\bm{x}^{\bm{W}}_{ik})$. Also set
temporary variables
$\bm{x}^{\prime}=(\bm{s}^{({ik})}_{\alpha_{ik}}+Q^{ik})\bm{x}^{\bm{\alpha}}_{ik}+\sum_{t\in\texttt{low}_{ik}^{\alpha}}\bm{s}^{({ik})}_{t}\bm{x}_{it}$
and $\bm{\gamma}^{\prime}=Y_{ik}\cdot h_{k}(\bm{x}^{\bm{W}}_{ik})$. Using
Assumption E on $\bm{x}^{\prime}$ and noticing
$P^{ik}=1-\bm{s}^{({ik})}_{\alpha_{ik}}-Q^{ik}$, observe that

$|\bm{\gamma}^{\bm{W}}_{ik}-\bm{\gamma}^{\prime}|\leq
BQ^{ik}\quad\text{and}\quad\bm{\gamma}^{\prime}\leq\bm{\gamma}^{\bm{\alpha}}_{ik}-c_{\bm{\alpha}}P^{ik}.$

Recall from (87) that, when $h_{k}$ are linear functions, $c_{\bm{\alpha}}$
can be chosen as

$\displaystyle
c_{\bm{\alpha}}:=\min_{i\in[n],k\in[K]}\min_{t\in\texttt{low}_{ik}^{\alpha}}-\gamma^{\textsl{gap}}_{ikt}>0.$

To summarize, applying Assumption E, we obtain the following score
inequalities

$\displaystyle\bm{\gamma}^{\bm{W}}_{ik}\leq\bm{\gamma}^{\bm{\alpha}}_{ik}-c_{\bm{\alpha}}P^{ik}+BQ^{ik},$
(91)
$\displaystyle|\bm{\gamma}^{\bm{W}}_{ik}-\bm{\gamma}^{\bm{\alpha}}_{ik}|\leq
L_{h}\|{\bm{x}^{\bm{W}}_{ik}-\bm{x}^{\bm{\alpha}}_{ik}}\|\leq
L_{h}\sum_{t\neq\alpha_{ik}}\bm{s}^{({ik})}_{t}\|{\bm{x}_{ikt}-\bm{x}^{\bm{\alpha}}_{ik}}\|\leq
B(1-\bm{s}^{({ik})}_{\alpha_{ik}}).$ (92)

We will use the $\bm{\gamma}^{\bm{W}}_{ik}-\bm{\gamma}^{\bm{\alpha}}_{ik}$
term in (91) to evaluate $\bm{W}$ against the reference loss (88). Let
${\bm{a}}^{({ik})}={\bm{X}}_{i}\bm{W}{\bm{z}}_{ik}$. Now since
$\bm{W}\in\texttt{cone}_{\varepsilon}({\bm{\alpha}})$, there exists
$t\in\texttt{low}_{ik}^{\alpha}$ obeying
${\bm{a}}^{({ik})}_{t}-\max_{\tau\in\texttt{high}_{ik}^{\alpha}}{\bm{a}}^{({ik})}_{\tau}\geq\varepsilon\|{\bm{W}}\|_{F}\geq\varepsilon_{\diamond}\|{\bm{W}}\|_{\diamond}$.
Denote $D^{ik}:=(\sum_{t\in[T]}e^{{\bm{a}}^{({ik})}_{t}})^{-1}$ to be the
softmax denominator i.e. sum of exponentials. We find that,

$\displaystyle
Q^{ik}=\sum_{\tau\in\texttt{high}_{ik}^{\alpha}}\bm{s}^{({ik})}_{\tau}=D^{ik}\sum_{\tau\in\texttt{high}_{ik}^{\alpha}}e^{{\bm{a}}^{({ik})}_{\tau}}\leq
D^{ik}Te^{{\bm{a}}^{({ik})}_{t}-\varepsilon\|{\bm{W}}\|_{F}}\leq
Te^{-\varepsilon_{\diamond}\|{\bm{W}}\|_{\diamond}}P^{ik}.$ (93)

Consequently, the score difference obeys

$\bm{\gamma}^{\bm{W}}_{ik}-\bm{\gamma}^{\bm{\alpha}}_{ik}\leq
BQ^{ik}-c_{\bm{\alpha}}P^{ik}\leq(BTe^{-\varepsilon_{\diamond}\|{\bm{W}}\|_{\diamond}}-c_{\bm{\alpha}})P^{ik}.$

Above, the right hand side is strictly negative as soon as
$\|{\bm{W}}\|_{\diamond}\geq
R_{0}:=\frac{1}{\varepsilon_{\diamond}}\log\frac{BT}{c_{\bm{\alpha}}}$. Note
that, this condition applies to all $(i,k)\in[n]\times[K]$ pairs uniformly for
the same $R_{0}$. Consequently, for any $\|{\bm{W}}\|_{\diamond}\geq R_{0}$,
for all $i,k$ and $\bm{W}\in\texttt{cone}_{\varepsilon}({\bm{\alpha}})$, we
have that $\bm{\gamma}^{\bm{W}}_{ik}<\bm{\gamma}^{\bm{\alpha}}_{ik}$.
Additionally, when ${\bm{\alpha}}=\texttt{opt}$, note that $Q^{ik}=0$ since
$\texttt{high}_{ik}^{\alpha}=\emptyset$. Thus, $R_{0}=0$ suffices to ensure
$\bm{\gamma}^{\bm{W}}_{ik}<\bm{\gamma}^{\bm{\alpha}}_{ik}$. Using the
strictly-decreasing nature of $\ell$, we conclude with the fact that for all
(finite) $\bm{W}\in\texttt{cone}_{\varepsilon}({\bm{\alpha}})$,

${\cal{L}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\sum_{k=1}^{K}\ell(\bm{\gamma}^{\bm{W}}_{ik})>{\cal{L}}_{\star}=\frac{1}{n}\sum_{i=1}^{n}\sum_{k=1}^{K}\ell(\bm{\gamma}^{\bm{\alpha}}_{ik}),$

which implies $\|{{\bm{\bar{W}}}(R)}\|_{\diamond}\rightarrow\infty$.

Step 2: To proceed, we show that ${\bm{\bar{W}}}(R)$ converges in direction to
${\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}$. Suppose this is not the case i.e.
convergence fails. We will obtain a contradiction by showing that
$\bar{\bm{W}}^{\textsl{mm}}_{R}=R\cdot\bar{\bm{W}}^{\textsl{mm}}$ achieves a
strictly superior loss compared to ${\bm{\bar{W}}}(R)$. Also define the
normalized parameter
${\bm{\bar{W}}}_{0}(R)=\frac{{\bm{\bar{W}}}(R)}{R\Xi_{\diamond}}$ and
$\bm{W}^{\prime}=\frac{{\bm{\bar{W}}}(R)}{\|{{\bm{\bar{W}}}(R)}\|_{\diamond}\Xi_{\diamond}}$.
Note that ${\bm{\bar{W}}}_{0}(R)$ is obtained by scaling down
$\bm{W}^{\prime}$ since $\|{{\bm{\bar{W}}}(R)}\|_{\diamond}\leq R$ and
$\bm{W}^{\prime}$ obeys
$\|{{\bm{\bar{W}}}_{0}(R)}\|_{\diamond}=\|{\bm{W}^{\textsl{mm}}_{\bm{\alpha}}}\|_{\diamond}$.

Since ${\bm{\bar{W}}}_{0}(R)$ fails to converge to
${\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}$, for some $\delta>0$, there exists
arbitrarily large $R>0$ such that
$\texttt{dist}\left({\bm{\bar{W}}}_{0}(R),{\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}\right)\geq\delta$.
This translates to the suboptimality in terms of margin constraints as
follows: First, distance with respect to the ${\diamond}$-norm obeys
$\texttt{dist}_{\diamond}\left({\bm{\bar{W}}}_{0}(R),{\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}\right)\geq\delta$
for some updated $\delta\leftarrow c_{\diamond}\delta$. Secondly, using
triangle inequality,

$\text{This implies that
either~{}~{}~{}}\|{{\bm{\bar{W}}}_{0}(R)}\|_{\diamond}\leq\|{\bm{W}^{\textsl{mm}}_{\bm{\alpha}}}\|_{\diamond}-\delta/2\text{~{}~{}~{}or~{}~{}~{}}\texttt{dist}_{\diamond}\left(\bm{W}^{\prime},{\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}\right)\geq\delta/2.$

In either scenario, ${\bm{\bar{W}}}_{0}(R)$ strictly violates one of the
margin constraints of (SAtt-SVM): If
$\|{{\bm{\bar{W}}}_{0}(R)}\|_{\diamond}\leq\|{\bm{W}^{\textsl{mm}}_{\bm{\alpha}}}\|_{\diamond}-\delta/2$,
then, since the optimal SVM objective is
$\|{\bm{W}^{\textsl{mm}}_{\bm{\alpha}}}\|_{\diamond}$, there exists a
constraint $(i,k)$ for which
$\left<{\bm{F}}^{\bm{\alpha}}_{{ik}}-{\bm{F}}_{{ikt}},{\bm{\bar{W}}}_{0}(R)\right>\leq
1-\frac{\delta}{2\|{\bm{W}^{\textsl{mm}}_{\bm{\alpha}}}\|_{\diamond}}$. If
$\texttt{dist}_{\diamond}\left(\bm{W}^{\prime},{\cal{W}}^{\textsl{mm}}_{{\bm{\alpha}}}\right)\geq\delta/2$,
then, $\bm{W}^{\prime}$ has same SVM objective but it is strictly bounded away
from the solution set. Thus, for some $\varepsilon:=\varepsilon(\delta)>0$,
$\bm{W}^{\prime}$ and its scaled down version ${\bm{\bar{W}}}_{0}(R)$ strictly
violate an SVM constraint achieving margin $\leq 1-\varepsilon$. Without
losing generality, suppose ${\bm{\bar{W}}}_{0}(R)$ violates the first
constraint. Thus, for a properly updated $\delta>0$ choice and for
$(i,k)=(1,1)$ and some support index $\tau\in\mathcal{T}_{11}$,

$\displaystyle\left<{\bm{F}}^{\bm{\alpha}}_{{11}}-{\bm{F}}_{{11}t},{\bm{\bar{W}}}_{0}(R)\right>\leq
1-\delta.$ (94)

Now, we will argue that this will lead to a contradiction by proving
${\cal{L}}(\bar{\bm{W}}^{\textsl{mm}}_{R})<{\cal{L}}({\bm{\bar{W}}}(R))$ for
sufficiently large $R$.

To obtain the result, we establish a refined softmax probability control as in
Step 1 by studying distance to ${\cal{L}}_{\star}$. Following (91), denote the
score function at ${\bm{\bar{W}}}(R)$ via
$\bm{\gamma}^{R}_{ik}=\bm{\gamma}_{ik}^{{\bm{\bar{W}}}(R)}$ as shorthand
notation. Similarly, let $\bm{s}^{R}_{ik}=\mathbb{S}({\bm{a}}^{R}_{ik})$ with
${\bm{a}}^{R}_{ik}={\bm{X}}_{i}{\bm{\bar{W}}}(R){\bm{z}}_{ik}$. Set the
corresponding notation for the reference parameter
$\bar{\bm{W}}^{\textsl{mm}}_{R}$ as
$\bm{\gamma}^{\star}_{ik},\bm{s}^{\star}_{ik},{\bm{a}}^{\star}_{ik}$.

Critically, recall the above inequalities (93) that applies to both
$\bm{W}\in\\{{\bm{\bar{W}}}(R),\bar{\bm{W}}^{\textsl{mm}}_{R}\\}\subset\texttt{cone}_{\varepsilon}({\bm{\alpha}})$
for an index $(i,k)$ and support index $t\in\mathcal{T}_{ik}$

$\displaystyle Q^{ik}$
$\displaystyle=\sum_{\tau\in\texttt{high}_{ik}^{\alpha}}\bm{s}_{ik\tau}=D^{ik}\sum_{\tau\in\texttt{high}_{ik}^{\alpha}}e^{{\bm{a}}_{{ik\tau}}}$
$\displaystyle\leq
D^{ik}Te^{{\bm{a}}_{ikt}-\varepsilon_{\diamond}\|{\bm{W}}\|_{\diamond}}\leq
Te^{-\varepsilon_{\diamond}\|{\bm{W}}\|_{\diamond}}P^{ik}\leq
Te^{-\varepsilon_{\diamond}\|{\bm{W}}\|_{\diamond}}(1-\bm{s}_{{ik}\alpha_{ik}}),$
(95)

where $P^{ik}=\sum_{\tau\in\texttt{low}_{ik}^{\alpha}}\bm{s}_{{ik\tau}}$ and
$P^{ik}+Q^{ik}=1-\bm{s}_{{ik}\alpha_{ik}}$.

Note that, setting
$R_{0}\geq{\cal{O}}(1/\varepsilon_{\diamond})={\cal{O}}(1/\varepsilon)$, we
guarantee that, for any $(i,k)\in[n]\times[K]$

$\displaystyle P^{ik}\geq Q^{ik}\implies P^{ik}\geq
0.5(1-\bm{s}_{{ik}\alpha_{ik}}).$ (96)

Additionally, when ${\bm{\alpha}}=\texttt{opt}$, note that $Q^{ik}=0$ since
$\texttt{high}_{ik}^{\alpha}=\emptyset$. Thus, $R_{0}=0$ suffices to ensure
(96).

To proceed, recall that $R\geq\|{{\bm{\bar{W}}}(R)}\|_{\diamond}\geq R_{0}$ by
definition since
${\bm{\bar{W}}}(R)\in\mathcal{C}_{\varepsilon,R_{0}}^{\diamond}$ and
$\Xi_{\diamond}:=1/\|{\bm{W}^{\textsl{mm}}_{\bm{\alpha}}}\|_{\diamond}$. On
top of above, we note the following softmax inequalities on the selected
tokens $\alpha_{ik}$

$\displaystyle\bm{s}^{\star}_{{ik}\alpha_{ik}}\geq\frac{1}{1+Te^{-R\Xi_{\diamond}}}\geq
1-Te^{-R\Xi_{\diamond}}\quad\text{for all}\quad(i,k)\in[n]\times[K],$ (97)
$\displaystyle
s^{R}_{{ik}\alpha_{ik}}\leq\frac{1}{1+e^{-(1-\delta)\|{{\bm{\bar{W}}}(R)}\|_{\diamond}\Xi_{\diamond}}}\leq\frac{1}{1+e^{-(1-\delta)R\Xi_{\diamond}}}\quad\text{for}\quad(i,k)=(1,1).$

The former inequality is thanks to $\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$
achieving $\geq 1$ margins on all tokens $[T]-\alpha_{ik}$ and the latter
arises from the $\delta$-margin violation of ${\bm{\bar{W}}}(R)$ at
$(i,k)=(1,1)$ i.e. Eq. (94). Since $\ell$ is strictly decreasing with
Lipschitz gradient and the scores are upper/lower bounded by an absolute
constant (as tokens are bounded, $(h_{k})_{k=1}^{K}$ are Lipschitz, and both
are fixed), we know that
$c_{\texttt{up}}\geq-\ell^{\prime}(\bm{\gamma}_{ik}^{\bm{W}})\geq
c_{\texttt{dn}}$ for some constants $c_{\texttt{up}}>c_{\texttt{dn}}>0$. Thus,
following Eq. (90) and the score decomposition (91), and using (95),(96),(97)
we can write

$\displaystyle{\cal{L}}({\bm{\bar{W}}}(R))-{\cal{L}}_{\star}$
$\displaystyle\geq\frac{1}{n}[\ell(\bm{\gamma}_{11}^{{\bm{\bar{W}}}(R)})-\ell(\bm{\gamma}^{\bm{\alpha}}_{11})]\geq\frac{c_{\texttt{dn}}}{n}(\bm{\gamma}^{\bm{\alpha}}_{{11}}-\bm{\gamma}_{11}^{{\bm{\bar{W}}}(R)})$
$\displaystyle\geq\frac{c_{\texttt{dn}}}{n}(c_{\bm{\alpha}}P^{11}_{{\bm{\bar{W}}}(R)}-BQ^{11}_{{\bm{\bar{W}}}(R)})$
(98)
$\displaystyle\geq\frac{c_{\texttt{dn}}}{n}(1-\bm{s}_{{11}\alpha_{11}}^{R})(0.5c_{\bm{\alpha}}-BTe^{-\varepsilon_{\diamond}\|{{\bm{\bar{W}}}(R)}\|_{\diamond}})$
$\displaystyle\geq\frac{c_{\texttt{dn}}}{n}\frac{1}{1+e^{(1-\delta)R\Xi_{\diamond}}}(0.5c_{\bm{\alpha}}-BTe^{-\varepsilon_{\diamond}R_{0}}).$

Above, recalling the choice
$R_{0}\geq{\cal{O}}(1/\varepsilon_{\diamond})={\cal{O}}(1/\varepsilon)$,
$R\geq R_{0}$ implies $BTe^{-\varepsilon_{\diamond}R_{0}}\leq
c_{\bm{\alpha}}/4$ to obtain

$\displaystyle{\cal{L}}({\bm{\bar{W}}}(R))-{\cal{L}}_{\star}\geq\frac{c_{\texttt{dn}}\cdot
c_{\bm{\alpha}}}{4n}\frac{1}{1+e^{(1-\delta)R\Xi_{\diamond}}}.$ (99)

Additionally when ${\bm{\alpha}}=\texttt{opt}$, since
$Q^{11}_{{\bm{\bar{W}}}(R)}=0$ in (98), the bound above holds with $R_{0}=0$
by directly using (98).

Conversely, we upper bound the difference between
${\cal{L}}(\bar{\bm{W}}^{\textsl{mm}}_{R})$ and ${\cal{L}}_{\star}$ as
follows. Define the worst-case loss difference for ${\bm{\bar{W}}}(R)$ as
$(i^{\prime},k^{\prime})=\arg\max_{i\in[n],k\in[K]}[\ell(\bm{\gamma}_{ik}^{\star})-\ell(\bm{\gamma}^{\bm{\alpha}}_{ik})]$.
Using (92)&(97), we write

$\displaystyle{\cal{L}}(\bar{\bm{W}}^{\textsl{mm}}_{R})-{\cal{L}}_{\star}$
$\displaystyle\leq\max_{i\in[n],k\in[K]}[\ell(\bm{\gamma}_{ik}^{\star})-\ell(\bm{\gamma}^{\bm{\alpha}}_{ik})]\leq
c_{\texttt{up}}\cdot(\bm{\gamma}^{\bm{\alpha}}_{i^{\prime}k^{\prime}}-\bm{\gamma}^{\star}_{i^{\prime}k^{\prime}})$
(100) $\displaystyle\leq
c_{\texttt{up}}\cdot(1-\bm{s}_{1\alpha_{j}}^{\star})B$ $\displaystyle\leq
c_{\texttt{up}}\cdot Te^{-R\Xi_{\diamond}}B.$

Combining the last inequality and (99), we conclude that
${\cal{L}}(\bar{\bm{W}}^{\textsl{mm}}_{R})<{\cal{L}}({\bm{\bar{W}}}(R))$
whenever

$c_{\texttt{up}}T\cdot e^{-R\Xi_{\diamond}}B<\frac{c_{\texttt{dn}}\cdot
c_{\bm{\alpha}}}{4n}\frac{1}{1+e^{(1-\delta)R\Xi_{\diamond}}}\iff\frac{e^{R\Xi_{\diamond}}}{1+e^{(1-\delta)R\Xi_{\diamond}}}>\frac{4c_{\texttt{up}}TnB}{c_{\texttt{dn}}c_{\bm{\alpha}}}.$

The left hand-side inequality holds for all sufficiently large $R$:
Specifically, as soon as $R$ is large enough to ensure
$e^{(1-\delta)R\Xi_{\diamond}}>1$ and
$R>\frac{1}{\delta\Xi_{\diamond}}\log(\frac{4CTnB}{c_{\bm{\alpha}}})$. This
completes the proof of the theorem via contradiction as we obtained
${\cal{L}}({\bm{\bar{W}}}(R))>{\cal{L}}(\bar{\bm{W}}^{\textsl{mm}}_{R})$.

###  D.3 Global Regularization Path and Proof of Theorem 2

The following result is a direct corollary of Theorem 10. Namely, we simply
restate the final line of this theorem that applies to optimal tokens.

######  Corollary 1 (Global Convergence of Regularization Path)

Suppose Assumptions A&E hold and the optimal indices
$\texttt{opt}_{ik}=\arg\max_{t\in[T]}\bm{\gamma}_{ikt}$ are unique. Consider
the global regularization path
${\bm{\bar{W}}}_{{\diamond},R}=\min_{\bm{W}\in\mathcal{R}_{m},\|{\bm{W}}\|_{\diamond}\leq
R}{\cal{L}}(\bm{W})$. Let ${\cal{W}}^{\textsl{mm}}_{\diamond}$ be the non-
empty solution set of (SAtt-SVM) with ${\bm{\alpha}}\leftarrow\texttt{opt}$
normalized to have unit ${\diamond}$-norm. Then

$\lim_{R\rightarrow\infty}\texttt{dist}\left(\frac{{\bm{\bar{W}}}_{{\diamond},R}}{R},{\cal{W}}^{\textsl{mm}}_{\diamond}\right)$

The next corollary directly targets application to (SERM-W) and (SERM-KQ).
This corollary is also a strict generalization of Theorem 2. Specifically, we
immediately recover Theorem 2 by specializing this to the single-output
setting $K\leftarrow 1$ and full-dimensional parameterization $m\leftarrow d$.

###### Corollary 2

Suppose Assumptions A&E hold and the optimal indices
$\texttt{opt}_{ik}=\arg\max_{t\in[T]}\bm{\gamma}_{ikt}$ are unique. Consider
the regularization paths associated to (SERM-W) and (SERM-KQ):

$\displaystyle{\bm{\bar{W}}}_{R}=\underset{\bm{W}\in\mathcal{R}_{m},\|{\bm{W}}\|_{F}\leq
R}{\arg\min}{\cal{L}}(\bm{W})\quad\text{and}\quad{\bm{\bar{K}}}_{R},{\bm{\bar{Q}}}_{R}=\underset{\|{{\bm{K}}}\|_{F}^{2}+\|{{\bm{Q}}}\|_{F}^{2}\leq
2R}{\arg\min}{\cal{L}}({\bm{K}},{\bm{Q}})$ (101)

Suppose (SAtt-SVM) is feasible for ${\bm{\alpha}}\leftarrow\texttt{opt}$. Let
$\bm{W}^{\textsl{mm}}$ be the unique solution of (SAtt-SVM) with Frobenius
norm and ${\cal{W}}^{\textsl{mm}}_{\star}$ be the solution set of (SAtt-SVM)
with nuclear norm and cost function
$\|{{\cal{W}}^{\textsl{mm}}_{\star}}\|_{\star}$. We have that

$\lim_{R\rightarrow\infty}\frac{{\bm{\bar{W}}}_{R}}{R}=\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}},\quad\lim_{R\rightarrow\infty}\texttt{dist}\left(\frac{{\bm{\bar{Q}}}_{R}{\bm{\bar{K}}}_{R}^{\top}}{R},\frac{{\cal{W}}^{\textsl{mm}}_{\star}}{\|{{\cal{W}}^{\textsl{mm}}_{\star}}\|_{\star}}\right)=0.$

Proof. We directly apply Corollary 1 with ${\diamond}=F$ and
${\diamond}=\star$ respectively. To obtain the result on ${\bm{\bar{W}}}_{R}$,
we note that $\bm{W}^{\textsl{mm}}$ is unique because Frobenius norm-squared
is strongly convex. To obtain the result on
$({\bm{\bar{Q}}}_{R},{\bm{\bar{K}}}_{R})$, we use Lemma 15 and observe that

${\bm{\bar{W}}}_{\star,R}:={\bm{\bar{Q}}}_{R}{\bm{\bar{K}}}_{R}^{\top}\in\underset{\bm{W}\in\mathcal{R}_{m},\|{\bm{W}}\|_{\star}\leq
R}{\arg\min}{\cal{L}}(\bm{W}).$

We then apply Corollary 1 with ${\diamond}=\star$ to conclude with the
convergence of the path ${\bm{\bar{W}}}_{\star,R}$.

##  Appendix E Proof of Theorem 8: Separability Under Mild Over-
Parameterization

Proof. We denote kronecker product of two matrices via $\otimes$.
Additionally, given $\bm{W}\in\mathbb{R}^{d\times d}$, let us denote its
vectorization $\bm{w}=\text{vec}(\bm{W})\in\mathbb{R}^{d^{2}}$. We first note
that separation is implied by the linear independence of the constraints.
Specifically, let
$\bm{f}_{it}=(\bm{x}^{\bm{\alpha}}_{i}-\bm{x}_{it})\otimes{\bm{z}}_{i}$. We
are interested in guaranteeing

$\left<\bm{w},\bm{f}_{it}\right>\geq 1\quad\text{for all}\quad
i\in[n],t\neq{\bm{\alpha}}_{i}.$

Note that, the constraints above are feasible as soon as $\bm{f}_{it}$’s are
linearly independent. Thus, we will instead prove linear independence of the
vectors $(\bm{f}_{it})_{i\in[n],t\neq\alpha_{i}}$. Also note that, since there
are finitely many ${\bm{\alpha}}$ choices, if we show almost sure separation
for a fixed but arbitrary ${\bm{\alpha}}$ choice, through union bound, we
recover the result for all ${\bm{\alpha}}$. Thus, we prove the result for a
fixed ${\bm{\alpha}}$ choice.

We will prove this result inductively. Let
${\bm{M}}_{n-1}\in\mathbb{R}^{(n-1)(T-1)\times d^{2}}$ denote the matrix whose
rows are given by the features $(\bm{f}_{it})_{i\in[n-1],t\neq\alpha_{i}}$.
Suppose the result is correct for $n-1$, thus, ${\bm{M}}_{n-1}$ is full row-
rank almost surely (post random gaussian perturbation). Now, fix
${\bm{M}}_{n-1}$ and, conditioned on ${\bm{M}}_{n-1}$ being full row-rank, let
us show that ${\bm{M}}_{n}$ is also full row-rank almost surely. To prove
this, consider the $n$’th example $({\bm{X}}_{n},{\bm{z}}_{n})$. Let
$(\bm{g}_{t})_{t=1}^{n},\bm{h}\in\mathbb{R}^{d}$ be random vectors with i.i.d.
$\mathcal{N}(0,\sigma^{2})$ entries. Consider the perturbed input
${\bm{X}}^{\prime}_{n}$ with tokens
$\bm{x}^{\prime}_{nt}=\bm{x}_{nt}+\bm{g}_{t}$ and
${\bm{z}}^{\prime}_{n}={\bm{z}}_{n}+\bm{h}$. Note that for self-attention, we
set ${\bm{z}}_{n}=\bm{x}_{n1}$ and $\bm{h}=\bm{g}_{1}$. From these, create the
matrix $\tilde{{\bm{M}}}_{n}$ with rows
$(\bm{f}^{\prime}_{nt})_{t\neq\alpha_{i}}$ where
$\bm{f}^{\prime}_{nt}=\bm{x}^{\prime}_{nt}\otimes{\bm{z}}^{\prime}_{n}$.
Observe that ${\bm{M}}_{n}=\begin{bmatrix}\tilde{{\bm{M}}}_{n}\\\
{\bm{M}}_{n-1}\end{bmatrix}$. To conclude with the result, we will apply Lemma
17. To apply this lemma, we have two claims.

Claim 1: Let $\bar{{\bm{z}}}_{n}$ be the projection of ${\bm{z}}^{\prime}_{n}$
on the orthogonal complement of $({\bm{z}}_{i})_{i=1}^{n-1}$. Consider the
matrix $\bar{{\bm{M}}}_{n}$ with rows
$\bar{\bm{f}}_{nt}=(\bm{x}^{\prime}_{n\alpha_{n}}-\bm{x}^{\prime}_{nt})\otimes\bar{{\bm{z}}}_{n}$
for $t\neq\alpha_{n}$. $\bar{{\bm{M}}}_{n}$ is rank $T-1$ almost surely
whenever $d\geq\max(T-1,n)$.

To see this claim, first denote the orthogonal complement of the span of the
vectors $({\bm{z}}_{i})_{i=1}^{n-1}$ by $Z_{n-1}$. The span is at most $n-1$
dimensional and, since $d\geq n$, $\text{dim}(Z_{n-1})\geq 1$. Consequently,
$\bar{{\bm{z}}}_{n}\neq 0$ almost surely because the gaussian variable
${\bm{z}}_{n}+\bm{h}$ will have nonzero projection on $Z_{n-1}$ almost surely.
Secondly, let $\bar{{\bm{X}}}\in\mathbb{R}^{(T-1)\times d}$ be the matrix
whose rows are equal to $\bm{x}^{\prime}_{n\alpha_{n}}-\bm{x}^{\prime}_{nt}$
for $t\neq\alpha_{n}$. $\bar{{\bm{X}}}$ is full row-rank almost surely, this
is because: Conditioned on $\bm{g}_{\alpha_{n}}$, the matrix $\bar{{\bm{X}}}$
is written as $\bar{{\bm{X}}}=\tilde{{\bm{X}}}+{\bm{G}}$ where
$\tilde{{\bm{X}}}$ is deterministic and ${\bm{G}}$ is i.i.d. gaussian. The
latter perturbation ensures full row-rank almost surely whenever $T-1\leq d$.
Finally, note that
$\bar{{\bm{M}}}_{n}=\bar{{\bm{X}}}\otimes\bar{{\bm{z}}}_{n}$. Since the rank
of the Kronecker product is multiplicative, we conclude with the claim.

Claim 2: Let $S_{n-1}\subset\mathbb{R}^{d^{2}}$ be the null space of
${\bm{M}}_{n-1}$. There exists a subspace $P\subseteq S_{n-1}$ such that rows
of $\bar{{\bm{M}}}_{n}$ are projections of the rows of ${\bm{M}}_{n}$ on $P$,
that is, $\bar{\bm{f}}_{nt}=\Pi_{P}(\bm{f}^{\prime}_{nt})$ where $\Pi$ denotes
set projection.

To show this claim, let us consider the matrix forms of the vectorized
features i.e. let us work with $\mathbb{R}^{d\times d}$ rather than
$R^{d^{2}}$. Denote the notation change as
$\bm{F}_{it}=(\bm{x}^{\bm{\alpha}}_{i}-\bm{x}_{it}){\bm{z}}_{i}^{\top}\leftrightarrow\bm{f}_{it}=(\bm{x}^{\bm{\alpha}}_{i}-\bm{x}_{it})\otimes{\bm{z}}_{i}$.
Recall that $Z_{n-1}$ denotes the orthogonal complement of
$({\bm{z}}_{i})_{i=1}^{n-1}$. Define $Q$ to be the set of matrices in
$\mathbb{R}^{d\times d}$ whose column space lies in $Z_{n-1}$ and $P$ to be
the vectorization of $Q$. We first show that $P$ is a subset of the null space
of $S_{n-1}$. To see this, fix any matrix ${\bm{A}}\in P$ and a row
$\bm{f}_{it}$ from ${\bm{M}}_{n-1}$. Matricized $\bm{f}_{it}$ can be written
as $\bm{F}_{it}={\bm{a}}{\bm{z}}_{i}^{\top}$ for ${\bm{z}}_{i}\in
Z_{n-1}^{\perp}$. Since ${\bm{A}}\in Q$, this implies
$\left<\bm{F}_{it},{\bm{A}}\right>={\bm{a}}^{\top}{\bm{A}}{\bm{z}}_{i}=0$ as
${\bm{A}}{\bm{z}}_{i}=0$. This holds for all $\bm{F}_{it}$, thus,
$\texttt{vectorized}({\bm{A}})\in\texttt{null}(S_{n-1})$.

Next, we need to show that $\bar{\bm{f}}_{nt}$ is the projection of
$\bm{f}^{\prime}_{nt}$ on $P$. To see this, we will show that
$\bar{\bm{f}}_{nt}\in P$ whereas $\bm{f}^{\prime}_{nt}-\bar{\bm{f}}_{nt}\in
P^{\perp}$ for all $t$. Write
$\bm{F}^{\prime}_{nt}=\texttt{matricized}(\bm{f}^{\prime}_{nt})={\bm{a}}{{\bm{z}}^{\prime}}_{n}^{\top}$.
We have that
$\bar{\bm{F}}_{nt}=\texttt{matricized}(\bar{\bm{f}}_{nt})={\bm{a}}\bar{{\bm{z}}}_{n}^{\top}$
where $\bar{{\bm{z}}}_{n}=\Pi_{Z_{n-1}}({\bm{z}}^{\prime}_{n})$. This implies
$\bar{\bm{F}}_{nt}\in Q$ and $\bar{\bm{f}}_{nt}\in P$. Similarly, since
${\bm{z}}^{\prime}_{n}-\bar{{\bm{z}}}_{n}\in Z_{n-1}^{\perp}$ which implies
$\bm{F}^{\prime}_{nt}-\bar{\bm{F}}_{nt}\in Q^{\perp}$.

To conclude with the proof, observe that, through Claims 1 and 2,
${\bm{M}}_{n}=\begin{bmatrix}\tilde{{\bm{M}}}_{n}\\\
{\bm{M}}_{n-1}\end{bmatrix}$ satisfies the requirements of Lemma 17 almost
surely, namely, projection of $\tilde{{\bm{M}}}_{n}$ onto a subset of the null
space of ${\bm{M}}_{n-1}$ being full rank. Thus, ${\bm{M}}_{n}$ is full rank
almost surely.

###### Lemma 17

Let ${\bm{A}}\in\mathbb{R}^{n\times p},{{\bm{B}}}\in\mathbb{R}^{m\times p}$.
Suppose $n+m\leq p$ and ${\bm{A}}$ is full row-rank. Denote the null space of
${\bm{A}}$ by $S_{A}^{\perp}$. Let $P$ be a subspace that is its subset i.e.
$P\subseteq S_{A}^{\perp}$. Let ${{\bm{B}}}^{\prime}$ be the matrix obtained
by projecting each of row of ${{\bm{B}}}$ on $P$ and suppose
${{\bm{B}}}^{\prime}$ is full rank. Then, the concatenation
${\bm{C}}=[{\bm{A}};{{\bm{B}}}]$ is full row-rank.

Proof. Let $({\bm{a}}_{i})_{i=1}^{n}$, $(\bm{b}_{i})_{i=1}^{m}$,
$(\bm{b}^{\prime}_{i})_{i=1}^{m}$ be the rows of
${\bm{A}},{{\bm{B}}},{{\bm{B}}}^{\prime}$ respectively, Suppose the set of
rows of ${\bm{A}}$ and ${{\bm{B}}}$ are linearly dependent. Then, for some
$(c_{i})_{i=1}^{n},(c^{\prime}_{i})_{i=1}^{m}$ (which are not all-zeros), we
have that

$\displaystyle\sum_{i=1}^{n}c_{i}{\bm{a}}_{i}+\sum_{i=1}^{m}c^{\prime}_{i}\bm{b}_{i}=0.$
(102)

We now rewrite this as follows to decouple $P$ and $P^{\perp}$:

$\sum_{i=1}^{n}c_{i}{\bm{a}}_{i}+\sum_{i=1}^{m}c^{\prime}_{i}\bm{b}^{\prime}_{i}+\sum_{i=1}^{m}c^{\prime}_{i}(\bm{b}^{\prime}_{i}-\bm{b}_{i})=0.$

The two terms on the left side lie in $P^{\perp}$. Projecting above inequality
to $P$, we find that $\sum_{i=1}^{m}c^{\prime}_{i}\bm{b}^{\prime}_{i}=0$.
Since $(\bm{b}^{\prime}_{i})_{i=1}^{m}$ are linearly independent, we find
$c^{\prime}_{i}=0$ for all $i\in[m]$. This implies
$\sum_{i=1}^{n}c_{i}{\bm{a}}_{i}=0$. Since $({\bm{a}}_{i})_{i=1}^{n}$ are
linearly independent, this implies $c_{i}=0$. Thus, (102) can only hold if all
coefficients are zero which is a contradiction.

##  Appendix F Supporting Experiments

Varying $m$$1-$correlation coefficient Figure 10: Convergence behavior of GD
when training attention weights $({\bm{K}},{\bm{Q}})\in\mathbb{R}^{d\times m}$
with random data and varying $m$. The misalignment between attention SVM and
GD,
$1-{\texttt{corr\\_coef}}(\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}},{\bm{K}}{\bm{Q}}^{\top})$,
is studied. $\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}}$ is from (Att-SVM⋆)
with GD tokens ${\bm{\alpha}}$ and $m=d$. Subfigures with fixed $n=5$ and
$T=5$ show that as $m$ approaches or exceeds $n$, ${\bm{K}}{\bm{Q}}^{\top}$
aligns more with $\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}}$.

Iterations$1-$correlation coefficient (a) Evolution of correlation under
varying $d$

Masked token threshold ($\Gamma$) (b) $\Gamma$ vs correlation coefficient

Figure 11: Behavior of GD with nonlinear nonconvex prediction head and multi-
token compositions. (a): Blue, green, red and teal curves represent the
evolution of $1-{\texttt{corr\\_coef}}(\bm{W},\bm{W}^{\texttt{SVMeq}})$ for
$d=4,6,8$ and $10$ respectively, which have been displayed in Figure 8(upper).
(b): Over the $500$ random instances as discussed in Figure 8, we filter
different instances by constructing masked set with tokens whose softmax
output $<\Gamma$ and vary $\Gamma$ from $10^{-16}$ to $10^{-6}$. The
corresponding results of
$1-{\texttt{corr\\_coef}}(\bm{W},\bm{W}^{\texttt{SVMeq}})$ are displayed in
blue, green, red and teal curves.

In this section, we introduce implementation details and additional
experiments. Code is available at

https://github.com/ucr-optml/max_margin_attention.

We create a 1-layer self-attention using PyTorch, training it with the SGD
optimizer and a learning rate of $\eta=0.1$. We apply normalized gradient
descent to ensure divergence of attention weights. The attention weight
$\bm{W}$ is then updated through

$\bm{W}(k+1)=\bm{W}(k)-\eta\frac{\nabla{\cal{L}}(\bm{W}(k))}{\|\nabla{\cal{L}}(\bm{W}(k))\|_{F}}.$

In the setting of $({\bm{K}},{\bm{Q}})$-parameterization, we noted that with
extended training iterations, the norm of the combined parameter
${\bm{K}}{\bm{Q}}^{\top}$ consistently rises, despite the gradient being
treated as zero due to computational limitations. To tackle this issue, we
introduce a minor regularization penalty to the loss function, ensuring that
the norms of ${\bm{K}}$ and ${\bm{Q}}$ remain within reasonable bounds. This
adjustment involves

$\widetilde{\cal{L}}({\bm{K}},{\bm{Q}})={\cal{L}}({\bm{K}},{\bm{Q}})+\lambda(\|{\bm{K}}\|^{2}_{F}+\|{\bm{Q}}\|^{2}_{F}).$

Here, we set $\lambda$ to be the the smallest representable number, e.g.
computed as $1+\lambda=1$ in Python, which is around $2.22\times 10^{-16}$.
Therefore, ${\bm{K}},{\bm{Q}}$ parameters are updated as follows.

${\bm{K}}(k+1)={\bm{K}}(k)-\eta\frac{\nabla\widetilde{\cal{L}}_{\bm{K}}({\bm{K}}(k),{\bm{Q}}(k))}{\|\nabla\widetilde{\cal{L}}_{\bm{K}}({\bm{K}}(k),{\bm{Q}}(k))\|_{F}},\qquad{\bm{Q}}(k+1)={\bm{Q}}(k)-\eta\frac{\nabla\widetilde{\cal{L}}_{\bm{Q}}({\bm{K}}(k),{\bm{Q}}(k))}{\|\nabla\widetilde{\cal{L}}_{\bm{Q}}({\bm{K}}(k),{\bm{Q}}(k))\|_{F}}.$

$\bullet$ As observed in previous work [TLZO23], and due to the exponential
expression of softmax nonlinearity and computation limitation, PyTorch has no
guarantee to select optimal tokens when the score gap is too small. Therefore
in Figures 2, 5 and 6, we generate random tokens making sure that
$\min_{i\in[n],t\neq\texttt{opt}_{i}}\bm{\gamma}_{i\texttt{opt}_{i}}-\bm{\gamma}_{it}\geq\underline{\gamma}$
and we choose $\underline{\gamma}=0.1$ in our experiments.

##### Rank sensitivity of $({\bm{K}},{\bm{Q}})$-parameterization (Figure 11).

In Figure 4 and Lemma 1, we have both theoretically and empirically
established that the rank of the SVM solution, denoted as
$\bm{W}^{\textsl{mm}}$ in (Att-SVM) or $\bm{W}^{\textsl{mm}}_{\star}$ in (Att-
SVM⋆), is at most rank $\max(n,d)$. Now, moving to Figure 11, we delve into GD
performance across various dimensions of
${\bm{K}},{\bm{Q}}\in\mathbb{R}^{d\times m}$ while keeping $d=20$ fixed and
varying $m$ from $1$ to $10$. In the upper subfigure, we maintain a constant
$n=5$ and vary $T$ within $\\{5,10,15\\}$, while in the lower subfigure, $T$
is fixed at $5$ and $n$ changes within $\\{5,10,15\\}$. Results are depicted
using blue, green, and red dashed curves, with both $y$-axes representing
$1-{\texttt{corr\\_coef}}(\bm{W},\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}})$,
where $\bm{W}$ represents the GD solution and
$\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}}$ is obtained from (Att-SVM⋆) by
employing token indices ${\bm{\alpha}}$ selected via GD and setting the rank
limit to $m=d$. Observing both subfigures, we note that a larger $n$
necessitates a larger $m$ for attention weights ${\bm{K}}{\bm{Q}}^{\top}$ to
accurately converge to the SVM solution (Figure 11(lower)). Meanwhile,
performances remain consistent across varying $T$ values (Figure 11(upper)).
This observation further validates Lemma 1. Furthermore, the results
demonstrate that $\bm{W}$ converges directionally towards
$\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}}$ as long as $m\gtrsim n$, thereby
confirming the assertion in our Theorem 7.

##### Behavior of GD with nonlinear nonconvex prediction head and multi-token
compositions (Figure 11).

To better investigate how correlation changes with data dimension $d$, we
collect the solid curves in Figure 8(upper) and construct as Figure 11(a).
Moreover, Figure 11(b) displays the average correlation of instances (refer to
scatters in Figure 8 (lower)), considering masked tokens with softmax
probability $<\Gamma$. Both findings highlight that higher $d$ enhances
alignment. For $d\geq 8$ or $\Gamma\leq 10^{-9}$, the GD solution $\bm{W}$
achieves a correlation of $>0.99$ with the SVM-equivalence
$\bm{W}^{\texttt{SVMeq}}$, defined in Section 6.

$\tau$$\lambda$ (a) $\tau$ and $\lambda$ parameters relationship

$\tau$# of selected tokens (b) $\tau$ and # of selected tokens relationship

Probabilities# of selected tokens (c) Distribution of # selected tokens over
varying $\tau$

Figure 12:  Behavior of GD when selecting multiple tokens.

##### Investigation of Lemma 6 over different $\tau$ selections (Figure 12).

Consider the setting of Section 6.1 and Lemma 6. Figure 9 explores the
influence of $\lambda$ on the count of tokens selected by GD-derived attention
weights. As $\lambda$ increases, the likelihood of selecting more tokens also
increases. Shifting focus to Figure 12, we examine the effect of $\tau$. For
each outcome, we generate random $\lambda$ values, retaining pairs
$(\lambda,{\bm{X}})$ satisfying $\tau$ constraints, with averages derived from
$100$ successful trials. The results indicate a positive correlation among
$\tau$, $\lambda$, and the number of selected tokens. Moreover, Figure 12(c)
provides a precise distribution of selected token counts across various $\tau$
values (specifically $\tau\in\\{3,5,7,9\\}$). The findings confirm that the
number of selected tokens remains within the limit of $\tau$, thus validating
the assertion made in Lemma 6.