# Transformers as Support Vector Machines

Davoud Ataee Tarzanagh1⋆ Yingcong Li2⋆ Christos Thrampoulidis3 Samet Oymak4†

###### Abstract

Since its inception in “Attention Is All You Need”, the transformer
architecture has led to revolutionary advancements in natural language
processing. The attention layer within the transformer admits a sequence of
input tokens ${\bm{X}}$ and makes them interact through pairwise similarities
computed as
$\texttt{softmax}({\bm{X}}{\bm{Q}}{\bm{K}}^{\top}{\bm{X}}^{\top})$, where
$({\bm{K}},{\bm{Q}})$ are the trainable key-query parameters. In this work, we
establish a formal equivalence between the optimization geometry of self-
attention and a hard-margin SVM problem that separates optimal input tokens
from non-optimal tokens using linear constraints on the outer-products of
token pairs. This formalism builds on [TLZO23] and allows us to characterize
the implicit bias of 1-layer transformers optimized with gradient descent, as
follows. (1) Optimizing the attention layer, parameterized by
$({\bm{K}},{\bm{Q}})$, with vanishing regularization, converges in direction
to an SVM solution minimizing the nuclear norm of the combined parameter
$\bm{W}:={\bm{K}}{\bm{Q}}^{\top}$. Instead, directly parameterizing by
$\bm{W}$ minimizes a Frobenius norm SVM objective. We characterize this
convergence, highlighting that it can occur in locally-optimal directions
rather than global ones. (2) Complementing this, for
$\bm{W}$-parameterization, we prove the local/global directional convergence
of gradient descent under suitable geometric conditions. Importantly, we show
that over-parameterization catalyzes global convergence by ensuring the
feasibility of the SVM problem and by guaranteeing a benign optimization
landscape devoid of stationary points. (3) While our theory applies primarily
to linear prediction heads, we propose a more general SVM equivalence that
predicts the implicit bias of 1-layer transformers with nonlinear heads/MLPs.
Our findings apply to general datasets, trivially extend to cross-attention
layer, and their practical validity is verified via thorough numerical
experiments. We also introduce open problems and future research directions.
We believe these findings inspire a new perspective, interpreting multilayer
transformers as a hierarchy of SVMs that separates and selects optimal tokens.

††footnotetext: 1 University of Pennsylvania<EMAIL_ADDRESS>2 University
of California, Riverside<EMAIL_ADDRESS>3 University of British Columbia,
<EMAIL_ADDRESS>4 University of Michigan<EMAIL_ADDRESS>⋆ Equal
contribution. † Corresponding author.

$\bm{W}$$({\bm{K}},{\bm{Q}})$ Figure 1: GD convergence during training of
cross-attention weight $\bm{W}$ or $({\bm{K}},{\bm{Q}})$ with data. Teal and
yellow markers represent tokens from ${\bm{X}}_{1}$ and ${\bm{X}}_{2}$, while
stars mark optimal tokens. Solid and dashed lines depict attention-SVM (Att-
SVM and Att-SVM⋆ respectively) directions mapped to ${\bm{z}}_{1}$ (red) and
${\bm{z}}_{2}$ (blue), with arrows illustrating GD trajectories converging
towards these SVM directions. Green circles denote GD $\leftrightarrow$ SVM
pairings.

PercentageVarying $d$Not LocalGlobalLocal Figure 2:  Percentage of different
convergence types when training cross-attention weights ($\bm{W}$) using GD
and varying dimension ($d$). Red and blue bars represent the percentages of
convergence to globally-optimal and locally-optimal (including global) SVM
solutions, respectively. Teal bars are complements of the blue bars. Larger
overparameterization ($d$) increases the likelihood of global convergence.

##  1 Introduction

Self-attention, the central component of the transformer architecture, has
revolutionized natural language processing (NLP) by empowering the model to
identify complex dependencies within input sequences [VSP+17]. By assessing
the relevance of each token to every other token, self-attention assigns
varying degrees of importance to different parts of the input sequence. This
mechanism has proven highly effective in capturing long-range dependencies,
which is essential for applications arising in NLP [KT19, BMR+20, RSR+20],
computer vision [FXM+21, LLC+21, TCD+21, CSL+23], and reinforcement learning
[JLL21, CLR+21, WWX+22]. Remarkable success of the self-attention mechanism
and transformers has paved the way for the development of sophisticated
language models such as GPT4 [Ope23], Bard [Goo23], LLaMA [TLI+23], and
ChatGPT [Ope22].

> Q: Can we characterize the optimization landscape and implicit bias of
> transformers?
>
> How does the attention layer select and compose tokens when trained with
> gradient descent?

We address these questions by rigorously connecting the optimization geometry
of the attention layer and a hard max-margin SVM problem, namely (Att-SVM),
that separates and selects the optimal tokens from each input sequence. This
formalism, which builds on the recent work [TLZO23], is practically meaningful
as demonstrated through experiments, and sheds light on the intricacies of
self-attention. Throughout, given input sequences
${\bm{X}},\bm{Z}\in\mathbb{R}^{T\times d}$ with length $T$ and embedding
dimension $d$, we study the core cross-attention and self-attention models:

$\displaystyle f_{\texttt{cross}}({\bm{X}},\bm{Z})$
$\displaystyle:=\mathbb{S}(\bm{Z}{\bm{Q}}\bm{K}^{\top}{\bm{X}}^{\top}){\bm{X}}\bm{V},$
(1a) $\displaystyle f_{\texttt{self}}({\bm{X}})$
$\displaystyle:=\mathbb{S}({\bm{X}}{\bm{Q}}\bm{K}^{\top}{\bm{X}}^{\top}){\bm{X}}\bm{V}.$
(1b)

Here, ${\bm{K}},{\bm{Q}}\in\mathbb{R}^{d\times m}$,
$\bm{V}\in\mathbb{R}^{d\times v}$ are the trainable key, query, value matrices
respectively; $\mathbb{S}(\cdot)$ denotes the softmax nonlinearity, which is
applied row-wise on ${\bm{X}}{\bm{Q}}\bm{K}^{\top}{\bm{X}}^{\top}$. Note that
self-attention (1b) is a special instance of the cross-attention (1a) by
setting $\bm{Z}\leftarrow{\bm{X}}$. To expose our main results, suppose the
first token of $\bm{Z}$ – denoted by ${\bm{z}}$ – is used for prediction.
Concretely, given a training dataset
$(Y_{i},{\bm{X}}_{i},{\bm{z}}_{i})_{i=1}^{n}$ with labels $Y_{i}\in\\{-1,1\\}$
and inputs ${\bm{X}}_{i}\in\mathbb{R}^{T\times
d},{\bm{z}}_{i}\in\mathbb{R}^{d}$, we consider the empirical risk minimization
with a decreasing loss function $\ell(\cdot):\mathbb{R}\rightarrow\mathbb{R}$,
represented as follows:

$\displaystyle{\cal{L}}(\bm{K},{\bm{Q}})=\frac{1}{n}\sum_{i=1}^{n}\ell\left(Y_{i}\cdot
f({\bm{X}}_{i},{\bm{z}}_{i})\right),~{}~{}~{}~{}\text{where}~{}~{}f({\bm{X}}_{i},{\bm{z}}_{i})=h\left({\bm{X}}^{\top}_{i}\mathbb{S}\left({\bm{X}}_{i}\bm{K}{\bm{Q}}^{\top}{\bm{z}}_{i}\right)\right).$
(2)

Here, $h(\cdot):\mathbb{R}^{d}\rightarrow\mathbb{R}$ is the prediction head
that subsumes the value weights ${\bm{V}}$. In this formulation, the model
$f(\cdot)$ precisely represents a one-layer transformer where an MLP follows
the attention layer. Note that, we recover the self-attention in (2) by
setting ${\bm{z}}_{i}\leftarrow\bm{x}_{i1}$, where $\bm{x}_{i1}$ denotes the
first token of the sequence ${\bm{X}}_{i}$ 111Note that for simplicity, we set
${\bm{z}}_{i}=\bm{x}_{i1}$, but it can be any other row of ${\bm{X}}_{i}$..
The softmax operation, due to its nonlinear nature, poses a significant
challenge when optimizing (2). The problem is nonconvex and nonlinear even
when the prediction head is fixed and linear. In this study, we focus on
optimizing the attention weights (${\bm{K}},{\bm{Q}}$ or $\bm{W}$) and
overcome such challenges to establish a fundamental SVM equivalence.222We fix
$h(\cdot)$ and only optimize the attention weights. This is partly to avoid
the degenerate case where $h(\cdot)$ can be used to achieve zero training loss
(e.g. via standard arguments like NTK [JGH18]) without providing any
meaningful insight into the functionality of the attention mechanism.

The paper’s main contributions are as follows:

  1. $\bullet$

Implicit bias of the attention layer (Secs. 2-3). Optimizing the attention
parameters $({\bm{K}},{\bm{Q}})$ with vanishing regularization converges in
direction towards a max-margin solution of (Att-SVM⋆) with the nuclear norm
objective of the combined parameter $\bm{W}:=\bm{K}{\bm{Q}}^{\top}$ (Thm 2).
In the case of directly parameterizing cross-attention by the combined
parameter $\bm{W}$, the regularization path (RP) directionally converges to
(Att-SVM) solution with the Frobenius norm objective. To our knowledge, this
is the first result to formally distinguish the optimization dynamics of
$\bm{W}$ vs $({\bm{K}},{\bm{Q}})$ parameterizations, revealing the low-rank
bias of the latter. Our theory clearly characterizes the _optimality_ of
selected tokens (Definition 1) and naturally extends to sequence-to-sequence
or causal classification settings (see SAtt-SVM and Theorem 10 in appendix).

  2. $\bullet$

Convergence of gradient descent (Secs. 4-5). The gradient descent (GD)
iterates for the combined key-query variable $\bm{W}$ converge in direction to
a _locally-optimal_ solution of (Att-SVM) with appropriate initialization and
a linear head $h(\cdot)$ (Sec. 5). For local optimality, selected tokens must
have higher scores than their neighboring tokens. Locally-optimal directions
are not necessarily unique and are characterized in terms of the problem
geometry following [TLZO23]. As a key contribution, we identify geometric
conditions that guarantee convergence to the globally-optimal direction (Sec.
4). These include: (i) optimal tokens being clearly distinguishable in terms
of their _scores_ , or, (ii) the initial gradient direction being aligned with
optimal tokens. Beyond these, we show that over-parameterization (i.e.
dimension $d$ being large, and equivalent conditions) catalyzes global
convergence by ensuring (1) feasibility of (Att-SVM), and, (2) benign
optimization landscape, in the sense that there are no stationary points and
no spurious locally-optimal directions (see Sec. 5.2). These are illustrated
in Figures 2 and 2.

  3. $\bullet$

Generality of SVM equivalence (Sec. 6). When optimizing with linear
$h(\cdot)$, the attention layer is inherently biased towards selecting a
single token from each sequence (a.k.a. hard attention). This is reflected in
(Att-SVM) and arises from output tokens being convex combinations of the input
tokens. In contrast, we show that nonlinear heads necessitate composing
multiple tokens, highlighting their importance in the transformer’s dynamics
(Sec 6.1). Using insights gathered from our theory, we propose a more general
SVM equivalence. Remarkably, we demonstrate that our proposal accurately
predicts the implicit bias of attention trained by gradient descent under
general scenarios not covered by theory (e.g. $h(\cdot)$ being an MLP).
Specifically, our general formulae decouple attention weights into two
components: A directional component governed by SVM which selects the tokens
by applying a 0-1 mask, and a finite component which dictates the precise
composition of the selected tokens by adjusting the softmax probabilities.

An important feature of these findings is that they apply to arbitrary
datasets (whenever SVM is feasible) and are numerically verifiable. We
extensively validate the max-margin equivalence and implicit bias of
transformers through enlightening experiments. We hold the view that these
findings aid in understanding transformers as hierarchical max-margin token-
selection mechanisms, and we hope that our outcomes will serve as a foundation
for upcoming studies concerning their optimization and generalization
dynamics.

Overview. The paper is structured as follows: Section 2 introduces
preliminaries on self-attention and optimization. Section 3 analyzes self-
attention’s optimization geometry, showing the RP of attention parameters
converges to a max-margin solution. Sections 4 and 5 present global and local
gradient descent analyses, respectively, demonstrating convergence of
$\bm{W}$, the key-query variable, towards the solution of (Att-SVM). Section 6
provides our results on nonlinear prediction heads and generalized SVM
equivalence. Section 7 extends the theory to sequential and causal
predictions. Section 8 discusses relevant literature. Finally, Section 9
concludes the paper with open problems and future research directions inspired
by our findings. All proofs are deferred to the appendix.

##  2 Preliminaries

Figure 3: Implicit biases of the attention layer and logistic regression.

Unveiling the Relationship Between Attention and Linear SVMs. For linear
classification, it is well-established that GD iterations on logistic loss and
separable datasets converge towards the hard margin SVM solution, which
effectively separates the two classes within the data [SHN+18, RZH03, ZY05].
The softmax nonlinearity employed by the attention layer exhibits an
exponentially-tailed behavior similar to the logistic loss; thus, attention
may also be biased towards margin-maximizing solutions. However, the attention
layer operates on tokens within an input sequence, rather than performing
classification directly. Therefore, its bias is towards an SVM, specifically
(Att-SVM), which aims to separate the tokens of input sequences by selecting
the relevant ones and suppressing the rest. Nonetheless, formalizing this
intuition is considerably more challenging: The presence of the highly
nonlinear and nonconvex softmax operation renders the analysis of standard GD
algorithms intricate. Additionally, the 1-layer transformer in (2) does not
inherently exhibit a singular bias towards a single (Att-SVM) problem, even
when using a linear head $h$. Instead, it can result in multiple locally
optimal directions induced by their associated SVMs. We emphasize that
[TLZO23] is the first work to make this attention$\leftrightarrow$SVM
connection. Here, we augment their framework to transformers by developing the
first guarantees for self/cross-attention layer, nonlinear prediction heads,
and global convergence.

Notation. For any integer $N\geq 1$, let $[N]:=\\{1,\dots,N\\}$. We use
lowercase and uppercase bold letters (e.g., ${\bm{a}}$ and ${\bm{A}}$) to
represent vectors and matrices, respectively. The entries of a vector
${\bm{a}}$ are denoted as ${\bm{a}}_{i}$. For a matrix ${\bm{A}}$,
$\|{{\bm{A}}}\|$ denotes the spectral norm, i.e. maximum singular value,
$\|{{\bm{A}}}\|_{\star}$ denotes the nuclear norm, i.e. summation of all
singular values, and
$\|{{\bm{A}}}\|_{F}:=\sqrt{{{\operatorname{trace}}}({\bm{A}}^{\top}{\bm{A}})}$
denotes the Frobenius norm. $\texttt{dist}\left(\cdot,\cdot\right)$ denotes
the Euclidean distance between two sets. The minimum of two numbers $a$ and
$b$ is denoted as $a\wedge b$, and the maximum as $a\vee b$. The big-O
notation $\mathcal{O}(\cdot)$ hides the universal constants.

Optimization problem definition. We use a linear head
$h(\bm{x})=\bm{v}^{\top}\bm{x}$ for most of our theoretical exposition. Given
dataset $(Y_{i},{\bm{X}}_{i},{\bm{z}}_{i})_{i=1}^{n}$, we minimize the
empirical risk of an 1-layer transformer using combined weights
$\bm{W}\in\mathbb{R}^{d\times d}$ or individual weights
${\bm{K}},{\bm{Q}}\in\mathbb{R}^{d\times m}$ for a fixed head and decreasing
loss function:

$\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),$
(W-ERM) $\displaystyle{\cal{L}}({\bm{K}},{\bm{Q}})$
$\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{K}}{\bm{Q}}^{\top}{\bm{z}}_{i})\right).$
(KQ-ERM)

We can recover the self-attention model by setting ${\bm{z}}_{i}$ to be the
first token of ${\bm{X}}_{i}$, i.e., ${\bm{z}}_{i}\leftarrow\bm{x}_{i1}$.
While the above formulation regresses a single label $Y_{i}$ per
$({\bm{X}}_{i},{\bm{z}}_{i})$, in Sections 7 & D, we show that our findings
gracefully extend to the sequence-to-sequence classification setting where we
output and classify $T$ tokens per inputs
${\bm{X}}_{i},\bm{Z}_{i}\in\mathbb{R}^{T\times d}$. The reader is referred to
Sections 6 & D for results on nonlinear prediction heads.

Optimization algorithms. Given a parameter $R>0$, we consider an
$\ell_{2}$-norm bound $R$, and define the regularized path solution associated
with Objectives (W-ERM) and (KQ-ERM), respectively as (W-RP) and (KQ-RP).
These update rules allow us to find the solution within a constrained region
defined by the norm bound. The RP illustrates the evolution of
${\bm{\bar{W}}}_{R}$ as $R$ increases, capturing the essence of GD where the
ridge constraint serves as an approximation for the number of iterations.
Previous studies, including [RZH03, SPR18, JDST20, TLZO23], have examined the
implicit bias of logistic regression and established a connection between the
directional convergence of the RP (i.e.,
$\lim_{R\rightarrow\infty}{\bm{\bar{W}}}_{R}/R$) and GD. In [TLZO23], the
concept of a local RP was also employed to investigate implicit bias along
local directions. For GD, with appropriate initialization and step size
$\eta>0$, we describe the optimization process associated with (W-ERM) and
(KQ-ERM) as (W-GD) and (KQ-GD), respectively.

Given $\bm{W}(0)\in\mathbb{R}^{d\times d}$, $\eta>0$, for $k\geq 0$ do:
$\bm{W}(k+1)=\bm{W}(k)-\eta\nabla{\cal{L}}(\bm{W}(k)).$ (W-GD) Given
${\bm{Q}}(0),{\bm{K}}(0)\in\mathbb{R}^{d\times m}$, $\eta>0$, for $k\geq 0$
do: $\begin{bmatrix}{\bm{K}}(k+1)\\\ {\bm{Q}}(k+1)\\\
\end{bmatrix}=\begin{bmatrix}{\bm{K}}(k)\\\ {\bm{Q}}(k)\\\
\end{bmatrix}-\eta\begin{bmatrix}\nabla{\cal{L}}_{{\bm{K}}}\left({\bm{K}}(k),{\bm{Q}}(k)\right)\\\
\nabla{\cal{L}}_{{\bm{Q}}}\left({\bm{K}}(k),{\bm{Q}}(k)\right)\end{bmatrix}.$
(KQ-GD)

Given $R>0$, find $d\times d$ matrix:
$\displaystyle{\bm{\bar{W}}}_{R}=\underset{\|{\bm{W}}\|_{F}\leq
R}{\arg\min}~{}{\cal{L}}(\bm{W}).$ (W-RP) Given $R>0$, find $d\times m$
matrices:
$\displaystyle({\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}}).$ (KQ-RP)

###  2.1 Optimal tokens and hard-margin SVM problem for cross attention

Given ${\bm{X}}_{i}\in\mathbb{R}^{T\times d},{\bm{z}}_{i}\in\mathbb{R}^{d}$,
we present a convex hard-margin SVM problem, denoted as (Att-SVM), that aims
to separate a specific token from the remaining tokens in the input sequence
${\bm{X}}_{i}$. This problem is jointly solved for all inputs, allowing us to
examine the optimization properties of cross-attention. To delve deeper, we
introduce the concept of optimal tokens, which are tokens that minimize the
training objective under the decreasing loss function $\ell(\cdot)$ as shown
in Lemma 2. This exploration will introduce the notions of token scores and
optimality, providing insights into the underlying principles of self-
attention mechanisms [TLZO23].

######  Definition 1 (Token Score and Optimality)

Given a prediction head $\bm{v}\in\mathbb{R}^{d}$, the score of a token
$\bm{x}_{it}$ of input ${\bm{X}}_{i}$ is defined as
$\bm{\gamma}_{it}=Y_{i}\cdot\bm{v}^{\top}\bm{x}_{it}$. The optimal token for
each input ${\bm{X}}_{i}$ is given by the index
$\texttt{opt}_{i}\in\arg\max_{t\in[T]}\bm{\gamma}_{it}$ for all $i\in[n]$.

By introducing token scores and identifying optimal tokens, we can better
understand the importance of individual tokens and their impact on the overall
objective. The token score quantifies the contribution of a token to the
prediction or classification task, while the optimal token represents the
token that exhibits the highest relevance within the corresponding input
sequence.

$\bullet$ Hard-margin SVM for $\bm{W}$-parameterization. Equipped with the set
of optimal indices $\texttt{opt}:=(\texttt{opt}_{i})_{i=1}^{n}$ as per
Definition 1, we introduce the following SVM formulation associated to
(W-ERM):

$\bm{W}^{\textsl{mm}}=\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].$ (Att-SVM)

The existence of matrix $\bm{W}^{\textsl{mm}}$ implies the separability of
tokens $(\texttt{opt}_{i})_{i=1}^{n}$ from the others. The terms
$a_{{it}}:=\bm{x}_{it}^{\top}\bm{W}{\bm{z}}_{i}$ represent the dot product
between the key-query features before applying the softmax nonlinearity. This
dot product is a crucial characteristic of self-attention and we can express
the SVM constraints in (Att-SVM) as $a_{i\texttt{opt}_{i}}\geq a_{it}+1$.
Thus, (Att-SVM) find the most efficient direction that ensures the optimal
token $\bm{x}_{i\texttt{opt}_{i}}$ achieves the highest similarity with query
${\bm{z}}_{i}$ among all key embeddings. Our first result shows that (Att-SVM)
is feasible under mild over-parameterization.

###### Theorem 1

Suppose $d\geq\max(T-1,n)$. Then, almost all datasets
$(Y_{i},{\bm{X}}_{i},{\bm{z}}_{i})_{i=1}^{n}$ – including the self-attention
setting with ${\bm{z}}_{i}\leftarrow\bm{x}_{i1}$ – obey the following: (Att-
SVM) is feasible i.e., $\bm{W}^{\textsl{mm}}$ separates the desired tokens
$\texttt{opt}=(\texttt{opt}_{i})_{i=1}^{n}$.

We note that the convex formulation (Att-SVM) does not fully capture the GD
geometry on (W-ERM). In a more general sense, GD can provably converge to an
SVM solution over locally-optimal tokens, as detailed in Section 5.1. For
deeper insight into (Att-SVM), consider that the attention layer’s output is a
convex mixture of input tokens. Thus, if minimizing the training loss involves
choosing the optimal token $\bm{x}_{i\texttt{opt}_{i}}$, the softmax
similarities should eventually converge to a one-hot vector, precisely
including $\bm{x}_{i\texttt{opt}_{i}}$ (assigned 1), while ignoring all other
tokens (assigned 0s). This convergence requires the attention weights $\bm{W}$
to diverge in norm to saturate the softmax probabilities. Due to the
exponential-tailed nature of the softmax function, the weights converge
directionally to the max-margin solution. This phenomenon resembles the
implicit bias of logistic regression on separable data [SHN+18, JT18]. Lemma 2
formalizes this intuition and rigorously motivates optimal tokens.

$\bullet$ Non-convex SVM problem for $(\mathbf{K},{\bm{Q}})$-parameterization.
The objective function (KQ-ERM) has an extra layer of nonconvexity compared to
(W-ERM) as $({\bm{K}},{\bm{Q}})$ corresponds to a matrix factorization of
$\bm{W}$. To study this, we introduce the following nonconvex SVM problem over
$({\bm{K}},{\bm{Q}})$ akin to (Att-SVM):

$\min_{{\bm{K}},{\bm{Q}}}~{}\frac{1}{2}\left(\|{\bm{K}}\|_{F}^{2}+\|{\bm{Q}}\|_{F}^{2}\right)\quad\text{subj.
to}\quad(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it})^{\top}{\bm{K}}{\bm{Q}}^{\top}{\bm{z}}_{i}\geq
1\quad\text{for all}\quad t\neq\texttt{opt}_{i},\quad i\in[n].$ (KQ-SVM)

Even if the direction of GD is biased towards the SVM solution, it does not
have to converge to the global minima of (KQ-SVM). Instead, it can potentially
converge towards a Karush–Kuhn–Tucker (KKT) point of the max-margin SVM. Such
KKT convergence has been studied by [LL19, JT20] in the context of other
nonconvex margin maximization problems. Fortunately, (KQ-ERM) may not be as
daunting as it may initially seem: Our experiments in Figures 2&7 reveal that
GD is indeed biased towards the global minima of (KQ-SVM). This global minima
is achieved by setting $\bm{W}:={\bm{K}}{\bm{Q}}^{\top}$ and finding the
factorization of $\bm{W}$ that minimizes the quadratic objective. This yields
the following $\bm{W}$-parameterized SVM with nuclear norm objective:

$\bm{W}^{\textsl{mm}}_{\star}\in\underset{\texttt{rank}(\bm{W})\leq
m}{\arg\min}\|{\bm{W}}\|_{\star}\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].$ (Att-SVM⋆)

Above, the nonconvex rank constraint arises from the fact that the rank of
$\bm{W}={\bm{K}}{\bm{Q}}^{\top}$ is at most $m$. However, under the condition
of the full parameterization where $m\geq d$, the rank constraint disappears,
leading to a convex nuclear norm minimization problem. Besides, the nuclear
norm objective inherently encourages a low-rank solution [RFP10, Faz02,
SRJ04]. Therefore, if the solution set of the unconstrained convex problem
satisfies $\texttt{rank}(\bm{W})\leq m$, the convex problem will recover the
solutions of (Att-SVM), even when $m<d$. Lemma 1, presented below,
demonstrates that this guarantee holds whenever $n\leq m$. This observation is
further supported by our experiments (see Fig. 4). Thus, it offers a
straightforward rationale for why setting $m<d$ is a reasonable practice,
particularly in scenarios involving limited data.

###### Lemma 1

Any optimal solution of (Att-SVM) or (Att-SVM⋆) is at most rank $n$. More
precisely, the row space of $\bm{W}^{\textsl{mm}}$ or
$\bm{W}^{\textsl{mm}}_{\star}$ lies within
$\texttt{span}(\\{{\bm{z}}_{i}\\}_{i=1}^{n})$.

Figures 4 illustrates rank range of solutions for (Att-SVM) and (Att-SVM⋆),
denoted as $\bm{W}^{\textsl{mm}}$ and $\bm{W}^{\textsl{mm}}_{\star}$, solved
using optimal tokens $(\texttt{opt}_{i})_{i=1}^{n}$ and setting $m=d$ (the
rank constraint is eliminated). Each result is averaged over 100 trials, and
for each trial, ${\bm{x}}_{it}$, ${{\bm{z}}}_{i}$, and linear head ${\bm{v}}$
are randomly sampled from the unit sphere. In Fig. 4(a), we fix $T=5$ and vary
$n$ across $\\{5,10,15\\}$. Conversely, in Fig. 4(b), we keep $n=5$ constant
and alter $T$ across $\\{5,10,15\\}$. Both figures confirm rank of
$\bm{W}^{\textsl{mm}}$ and $\bm{W}^{\textsl{mm}}_{\star}$ are bounded by
$\max(n,d)$, validating Lemma 1.

$\bm{W}^{\textsl{mm}}$$\bm{W}^{\textsl{mm}}_{\star}$Varying $d$Rank of SVM
solution (a) Rank of attention SVM solutions with fixed $T=5$

$\bm{W}^{\textsl{mm}}$$\bm{W}^{\textsl{mm}}_{\star}$Varying $d$Rank of SVM
solution (b) Rank of attention SVM solutions with fixed $n=5$

Figure 4: Rank range of solutions for (Att-SVM) and (Att-SVM⋆), denoted as
$\bm{W}^{\textsl{mm}}$ and $\bm{W}^{\textsl{mm}}_{\star}$, solved using
optimal tokens $(\texttt{opt}_{i})_{i=1}^{n}$ and setting $m=d$ (the rank
constraint is eliminated). Both figures confirm ranks of
$\bm{W}^{\textsl{mm}}$ and $\bm{W}^{\textsl{mm}}_{\star}$ are bounded by
$\max(n,d)$, validating Lemma 1.

##  3 Understanding Implicit Bias of Self-Attention

We start by motivating the optimal token definition and establishing the
global convergence of RPs which shed light on the implicit bias of attention
parameterizations. Throughout, we maintain the following assumption regarding
the loss function.

###### Assumption A

Over any bounded interval:  (i) $\ell:\mathbb{R}\rightarrow\mathbb{R}$is
strictly decreasing and bounded from below;  (ii) The derivative of $\ell$,
denoted as $\ell^{\prime}$, is $M_{0}$-Lipschitz continuous; and  (iii)
$|\ell^{\prime}(u)|\leq M_{1}$.

Assumption A encompasses many commonly used loss functions, especially for
classification, such as the logistic loss
$\ell\left(u\right)=\log\left(1+e^{-u}\right)$, exponential loss
$\ell\left(u\right)=e^{-u}$, and correlation loss $\ell(u)=-u$.

######  Lemma 2 (Optimal Tokens Minimize Training Loss)

Suppose Assumption A (i) and (iii) hold, and not all tokens are optimal per
Definition 1. Then, training risk obeys
${\cal{L}}(\bm{W})>{\cal{L}}_{\star}:=\frac{1}{n}\sum_{i=1}^{n}\ell(\bm{\gamma}_{i\texttt{opt}_{i}})$.
Additionally, suppose there are optimal indices $(\texttt{opt}_{i})_{i=1}^{n}$
for which (Att-SVM) is feasible, i.e. there exists a $\bm{W}$ separating
optimal tokens. This $\bm{W}$ choice obeys
$\lim_{R\rightarrow\infty}{\cal{L}}(R\cdot\bm{W})={\cal{L}}_{\star}$.

The result presented in Lemma 2 originates from the observation that the
output tokens of the attention layer constitute a convex combination of the
input tokens. Consequently, when subjected to a strictly decreasing loss
function, attention optimization inherently leans towards the selection of a
singular token, specifically, the optimal token
$(\texttt{opt}_{i})_{i=1}^{n}$. Our following theorem unveils the implicit
bias ingrained within attention parameterizations through RP analysis.

###### Theorem 2

Suppose Assumptions A (i) and (iii) hold, optimal indices
$(\texttt{opt}_{i})_{i=1}^{n}$ are unique, and (Att-SVM) is feasible. Let
$\bm{W}^{\textsl{mm}}$ be the unique solution of (Att-SVM), and let
${\cal{W}}^{\textsl{mm}}_{\star}$ be the solution set of (Att-SVM⋆) with
nuclear norm achieving objective $C_{\star}$. Then, Algorithms W-RP and KQ-RP,
respectively, satisfy:

  * •

$\bm{W}$-parameterization has Frobenius norm bias:
$\underset{R\rightarrow\infty}{\lim}\frac{{\bm{\bar{W}}}_{R}}{R}=\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}$.

  * •

$({\bm{K}},{\bm{Q}})$-parameterization has nuclear norm bias:
$\underset{R\rightarrow\infty}{\lim}\texttt{dist}\left(\frac{{\bm{\bar{K}}}_{R}{\bm{\bar{Q}}}_{R}^{\top}}{R},\frac{{\cal{W}}^{\textsl{mm}}_{\star}}{C_{\star}}\right)=0$.

    * –

Setting $m=d$: (Att-SVM⋆) is a convex problem without rank constraints.

Theorem 2 demonstrates that the RP of the
$({\bm{K}},{\bm{Q}})$-parameterization converges to a max-margin solution of
(Att-SVM⋆) with nuclear norm objective on $\bm{W}={\bm{K}}{\bm{Q}}^{\top}$.
When self-attention is directly parameterized by $\bm{W}$, the RP converges to
the solution of (Att-SVM) with a Frobenius norm objective. This result is the
first to distinguish the optimization dynamics of $\bm{W}$ and
$({\bm{K}},{\bm{Q}})$ parameterizations, revealing the low-rank bias of the
latter. These findings also provide a clear characterization of token
optimality (Definition 1) and extend naturally to the setting with multiple
optimal tokens per sequence (Theorem 10 in appendix). By definition, the RP
captures the global geometry and cannot be used for the implicit bias of GD
towards locally-optimal directions. Sections 4 and 5 accomplish this goal
through gradient-descent and localized RP analysis to obtain locally-
applicable SVM equivalences. Note that, this theorem requires each input
sequence has a unique optimal token per Definition 1. Fortunately, this is a
very mild condition as it holds for almost all datasets, namely, as soon as
input features are slightly perturbed.

Theorem 2 establishes the implicit bias of attention from the perspective of
RP analysis. This leads to the question: To what extent is this RP theory
predictive of the implicit bias exhibited by GD? To delve into this, we
examine the gradient paths of $\bm{W}(k)$ or $({\bm{K}}(k),{\bm{Q}}(k))$ and
present the findings in Figure 2. We consider a scenario where $n=d=2$ and
$T=3$, and employ cross-attention, where tokens $({\bm{z}}_{1},{\bm{z}}_{2})$
are generated independently of the inputs $({\bm{X}}_{1},{\bm{X}}_{2})$. In
Figure 2’s lower right corner, the teal and yellow markers correspond to
tokens from ${\bm{X}}_{1}$ and ${\bm{X}}_{2}$, respectively. The stars
indicate the optimal token for each input. To provide a clearer view of the
gradient convergence path, we illustrate the outcomes of training the
attention weight $\bm{W}$ or $({\bm{K}},{\bm{Q}})$ in the form of
$\bm{W}{\bm{z}}_{i}$ or ${\bm{K}}{\bm{Q}}^{\top}{\bm{z}}_{i}$, where
$i={1,2}$. With reference to Equations (Att-SVM) and (Att-SVM⋆), the red and
blue solid lines delineate the directions of
$\bm{W}^{\textsl{mm}}{\bm{z}}_{1}$ and $\bm{W}^{\textsl{mm}}{\bm{z}}_{2}$,
correspondingly. Conversely, the red and blue dashed lines show the directions
of $\bm{W}^{\textsl{mm}}_{\star}{\bm{z}}_{1}$ and
$\bm{W}^{\textsl{mm}}_{\star}{\bm{z}}_{2}$. The red/blue solid/dashed arrows
denote the corresponding directions of gradient evolution. Figure 2 provides a
clear depiction of the incremental alignment of $\bm{W}(k)$ and
${\bm{K}}(k){\bm{Q}}(k)^{\top}$ with their respective attention SVM solutions
as $k$ increases. This strongly supports the assertions of Theorem 2.

It is worth noting that (Att-SVM⋆) imposes a nonconvex rank constraint, i.e.,
$\texttt{rank}(\bm{W})\leq m$. Nevertheless, this constraint becomes
inconsequential if the unconstrained problem, with $m$ set to be greater than
or equal to $d$, admits a low-rank solution, as demonstrated in Lemma 1.
Consequently, in our experimental endeavors, we have the flexibility to employ
the unconstrained attention SVM for predicting the implicit bias. This concept
is succinctly summarized by the following lemma.

###### Lemma 3

Let ${\cal{W}}^{\textsl{mm}}_{\star}$ be the solution set of (Att-SVM⋆) with
nuclear norm achieving objective $C_{\star}$. Further let
${\cal{W}}^{\textsl{mm}}_{\texttt{cvx}}$ be the solution set of (Att-SVM⋆)
with $m=d$ achieving objective $C_{\texttt{cvx}}$. If
${\cal{W}}^{\textsl{mm}}_{\star}\cap{\cal{W}}^{\textsl{mm}}_{\texttt{cvx}}\neq\emptyset$,
then $C_{\star}=C_{\texttt{cvx}}$ and
${\cal{W}}^{\textsl{mm}}_{\star}\subseteq{\cal{W}}^{\textsl{mm}}_{\texttt{cvx}}$.
Also, if the elements of ${\cal{W}}^{\textsl{mm}}_{\texttt{cvx}}$ have rank at
most $m$, then,
${\cal{W}}^{\textsl{mm}}_{\star}={\cal{W}}^{\textsl{mm}}_{\texttt{cvx}}$.

##  4 Global Convergence of Gradient Descent

In this section, we will establish conditions that guarantee the global
convergence of GD. Concretely, we will investigate when GD solution selects
the _optimal token within each input sequence_ through the softmax
nonlinearity and coincides with the solution of the RP. Section 5 will
complement this with showing that self-attention can more generally converge
to locally-optimal max-margin directions. We identify the following conditions
as provable catalysts for global convergence: (1) Optimal tokens have
relatively large scores, (2) Initial gradient direction is favorable, (3)
Overparameterization, i.e. $d$ is appropriately large.

###  4.1 Properties of optimization landscape

We start by establishing some fundamental properties of Objectives (W-ERM) and
(KQ-ERM).

###### Lemma 4

Under Assumption A, $\nabla{\cal{L}}(\bm{W})$,
$\nabla_{{\bm{K}}}{\cal{L}}({\bm{K}},{\bm{Q}})$, and
$\nabla_{{\bm{Q}}}{\cal{L}}({\bm{K}},{\bm{Q}})$ are $L_{\bm{W}}$,
$L_{{\bm{K}}}$, $L_{{\bm{Q}}}$–Lipschitz continuous, respectively, where
$a_{i}=\|\bm{v}\|~{}\|{\bm{z}}_{i}\|^{2}\|{\bm{X}}_{i}\|^{3}$,
$b_{i}=M_{0}\|\bm{v}\|\|{\bm{X}}_{i}\|+3M_{1}$ for all $i\in[n]$,

$\displaystyle L_{\bm{W}}:=\frac{1}{n}\sum_{i=1}^{n}a_{i}b_{i},\qquad
L_{{\bm{K}}}:=\|{\bm{Q}}\|L_{\bm{W}},\quad~{}~{}\textnormal{and}\quad
L_{{\bm{Q}}}:=\|{\bm{K}}\|L_{\bm{W}}.$ (3)

The next assumption will play an important role for ensuring the attention
layer has a benign optimization landscape.

###### Assumption B

Optimal tokens’ indices $(\texttt{opt}_{i})_{i=1}^{n}$ are unique and one of
the following conditions on the tokens holds:

  1. B.1

All tokens are support vectors, i.e.,
$(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it})^{\top}\bm{W}^{\textsl{mm}}{\bm{z}}_{i}=1$
for all $t\neq\texttt{opt}_{i}$ and $i\in[n]$.

  2. B.2

The tokens’ scores, as defined in Definition 1, satisfy
$\bm{\gamma}_{it}=\bm{\gamma}_{i\tau}<\bm{\gamma}_{i\texttt{opt}_{i}}$, for
all $t,\tau\neq\texttt{opt}_{i}$ and $i\in[n]$.

PercentageVarying $d$Not LocalGlobalLocalAssum B.1 (a) Convergence behaviour
of GD for $\bm{W}$-parameterization

PercentageVarying $d$Not LocalGlobalLocalAssum B.1 (b) Convergence behaviour
of GD for $({\bm{K}},{\bm{Q}})$-parameterization

Figure 5: Percentage of different convergence types of GD when training cross-
attention weights (a): $\bm{W}$ or (b): $({\bm{K}},{\bm{Q}})$ with varying
$d$. In both figures, red, blue, and teal bars represent the percentages of
Global, Local (including Global), and Not Local convergence, respectively. The
green bar corresponds to Assumption B.1 where all tokens act as support
vectors. Larger overparameterization ($d$) relates to a higher percentage of
globally-optimal SVM convergence.

Assumption B.1 is directly linked to overparameterization and holds practical
significance. In scenarios such as classical SVM classification, where the
goal is to separate labels, overparameterization leads to the situation where
_all training points become support vectors_. Consequently, the SVM solution
aligns with the least-squares minimum-norm interpolation, a concept
established in [MNS+21, HMX21] under broad statistical contexts. Assumption
B.1 represents an equivalent manifestation of this condition. Therefore, in
cases involving realistic data distributions with sufficiently large $d$, we
expect the same phenomena to persist, causing the SVM solution
$\bm{W}^{\textsl{mm}}$ to coincide with (Att-SVM).

Drawing on insights from [HMX21, Theorem 1] and our Theorem 1, we expect that
the necessary degree of overparameterization remains moderate. Specifically,
in instances where input sequences follow an independent and identically
distributed (IID) pattern and tokens exhibit IID isotropic distributions, we
posit that $d\gtrsim(T+n)\log(T+n)$ will suffice. More generally, the extent
of required overparameterization will be contingent on the covariance of
tokens [BLLT20, MNS+21] and the distribution characteristics of input
sequences [WT22].

Assumption B.2 stipulates that non-optimal tokens possess identical scores
which constitutes a relatively stringent assumption that we will subsequently
relax. Under Assumption B, we establish that optimization problems (W-ERM) and
(KQ-ERM) lack stationary points, and when trained using GD, the norm of
parameters will diverge.

###### Theorem 3

Suppose Assumption A on the loss function $\ell$ and Assumption B on the
tokens hold. Then,

  * •

There is no $\bm{W}\in\mathbb{R}^{d\times d}$ satisfying
$\nabla{\cal{L}}(\bm{W})=0$.

  * •

Algorithm W-GD with the step size
$\eta\leq\mathcal{O}\left(1/L_{\bm{W}}\right)$ and any starting point
$\bm{W}(0)$ satisfies
$\lim_{k\rightarrow\infty}\left\|\nabla\mathcal{L}\left(\bm{W}\left(k\right)\right)\right\|_{F}=0$,
and $\lim_{k\rightarrow\infty}\|{\bm{W}(k)}\|_{F}=\infty$.

The feasibility of SVM (per Theorem 1) is a necessary condition for
convergence of GD to the $\bm{W}^{\textsl{mm}}$ direction. However, it does
not inform us about the optimization landscape. Two additional criteria are
essential for convergence: the absence of stationary points
$\nabla{\cal{L}}(\bm{W})=0$ and divergence of parameter norm to infinity.
Theorem 3 above precisely guarantees both of these criteria under Assumption
B.

###  4.2 Provable global convergence of 1-layer transformer

Here, we provide three conditions for achieving global convergence towards the
max-margin direction $\bm{W}^{\textsl{mm}}$ based on: (I) token score
constraints, (II) initial gradient direction, and (III) over-parameterization.
For the first two, we provide precise theoretical guarantees. For the third,
we provide strong empirical evidence, supported by Theorem 3, and a formal
conjecture described in Section 5.2. We remind that, we optimize the attention
weights $\bm{W}$ while fixing the linear prediction head
$h(\bm{x})=\bm{v}^{\top}\bm{x}$. This avoids trivial convergence guarantees
where an over-parameterized $h(\cdot)$ – whether it is a linear model or MLP –
can be used to achieve zero training loss without providing any meaningful
insights into the functionality of the attention mechanism.

##### (I) Global convergence under score constraints.

Our next result establishes the global convergence of GD to the max-margin
direction $\bm{W}^{\textsl{mm}}$ under Assumption B.2 that non-optimal tokens
have identical scores but lower than the score of the optimal token.

###### Theorem 4

Suppose Assumption A on the loss function $\ell$ and Assumption B.2 on the
tokens’ score hold. Then, Algorithm W-GD with the step size
$\eta\leq\mathcal{O}\left(1/L_{\bm{W}}\right)$ and any starting point
$\bm{W}(0)$ satisfies
$\lim_{k\rightarrow\infty}\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}}=\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}$.

To proceed, we also provide a relaxation of Assumption B.2 and, show that,
global convergence still happens under approximately equal scores.

###### Theorem 5

Suppose Assumption A on the loss function $\ell$ hold. For any initialization
$\bm{W}(0)$, there exists a dataset-dependent sufficiently small $\delta>0$
such that the following holds: Suppose non-optimal scores obey
$|\bm{\gamma}_{it}-\bm{\gamma}_{i\tau}|\leq\delta$ for all
$t,\tau\neq\texttt{opt}_{i},~{}i\in[n]$. Then, Algorithm W-GD, with
$\eta\leq\mathcal{O}\left(1/L_{\bm{W}}\right)$ obeys
$\lim_{k\rightarrow\infty}\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}}=\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}$.

##### (II) Global convergence under good initial gradient.

To ensure global convergence, we identify an assumption that prevents GD from
getting trapped at suboptimal tokens that offer no scoring advantage compared
to other choices.

######  Assumption C (First GD step is separating)

For all non-optimal tokens $t\neq\texttt{opt}_{i},~{}i\in[n]$:
$(\bm{x}_{it}-\bm{x}_{i\texttt{opt}_{i}})^{\top}\nabla{\cal{L}}(0){\bm{z}}_{i}>0$.

###### Theorem 6

Suppose Assumption A on the loss function $\ell$ and Assumption C on the
initial gradient holds. Consider GD iterates
$\bm{W}(k+1)=\bm{W}(k)-\eta(k)\nabla{\cal{L}}(\bm{W}(k))$ on (W-ERM) with
$\bm{W}(0)=0$ and step sizes
$\eta(0)\leq\mathcal{O}(1/\|\bm{W}^{\textsl{mm}}\|)$ and
$\eta(k)\leq\mathcal{O}\left(1/L_{\bm{W}}\right)$ for all $k\geq 1$. Then,
$\lim_{k\rightarrow\infty}\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}}=\frac{\bm{W}^{\textsl{mm}}}{\|{\bm{W}^{\textsl{mm}}}\|_{F}}$.

Varying $d$Prob. of global convergence$\bm{W}$$({\bm{K}},{\bm{Q}})$ (a) Global
convergence for varying $n,d$

Varying $d$Prob. of global convergence$\bm{W}$$({\bm{K}},{\bm{Q}})$ (b) Global
convergence for varying $T,d$

Figure 6: Global convergence behavior of GD when training cross-attention
weights $\bm{W}$ (solid) or $({\bm{K}},{\bm{Q}})$ (dashed) with random data.
The blue, green, and red curves represent the probabilities of global
convergence for (a): fixing $T=5$ and varying $n\in\\{5,10,20\\}$ and (b):
fixing $n=5$ and varying $T\in\\{5,10,20\\}$. Results demonstrate that for
both attention models, as $d$ increases (due to over-parameterization),
attention weights tend to select optimal tokens
$(\texttt{opt}_{i})_{i=1}^{n}$.

##### (III) Global Convergence via overparameterization.

In the context of standard neural networks, overparameterization has been
recognized as a pivotal factor for the global convergence of GD [DZPS18,
AZLS19, LL18, OS19]. However, conventional approaches like the neural tangent
kernel [JGH18] do not seamlessly apply to our scenario, given our assumption
on fixed $h$ and the avoidance of achieving zero loss by trivially fitting
$h$. Furthermore, even when we train $h$ to achieve zero loss, it doesn’t
provide substantial insights into the implicit bias of the attention weights
$\bm{W}$. Conversely, Theorem 3 illustrates the benefits of over-
parameterization in terms of convergence. Considering that Assumption B.1 is
anticipated to hold as the dimension $d$ increases, the norm of the GD
solution is bound to diverge to infinity. This satisfies a prerequisite for
converging towards the globally-optimal SVM direction $\bm{W}^{\textsl{mm}}$.

The trend depicted in Figure 5, where the percentage of global convergence
(red bars) approaches $100\%$ and Assumption B.1 holds with higher probability
(green bars) as $d$ grows, reinforces this insight. Specifically, Fig. 5(a) is
similar to Figure 2 but with additional green bars representing the percentage
of the scenarios where almost all tokens act as support vectors (Assumption
B.1), and Fig. 5(b) displays the same evaluation over
$({\bm{K}},{\bm{Q}})$-parameterization setting. In both experiments, and for
each chosen $d$ value, a total of $500$ random instances are conducted under
the conditions of $n=T=5$. The outcomes are reported in terms of the
percentages of Not Local, Local, and Global convergence, represented by the
teal, blue, and red bars, respectively. We validate Assumption B.1 as follows:
Given a problem instance, we compute the average margin over all non-optimal
tokens of all inputs and declare that problem satisfies Assumption B.1, if the
average margin is below 1.1 (where 1 is the minimum). Here, recall that margin
of a non-optimal token is defined as
$(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it})^{\top}\bm{W}^{\textsl{mm}}{\bm{z}}_{i}$
or
$(\bm{x}_{i\texttt{opt}_{i}}-\bm{x}_{it})^{\top}\bm{W}^{\textsl{mm}}_{\star}{\bm{z}}_{i}$
for $t\neq\texttt{opt}_{i}$.

Furthermore, the observations in Figure 6 regarding the percentages of
achieving global convergence reaching 100 with larger $d$ reaffirm that
overparameterization leads the attention weights to converge directionally
towards the optimal max-margin direction outlined by (Att-SVM) and (Att-SVM⋆).

In the upcoming section, we will introduce locally-optimal directions, to
which GD can be proven to converge when appropriately initialized. We will
then establish a condition that ensures the _globally-optimal direction is the
sole viable locally-optimal direction_. This culmination will result in a
formal conjecture detailed in Section 5.2.

##  5 Understanding Local Convergence of the Attention Layer

So far, we have primarily focused on the convergence to the global direction
dictated by (Att-SVM). In this section, we investigate and establish the local
directional convergence of GD as well as RP.

###  5.1 Local convergence of gradient descent

To proceed, we introduce locally-optimal directions by adapting Definition 2
of [TLZO23].

######  Definition 2 (Support indices and locally-optimal direction)

Fix token indices ${\bm{\alpha}}=(\alpha_{i})_{i=1}^{n}$. Solve (Att-SVM) with
$(\texttt{opt}_{i})_{i=1}^{n}$ replaced with
$\bm{\alpha}=(\alpha_{i})_{i=1}^{n}$ to obtain
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$. Consider the set
$\mathcal{T}_{i}\subset[T]$ such that
$(\bm{x}_{i\alpha_{i}}-\bm{x}_{it})^{\top}\bm{W}^{\textsl{mm}}_{\bm{\alpha}}{\bm{z}}_{i}=1$
for all $t\in\mathcal{T}_{i}$. We refer to $(\mathcal{T}_{i})_{i=1}^{n}$ as
the support indices of ${\bm{\alpha}}$. Additionally, if for all $i\in[n]$ and
$t\in\mathcal{T}_{i}$ scores per Definition 1 obey
$\bm{\gamma}_{i\alpha_{i}}>\bm{\gamma}_{it}$, indices
${\bm{\alpha}}=(\alpha_{i})_{i=1}^{n}$ are called _locally-optimal_ and
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$ is called a _locally-optimal direction_.

In words, the concept of local optimality requires that the selected tokens
denoted as ${\bm{\alpha}}$ should have scores that are higher than the scores
of their neighboring tokens referred to as support indices. It is important to
observe that the tokens defined as
$\texttt{opt}=(\texttt{opt}_{i})_{i=1}^{n}$, which we term as the optimal
tokens, inherently satisfy the condition of local optimality. Moving forward,
we will provide Theorem 7 which establishes that when the process of GD is
initiated along a direction that is locally optimal, it gradually converges in
that particular direction, eventually aligning itself with
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$. This theorem immediately underscores the
fact that if there exists a direction of local optimality (apart from the
globally optimal direction $\bm{W}^{\textsl{mm}}$), then when GD commences
from any arbitrary starting point, it does not achieve global convergence
towards $\bm{W}^{\textsl{mm}}$. In the context of Theorem 4, this insight
further demonstrates the general necessity of Assumption B.2. This assumption
is pivotal because without it, scenarios in which token scores are not equal
might potentially lead to the emergence of directions that are locally optimal
(as specified in Definition 2), thereby impeding the possibility of achieving
global convergence.

To provide a basis for discussing local convergence of GD, we establish a cone
centered around $\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$ using the following
construction. For a fixed $\bar{\bm{V}}\in\mathbb{R}^{d\times d}$, parameters
$\mu\in(0,1)$ and $R>0$, we define $\mathcal{C}_{\mu,R}(\bar{\bm{V}})$ as the
set of matrices $\bm{W}\in\mathbb{R}^{d\times d}$ such that
$\|{\bm{W}}\|_{F}\geq R$ and the correlation coefficient between $\bm{W}$ and
$\bar{\bm{V}}$ is at least $1-\mu$:

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

###### Theorem 7

Suppose Assumption A on the loss $\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{\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}$. Then,

  1. T1.

There exist parameters $\mu=\mu({\bm{\alpha}})\in(0,1)$ and $R>0$ such that
$\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}}_{\bm{\alpha}})$ does not contain any
stationary points.

  2. T2.

For any starting point $\bm{W}(0)$ in the confidence region
$\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}}_{\bm{\alpha}})$, consider Algorithm
W-GD with stepsize
$\eta\leq\mathcal{O}\left(\min\left(\frac{1}{L_{\bm{W}}},\zeta(\mu)\right)\right)$,
where $\zeta(\mu)$ is a correlation dependent parameter. Then, we have
$\lim_{k\rightarrow\infty}\|{\bm{W}(k)}\|_{F}=\infty$ and
$\lim_{k\rightarrow\infty}\frac{\bm{W}(k)}{\|{\bm{W}(k)}\|_{F}}=\frac{\bm{W}^{\textsl{mm}}_{\bm{\alpha}}}{\|{\bm{W}^{\textsl{mm}}_{\bm{\alpha}}}\|_{F}}$.

This theorem establishes the existence of positive parameters
$\mu=\mu({\bm{\alpha}})>0$ and $R>0$ such that there are no stationary points
within $\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}}_{\bm{\alpha}})$. Furthermore,
if GD is initiated within
$\mathcal{C}_{\mu,R}(\bm{W}^{\textsl{mm}}_{\bm{\alpha}})$, it will converge in
the direction of
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}/\|{\bm{W}^{\textsl{mm}}_{\bm{\alpha}}}\|_{F}$.
It is important to note that Theorem 7 does not make any assumptions on the
tokens as opposed to Theorem 4. However, it does require a proper
initialization that exhibits a correlation between the initial direction and
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$. Additionally, we observe that the
stepsize is adjusted based on the Lipschitz constant of the gradient of the
objective function, the correlation of the initial direction with
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$, which is parameterized by $\mu$ and the
correlation of the negative gradient at $\bm{W}(k)$ and
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$, parameterized by $\rho$.

$\bm{W}$$({\bm{K}},{\bm{Q}})$IterationsSoftmax probability (a) Evolution of
softmax probabilities

$\bm{W}$$({\bm{K}},{\bm{Q}})$IterationsCorrelation coefficient (b) Corr.
coeff. of GD and $\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$

$\bm{W}$$({\bm{K}},{\bm{Q}})$IterationsCorrelation coefficient (c) Corr.
coeff. of GD and $\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}}$

Figure 7: Local convergence behaviour of GD when training cross-attention
weights $\bm{W}$ (blue) or $({\bm{K}},{\bm{Q}})$ (red) with random data:  (a)
displays the largest entry of the softmax outputs averaged over the dataset;
(b&c) display the Pearson correlation coefficients of GD trajectories and the
SVM solutions (b) with the Frobenius norm objective
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$ (solution of (Att-SVM)) and (c) with the
nuclear norm objective $\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}}$ (solution
of (Att-SVM⋆)). These demonstrate the Frobenius norm bias of $\bm{W}(k)$ and
the nuclear norm bias of ${\bm{K}}(k){\bm{Q}}(k)^{\top}$.

In Figure 7 we consider setting where $n=6$, $T=8$, and $d=10$. The displayed
results are averaged from $100$ random trials. We train cross-attention models
with $\bm{x}_{it},{\bm{z}}_{i},\bm{v}\in\mathbb{R}^{d}$ randomly sampled from
unit sphere, and apply the normalized GD approach with fixed step size
$\eta=0.1$. In Figure 7(a) we calculate the softmax probability via
$\frac{1}{n}\sum_{i=1}^{n}\max_{t\in[T]}\mathbb{S}({\bm{X}}_{i}\tilde{\bm{W}}(k){\bm{z}}_{i})_{t}$
for either $\tilde{\bm{W}}=\bm{W}$ or ${\bm{K}}{\bm{Q}}^{\top}$ at each
iteration. Both scenarios result in probability $1$, which indicates that
attention weights succeed in selecting one token per input. Then following
Def. 2 let ${\bm{\alpha}}=(\alpha_{i})_{i=1}^{n}$ be the token indices
selected by GD and denote $\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}}$ as the
corresponding SVM solution of (Att-SVM⋆). Define the correlation coefficient
of two matrices as
$\texttt{corr\\_coef}(\bm{W}_{1},\bm{W}_{2}):=\left<\bm{W}_{1},\bm{W}_{2}\right>/\|\bm{W}_{1}\|_{F}\|\bm{W}_{2}\|_{F}$.
Figures 7(b) and 7(c) illustrate the correlation coefficients of attention
weights ($\bm{W}(k)$ and ${\bm{K}}(k){\bm{Q}}(k)^{\top}$) with respect to
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$ and
$\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}}$. The results demonstrate that
$\bm{W}$ (${\bm{K}}{\bm{Q}}^{\top}$) ultimately reaches a $1$ correlation with
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$
($\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}}$), which suggests that $\bm{W}$
(${\bm{K}}{\bm{Q}}^{\top}$) converges in the direction of
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$
($\bm{W}^{\textsl{mm}}_{\star,{\bm{\alpha}}}$). This further validates Theorem
7.

###  5.2 Overparameterization conjecture: When local-optimal directions
disappear

In Section 4 we demonstrated that larger $d$ serves as a catalyst for global
convergence to select the optimal indices
$\texttt{opt}=(\texttt{opt}_{i})_{i=1}^{n}$. However, Section 5.1 shows that
the convergence can be towards locally-optimal directions rather than global
ones. How do we reconcile these? Under what precise conditions, can we expect
global convergence?

The aim of this section is gathering these intuitions and stating a concrete
conjecture on the global convergence of the attention layer under geometric
assumptions related to overparameterization. To recap, Theorem 1 characterizes
when (Att-SVM) is feasible and Theorem 3 characterizes when the parameter norm
provably diverges to infinity, i.e. whenever all tokens are support vectors of
(Att-SVM) (Assumption B.1 holds). On the other hand, this is not sufficient
for global convergence, as GD can converge in direction to locally-optimal
directions per Section 5.1. Thus, to guarantee global convergence, we need to
ensure that globally-optimal direction is the only viable one. Our next
assumption is a fully-geometric condition that precisely accomplishes this.

######  Assumption D (There is always an optimal support index)

For any choice of ${\bm{\alpha}}=(\alpha_{i})_{i=1}^{n}$ with
${\bm{\alpha}}\neq\texttt{opt}$ when solving (Att-SVM) with
$\texttt{opt}\leftarrow{\bm{\alpha}}$, there exists $i\in[n]$ such that
$\alpha_{i}\neq\texttt{opt}_{i}$ and $\texttt{opt}_{i}\in\mathcal{T}_{i}$.

This guarantees that no ${\bm{\alpha}}\neq\texttt{opt}$ can be locally-optimal
because it has a support index with higher score at the input $i$ with
$\alpha_{i}\neq\texttt{opt}_{i}$. Thus, this ensures that global direction
$\bm{W}^{\textsl{mm}}$ is the unique locally-optimal direction obeying Def. 2.
Finally, note that local-optimality in Def. 2 is one-sided: GD can provably
converge to locally-optimal directions, while we do not provably exclude the
existence of other directions. Yet, Theorem 4 of [TLZO23] shows that local RPs
(see Section 5.3 for details) can only converge to locally-optimal directions
for almost all datasets333To be precise, they prove this for their version of
Def. 2, which is stated for an attention model
$f({\bm{p}})=\bm{v}^{\top}{\bm{X}}^{\top}\mathbb{S}({\bm{X}}\bm{W}{\bm{p}})$
admitting an analogous SVM formulation. . This and Figure 5 provide strong
evidence that Def. 2 captures all possible convergence directions of GD, and
as a consequence, that Assumption D guarantees that $\bm{W}^{\textsl{mm}}$ is
the only viable direction to converge.

➡ Integrating the results and global convergence conjecture. Combining
Assumptions B.1 and D, we have concluded that gradient norm diverges and
$\bm{W}^{\textsl{mm}}$ is the only viable direction to converge. Thus, we
conclude this section with the following conjecture: Suppose
$\texttt{opt}=(\texttt{opt}_{i})_{i=1}^{n}$ are the unique optimal indices
with strictly highest score per sequence and Assumptions B.1 and D hold. Then,
for almost all datasets (e.g. add small IID gaussian noise to input features),
GD with a proper constant learning rate converges to $\bm{W}^{\textsl{mm}}$ of
(Att-SVM) in direction.

To gain some intuition, consider Figure 5: Here, red bars denote the frequency
of global convergence whereas green bars denote the frequency of Assumption
B.1 holding over random problem instances. In short, this suggests that
Assumption B.1 occurs less frequently than global convergence, which is
consistent with our conjecture. On the other hand, verifying Assumption D is
more challenging due to its combinatorial nature. A stronger condition that
implies Assumption D is when _all optimal indices
$(\texttt{opt}_{i})_{i=1}^{n}$ are support vectors of the SVM. That is, either
$\texttt{opt}_{i}=\alpha_{i}$ or $\texttt{opt}_{i}\in\mathcal{T}_{i}$,
$\forall~{}i\in[n]$._ When the data follows a statistical model, this stronger
condition could be verified through probabilistic tools building on our
earlier “all training points are support vectors” discussion [MNS+21]. More
generally, we believe a thorough statistical investigation of (Att-SVM) is a
promising direction for future work.

###  5.3 Guarantees on local regularization path

In this section, we provide a _localized_ regularization path analysis for
general objective functions. As we shall see, this will also allow us to
predict the local solutions of gradient descent described in Section 5.1. Let
${\diamond}$ denote a general norm objective. Given indices
${\bm{\alpha}}=(\alpha_{i})_{i=1}^{n}$, consider the formulation

$\bm{W}^{\textsl{mm}}_{{\diamond},\alpha}=\underset{\texttt{rank}(W)\leq
m}{\arg\min}\|\bm{W}\|_{\diamond}\quad\text{subj.
to}\quad\min_{t\neq\alpha_{ik}}(\bm{x}_{i\alpha_{ik}}-\bm{x}_{it})^{\top}\bm{W}{\bm{z}}_{i}\geq
1,\quad\textnormal{for all}\quad i\in[n],k\in[K].$ (${\diamond}$-SVM)

In this section, since ${\diamond}$ is clear from the context, we will use the
shorthand
$\bm{W}^{\textsl{mm}}_{\alpha}:=\bm{W}^{\textsl{mm}}_{{\diamond},\alpha}$ and
denote the optimal solution set of (${\diamond}$-SVM) as
${\cal{W}}^{\textsl{mm}}:={\cal{W}}^{\textsl{mm}}_{{\diamond},\alpha}$. It is
important to note that if the ${\diamond}$-norm is not strongly convex,
${\cal{W}}^{\textsl{mm}}$ may not be a singleton. Additionally, when $m=d$,
the rank constraint becomes vacuous, and the problem becomes convex. The
following result is a slight generalization of Theorem 1 and demonstrates that
choosing a large $d$ ensures the feasibility of (${\diamond}$-SVM) uniformly
over all choices of ${\bm{\alpha}}$.

###### Theorem 8

Suppose $d\geq\max(T-1,n)$ and $m=d$. Then, almost all datasets444Here,
_“almost all datasets”_ means that adding i.i.d. gaussian noise, with
arbitrary nonzero variance, to the input features will almost surely result in
SVM’s feasibility. $(Y_{i},{\bm{X}}_{i},{\bm{z}}_{i})_{i=1}^{n}$ – including
the self-attention setting with ${\bm{z}}_{i}\leftarrow\bm{x}_{i1}$ – obey the
following: For any choice of indices
${\bm{\alpha}}=(\alpha_{i})_{i=1}^{n}\subset[T]$, (${\diamond}$-SVM) is
feasible, i.e. the attention layer can separate and select indices
${\bm{\alpha}}$.

To proceed, we define the _local regularization path_ , which is obtained by
solving the ${\diamond}$-norm-constrained problem over a
${\bm{\alpha}}$-dependent cone denoted as
$\texttt{cone}_{\varepsilon}({\bm{\alpha}})$. This cone has a simple
interpretation: it prioritizes tokens with a lower score than ${\bm{\alpha}}$
over tokens with a higher score than ${\bm{\alpha}}$. This interpretation
sheds light on the convergence towards locally optimal directions: lower-score
tokens create a barrier for ${\bm{\alpha}}$ and prevent optimization from
moving towards higher-score tokens.

######  Definition 3 (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}\right\\},\quad\texttt{high}^{\alpha}({\bm{X}})=\left\\{t\in[T]-\\{\alpha\\}{~{}\big{|}~{}}\bm{\gamma}_{t}\geq\bm{\gamma}_{\alpha}\right\\}.$

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

$\displaystyle\texttt{cone}_{\varepsilon}({\bm{\alpha}})=\left\\{\texttt{rank}(W)\leq
m{~{}\big{|}~{}}\min_{i\in[n]}\max_{t\in\texttt{low}^{\alpha}_{i}}\min_{\tau\in\texttt{high}^{\alpha}_{i}}(\bm{x}_{it}-\bm{x}_{i\tau})^{\top}\bm{W}{\bm{z}}_{i}\geq\varepsilon\|{\bm{W}}\|_{F}\right\\}.$
(5)

Our next lemma relates this cone definition to locally-optimal directions of
Definition 2.

###### Lemma 5

Suppose (${\diamond}$-SVM) is feasible. 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}_{i}\in\arg\max_{t\in[T]}\bm{\gamma}_{it}$ are unique. Then,
$\texttt{cone}_{\varepsilon}(\texttt{opt})$ is the set of all rank-$\leq$$m$
matrices (i.e. global set).

Lemma 5 can be understood as follows: Among the SVM solutions
$\bm{W}^{\textsl{mm}}_{\bm{\alpha}}$, only those that are locally optimal
demonstrate a barrier of low-score tokens, effectively acting as a protective
shield against higher-score tokens. Moreover, in the case of globally optimal
tokens (with the highest scores), the global set
$\texttt{cone}_{\varepsilon}(\texttt{opt})$ can be chosen, as they inherently
do not require protective measures. The subsequent result introduces our
principal theorem, which pertains to the regularization path converging
towards the locally-optimal direction over
$\texttt{cone}_{\varepsilon}({\bm{\alpha}})$ whenever ${\bm{\alpha}}$ is
locally optimal.

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

Suppose Assumption A (i), (iii), and E hold. Fix locally-optimal token indices
${\bm{\alpha}}=(\alpha_{i})_{i=1}^{n}$ and $R_{0},\varepsilon>0$. Consider the
norm-constrained variation of (5) defined as

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

Define local RP as
${\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}}$ be the set of minima for
(${\diamond}$-SVM) and $\Xi_{\diamond}>0$ be the associated margin i.e.
$\Xi_{\diamond}=1/\|{\bm{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}}\right)=0$.
Additionally, suppose optimal indices
$\texttt{opt}=(\texttt{opt}_{i})_{i=1}^{n}$ are unique and set
${\bm{\alpha}}\leftarrow\texttt{opt}$. Then, the same convergence guarantee on
regularization path holds by setting
$\mathcal{C}_{\varepsilon,R_{0}}^{\diamond}$ as the set of rank-$\leq$$m$
matrices.

Note that when setting $m=d$, the rank constraint is eliminated. Consequently,
specializing this theorem to the Frobenius norm aligns it with Theorem 7. On
the other hand, by assigning ${\diamond}$ as the nuclear norm and
${\bm{\alpha}}\leftarrow\texttt{opt}$, the global inductive bias of the
nuclear norm is recovered, as stated in Theorem 2.

We would like to emphasize that both this theorem and Theorem 2 are specific
instances of Theorem 10 found in Appendix D. It is worth noting that, within
this appendix, we establish all regularization path results for sequence-to-
sequence classification, along with a general class of _monotonicity-
preserving_ prediction heads outlined in Assumption E. The latter
significantly generalizes linear heads, highlighting the versatility of our
theory. The following section presents our discoveries concerning general
nonlinear heads.

##  6 Understanding Multi-token Compositions: Toward A More General Max-Margin
and Directional Convergence Theory

So far, our theory has focused on the setting where the attention layer
selects a single optimal token within each sequence. As we have discussed,
this is theoretically well-justified under linear head assumption and certain
nonlinear generalizations. On the other hand, for arbitrary nonconvex
$h(\cdot)$ or multilayer transformer architectures, it is expected that
attention will select multiple tokens per sequence. This motivates us to ask:

> Q: What is the implicit bias and the form of $\bm{W}(k)$ when the GD
> solution is composed by multiple tokens?

In this section, our goal is to derive and verify the generalized behavior of
GD. Let $\bm{o}_{i}={\bm{X}}_{i}^{\top}\bm{s}^{\bm{W}}_{i}$ denote the token
generated by the attention layer where
$\bm{s}^{\bm{W}}_{i}=\mathbb{S}({\bm{X}}_{i}\bm{W}{\bm{z}}_{i})$ are the
softmax probabilities. Suppose GD trajectory converges to achieve the risk
${\cal{L}}_{\star}=\min_{\bm{W}}{\cal{L}}(\bm{W})$. Suppose the eventual token
composition achieving ${\cal{L}}_{\star}$ is given by

$\bm{o}^{\star}_{i}={\bm{X}}_{i}^{\top}\bm{s}^{\star}_{i},$

where $\bm{s}^{\star}_{i}$ are the eventual softmax probability vectors that
dictate the token composition. Since attention maps are sparse in practice, we
are interested in the scenario where $\bm{s}^{\star}_{i}$ is sparse i.e. it
contains some zero entries. This can only be accomplished by letting
$\|{\bm{W}}\|_{F}\rightarrow\infty$. However, unlike the earlier sections, we
wish to allow for arbitrary $\bm{s}^{\star}_{i}$ rather than a one-hot vector
which selects a single token.

To proceed, we aim to understand the form of GD solution $\bm{W}(k)$
responsible for composing $\bm{o}^{\star}_{i}$ via the softmax map
$\bm{s}^{\star}_{i}$ as $R\rightarrow\infty$. Intuitively, $\bm{W}(k)$ should
be decomposed into two components via

$\displaystyle\bm{W}(k)\approx{\bm{W}^{\textsl{fin}}}+\|{\bm{W}(k)}\|_{F}\cdot\bar{\bm{W}}^{\textsl{mm}}.$
(6)

where ${\bm{W}^{\textsl{fin}}}$ is the finite component and
$\bar{\bm{W}}^{\textsl{mm}}$ is the directional component with
$\|{\bar{\bm{W}}^{\textsl{mm}}}\|_{F}=1$. Define the selected set
$\mathcal{O}_{i}\subseteq[T]$ to be the indices $s^{\star}_{it}\neq 0$ and the
masked (i.e. suppressed) set as $\bar{\mathcal{O}}_{i}=[T]-\mathcal{O}_{i}$
where softmax entries are zero. In the context of earlier sections, we could
also call these the _optimal set_ and the _non-optimal set_ , respectively.

  * $\bullet$

Finite component: The job of ${\bm{W}^{\textsl{fin}}}$ is to assign nonzero
softmax probabilities within each $\bm{s}^{\star}_{i}$. This is accomplished
by ensuring that, ${\bm{W}^{\textsl{fin}}}$ induces the probabilities of
$\bm{s}^{\star}_{i}$ over $\mathcal{O}_{i}$ by satisfying the softmax
equations

$\frac{e^{\bm{x}_{it}^{\top}{\bm{W}^{\textsl{fin}}}{\bm{z}}_{i}}}{e^{\bm{x}_{i\tau}^{\top}{\bm{W}^{\textsl{fin}}}{\bm{z}}_{i}}}=e^{(\bm{x}_{it}-\bm{x}_{i\tau})^{\top}{\bm{W}^{\textsl{fin}}}{\bm{z}}_{i}}=s^{\star}_{it}/s^{\star}_{i\tau}.$

for $t,\tau\in\mathcal{O}_{i}$. Consequently, this ${\bm{W}^{\textsl{fin}}}$
should satisfy the following linear constraints

$(\bm{x}_{it}-\bm{x}_{i\tau})^{\top}{\bm{W}^{\textsl{fin}}}{\bm{z}}_{i}=\log(s^{\star}_{it}/s^{\star}_{i\tau})\quad\text{for
all}\quad t,\tau\in\mathcal{O}_{i},~{}i\in[n].$ (7)

  * $\bullet$

Directional component: While ${\bm{W}^{\textsl{fin}}}$ creates the composition
by allocating the nonzero softmax probabilities, it does not explain sparsity
of attention map. This is the role of $\bar{\bm{W}}^{\textsl{mm}}$, which is
responsible for selecting the selected tokens $\mathcal{O}_{i}$ and
suppressing the masked ones $\bar{\mathcal{O}}_{i}$ by assigning zero softmax
probability to them. To predict direction component, we build on the theory
developed in earlier sections. Concretely, there are two constraints
$\bar{\bm{W}}^{\textsl{mm}}$ should satisfy

    1. 1.

Equal similarity over selected tokens: For all $t,\tau\in\mathcal{O}_{i}$, we
have that $(\bm{x}_{it}-\bm{x}_{i\tau})^{\top}\bm{W}{\bm{z}}_{i}=0$. This way,
softmax scores assigned by ${\bm{W}^{\textsl{fin}}}$ are not disturbed by the
directional component and
${\bm{W}^{\textsl{fin}}}+R\bar{\bm{W}}^{\textsl{mm}}$ will still satisfy the
softmax equations (7).

    2. 2.

Max-margin against masked tokens: For all
$t\in\mathcal{O}_{i},\tau\in\bar{\mathcal{O}}_{i}$, enforce the margin
constraint $(\bm{x}_{it}-\bm{x}_{i\tau})^{\top}\bm{W}{\bm{z}}_{i}\geq 1$
subject to minimum norm $\|{\bm{W}}\|_{F}$.

Combining these, we obtain the following convex generalized SVM formulation

$\displaystyle\bm{W}^{\textsl{mm}}=\arg\min_{\bm{W}}\|{\bm{W}}\|_{F}\quad\text{subj.
to}\quad\begin{cases}\forall~{}t\in\mathcal{O}_{i},\tau\in\bar{\mathcal{O}}_{i}:~{}(\bm{x}_{it}-\bm{x}_{i\tau})^{\top}\bm{W}{\bm{z}}_{i}\geq
1,\\\
\forall~{}t,\tau\in\mathcal{O}_{i}:~{}\quad\quad(\bm{x}_{it}-\bm{x}_{i\tau})^{\top}\bm{W}{\bm{z}}_{i}=0,\end{cases}\quad\forall
1\leq i\leq n.$ (Gen-SVM)

and set the normalized direction in (6) to
$\bar{\bm{W}}^{\textsl{mm}}=\bm{W}^{\textsl{mm}}/\|{\bm{W}^{\textsl{mm}}}\|_{F}$.

It is important to note that (Gen-SVM) offers a substantial generalization
beyond the scope of the previous sections, where the focus was on selecting a
single token from each sequence, as described in the main formulation (Att-
SVM). This broader solution class introduces a more flexible approach to the
problem.

Iterations$d=4$Iterations$d=6$Iterations$d=8$Iterations$d=10$$1-$correlation
coefficient$\bm{W}^{\texttt{SVMeq}}$$\bm{W}^{\textsl{mm}}$$\bm{W}^{\texttt{1token}}$$\bm{W}^{\texttt{SVMeq}}$$\bm{W}^{\textsl{mm}}$$\bm{W}^{\texttt{1token}}$$\bm{W}^{\texttt{SVMeq}}$$\bm{W}^{\textsl{mm}}$$\bm{W}^{\texttt{1token}}$$\bm{W}^{\texttt{SVMeq}}$$\bm{W}^{\textsl{mm}}$$\bm{W}^{\texttt{1token}}$

$\max_{i,\tau}s_{i\tau},~{}\tau\in\bar{\mathcal{O}}_{i}$$\max_{i,\tau}s_{i\tau},~{}\tau\in\bar{\mathcal{O}}_{i}$$\max_{i,\tau}s_{i\tau},~{}\tau\in\bar{\mathcal{O}}_{i}$$\max_{i,\tau}s_{i\tau},~{}\tau\in\bar{\mathcal{O}}_{i}$$1-$correlation
coefficient

Figure 8:  Behavior of GD with nonlinear nonconvex prediction head and multi-
token compositions. Upper: The correlation between GD solution and three
distinct baselines: ($\cdots$) $\bm{W}^{\textsl{mm}}$ obtained from (Gen-SVM);
(—) $\bm{W}^{\texttt{SVMeq}}$ obtained by calculating
${\bm{W}^{\textsl{fin}}}$ and determining the best linear combination
${\bm{W}^{\textsl{fin}}}+\gamma\bar{\bm{W}}^{\textsl{mm}}$ that maximizes
correlation with the GD solution; and (\- -) $\bm{W}^{\texttt{1token}}$
obtained by solving (Att-SVM) and selecting the highest probability token from
the GD solution. Lower: Scatterplot of the largest softmax probability over
masked tokens (per our $s_{i\tau}\leq 10^{-6}$ criteria) vs correlation
coefficient.

We present experiments showcasing the predictive power of the (Gen-SVM)
equivalence in nonlinear scenarios. We conducted these experiments on random
instances using an MLP denoted as $h(\cdot)$, which takes the form of
${\mathbf{1}}^{\top}\texttt{ReLU}(\bm{x})$. We begin by detailing the
preprocessing step and our setup. For the attention SVM equivalence analytical
prediction, clear definitions of the selected and masked sets are crucial.
These sets include token indices with nonzero and zero softmax outputs,
respectively. However, practically, reaching a precisely zero output is not
feasible. Hence, we define the selected set as tokens with softmax outputs
exceeding $10^{-3}$, and the masked set as tokens with softmax outputs below
$10^{-6}$. We also excluded instances with softmax outputs falling between
$10^{-6}$ and $10^{-3}$ to distinctly separate the concepts of _selected_ and
_masked_ sets, thereby enhancing the predictive accuracy of the attention SVM
equivalence. In addition to the filtering process, we focus on scenarios where
the label $Y=-1$ exists to enforce _non-convexity_ of prediction head
$Y_{i}\cdot h(\cdot)$. It is worth mentioning that when all labels are $1$,
due to the convexity of $Y_{i}\cdot h(\cdot)$, GD tends to select one token
per input, and Equations (Gen-SVM) and (Att-SVM) yield the same solutions. The
results are displayed in Figure 8, where $n=3$, $T=4$, and $d$ varies within
${4,6,8,10}$. We conduct 500 random trials for different choices of $d$, each
involving ${\bm{x}}_{it}$, ${{\bm{z}}}_{i}$, and ${\bm{v}}$ randomly sampled
from the unit sphere. We apply normalized GD with a step size $\eta=0.1$ and
run $2000$ iterations for each trial.

  1. $\bullet$

Figure 8 (upper) illustrates the correlation evolution between the GD solution
and three distinctive baselines: ($\cdots$) $\bm{W}^{\textsl{mm}}$ obtained
from (Gen-SVM); (—) $\bm{W}^{\texttt{SVMeq}}$ obtained by calculating
${\bm{W}^{\textsl{fin}}}$ and determining the best linear combination
${\bm{W}^{\textsl{fin}}}+\gamma\bar{\bm{W}}^{\textsl{mm}}$ that maximizes
correlation with the GD solution; and (\- -) $\bm{W}^{\texttt{1token}}$
obtained by solving (Att-SVM) and selecting the highest probability token from
the GD solution. For clearer visualization, the logarithmic scale of
correlation misalignment is presented in Figure 8. In essence, our findings
show that $\bm{W}^{\texttt{1token}}$ yields unsatisfactory outcomes, whereas
$\bm{W}^{\textsl{mm}}$ attains a significant correlation coefficient in
alignment with our expectations. Ultimately, our comprehensive SVM-equivalence
$\bm{W}^{\texttt{SVMeq}}$ further enhances correlation, lending support to our
analytical formulas. It’s noteworthy that SVM-equivalence displays higher
predictability in a larger $d$ regime (with an average correlation exceeding
$0.99$). This phenomenon might be attributed to more frequent directional
convergence in higher dimensions, with overparameterization contributing to a
smoother loss landscape, thereby expediting optimization.

  2. $\bullet$

Figure 8 (lower) offers a scatterplot overview of the $500$ random problem
instances that were solved. The $x$-axis represents the largest softmax
probability over the masked set, denoted as $\max_{i,\tau}s_{i\tau}$ where
$\tau\in\bar{\mathcal{O}}_{i}$. Meanwhile, the $y$-axis indicates the
predictivity of the SVM-equivalence, quantified as
$1-\texttt{corr\\_coef}(\bm{W},\bm{W}^{\texttt{SVMeq}})$. From this analysis,
two significant observations arise. Primarily, there exists an inverse
correlation between softmax probability and SVM-predictivity. This correlation
is intuitive, as higher softmax probabilities signify a stronger divergence
from our desired _masked set_ state (ideally set to $0$). Secondly, as
dimensionality ($d$) increases, softmax probabilities over the masked set tend
to converge towards the range of $10^{-15}$ (effectively zero).
Simultaneously, attention SVM-predictivity improves, creating a noteworthy
correlation.

###  6.1 When does attention select multiple tokens?

In this section, we provide a concrete example where the optimal solution
indeed requires combining multiple tokens in a nontrivial fashion. Here, by
nontrivial we mean that, we select more than 1 tokens from an input sequence
but we don’t select all of its tokens. Recall that, for linear prediction
head, attention will ideally select the single token with largest score for
almost all datasets. Perhaps not surprisingly, this behavior will not persist
for nonlinear prediction heads. For instance in Figure 8, the GD output
$\bm{W}$ aligned better in direction with $\bm{W}^{\textsl{mm}}$ than
$\bm{W}^{\texttt{1token}}$. Specifically, here we prove that if we make the
function $h_{Y}(\bm{x}):=Y\cdot h(\bm{x})$ concave, then optimal softmax map
can select multiple tokens in a controllable fashion. $h_{Y}(\bm{x})$ can be
viewed as generalization of the linear score function
$Y\cdot\bm{v}^{\top}\bm{x}$. In the example below, we induce concavity by
incorporating a small $-\lambda\|{\bm{x}}\|^{2}$ term within a linear
prediction head and setting
$h(\bm{x})=\bm{v}^{\top}\bm{x}-\lambda\|{\bm{x}}\|^{2}$ with $Y=1$.

###### Lemma 6

Given $\bm{v}\in\mathbb{R}^{d}$, recall the score vector
$\bm{\gamma}={\bm{X}}\bm{v}$. Without losing generality, assume $\bm{\gamma}$
is non-increasing. Define the vector of score gaps
$\bm{\gamma}^{\textsl{gap}}\in\mathbb{R}^{T-1}$ with entries
$\bm{\gamma}^{\textsl{gap}}_{t}=\bm{\gamma}_{t}-\bm{\gamma}_{t+1}$. Suppose
all tokens within the input sequence are orthonormal and for some $\tau\geq
2$, we have that

$\displaystyle\tau\bm{\gamma}^{\textsl{gap}}_{\tau}/2>\bm{\gamma}^{\textsl{gap}}_{1}.$
(8)

Set $h(\bm{x})=\bm{v}^{\top}\bm{x}-\lambda\|{\bm{x}}\|^{2}$ where
$\tau\bm{\gamma}^{\textsl{gap}}_{\tau}/2>\lambda>\bm{\gamma}^{\textsl{gap}}_{1}$,
$\ell(x)=-x$, and $Y=1$. Let ${\bm{\Delta}}_{T}$ denote the $T$-dimensional
simplex. Define the unconstrained softmax optimization associated to the
objective $h$ where we make $\bm{s}:=\mathbb{S}({\bm{X}}\bm{W}{\bm{z}})$ a
free variable, namely,

$\displaystyle\min_{\bm{s}\in{\bm{\Delta}}_{T}}\ell(h({\bm{X}}\bm{s}))=\min_{\bm{s}\in{\bm{\Delta}}_{T}}\lambda\|{{\bm{X}}^{\top}\bm{s}}\|^{2}-\bm{v}^{\top}{\bm{X}}^{\top}\bm{s}.$
(9)

Then, the optimal solution $\bm{s}^{\star}$ contains at least $2$ and at most
$\tau$ nonzero entries.

$\lambda$# of selected tokens (a) $\lambda$ vs # selected tokens

$1-$correlation coefficientIterations (b) $\lambda$ vs correlation coefficient

Iterations$1-$correlation coefficient$\text{\\# selected}=1$$\text{\\#
selected}=2$$\text{\\# selected}=5$$\text{\\# selected}=9$ (c) # selected
tokens vs correlation coefficient

Figure 9:  Behavior of GD when selecting multiple tokens. (a) The number of
selected tokens increases with $\lambda$. (b) Predictivity of attention SVM
solutions for varying $\lambda$; Dotted curves depict the correlation
corresponding to $\bm{W}^{\textsl{mm}}$ calculated via (Gen-SVM) and solid
curves represent the correlation to $\bm{W}^{\texttt{SVMeq}}$, which
incorporates the ${\bm{W}^{\textsl{fin}}}$ correction. (c) Similar to (b), but
evaluating correlations over different numbers of selected tokens.

Figure 9 presents experimental findings concerning Lemma 6 across random
problem instances. For this experiment, we set $n=1$, $T=10$, and $d=10$. The
results are averaged over $100$ random trials, with each trial involving the
generation of randomly orthonormal vectors $\bm{x}_{1t}$ and the random
sampling of vector $\bm{v}$ from the unit sphere. Similar to the processing
step in Figure 8, and following Figure 8 (lower) which illustrates that
smaller softmax outputs over masked sets correspond to higher correlation
coefficients, we define the selected and masked token sets. Specifically,
tokens with softmax outputs $>10^{-3}$ are considered selected, while tokens
with softmax outputs $<10^{-8}$ are masked. Instances with softmax outputs
between $10^{-8}$ and $10^{-3}$ are filtered out.

Figure 9(a) shows that the number of selected tokens grows alongside
$\lambda$, a prediction consistent with Lemma 6. When $\lambda=0$, the head
$h(\bm{x})=\bm{v}^{\top}\bm{x}$ is linear, resulting in the selection of only
one token per input. Conversely, as $\lambda$ exceeds a certain threshold
(e.g., $\lambda>2.0$ based on our criteria), the optimization consistently
selects all tokens. Figure 9(b) and 9(c) delve into the predictivity of
attention SVM solutions for varying $\lambda$ and different numbers of
selected tokens. The dotted curves in both figures represent
$1-{\texttt{corr\\_coef}}(\bm{W},\bm{W}^{\textsl{mm}})$, while solid curves
indicate $1-{\texttt{corr\\_coef}}(\bm{W},\bm{W}^{\texttt{SVMeq}})$, where
$\bm{W}$ denotes the GD solution. Overall, the SVM-equivalence demonstrates a
strong correlation with the GD solution (consistently above $0.95$). However,
selecting more tokens (aligned with larger $\lambda$ values) leads to reduced
predictivity.

To sum up, we have showcased the predictive capacity of the generalized SVM
equivalence regarding the inductive bias of 1-layer transformers with
nonlinear heads. Nevertheless, it’s important to acknowledge that this section
represents an initial approach to a complex problem, with certain caveats
requiring further investigation (e.g., the use of filtering in Figures 8 and
9, and the presence of imperfect correlations). We aspire to conduct a more
comprehensive investigation, both theoretically and empirically, in
forthcoming work.

##  7 Extending the Theory to Sequential and Causal Predictions

While our formulations (W-ERM & KQ-ERM) regress a single label $Y_{i}$ per
$({\bm{X}}_{i},{\bm{z}}_{i})$, in Appendix D, we show that our findings
gracefully extend to the sequence-to-sequence classification setting where we
output and classify all $T$ tokens per input
${\bm{X}}_{i},\bm{Z}_{i}\in\mathbb{R}^{T\times d}$. In this scenario, we prove
that all of our RP guarantees remain intact after introducing a slight
generalization of (Att-SVM). To further elaborate, consider the following ERM
problems for sequential and causal settings:

$\displaystyle{\cal{L}}^{\texttt{seq}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\sum_{k=1}^{T}\ell(Y_{ik}\cdot
h({\bm{X}}_{i}^{\top}\mathbb{S}({\bm{X}}_{i}\bm{W}{\bm{z}}_{ik}))),\quad\text{and}\quad{\cal{L}}^{\texttt{csl}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\sum_{k=1}^{T}\ell(Y_{ik}\cdot
h({\bm{X}}_{i}^{\top}\mathbb{S}_{k}({\bm{X}}_{i}\bm{W}{\bm{z}}_{ik}))).$ (10)

Both equations train $T$ tokens per input $({\bm{X}}_{i},\bm{Z}_{i})$ and, as
usual, we recover self-attention via $\bm{Z}_{i}\leftarrow{\bm{X}}_{i}$. For
the causal setting, we use the masked attention $\mathbb{S}_{k}(\cdot)$ which
calculates the softmax probabilities over the first $k$ entries of its input
and sets the remaining $T-k$ entries to zero. This way, the $k$’th prediction
of the transformer only utilizes tokens from $1$ to $k$ and not the future
tokens.

Let ${\bm{\alpha}}=(\alpha_{ik})_{ik=(1,1)}^{(n,k)}$ be tokens to be selected
by attention (e.g. locally-optimal indices, see Def. 4). Then, sequential
generalization of (Att-SVM) corresponding to
${\cal{L}}^{\texttt{seq}}(\bm{W})$ is given by

$\min_{\bm{W}}\|{\bm{W}}\|_{F}\quad\text{subj.
to}\quad(\bm{x}_{i\alpha_{ik}}-\bm{x}_{it})^{\top}\bm{W}{\bm{z}}_{ik}\geq
1\quad\textnormal{for
all}\quad{t\neq\alpha_{ik}},~{}~{}k\in[T],~{}~{}i\in[n]~{}~{}.$ (11)

We refer the reader to Appendix D which rigorously establishes all RP results
for this sequential classification setting. On the other hand, for the causal
inference setting SVM should reflect the fact that the model is not allowed to
make use of future tokens. Note that, SVM constraints directly arise from
softmax calculations. Thus, since attention is masked over the indices $t\leq
k$ and $k\in[T]$, the SVM constraints should apply over the same mask. Thus,
we can consider the straightforward generalization of global-optimality where
$\texttt{opt}_{ik}$ is the token index with highest score over indices $t\leq
k$ and introduce the similar definition for local-optimality. This leads to
the following variation of (11) where we aim to select indices
$\alpha_{ik}\in[k]$ from the first $k$ tokens

$\min_{\bm{W}}\|{\bm{W}}\|_{F}\quad\text{subj.
to}\quad(\bm{x}_{i\alpha_{ik}}-\bm{x}_{it})^{\top}\bm{W}{\bm{z}}_{ik}\geq
1\quad\textnormal{for all}\quad t\neq\alpha_{ik},~{}~{}t\leq
k,~{}~{}k\in[T],~{}~{}i\in[n].$

Causal attention is a special case of a general attention mask which can
restrict softmax to arbitrary subset of the tokens. Such general masking can
be handled similar to (7) by enforcing SVM constraints over the nonzero
support of the mask. Finally, the discussion so far extends our main
theoretical results and focuses on selecting single token per sequence. It can
further be enriched by the generalized SVM equivalence developed in Section 6
to select and compose multiple tokens by generalizing (Gen-SVM).

##  8 Related work

###  8.1 Implicit Regularization, Matrix Factorization, Sparsity

Extensive research has delved into gradient descent’s implicit bias in
separable classification tasks, often using logistic or exponentially-tailed
losses for margin maximization [SHN+18, GLSS18, NLG+19, JT21, KPOT21, MWG+20,
JT20]. The findings have also been extended to non-separable data using
gradient-based techniques [JT18, JT19, JDST20]. Implicit bias in regression
problems and losses has been investigated, utilizing methods like mirror
descent [WGL+20, GLSS18, YKM20, VKR19, AW20a, AW20b, ALH21, SATA22].
Stochastic gradient descent has also been a subject of interest regarding its
implicit bias [LWM19, BGVV20, LR20, HWLM20, LWA22, DML21, ZWB+21]. This
extends to the implicit bias of adaptive and momentum-based methods [QQ19,
WMZ+21, WMCL21, JST21].

In linear classification, GD iterations on logistic loss and separable
datasets converge to the hard margin SVM solution [SHN+18, RZH03, ZY05]. The
attention layer’s softmax nonlinearity behaves similarly, potentially favoring
margin-maximizing solutions. Yet, the layer operates on tokens in input
sequences, not for direct classification. Its bias leans toward an (Att-SVM),
selecting relevant tokens while suppressing others. However, formalizing this
intuition presents significant challenges: Firstly, our problem is nonconvex
(even in terms of the $\bm{W}$-parameterization), introducing new challenges
and complexities. Secondly, it requires the introduction of novel concepts
such as locally-optimal tokens, demanding a tailored analysis focused on the
cones surrounding them. Our findings on the implicit bias of
$({\bm{K}},{\bm{Q}})$-parameterization share conceptual similarities with
[SRJ04], which proposes and analyzes a max-margin matrix factorization
problem. Similar problems have also been studied more recently in the context
of neural-collapse phenomena [PHD20] through an analysis of the implicit bias
and regularization path of the unconstrained features model with cross-entropy
loss [TKVB22]. However, a fundamental distinction from these works lies in the
fact that attention solves a different max-margin problem that separate
tokens. Moreover, our results on $({\bm{K}},{\bm{Q}})$-parameterization are
inherently connected to the rich literature on low-rank factorization [GWB+17,
ACHL19, TVS23, TBS+16, SS21], stimulating further research. [TLZO23] is the
first work to establish the connection between attention and SVM, which is
closest to our work. Here, we augment their framework, initially developed for
a simpler attention model, to transformers by providing the first guarantees
for self/cross-attention layers, nonlinear prediction heads, and realistic
global convergence guarantees. While our Assumption B.2 and local-convergence
analysis align with [TLZO23], our contributions in global convergence
analysis, benefits of overparameterization, and the generalized SVM-
equivalence in Section 6 are unique to this work.

It is well-known that attention map (i.e. softmax outputs) act as a feature
selection mechanism and reveal the tokens that are relevant to classification.
On the other hand, sparsity and lasso regression (i.e. $\ell_{1}$
penalization) [Don06, Tib96, TG07, CDS01, CRT06] have been pivotal tools in
the statistics literature for feature selection. Softmax and lasso regression
exhibit interesting parallels: The Softmax output
$\bm{s}=\mathbb{S}({\bm{X}}\bm{W}{\bm{z}})$ obeys $\|\bm{s}\|_{\ell_{1}}=1$ by
design. Softmax is also highly receptive to being sparse because decreasing
the temperature (i.e. scaling up the weights $\bm{W}$) eventually leads to a
one-hot vector unless all logits are equal. We (also, [TLZO23]) have used
these intuitions to formalize attention as a _token selection mechanism_. This
aspect is clearly visible in our primary SVM formulation (Att-SVM) which
selects precisely one token from each input sequence (i.e. hard attention).
Section 6 has also demonstrated how (Gen-SVM) can explain more general
sparsity patterns by precisely selecting desired tokens and suppressing
others. We hope that this SVM-based token-selection viewpoint will motivate
future work and deeper connections to the broader feature-selection and
compressed sensing literature.

###  8.2 Attention Mechanism and Transformers

Transformers, as highlighted by [VSP+17], revolutionized the domains of NLP
and machine translation. Prior work on self-attention [CDL16, PTDU16, PXS18,
LFS+17] laid the foundation for this transformative paradigm. In contrast to
conventional models like MLPs and CNNs, self-attention models employ global
interactions to capture feature representations, resulting in exceptional
empirical performance.

Despite their achievements, the mechanisms and learning processes of attention
layers remain enigmatic. Recent investigations [EGKZ22, SEO+22, ENM22, BV22,
DCL21] have concentrated on specific aspects such as sparse function
representation, convex relaxations, and expressive power. Expressivity
discussions concerning hard-attention [Hah20] or attention-only architectures
[DCL21] are connected to our findings when $h(\cdot)$ is linear. In fact, our
work reveals how linear $h$ results in attention’s optimization dynamics to
collapse on a single token whereas nonlinear $h$ provably requires attention
to select and compose multiple tokens. This supports the benefits of the MLP
layer for expressivity of transformers. There is also a growing body of
research aimed at a theoretical comprehension of in-context learning and the
role played by the attention mechanism [ASA+22, LIPO23, ACDS23, ZFB23, BCW+23,
GRS+23]. [SEO+22] investigate self-attention with linear activation instead of
softmax, while [ENM22] approximate softmax using a linear operation with unit
simplex constraints. Their primary goal is to derive convex reformulations for
training problems grounded in empirical risk minimization (ERM). In contrast,
our methodologies, detailed in equations (W-ERM) and (KQ-ERM), delve into the
nonconvex domain.

[MRG+20, BALA+23] offer insights into the implicit bias of optimizing
transformers. Specifically, [MRG+20] provide empirical evidence that an
increase in attention weights results in a sparser softmax, which aligns with
our theoretical framework. [BALA+23] study incremental learning and furnish
both theory and numerical evidence that increments of the softmax attention
weights (${\bm{K}}{\bm{Q}}^{\top}$) are low-rank. Our theory aligns with this
concept, as the SVM formulation (KQ-SVM) of $({\bm{K}},{\bm{Q}})$
parameterization inherently exhibits low-rank properties through the nuclear
norm objective, rank-$m$ constraint, and implicit constraint induced by Lemma
1.

Several parallel recent works [JSL22, LWLC23, TWCD23, NLL+23, ORST23, NNH+23]
aim to delineate the optimization and generalization dynamics of transformers.
However, their findings usually apply under strict statistical assumptions
about the data, while our study offers a comprehensive optimization-theoretic
analysis of the attention model, establishing a formal linkage to max-margin
problems and SVM geometry. This allows our findings to encompass the problem
geometry and apply to diverse datasets. Overall, the max-margin equivalence
provides a fundamental comprehension of the optimization geometry of
transformers, offering a framework for prospective research endeavors, as
outlined in the subsequent section.

##  9 Discussion, Future Directions, and Open Problems

Our optimization-theoretic characterization of the self-attention model
provides a comprehensive understanding of its underlying principles. The
developed framework, along with the research presented in [TLZO23], introduces
new avenues for studying transformers and language models. The key findings
include:

  1. $\checkmark$

The optimization geometry of self-attention exhibits a fascinating connection
to hard-margin SVM problems. By leveraging linear constraints formed through
outer products of token pairs, optimal input tokens can be effectively
separated from non-optimal ones.

  2. $\checkmark$

When gradient descent is employed without early-stopping, implicit
regularization and convergence of self-attention naturally occur. This
convergence leads to the maximum margin solution when minimizing specific
requirements using logistic loss, exp-loss, or other smooth decreasing loss
functions. Moreover, this implicit bias is unaffected by the step size, as
long as it is sufficiently small for convergence, and remains independent of
the initialization process.

The fact that gradient descent leads to a maximum margin solution may not be
surprising to those who are familiar with the relationship between
regularization path and gradient descent in linear and nonlinear neural
networks [SHN+18, GLSS18, NLG+19, JT21, MWG+20, JT20]. However, there is a
lack of prior research or discussion regarding this connection to the
attention mechanism. Moreover, there has been no rigorous analysis or
investigation into the exactness and independence of this bias with respect to
the initialization and step size. Thus, we believe our findings and insights
deepen our understanding of transformers and language models, paving the way
for further research in this domain. Below, we discuss some notable directions
and highlight open problems that are not resolved by the existing theory.

  * •

Convergence Rates: The current paper establishes asymptotic convergence of
gradient descent; nonetheless, there is room for further exploration to
characterize non-asymptotic convergence rates. Indeed, such an exploration can
also provide valuable insights into the choice of learning rate,
initialization, and the optimization method.

  * •

Gradient descent on $({\bm{K}},{\bm{Q}})$ parameterization: We find it
remarkable that regularization path analysis was able to predict the implicit
bias of gradient descent. Complete analysis of gradient descent is inherently
connected to the fundamental question of low-rank factorization [GWB+17,
LMZ18]. We believe formalizing the implicit bias of gradient descent under
margin constraints presents an exciting open research direction for further
research.

  * •

Generalization analysis: An important direction is the generalization
guarantees for gradient-based algorithms. The established connection to hard-
margin SVM can facilitate this because the SVM problem is amenable to
statistical analysis. This would be akin to how kernel/NTK analysis for deep
nets enabled a rich literature on generalization analysis for traditional deep
learning.

  * •

Global convergence of gradient descent: We lack a complete characterization of
the directional convergence of gradient descent. We ask: _Where does gradient
descent directionally-converge from arbitrary initialization for 1-layer self-
attention?_ The role of over-parameterization as conjectured in Section 5.2
and the notion of locally-optimal directions discussed in Section 5 constitute
important pieces of this puzzle (also see the discussion in [TLZO23]).

  * •

Realistic architectures: Naturally, we wish to explore whether max-margin
equivalence can be extended to more realistic settings: Can the theory be
expanded to handle multi-head attention, multi-layer architectures, and MLP
nonlinearities? We believe the results in Section 6 take an important step
towards this direction by including analytical formulae for the implicit bias
of the attention layer under nonlinear prediction heads.

  * •

Jointly optimizing attention and prediction head: It would be interesting to
study the joint optimization dynamics of attention weights and prediction head
$h(\cdot)$. This problem can be viewed as a novel low-rank factorization type
problem where $h(\cdot)$ and $\bm{W}$ are factors of the optimization problem,
only, here, $\bm{W}$ passes through the softmax nonlinearity. To this aim,
[TLZO23] provides a preliminary geometric characterization of the implicit
bias for a simpler attention model using regularization path analysis. Such
findings can potentially be generalized to the analysis of gradient methods
and full transformer block.

## Acknowledgements

This work was supported by the NSF grants CCF-2046816 and CCF-2212426, Google
Research Scholar award, and Army Research Office grant W911NF2110312. The
authors thank Xuechen Zhang, Ankit Singh Rawat, Mahdi Soltanolkotabi, Jason
Lee, Arkadas Ozakin, Ramya Korlakai Vinayak, and Babak Hassibi for helpful
suggestions and discussion.

## References

  * [ACDS23] Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra.  Transformers learn to implement preconditioned gradient descent for in-context learning.  arXiv preprint arXiv:2306.00297, 2023.
  * [ACHL19] Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo.  Implicit regularization in deep matrix factorization.  Advances in Neural Information Processing Systems, 32, 2019.
  * [ALH21] Navid Azizan, Sahin Lale, and Babak Hassibi.  Stochastic mirror descent on overparameterized nonlinear models.  IEEE Transactions on Neural Networks and Learning Systems, 33(12):7717–7727, 2021.
  * [ASA+22] Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou.  What learning algorithm is in-context learning? investigations with linear models.  arXiv:2211.15661, 2022.
  * [AW20a] Ehsan Amid and Manfred K Warmuth.  Winnowing with gradient descent.  In Conference on Learning Theory, pages 163–182. PMLR, 2020.
  * [AW20b] Ehsan Amid and Manfred KK Warmuth.  Reparameterizing mirror descent as gradient descent.  Advances in Neural Information Processing Systems, 33:8430–8439, 2020.
  * [AZLS19] Zeyuan Allen-Zhu, Yuanzhi Li, and Zhao Song.  A convergence theory for deep learning via over-parameterization.  In International Conference on Machine Learning, pages 242–252. PMLR, 2019.
  * [BALA+23] Enric Boix-Adsera, Etai Littwin, Emmanuel Abbe, Samy Bengio, and Joshua Susskind.  Transformers learn through gradual rank increase.  arXiv preprint arXiv:2306.07042, 2023.
  * [BCW+23] Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei.  Transformers as statisticians: Provable in-context learning with in-context algorithm selection.  arXiv preprint arXiv:2306.04637, 2023.
  * [BGVV20] Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant.  Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process.  In Conference on learning theory, pages 483–513. PMLR, 2020.
  * [BLLT20] Peter L Bartlett, Philip M Long, Gábor Lugosi, and Alexander Tsigler.  Benign overfitting in linear regression.  Proceedings of the National Academy of Sciences, 117(48):30063–30070, 2020.
  * [BMR+20] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, and et al.  Language models are few-shot learners.  In Advances in neural information processing systems, volume 33, pages 1877–1901, 2020.
  * [BV22] Pierre Baldi and Roman Vershynin.  The quarks of attention.  arXiv preprint arXiv:2202.08371, 2022.
  * [Car21] Marcus Carlsson.  von neumann’s trace inequality for hilbert–schmidt operators.  Expositiones Mathematicae, 39(1):149–157, 2021.
  * [CDL16] Jianpeng Cheng, Li Dong, and Mirella Lapata.  Long short-term memory-networks for machine reading.  In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 551–561, Austin, Texas, November 2016. Association for Computational Linguistics.
  * [CDS01] Scott Shaobing Chen, David L Donoho, and Michael A Saunders.  Atomic decomposition by basis pursuit.  SIAM review, 43(1):129–159, 2001.
  * [CLR+21] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch.  Decision transformer: Reinforcement learning via sequence modeling.  In Advances in Neural Information Processing Systems, volume 34, pages 15084–15097, 2021.
  * [CRT06] Emmanuel J Candès, Justin Romberg, and Terence Tao.  Robust uncertainty principles: Exact signal reconstruction from highly incomplete frequency information.  IEEE Transactions on information theory, 52(2):489–509, 2006.
  * [CSL+23] Yingyi Chen, Xi Shen, Yahui Liu, Qinghua Tao, and Johan AK Suykens.  Jigsaw-vit: Learning jigsaw puzzles in vision transformer.  Pattern Recognition Letters, 166:53–60, 2023.
  * [DCL21] Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas.  Attention is not all you need: Pure attention loses rank doubly exponentially with depth.  In International Conference on Machine Learning, pages 2793–2803. PMLR, 2021.
  * [DML21] Alex Damian, Tengyu Ma, and Jason Lee.  Label noise sgd provably prefers flat global minimizers.  arXiv preprint arXiv:2106.06530, 2021.
  * [Don06] David L Donoho.  Compressed sensing.  IEEE Transactions on information theory, 52(4):1289–1306, 2006\.
  * [DZPS18] Simon S Du, Xiyu Zhai, Barnabas Poczos, and Aarti Singh.  Gradient descent provably optimizes over-parameterized neural networks.  arXiv preprint arXiv:1810.02054, 2018.
  * [EGKZ22] Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang.  Inductive biases and variable creation in self-attention mechanisms.  In International Conference on Machine Learning, pages 5793–5831. PMLR, 2022.
  * [ENM22] Tolga Ergen, Behnam Neyshabur, and Harsh Mehta.  Convexifying transformers: Improving optimization and understanding of transformer networks.  arXiv:2211.11052, 2022.
  * [Faz02] Maryam Fazel.  Matrix rank minimization with applications.  PhD thesis, PhD thesis, Stanford University, 2002.
  * [FXM+21] Haoqi Fan, Bo Xiong, Karttikeya Mangalam, Yanghao Li, Zhicheng Yan, Jitendra Malik, and Christoph Feichtenhofer.  Multiscale vision transformers.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 6824–6835, 2021.
  * [GLSS18] Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro.  Characterizing implicit bias in terms of optimization geometry.  In International Conference on Machine Learning, pages 1832–1841. PMLR, 2018.
  * [Goo23] Google.  Try bard, an ai experiment by google.  https://bard.google.com, 2023.
  * [GRS+23] Angeliki Giannou, Shashank Rajput, Jy-yong Sohn, Kangwook Lee, Jason D Lee, and Dimitris Papailiopoulos.  Looped transformers as programmable computers.  arXiv:2301.13196, 2023.
  * [GWB+17] Suriya Gunasekar, Blake E Woodworth, Srinadh Bhojanapalli, Behnam Neyshabur, and Nati Srebro.  Implicit regularization in matrix factorization.  Advances in neural information processing systems, 30, 2017.
  * [Hah20] Michael Hahn.  Theoretical limitations of self-attention in neural sequence models.  Transactions of the Association for Computational Linguistics, 8:156–171, 2020.
  * [HMX21] Daniel Hsu, Vidya Muthukumar, and Ji Xu.  On the proliferation of support vectors in high dimensions.  In International Conference on Artificial Intelligence and Statistics, pages 91–99. PMLR, 2021.
  * [HWLM20] Jeff Z HaoChen, Colin Wei, Jason D Lee, and Tengyu Ma.  Shape matters: Understanding the implicit bias of the noise covariance.  arXiv preprint arXiv:2006.08680, 2020.
  * [JDST20] Ziwei Ji, Miroslav Dudík, Robert E Schapire, and Matus Telgarsky.  Gradient descent follows the regularization path for general losses.  In Conference on Learning Theory, pages 2109–2136. PMLR, 2020.
  * [JGH18] Arthur Jacot, Franck Gabriel, and Clément Hongler.  Neural tangent kernel: Convergence and generalization in neural networks.  arXiv preprint arXiv:1806.07572, 2018.
  * [JLL21] Michael Janner, Qiyang Li, and Sergey Levine.  Offline reinforcement learning as one big sequence modeling problem.  Advances in neural information processing systems, 34:1273–1286, 2021.
  * [JSL22] Samy Jelassi, Michael Eli Sander, and Yuanzhi Li.  Vision transformers provably learn spatial structure.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022.
  * [JST21] Ziwei Ji, Nathan Srebro, and Matus Telgarsky.  Fast margin maximization via dual acceleration.  In International Conference on Machine Learning, pages 4860–4869. PMLR, 2021.
  * [JT18] Ziwei Ji and Matus Telgarsky.  Risk and parameter convergence of logistic regression.  arXiv preprint arXiv:1803.07300, 2018.
  * [JT19] Ziwei Ji and Matus Telgarsky.  The implicit bias of gradient descent on nonseparable data.  In Conference on Learning Theory, pages 1772–1798. PMLR, 2019.
  * [JT20] Ziwei Ji and Matus Telgarsky.  Directional convergence and alignment in deep learning.  In H. Larochelle, M. Ranzato, R. Hadsell, M. F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 17176–17186. Curran Associates, Inc., 2020.
  * [JT21] Ziwei Ji and Matus Telgarsky.  Characterizing the implicit bias via a primal-dual analysis.  In Algorithmic Learning Theory, pages 772–804. PMLR, 2021.
  * [KPOT21] Ganesh Ramachandra Kini, Orestis Paraskevas, Samet Oymak, and Christos Thrampoulidis.  Label-imbalanced and group-sensitive classification under overparameterization.  Advances in Neural Information Processing Systems, 34:18970–18983, 2021.
  * [KT19] Jacob Devlin Ming-Wei Chang Kenton and Lee Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  In Proceedings of NAACL-HLT, pages 4171–4186, 2019.
  * [LFS+17] Zhouhan Lin, Minwei Feng, Cicero Nogueira dos Santos, Mo Yu, Bing Xiang, Bowen Zhou, and Yoshua Bengio.  A structured self-attentive sentence embedding.  In International Conference on Learning Representations, 2017.
  * [LIPO23] Yingcong Li, M Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak.  Transformers as algorithms: Generalization and stability in in-context learning.  In International Conference on Machine Learning, 2023.
  * [LL18] Yuanzhi Li and Yingyu Liang.  Learning overparameterized neural networks via stochastic gradient descent on structured data.  Advances in neural information processing systems, 31, 2018.
  * [LL19] Kaifeng Lyu and Jian Li.  Gradient descent maximizes the margin of homogeneous neural networks.  arXiv preprint arXiv:1906.05890, 2019.
  * [LLC+21] Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo.  Swin transformer: Hierarchical vision transformer using shifted windows.  In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages 10012–10022, 2021.
  * [LMZ18] Yuanzhi Li, Tengyu Ma, and Hongyang Zhang.  Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations.  In Conference On Learning Theory, pages 2–47. PMLR, 2018.
  * [LR20] TENGYUAN LIANG and ALEXANDER RAKHLIN.  Just interpolate: Kernel “ridgeless” regression can generalize.  The Annals of Statistics, 48(3):1329–1347, 2020.
  * [LWA22] Zhiyuan Li, Tianhao Wang, and Sanjeev Arora.  What happens after SGD reaches zero loss? –a mathematical framework.  In International Conference on Learning Representations, 2022.
  * [LWLC23] Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen.  A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity.  arXiv preprint arXiv:2302.06015, 2023.
  * [LWM19] Yuanzhi Li, Colin Wei, and Tengyu Ma.  Towards explaining the regularization effect of initial large learning rate in training neural networks.  arXiv preprint arXiv:1907.04595, 2019.
  * [MNS+21] Vidya Muthukumar, Adhyyan Narang, Vignesh Subramanian, Mikhail Belkin, Daniel Hsu, and Anant Sahai.  Classification vs regression in overparameterized regimes: Does the loss function matter?  The Journal of Machine Learning Research, 22(1):10104–10172, 2021\.
  * [MRG+20] William Merrill, Vivek Ramanujan, Yoav Goldberg, Roy Schwartz, and Noah Smith.  Effects of parameter norm growth during transformer training: Inductive bias from gradient descent.  arXiv preprint arXiv:2010.09697, 2020.
  * [MWG+20] Edward Moroshko, Blake E Woodworth, Suriya Gunasekar, Jason D Lee, Nati Srebro, and Daniel Soudry.  Implicit bias in deep linear classification: Initialization scale vs training accuracy.  Advances in neural information processing systems, 33:22182–22193, 2020.
  * [NLG+19] Mor Shpigel Nacson, Jason Lee, Suriya Gunasekar, Pedro Henrique Pamplona Savarese, Nathan Srebro, and Daniel Soudry.  Convergence of gradient descent on separable data.  In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3420–3428. PMLR, 2019.
  * [NLL+23] Lorenzo Noci, Chuning Li, Mufan Bill Li, Bobby He, Thomas Hofmann, Chris Maddison, and Daniel M Roy.  The shaped transformer: Attention models in the infinite depth-and-width limit.  arXiv preprint arXiv:2306.17759, 2023.
  * [NNH+23] Tan Minh Nguyen, Tam Minh Nguyen, Nhat Ho, Andrea L Bertozzi, Richard Baraniuk, and Stanley Osher.  A primal-dual framework for transformers and neural networks.  In The Eleventh International Conference on Learning Representations, 2023.
  * [OH10] Samet Oymak and Babak Hassibi.  New null space results and recovery thresholds for matrix rank minimization.  arXiv preprint arXiv:1011.6326, 2010.
  * [OMFH11] Samet Oymak, Karthik Mohan, Maryam Fazel, and Babak Hassibi.  A simplified approach to recovery conditions for low rank matrices.  In 2011 IEEE International Symposium on Information Theory Proceedings, pages 2318–2322. IEEE, 2011.
  * [Ope22] OpenAI.  OpenAI: Introducing ChatGPT.  https://openai.com/blog/chatgpt, 2022.
  * [Ope23] OpenAI.  Gpt-4 technical report.  arXiv preprint arXiv:2303.08774, 2023.
  * [ORST23] Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis.  On the role of attention in prompt-tuning.  In International Conference on Machine Learning, 2023.
  * [OS19] Samet Oymak and Mahdi Soltanolkotabi.  Overparameterized nonlinear learning: Gradient descent takes the shortest path?  In International Conference on Machine Learning, pages 4951–4960. PMLR, 2019.
  * [PHD20] Vardan Papyan, XY Han, and David L Donoho.  Prevalence of neural collapse during the terminal phase of deep learning training.  Proceedings of the National Academy of Sciences, 117(40):24652–24663, 2020.
  * [PTDU16] Ankur Parikh, Oscar Täckström, Dipanjan Das, and Jakob Uszkoreit.  A decomposable attention model for natural language inference.  In Proceedings of the 2016 Conference on Empirical Methods in Natural Language Processing, pages 2249–2255, Austin, Texas, November 2016. Association for Computational Linguistics.
  * [PXS18] Romain Paulus, Caiming Xiong, and Richard Socher.  A deep reinforced model for abstractive summarization.  In International Conference on Learning Representations, 2018.
  * [QQ19] Qian Qian and Xiaoyuan Qian.  The implicit bias of adagrad on separable data.  Advances in Neural Information Processing Systems, 32, 2019.
  * [RFP10] Benjamin Recht, Maryam Fazel, and Pablo A Parrilo.  Guaranteed minimum-rank solutions of linear matrix equations via nuclear norm minimization.  SIAM review, 52(3):471–501, 2010.
  * [RSR+20] Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J Liu.  Exploring the limits of transfer learning with a unified text-to-text transformer.  Journal of Machine Learning Research, 21(1):5485–5551, 2020.
  * [RXH11] Benjamin Recht, Weiyu Xu, and Babak Hassibi.  Null space conditions and thresholds for rank minimization.  Mathematical programming, 127:175–202, 2011.
  * [RZH03] Saharon Rosset, Ji Zhu, and Trevor Hastie.  Margin maximizing loss functions.  Advances in neural information processing systems, 16, 2003.
  * [SATA22] Haoyuan Sun, Kwangjun Ahn, Christos Thrampoulidis, and Navid Azizan.  Mirror descent maximizes generalized margin and can be implemented efficiently.  Advances in Neural Information Processing Systems, 35:31089–31101, 2022.
  * [SEO+22] Arda Sahiner, Tolga Ergen, Batu Ozturkler, John Pauly, Morteza Mardani, and Mert Pilanci.  Unraveling attention via convex duality: Analysis and interpretations of vision transformers.  In International Conference on Machine Learning, pages 19050–19088. PMLR, 2022.
  * [SHN+18] Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro.  The implicit bias of gradient descent on separable data.  The Journal of Machine Learning Research, 19(1):2822–2878, 2018\.
  * [SPR18] Arun Suggala, Adarsh Prasad, and Pradeep K Ravikumar.  Connecting optimization and regularization paths.  Advances in Neural Information Processing Systems, 31, 2018.
  * [SRJ04] Nathan Srebro, Jason Rennie, and Tommi Jaakkola.  Maximum-margin matrix factorization.  Advances in neural information processing systems, 17, 2004.
  * [SS21] Dominik Stöger and Mahdi Soltanolkotabi.  Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparameterized low-rank matrix reconstruction.  Advances in Neural Information Processing Systems, 34:23831–23843, 2021.
  * [TBS+16] Stephen Tu, Ross Boczar, Max Simchowitz, Mahdi Soltanolkotabi, and Ben Recht.  Low-rank solutions of linear matrix equations via procrustes flow.  In International Conference on Machine Learning, pages 964–973. PMLR, 2016.
  * [TCD+21] Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou.  Training data-efficient image transformers & distillation through attention.  In International Conference on Machine Learning, pages 10347–10357. PMLR, 2021.
  * [TG07] Joel A Tropp and Anna C Gilbert.  Signal recovery from random measurements via orthogonal matching pursuit.  IEEE Transactions on information theory, 53(12):4655–4666, 2007\.
  * [Tib96] Robert Tibshirani.  Regression shrinkage and selection via the lasso.  Journal of the Royal Statistical Society Series B: Statistical Methodology, 58(1):267–288, 1996.
  * [TKVB22] Christos Thrampoulidis, Ganesh Ramachandra Kini, Vala Vakilian, and Tina Behnia.  Imbalance trouble: Revisiting neural-collapse geometry.  Advances in Neural Information Processing Systems, 35:27225–27238, 2022.
  * [TLI+23] Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  arXiv preprint arXiv:2302.13971, 2023.
  * [TLZO23] Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak.  Margin maximization in attention mechanism.  arXiv preprint arXiv:2306.13596, 2023.
  * [TVS23] Nadav Timor, Gal Vardi, and Ohad Shamir.  Implicit regularization towards rank minimization in relu networks.  In International Conference on Algorithmic Learning Theory, pages 1429–1459. PMLR, 2023.
  * [TWCD23] Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du.  Scan and snap: Understanding training dynamics and token composition in 1-layer transformer.  arXiv:2305.16380, 2023.
  * [VKR19] Tomas Vaskevicius, Varun Kanade, and Patrick Rebeschini.  Implicit regularization for optimal sparse recovery.  Advances in Neural Information Processing Systems, 32:2972–2983, 2019.
  * [VSP+17] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  Advances in neural information processing systems, 30, 2017.
  * [WGL+20] Blake Woodworth, Suriya Gunasekar, Jason D Lee, Edward Moroshko, Pedro Savarese, Itay Golan, Daniel Soudry, and Nathan Srebro.  Kernel and rich regimes in overparametrized models.  In Conference on Learning Theory, pages 3635–3673. PMLR, 2020.
  * [WMCL21] Bohan Wang, Qi Meng, Wei Chen, and Tie-Yan Liu.  The implicit bias for adaptive optimization algorithms on homogeneous neural networks.  In International Conference on Machine Learning, pages 10849–10858. PMLR, 2021.
  * [WMZ+21] Bohan Wang, Qi Meng, Huishuai Zhang, Ruoyu Sun, Wei Chen, and Zhi-Ming Ma.  Momentum doesn’t change the implicit bias.  arXiv preprint arXiv:2110.03891, 2021.
  * [WT22] Ke Wang and Christos Thrampoulidis.  Binary classification of gaussian mixtures: Abundance of support vectors, benign overfitting, and regularization.  SIAM Journal on Mathematics of Data Science, 4(1):260–284, 2022\.
  * [WWX+22] Haixu Wu, Jialong Wu, Jiehui Xu, Jianmin Wang, and Mingsheng Long.  Flowformer: Linearizing transformers with conservation flows.  In International Conference on Machine Learning, pages 24226–24242, 2022.
  * [YKM20] Chulhee Yun, Shankar Krishnan, and Hossein Mobahi.  A unifying view on implicit bias in training linear neural networks.  arXiv preprint arXiv:2010.02501, 2020.
  * [ZFB23] Ruiqi Zhang, Spencer Frei, and Peter L Bartlett.  Trained transformers learn linear models in-context.  arXiv preprint arXiv:2306.09927, 2023.
  * [ZWB+21] Difan Zou, Jingfeng Wu, Vladimir Braverman, Quanquan Gu, Dean P Foster, and Sham Kakade.  The benefits of implicit regularization from sgd in least squares problems.  Advances in Neural Information Processing Systems, 34:5456–5468, 2021.
  * [ZY05] Tong Zhang and Bin Yu.  Boosting with early stopping: Convergence and consistency.  Annals of Statistics, page 1538, 2005.

##### Roadmap.

The appendix is organized as follows:

  * •

Appendix A provides auxiliary lemmas about the training risk.

  * •

Appendix B presents the proofs for the global convergence of gradient descent
(Section 4).

  * •

Appendix C presents the proofs for the local convergence of gradient descent
(Section 5).

  * •

Appendix D provides a general regularization path analysis. This analysis
addresses the inductive bias of the attention layer for general norm
objectives and beyond-linear prediction heads (see Def E) under a sequence-to-
sequence classification model. The seq2seq aspect also goes beyond our results
in the main body where we predict using single output token (Sections 3 and
5.3).

  * •

Appendix E provides the proof of Theorem 8.

  * •

Appendix F provides additional experiments and their discussion.

###### Contents

  1. 1 Introduction
  2. 2 Preliminaries
    1. 2.1 Optimal tokens and hard-margin SVM problem for cross attention
  3. 3 Understanding Implicit Bias of Self-Attention
  4. 4 Global Convergence of Gradient Descent
    1. 4.1 Properties of optimization landscape
    2. 4.2 Provable global convergence of 1-layer transformer
  5. 5 Understanding Local Convergence of the Attention Layer
    1. 5.1 Local convergence of gradient descent
    2. 5.2 Overparameterization conjecture: When local-optimal directions disappear
    3. 5.3 Guarantees on local regularization path
  6. 6 Understanding Multi-token Compositions: Toward A More General Max-Margin and Directional Convergence Theory
    1. 6.1 When does attention select multiple tokens?
  7. 7 Extending the Theory to Sequential and Causal Predictions
  8. 8 Related work
    1. 8.1 Implicit Regularization, Matrix Factorization, Sparsity
    2. 8.2 Attention Mechanism and Transformers
  9. 9 Discussion, Future Directions, and Open Problems
  10. A Auxiliary Lemmas
    1. A.1 Proof of Lemma 1
    2. A.2 Proof of Lemma 2
    3. A.3 Proof of Lemma 4
    4. A.4 A Useful Lemma for Gradient Descent Analysis
  11. B Global Convergence of Gradient Descent
    1. B.1 Divergence of $\|W(k)\|_{F}$
      1. B.1.1 Proof of Theorem 3
    2. B.2  Global Convergence: Equal Score
      1. B.2.1 Proof of Theorem 4
    3. B.3 Global Convergence: Non-Equal Scores
      1. B.3.1 Proof of Theorem 5
    4. B.4 Global Convergence: Zero Initialization
      1. B.4.1 Proof of Theorem 6
  12. C Local Convergence of Gradient Descent
    1. C.1 Proof of Theorem 7
    2. C.2 Proof of Lemma 6
  13. D Convergence of Regularization Path for Sequence-to-Sequence Setting
    1. D.1 Regularization Path Analysis
    2. D.2 Local Regularization Path and Proof of Theorem 9
    3. D.3 Global Regularization Path and Proof of Theorem 2
  14. E Proof of Theorem 8: Separability Under Mild Over-Parameterization
  15. F Supporting Experiments

##  Appendix A Auxiliary Lemmas

###  A.1 Proof of Lemma 1

Suppose the claim is wrong and row space of $\bm{W}^{\textsl{mm}}_{\diamond}$
does not lie within $\mathcal{S}=\texttt{span}(\\{{\bm{z}}_{i}\\}_{i=1}^{n})$.
Let $\bm{W}=\Pi_{\mathcal{S}}(\bm{W}^{\textsl{mm}}_{\diamond})$ denote the
matrix obtained by projecting the rows of $\bm{W}^{\textsl{mm}}_{\diamond}$ on
$\mathcal{S}$. Observe that $\bm{W}$ satisfies all SVM constraints since
$\bm{W}{\bm{z}}_{i}=\bm{W}^{\textsl{mm}}_{\diamond}{\bm{z}}_{i}$ for all
$i\in[n]$. For Frobenius norm, using
$\bm{W}^{\textsl{mm}}_{\diamond}\neq\bm{W}$, we obtain a contradiction via
$\|{\bm{W}^{\textsl{mm}}_{\diamond}}\|_{F}^{2}=\|{\bm{W}}\|_{F}^{2}+\|{\bm{W}^{\textsl{mm}}_{\diamond}-\bm{W}}\|_{F}^{2}>\|{\bm{W}}\|_{F}^{2}$.
For nuclear norm, we can write $\bm{W}={\bm{U}}{\bm{{\Sigma}}}{\bm{V}}^{\top}$
with ${\bm{{\Sigma}}}\in\mathbb{R}^{r\times r}$ where $r$ is dimension of
$\mathcal{S}$ and $\texttt{column\\_span}({\bm{V}})=\mathcal{S}$.
To proceed, we split the problem into two scenarios.

Scenario 1: Let ${\bm{U}}_{\perp},{\bm{V}}_{\perp}$ be orthogonal complements
of ${\bm{U}},{\bm{V}}$ – viewing matrices with orthonormal columns as
subspaces. Suppose
${\bm{U}}_{\perp}^{\top}\bm{W}^{\textsl{mm}}_{\diamond}{\bm{V}}_{\perp}\neq
0$. Then, singular value inequalities (which were also used in earlier works
on nuclear norm analysis [RXH11, OH10, OMFH11]) guarantee that
$\|{\bm{W}^{\textsl{mm}}_{\diamond}}\|_{\star}\geq\|{{\bm{U}}^{\top}\bm{W}^{\textsl{mm}}_{\diamond}{\bm{V}}}\|_{\star}+\|{{\bm{U}}_{\perp}^{\top}\bm{W}^{\textsl{mm}}_{\diamond}{\bm{V}}_{\perp}}\|_{\star}>\|{\bm{W}}\|_{\star}$.

Scenario 2: Now suppose
${\bm{U}}_{\perp}^{\top}\bm{W}^{\textsl{mm}}_{\diamond}{\bm{V}}_{\perp}=0$.
Since $\bm{W}^{\textsl{mm}}_{\diamond}{\bm{V}}_{\perp}\neq 0$, this implies
${\bm{U}}^{\top}\bm{W}^{\textsl{mm}}_{\diamond}{\bm{V}}_{\perp}\neq 0$. Let
$\bm{W}^{\prime}={\bm{U}}{\bm{U}}^{\top}\bm{W}^{\textsl{mm}}_{\diamond}$ which
is a rank-$r$ matrix. Since $\bm{W}^{\prime}$ is a subspace projection, we
have
$\|{\bm{W}^{\prime}}\|_{\star}\leq\|{\bm{W}^{\textsl{mm}}_{\diamond}}\|_{\star}$.
Next, observe that
$\|{\bm{W}}\|_{\star}=\texttt{trace}({\bm{U}}^{\top}\bm{W}{\bm{V}})=\texttt{trace}({\bm{U}}^{\top}\bm{W}^{\prime}{\bm{V}})$.
On the other hand,
$\texttt{trace}({\bm{U}}^{\top}\bm{W}^{\prime}{\bm{V}})<\|{\bm{W}^{\prime}}\|_{\star}$
because the equality in _von Neumann’s trace inequality_ happens if and only
if the two matrices we are inner-producting, namely
$(\bm{W}^{\prime},{\bm{U}}{\bm{V}}^{\top})$, share a joint set of singular
vectors [Car21]. However, this is not true as the row space of
$\bm{W}^{\textsl{mm}}_{\diamond}$ does not lie within $\mathcal{S}$. Thus, we
obtain
$\|{\bm{W}}\|_{\star}<\|{\bm{W}^{\prime}}\|_{\star}\leq\|{\bm{W}^{\textsl{mm}}_{\diamond}}\|_{\star}$
concluding the proof via contradiction. $\hfill\blacksquare$

###  A.2 Proof of Lemma 2

The token at the output of the attention layer is given by
${\bm{a}}_{i}={\bm{X}}_{i}^{\top}\mathbb{S}({\bm{X}}_{i}\bm{W}{\bm{z}}_{i})$.
Here, ${\bm{a}}_{i}$ can be written as
${\bm{a}}_{i}=\sum_{t\in[T]}c_{it}\bm{x}_{it}$ where $c_{it}\geq 0$ and
$\sum_{t\in[T]}c_{it}=1$. Note that, for any finite $\bm{W}$, $c_{it}$ as
softmax probabilities are strictly positive. To proceed, using the linearity
of $h(\bm{x})=\bm{v}^{\top}\bm{x}$ and strictly-decreasing nature of the loss
$\ell$, we find that

$\displaystyle{\cal{L}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\ell(Y_{i}\cdot
h({\bm{a}}_{i}))=\frac{1}{n}\sum_{i=1}^{n}\ell(Y_{i}\cdot\sum_{t\in[T]}c_{it}h(\bm{x}_{it}))\geq\frac{1}{n}\sum_{i=1}^{n}\ell(Y_{i}\cdot
h(\bm{x}_{i\texttt{opt}_{i}}))={\cal{L}}_{\star},$

which implies that ${\cal{L}}(\bm{W})\geq{\cal{L}}_{\star}$ for any $\bm{W}$.

On the other hand, since not all tokens are optimal, there exists a token
index $(i,t)$ for which $Y_{i}\cdot h(\bm{x}_{it})<Y_{i}\cdot
h(\bm{x}_{i\texttt{opt}_{i}})$. Since all softmax entries obey $c_{it}>0$ for
finite $\bm{W}$, this implies the strict inequality $\ell(Y_{i}\cdot
h({\bm{a}}_{i}))>\ell(Y_{i}\cdot h(\bm{x}_{i\texttt{opt}_{i}}))$. This leads
to the desired conclusion ${\cal{L}}(\bm{W})>{\cal{L}}_{\star}$.

Secondly, suppose (Att-SVM) is feasible i.e. there exists a $\bm{W}$
separating some optimal indices $(\texttt{opt}_{i})_{i=1}^{n}$ from the other
tokens. Note that, this does not exclude the existence of other optimal
indices. This implies that, letting
$\lim_{R\rightarrow\infty}\mathbb{S}({\bm{X}}_{i}(R\cdot\bm{W}){\bm{z}}_{i})$
saturates the softmax and will be equal to the indicator function at
$\texttt{opt}_{i}$ for all inputs $i\in[n]$. Thus, $c_{it}\rightarrow 0$ for
$t\neq\texttt{opt}_{i}$ and $c_{it}\rightarrow 1$ for $t=\texttt{opt}_{i}$.
Using $M_{1}$-Lipschitzness of $\ell$, we can write

$\left|\ell(Y_{i}\cdot h(\bm{x}_{i\texttt{opt}_{i}}))-\ell(Y_{i}\cdot
h({\bm{a}}_{i}))\right|\leq
M_{1}\left|h({\bm{a}}_{i})-h(\bm{x}_{i\texttt{opt}_{i}})\right|.$

Since $h$ is linear, it is $\|{\bm{v}}\|$-Lipschitz implying

$\displaystyle\left|\ell(Y_{i}\cdot
h(\bm{x}_{i\texttt{opt}_{i}}))-\ell(Y_{i}\cdot h({\bm{a}}_{i}))\right|\leq
M_{1}\|{\bm{v}}\|\cdot\|{{\bm{a}}_{i}-\bm{x}_{i\texttt{opt}_{i}}}\|.$

Since ${\bm{a}}_{i}\rightarrow\bm{x}_{i\texttt{opt}_{i}}$ as
$R\rightarrow\infty$, we conclude with the advertised result.
$\hfill\blacksquare$

###  A.3 Proof of Lemma 4

Proof. Let

$\bm{\gamma}_{i}=Y_{i}\cdot{\bm{X}}_{i}\bm{v},\quad\bm{h}_{i}={\bm{X}}_{i}\bm{W}{\bm{z}}_{i}.$
(12)

From Assumption A, $\ell:\mathbb{R}\rightarrow\mathbb{R}$ is differentiable.
Hence, the gradient evaluated at $\bm{W}$ is given by

$\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},$
(13)

where

$\mathbb{S}^{\prime}(\bm{h})=\text{diag}\left(\mathbb{S}(\bm{h})\right)-\mathbb{S}(\bm{h})\mathbb{S}(\bm{h})^{\top}\in\mathbb{R}^{T\times
T}.$ (14)

Note that

$\|\mathbb{S}^{\prime}(\bm{h})\|\leq\|\mathbb{S}^{\prime}(\bm{h})\|_{F}\leq
1.$ (15)

Hence, for any $\bm{W},\dot{\bm{W}}\in\mathbb{R}^{d\times d}$, $i\in[n]$, we
have

$\displaystyle\left\|\mathbb{S}(\bm{h}_{i})-\mathbb{S}(\dot{\bm{h}}_{i})\right\|\leq\left\|\bm{h}_{i}-\dot{\bm{h}}_{i}\right\|\leq\|{\bm{X}}_{i}\|~{}\|{\bm{z}}_{i}\|~{}\left\|\bm{W}-\dot{\bm{W}}\right\|_{F},$
(16a) where $\dot{\bm{h}}_{i}={\bm{X}}_{i}\dot{\bm{W}}{\bm{z}}_{i}$.

Similarly,

$\displaystyle\left\|\mathbb{S}^{\prime}(\bm{h}_{i})-\mathbb{S}^{\prime}(\dot{\bm{h}}_{i})\right\|_{F}$
$\displaystyle\leq\left\|\mathbb{S}(\bm{h}_{i})-\mathbb{S}(\dot{\bm{h}_{i}})\right\|+\left\|\mathbb{S}(\bm{h}_{i})\mathbb{S}(\bm{h}_{i})^{\top}-\mathbb{S}(\dot{\bm{h}_{i}})\mathbb{S}(\dot{\bm{h}_{i}})^{\top}\right\|_{F}$
$\displaystyle\leq
3\|{\bm{X}}_{i}\|~{}\|{\bm{z}}_{i}\|~{}\left\|\bm{W}-\dot{\bm{W}}\right\|_{F}.$
(16b)

Next, for any $\bm{W},\dot{\bm{W}}\in\mathbb{R}^{d\times d}$, we get

$\displaystyle\left\|\nabla\mathcal{L}(\bm{W})-\nabla\mathcal{L}(\dot{\bm{W}})\right\|_{F}$