[B]bvmagenta [P]pdorange

# Implicit Bias and Fast Convergence Rates for Self-attention

Bhavya Vasudeva1† Puneesh Deora2† Christos Thrampoulidis2

1University of Southern California 2University of British Columbia

###### Abstract

Self-attention, the core mechanism of transformers, distinguishes them from
traditional neural networks and drives their outstanding performance. Towards
developing the fundamental optimization principles of self-attention, we
investigate the implicit bias of gradient descent (GD) in training a self-
attention layer with fixed linear decoder in binary classification. Drawing
inspiration from the study of GD in linear logistic regression over separable
data, recent work demonstrates that as the number of iterations $t$ approaches
_infinity_ , the key-query matrix $\bm{W}_{t}$ converges _locally_ (with
respect to the initialization direction) to a hard-margin SVM solution
$\bm{W}_{\text{mm}}$. Our work enhances this result in four aspects. Firstly,
we identify non-trivial data settings for which convergence is provably
_global_ , thus shedding light on the optimization landscape. Secondly, we
provide the first _finite-time_ convergence rate for $\bm{W}_{t}$ to
$\bm{W}_{\text{mm}}$, along with quantifying the rate of sparsification in the
attention map. Thirdly, through an analysis of normalized GD and Polyak step-
size, we demonstrate analytically that _adaptive_ step-size rules can
accelerate the convergence of self-attention. Additionally, we remove the
restriction of prior work on a fixed linear decoder. Our results reinforce the
implicit-bias perspective of self-attention and strengthen its connections to
implicit-bias in linear logistic regression, despite the intricate non-convex
nature of the former.

00footnotetext:  †Equal contribution.
Emails<EMAIL_ADDRESS>{puneeshdeora<EMAIL_ADDRESS>This work is
supported by NSERC Discovery Grant RGPIN-2021-03677 and NSF CAREER Award
CCF-2239265. The authors acknowledge use of the Discovery cluster by USC
Advanced Research Computing.

##  1 Introduction

Self-attention serves as the fundamental building block of transformers,
distinguishing them from traditional neural networks (Vaswani et al.,, 2017)
and driving their outstanding performance across various applications,
including natural language processing and generation (Devlin et al.,, 2019;
Brown et al.,, 2020; Raffel et al.,, 2020), as well as computer vision
(Dosovitskiy et al.,, 2021; Radford et al.,, 2021; Touvron et al.,, 2021).
With transformers establishing themselves as the de-facto deep-learning
architecture, driving advancements in applications seamlessly integrated into
society’s daily life at an unprecedented pace (OpenAI,, 2022), there has been
a surge of recent interest in the mathematical study of the fundamental
optimization and statistical principles of the self-attention mechanism; see
Section 6 on related work for an overview.

In pursuit of this objective, Tarzanagh et al., 2023b ; Tarzanagh et al.,
2023a  have initiated an investigation into the _implicit bias_ of gradient
descent (GD) in training a self-attention layer with fixed linear decoder in a
binary classification task. Concretely, the study paradigm of implicit bias
seeks to characterize structural properties of the weights learned by GD when
the training objective has multiple solutions. The prototypical instance of
this paradigm is GD training of linear logistic regression on separable data:
among infinitely many possible solutions to logistic-loss minimization (each
linear separator defines one such solution), GD learns weights that converge
in direction to the (unique) max-margin class separator (Soudry et al.,, 2018;
Ji and Telgarsky,, 2018). Notably, convergence is global, holding irrespective
of the initial weights’ direction, and comes with explicit rates that
characterize its speed with respect to the number of iterations. Drawing an
analogy to this prototypical instance, when training self-attention with
linear decoder in a binary classification task, Tarzanagh et al., 2023a
defines a hard-margin SVM problem (W-SVM) that separates, with maximal margin,
_optimal_ input tokens from non-optimal ones based on their respective softmax
logits. For this, they show that the key-query weights $\bm{W}_{t}$ found by
GD converge locally to the solution $\bm{W}_{\text{mm}}$ of (W-SVM) as the
number of iterations $t$ grows to infinity.

Figure 1: Comparison of train and test dynamics of various optimizers—SGD,
stochastic normalized GD (SNGD), stochastic Polyak step (SPS), and Adam—while
fine-tuning a pre-trained BERT model on the MNLI dataset; see App. C for
details. SNGD and SPS, employing adaptive step-size rules, demonstrate
significantly faster training, closely resembling the performance of Adam.
Motivated in part by this observation, our work establishes fast convergence
rates for NGD and PS for single-layer self-attention.

Despite the intriguing analogy between the two settings, the findings of
Tarzanagh et al., 2023a  are highly non-trivial not only because the nature of
the max-margin solution differs, but also because of the intricate non-convex
optimization landscape introduced by the presence of self-attention. The non-
convexity, induced by the softmax operator in the self-attention layer,
complicates the analysis and is the reason why the convergence result of
Tarzanagh et al., 2023a  is: (i) _local_ , holding only for an appropriate
initialization direction, and (ii) _asymptotic_ , applicable only as
iterations $t$ approach infinity.

Identifying these limitations, this work is motivated by the following
questions:

  * Q1:

_Are there settings under which GD iterates converge _globally_ to the
solution $\bm{W}_{\text{mm}}$ of (W-SVM)?_

  * Q2:

_Is it possible to obtain _finite-time_ rates of convergence to
$\bm{W}_{\text{mm}}$?_

Additionally, motivated by the practical benefit of using adaptive learning
rates in transformer optimization (see for example Fig. 1), we pose an
additional open question:

  * Q3:

_Can using _adaptive learning rates_ in self-attention optimization accelerate
the convergence to $\bm{W}_{\text{mm}}$?_

Contributions. Our work addresses the above open questions, thus contributing
fundamental insights on the optimization properties of the self-attention
mechanism.

Concretely, we study a single-layer self-attention model
$\Phi({\bm{X}},\bm{\theta})=\bm{u}^{\top}{\bm{X}}^{\top}\bm{\varphi}({\bm{X}}\bm{W}\bm{x}_{1})$,
where, $\bm{\varphi}(\cdot)$ is the softmax nonlinearity,
${\bm{X}}=[\bm{x}_{1},\cdots,\bm{x}_{\mathrm{T}}]$ is the sequence of input
tokens, $\bm{W}$ is the key-query matrix, and $\bm{u}$ is a linear decoder.
Following the setup of Tarzanagh et al., 2023a , for each sequence ${\bm{X}}$,
we associate each token $\bm{x}_{\tau}$ with a score
$\gamma_{\tau}:=y\bm{u}^{\top}\bm{x}_{\tau}$ and let
$\texttt{opt}\in[\mathrm{T}]$ be the index of the token with the largest
score. Given a training set of $n$ sequences ${\bm{X}}_{i}$, this defines a
hard-margin SVM problem with solution $\bm{W}_{\text{mm}}$ that separates the
optimal tokens with maximum margin; see Section 2 for details.

Motivated by question Q2:, we study here the optimization properties of
training $\Phi({\bm{X}};\bm{\theta})$ with exponential loss using _normalized
GD_ , which sets the learning rate $\eta_{t}$ at iteration $t$ adaptively as
$\eta_{t}=\frac{\eta}{\lVert\nabla_{{\bm{\theta}}}\widehat{L}({\bm{\theta}}_{t})\rVert}$,
for some constant $\eta>0$. Our results also extend to Polyak step-size,
another adaptive step-size rule.

Our first set of results (Section 3), for fixed decoder $\bm{u}$ (similar to
Tarzanagh et al., 2023a ), answer questions Q1:-Q2: as follows. In response to
Q1:, we begin by identifying sufficient conditions on the input data under
which GD converges _globally_ , in direction, to $\bm{W}_{\text{mm}}$. Then,
simultaneously addressing Q1: and Q2:, we establish fast finite-time rates of
convergence for normalized GD by proving that the iterates $\bm{W}_{t}$, at
any time $t$, satisfy
$\left\|\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}-\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\|\leq\mathcal{O}\left(t^{-3/4}\right)$.
We identify two key ingredients towards establishing these results. First, we
are able to show that the Euclidean norm $\|\bm{W}_{t}\|$ of the iterates
grows at a rate $\Theta(t)$. Second, and more importantly, we demonstrate that
if the iterate $W_{t}$ at any time $t$ violates the constraints of (W-SVM)
corresponding to any training sample, then the softmax score of the optimal
token for that sample must be non-decreasing. In turn, this enables us to
establish that the iterates, initialized in any direction, remain inside a
cone around $\bm{W}_{mm}$. Our convergence results also imply an explicit rate
at which softmax-attention gets sparsified as softmax scores of optimal tokens
converge to $1$ at an exponential rate $\mathcal{O}(\exp(-\eta t))$.

Our second set of results (Section 4), raises the limitation of fixed linear
decoder of prior work and applies to joint training of $\bm{u}$ and of the
attention weights $\bm{W}$. In response to Q1:, we construct a representative
data model with Gaussian-distributed tokens and prove that GD converges
_globally_ , in direction, to $\bm{W}_{\text{mm}}$. To address Q1:, we show
that for normalized GD with an aggressive step-size, the iterates
$\bm{W}_{t}$, at any time $t$, satisfy
$\left\|\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}-\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\|\leq\mathcal{O}\left(1/\log
t\right)$111Note the rate here is slower compared to the fixed-decoder case.
As detailed in Section 4, this is due to the additional $(t+1)^{-1}$ factor in
the step-size, which results in a slow down of the rate of growth of
$\lVert\bm{W}_{t}\rVert$. Our proof requires this additional factor to account
for the non-smooth objective.. Further, we prove that the linear decoder
$\bm{u}$ converges, in direction, to $\bm{u}_{\text{mm}}$, the solution of the
hard-margin SVM problem (u-SVM) that separates the examples with maximal
margin, using only the _optimal_ input tokens at a $\mathcal{O}(t^{-\eta})$
rate. Finally, to completely characterize the training dynamics, we show fast
train loss convergence at a $\mathcal{O}(\exp(t^{1/3}))$ rate. We highlight
three key technical contributions towards proving these results: (i) a growing
token score gap $\gamma_{\texttt{opt},t}-\gamma_{\tau,t}>0$, between optimal
opt and non-optimal $\tau\neq\texttt{opt}$ tokens, for any time $t>0$; (ii)
non-decreasing softmax scores of the optimal tokens with time; and (iii) a
constant loss ratio across all training sequences valid for any time. These
properties pave way for a PL-like inequality which is crucial to show the loss
convergence rates.

Throughout, we validate our findings on both synthetic and real datasets (see
Section 5). We end with Section 7 discussing the implications of our results
and the open questions they raise.

##  2 Preliminaries

##### Notation.

Lowercase and uppercase bold letters (e.g., $\bm{a}$ and $\bm{A}$) represent
vectors and matrices, respectively. $\lVert\bm{a}\rVert$ and
$\lVert\bm{A}\rVert$ denote Euclidean norms. $\lVert\bm{A}\rVert_{2}$ denotes
the spectral norm and $\lambda_{\min}(\bm{A})$ denotes the minimum eigenvalue.
Also,
$\lVert\bm{A}\rVert_{2,\infty}=\underset{i}{\max}\lVert\bm{a}_{i}\rVert$,
where $\bm{a}_{i}$ is the $i^{\text{th}}$ row of $\bm{A}$. We use
$[\bm{a}_{1},\bm{a}_{2}]=\\{\bm{a}:\bm{a}=\beta\bm{a}_{1}+(1-\beta)\bm{a}_{2},\beta\in[0,1]\\}$
to denote the line segment between $\bm{a}_{1}$ and $\bm{a}_{2}$, and
$\operatorname{concat}(\cdot,\cdot)$ to denote the concatenation operation.
$a\,\wedge\,b$ denotes the minimum of numbers $a$ and $b$. $a\,\vee\,b$
denotes their maximum. We also use standard notation $\mathcal{O}$, $\Omega$
to hide absolute constants, and $\widetilde{\mathcal{O}}$,
$\widetilde{\Omega}$ to hide poly-logarithmic factors. All logarithms are
natural logarithms (base $e$).

##### Single-head Self-attention Model.

The output of a single-head self-attention layer, parameterized by key, query
and value matrices $\bm{W}_{Q},\bm{W}_{K}\in\mathbb{R}^{d\times d_{1}}$ and
$\bm{W}_{V}\in\mathbb{R}^{d\times d_{2}}$, is given by

$\displaystyle\text{ATTN}({\bm{X}};\bm{W}_{Q},\bm{W}_{K},\bm{W}_{V}):=\bm{\varphi}({\bm{X}}\bm{W}_{Q}\bm{W}_{K}^{\top}{\bm{X}}^{\top}){\bm{X}}\bm{W}_{V},$

where ${\bm{X}}\in\mathbb{R}^{\mathrm{T}\times d}$ is an input sequence of
$\mathrm{T}$ tokens with dimension $d$, and
$\bm{\varphi}({\bm{X}}\bm{W}_{Q}\bm{W}_{K}^{\top}{\bm{X}}^{\top})\in\mathbb{R}^{\mathrm{T}\times\mathrm{T}}$
is the attention map, where the softmax map
$\bm{\varphi}(\cdot):\mathbb{R}^{\mathrm{T}}\rightarrow\mathbb{R}^{\mathrm{T}}$
is applied row-wise. We can compose the output of the attention layer with a
linear projection head to obtain the final prediction as

$\displaystyle\Phi({\bm{X}},{\bm{\theta}}):=\bm{u}^{\top}{\bm{X}}^{\top}\bm{\varphi}({\bm{X}}\bm{W}{\bm{x}}_{1}),$
(1)

where ${\bm{\theta}}:=\operatorname{concat}(\bm{u},\bm{W})$ denotes the set of
trainable parameters. Note that here, we reparameterize222Similar model has
been studied in Tarzanagh et al., 2023a ; Deora et al., (2023); Tian et al.,
2023a . the key-query product matrix as
$\bm{W}:=\bm{W}_{Q}\bm{W}_{K}^{\top}\in\mathbb{R}^{d\times d}$, use the first
token for prediction333This is without loss of generality as our results hold
for any token $\tau\in[\mathrm{T}]$. and subsume the value weights
$\bm{W}_{V}$ within the prediction head $\bm{u}\in\mathbb{R}^{d}$.

Let
$\bm{\varphi}^{\prime}({\bm{v}}):=\nabla\bm{\varphi}({\bm{v}})=\operatorname{diag}(\bm{\varphi}({\bm{v}}))-\bm{\varphi}({\bm{v}})\bm{\varphi}({\bm{v}})^{\top}$
denote the Jacobian of the softmax vector $\bm{\varphi}({\bm{v}})$ at
${\bm{v}}\in\mathbb{R}^{\mathrm{T}}$. Let
$\bm{a}(\bm{W}):={\bm{X}}\bm{W}{\bm{x}}_{1}$ denote the vector of softmax
logits, then
$\bm{\varphi}^{\prime}(\bm{a}(\bm{W}))\in\mathbb{R}^{\mathrm{T}\times\mathrm{T}}$
denotes the Jacobian at $\bm{a}(\bm{W})$.

##### Gradient-descent Training.

Given a training dataset $\\{({\bm{X}}_{i},y_{i})\\}_{i=1}^{n}$ with labels
$y_{i}\in\\{-1,1\\}$ and input sequence
${\bm{X}}_{i}:=[{\bm{x}}_{i,1},{\bm{x}}_{i,2},\dots,{\bm{x}}_{i,\mathrm{T}}]^{\top}$,
we seek to minimize the standard empirical risk objective

$\displaystyle\widehat{L}({\bm{\theta}}):=\frac{1}{n}\sum\limits_{i=1}^{n}\ell(y_{i}\Phi({\bm{X}}_{i};{\bm{\theta}})),$

where $\ell:\mathbb{R}\rightarrow\mathbb{R}$ is a decreasing loss function. We
focus on gradient descent (GD) optimization with adaptive time-dependent step-
size $\eta_{t}$. Concretely, we study two variants: normalized gradient
descent (NGD) (Hazan et al.,, 2015; Nacson et al.,, 2019) and Polyak step-size
(Loizou et al.,, 2021). NGD sets
$\eta_{t}=\frac{\eta}{\lVert\nabla_{{\bm{\theta}}}\widehat{L}({\bm{\theta}}_{t})\rVert}$,
for some $\eta>0$ giving

${\bm{\theta}}_{t+1}={\bm{\theta}}_{t}-\eta\frac{\nabla_{{\bm{\theta}}}\widehat{L}({\bm{\theta}}_{t})}{\lVert\nabla_{{\bm{\theta}}}\widehat{L}({\bm{\theta}}_{t})\rVert}.$
(2)

Accordingly, Polyak’s step-size rule sets
$\eta_{t}=\frac{\widehat{L}({\bm{\theta}}_{t})-\widehat{L}^{*}}{2\lVert\nabla_{\bm{\theta}}\widehat{L}({\bm{\theta}}_{t})\rVert^{2}}$
with
$\widehat{L}^{*}=\underset{{\bm{\theta}}}{\min}\,\widehat{L}({\bm{\theta}})$,
giving

$\displaystyle{\bm{\theta}}_{t+1}={\bm{\theta}}_{t}-\frac{\widehat{L}({\bm{\theta}}_{t})-\widehat{L}^{*}}{2\lVert\nabla_{\bm{\theta}}\widehat{L}({\bm{\theta}}_{t})\rVert^{2}}\nabla_{{\bm{\theta}}}\widehat{L}({\bm{\theta}}_{t}).$
(3)

Our results are applicable to both step-size rules. For specificity, we
present them for NGD and include remarks for Polyak step-size.

##### Optimal Tokens and Hard-margin SVM.

Following the setup in Tarzanagh et al., 2023b ; Tarzanagh et al., 2023a , we
define token scores as follows.

######  Definition 1 (Token scores and Optimality).

Given a fixed prediction head $\bm{u}_{*}\in\mathbb{R}^{d}$, the token score
vector for a sample $(y,{\bm{X}})$ is given by
$\bm{\gamma}=y{\bm{X}}\bm{u}_{*}$. The optimal token index444Similar to
Tarzanagh et al., 2023a , we assume unique optimal token which holds for
almost all datasets. is
$\texttt{opt}=\underset{\tau\in[\mathrm{T}]}{\operatorname*{argmax}}\,\gamma_{\tau}$,
where $\gamma_{\tau}=y{\bm{x}}_{\tau}^{\top}\bm{u}_{*}$ denotes the token
score for the token ${\bm{x}}_{\tau}$.

Similarly, $\texttt{opt}_{i}$ denotes the optimal token index for a sample
$(y_{i},{\bm{X}}_{i})$, for any $i\in[n]$. Intuitively, these are the tokens
that minimize the training loss upon selection (Lemma 2 in Tarzanagh et al.,
2023a ). Given a set of optimal token indices
${\texttt{OPT}}:=\\{\texttt{opt}_{i}\\}_{i=1}^{n}$, define the following hard-
margin SVM problem, which aims to separate, with maximal margin, optimal
tokens from the rest of the tokens for every input sequence:

$\displaystyle\bm{W}_{\text{mm}}=\underset{\bm{W}}{\operatorname*{argmin}}~{}\lVert\bm{W}\rVert\quad\text{subj.
to}\quad({\bm{x}}_{i,\texttt{opt}_{i}}-{\bm{x}}_{i,\tau})^{\top}\bm{W}{\bm{x}}_{i,1}\geq
1\,\,\forall i\in[n],\tau\in[\mathrm{T}]\setminus\\{\texttt{opt}_{i}\\}.$
(W-SVM)

Throughout, we assume that (W-SVM) is feasible, i.e. softmax logits
${\bm{x}}_{i,\texttt{opt}_{i}}^{\top}\bm{W}{\bm{x}}_{i,1}$ of optimal tokens
can be separated from softmax logits
${\bm{x}}_{i,\tau}^{\top}\bm{W}{\bm{x}}_{i,1}$ of the rest of the tokens
$\tau\neq\texttt{opt}_{i}$.

##  3 Training Dynamics of $\bm{W}$

This section considers fixed prediction head $\bm{u}=\bm{u}_{*}$. This allows
focusing first on understanding the dynamics of the token selection mechanism
iduced by training key-query weight matrix $\bm{W}$. Later, in Section 4, we
extend the analysis to joint optimization of $\bm{u}$ and $\bm{W}$. We start
by setting up (mild) assumptions on the data and initialization, and then
present our main results.

###  3.1 Setup

Let $\Lambda:=\lVert\bm{W}_{\text{mm}}\rVert$,
$B:=\underset{i}{\max}\lVert{\bm{X}}_{i}\rVert_{2,\infty}$,
$\Upsilon:=\frac{\underset{i}{\max}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})}{\underset{i}{\min}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})}\frac{\underset{i}{\min}\exp(\gamma_{i,\texttt{opt}_{i}})}{\underset{i}{\max}\exp(\gamma_{i})}$.
We only train $\bm{W}$ (thus, ${\bm{\theta}}=\bm{W}$) and using Eq. (2), the
updates at any $t>0$ become

$\bm{W}_{t+1}=\bm{W}_{t}-\eta\frac{\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})}{\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})\rVert}.$
(4)

For the data, we require that the tokens are nearly orthogonal. This is often
the case in high-dimensional settings; see Example 1 and references below.
555Our results also extend to cases where the OPT tokens are antipodal, while
the remaining tokens are nearly orthogonal.

######  Assumption 1 (Nearly-orthogonal Tokens).

For any $i,j,k\in[n]$, and for any $\tau,\tau^{\prime}\in[\mathrm{T}]$,

$\displaystyle\lVert{\bm{x}}_{i,\tau}\rVert^{2}$ $\displaystyle\geq
4n\Upsilon\mathrm{T}|\langle{\bm{x}}_{i,\tau},{\bm{x}}_{j,\tau^{\prime}}\rangle|,\quad
j\neq i,\tau^{\prime}\neq\texttt{opt}_{j},$
$\displaystyle\langle{\bm{x}}_{i,\texttt{opt}_{i}},{\bm{x}}_{j,\texttt{opt}_{j}}\rangle$
$\displaystyle\geq
4n\Upsilon\mathrm{T}|\langle{\bm{x}}_{i,\tau},{\bm{x}}_{k,\tau^{\prime}}\rangle|,\quad
y_{i}=y_{j}\neq y_{k}.$

The following example illustrates the above assumption.

######  Example 1.

Let $\bm{\mu}_{\pm}\in\mathbb{R}^{d}$, $\bm{\mu}_{+}\perp\bm{\mu}_{-}$,
$\lVert\bm{\mu}_{+}\rVert=\lVert\bm{\mu}_{-}\rVert=U$, and data generated as
follows:

$\displaystyle y\sim\operatorname{Unif}(\\{\pm
1\\}),\quad\texttt{opt}\sim\operatorname{Unif}([\mathrm{T}]),\quad\bm{\nu}\sim\mathcal{N}(\mathbf{0},\sigma^{2}\mathbf{\Sigma}),$
$\displaystyle\quad{\bm{x}}_{\texttt{opt}}=\bm{\nu}+\begin{cases}\bm{\mu}_{+}&\text{,\,$y=1$}\\\
\bm{\mu}_{-}&\text{,\,$y=-1$}\\\
\end{cases},\quad{\bm{x}}_{\tau}\sim\mathcal{N}(\mathbf{0},\rho^{2}\mathbf{\Sigma}),\,\forall\tau\in[\mathrm{T}]\setminus\texttt{opt},$

where
$\mathbf{\Sigma}:={\bm{I}}_{d}-U^{-2}\bm{\mu}_{+}\bm{\mu}_{+}^{\top}-U^{-2}\bm{\mu}_{-}\bm{\mu}_{-}^{\top}$.
It is easy to show that when $d=\widetilde{\Omega}(n^{2}\mathrm{T}^{2})$, the
tokens are nearly orthogonal with high probability. Further let high enough
signal-to-noise ratio
$U^{2}\geq\widetilde{\Omega}((\sigma\,\vee\,\rho)^{2}n\mathrm{T}\sqrt{d})$ and
appropriate $\bm{u}_{*}$ such that the token scores satisfy
$\Upsilon=\mathcal{O}(1)$, then, the data model satisfies Assumption 1 (see
Lemma 8 in the Appendix for details). We remark that similar data models have
been considered in prior works on the analysis of linear overparameterized
models (Muthukumar et al.,, 2019; Chatterji and Long,, 2021; Wang and
Thrampoulidis,, 2022; Cao and Gu,, 2019; Wang et al.,, 2021), NNs with Leaky
ReLU activation (Frei et al., 2022a, ; Frei et al.,, 2023), CNNs (Cao et al.,,
2022; Kou et al., 2023b, ), and self-attention models (Deora et al.,, 2023; Li
et al., 2023a, ).

Further, similar to Tarzanagh et al., 2023a , we make the following assumption
on token scores.

######  Assumption 2.

For any $i\in[n]$, and any $\tau\neq\texttt{opt}_{i}\in[\mathrm{T}]$, the
token scores satisfy $\gamma_{i,\tau}=\gamma_{i}$.

Note that by definition, $\gamma_{i,\texttt{opt}_{i}}>\gamma_{i}$. In Example
1, Assumption 2 is satisfied by choosing an appropriate $\bm{u}_{*}$ (see
Lemma 8 in the Appendix for details). Under Assumption 2, it can be shown that
$\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})=\mathbf{0}$ iff the softmax scores are
saturated, i.e., $\varphi_{i,\texttt{opt}_{i}}$ is $0$ or $1$ for every
$i\in[n]$. To prevent initialization at a stationary point, we incorporate an
assumption of a small initialization scale for $\bm{W}_{0}$, as formalized in
Lemma 4 in the Appendix. Importantly, this does _not_ impose any constraints
on the direction of initialization, which is necessary to demonstrate global
convergence.

As mentioned in the Introduction, our goal is to demonstrate that $\bm{W}_{t}$
can globally converge to the SVM solution $\bm{W}_{\text{mm}}$ and,
additionally, establish explicit rates of convergence. It is known, as per
Tarzanagh et al., 2023a , that the optimization landscape of self-attention
can cause GD iterates to converge to local SVM directions (different than
$\bm{W}_{\text{mm}}$). Therefore, achieving global convergence necessitates
specific conditions to be satisfied. Assumptions 1 and 2, as outlined above,
serve this purpose. While being only sufficient conditions, they lead to the
first-known global convergence result with finite-time guarantees, contrasting
with the local and asymptotic convergence in Tarzanagh et al., 2023a .

###  3.2 Main Results

We now present and discuss our main results. For ease of exposition, we use
here exponential loss in our proofs, but our results directly extend to any
continuously differentiable, decreasing loss function.

The following theorem establishes the directional rate of convergence of
$\bm{W}_{t}$ to $\bm{W}_{\text{mm}}$. A proof sketch is provided in Section
3.3.

Figure 2: Training dynamics of a single-head self-attention model (Eq. (1))
when optimizing only $\bm{W}$ on synthetic data with nearly orthogonal tokens
(Example 1 with $\sigma=0$). The observed softmax score saturation, norm
growth of $\bm{W}$ and directional alignment with $\bm{W}_{\text{mm}}$ closely
match with our theoretical results (Lemma 1, Eq. (12) and Theorem 1,
respectively).

######  Theorem 1 (IB Rate, Informal).

Under small initialization scale (Lemma 4 in the Appendix) and Assumptions
1-2, using the NGD updates in Eq. (4) with
$\eta=\widetilde{\mathcal{O}}(B^{-2})$, it holds for any $t\geq
t_{0}=\operatorname{poly}\left(B,\Lambda,\mathrm{T},n\Upsilon\right)$ and for
$C(\eta,B,\Lambda,t_{0})=\operatorname{poly}\left(\eta,B,\Lambda,t_{0}\right)$,
that

$\displaystyle\left\langle\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq
1-\frac{C(\eta,B,\Lambda,t_{0})}{t^{3/4}}.$

In particular, assuming $B=\mathcal{O}(1),T=\mathcal{O}(1)$ and
$\Upsilon=\mathcal{O}(1)$, we have $\Lambda=\mathcal{O}(1)$. Thus, for any
$t\geq t_{0}=\Omega(n)$, it holds

$\displaystyle\left\|\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}-\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\|\leq\mathcal{O}\left(t^{-3/4}\right).$
(5)

Theorem 1 establishes that normalized iterates
$\bm{W}_{t}/\lVert\bm{W}_{t}\rVert$ globally converge to
$\bm{W}_{\text{mm}}/\lVert\bm{W}_{\text{mm}}\rVert$, starting from any
initialization direction. Furthermore, it provides a rate of convergence of
order $\mathcal{O}(t^{-3/4})$. In Fig. 2, we validate the predictions of the
theorem using synthetic data generated as per Example 1 with $\sigma=0$.
Specifically, the figure shows the train loss, the norm growth for iterates
$\bm{W}_{t}$, the softmax score $\varphi_{i,\texttt{opt}_{i}}$ for the opt
token (averaged over $i\in[n]$), and the alignment of $\bm{W}_{t}$ with
$\bm{W}_{\text{mm}}$. Note that since $\bm{u}_{*}$ is fixed, the train loss
does not converge to $0$. Observe that the directional alignment for NGD is
closely predicted by our result in Theorem 1. Moreover, note that NGD is
significantly faster than standard GD.

From the definition of (W-SVM), recall that $\bm{W}_{\text{mm}}$ separates
optimal tokens from the rest. Also, by Theorem 1, $\bm{W}_{t}$ converges
directionally to $\bm{W}_{\text{mm}}$, i.e.
$\bm{W}_{t}\approx\frac{\|\bm{W}_{t}\|}{\|\bm{W}_{\text{mm}}\|}\bm{W}_{\text{mm}}$.
In the proof of the theorem (see Eq. (12)), we show that
$\lVert\bm{W}_{t}\rVert=\Theta(t)$ grows with iterations. Combining these
observations implies that softmax scores of optimal tokens
$\varphi_{i,\texttt{opt}_{i}}^{t}$ approach $1$, while the scores of the rest
of the tokens approach $0$. To see this, note that for
$\bm{W}_{t}\approx\frac{\|\bm{W}_{t}\|}{\|\bm{W}_{\text{mm}}\|}\bm{W}_{\text{mm}}$,
the opt softmax score
$\varphi_{i,\texttt{opt}_{i}}^{t}\geq\frac{1}{1+\sum\limits_{\tau\neq\texttt{opt}_{i}}\exp(-\frac{\|\bm{W}_{t}\|}{\|\bm{W}_{\text{mm}}\|})}$
approaches $1$ as $\lVert\bm{W}_{t}\rVert$ grows. Thus, during training, the
attention map becomes sparse. We formalize these observations in Lemma 1
below, where we also derive an explicit rate describing the sparsification of
the attention map.

######  Lemma 1 (Softmax score rate).

Under small initialization scale and Assumptions 1 and 2, using the updates in
Eq. (4), for any $i\in[n]$ and any
$t\geq\operatorname{poly}\left(B,\Lambda,\mathrm{T}\right)t_{0}^{4/3}$,

$\displaystyle\varphi_{i,\texttt{opt}_{i}}^{t}\geq$
$\displaystyle\,\,\frac{1}{1+(T-1)\exp\left(-\eta(8B^{2}\Lambda^{2})^{-1}t\right)}.$

Note that the saturation rate of the softmax-score derived in Lemma 1 closely
aligns with the empirical rate depicted in Fig. 2. Together, Theorem 1 and
Lemma 1 completely characterize the implicit bias of GD and global convergence
to $\bm{W}_{\text{mm}}$.

######  Remark 1 (Comparison to Tarzanagh et al., 2023a .).

Theorem 1 establishes _finite-time_ convergence of NGD to $\bm{W}_{\text{mm}}$
starting from _any_ initialization direction. This marks a clear improvement
over prior work Tarzanagh et al., 2023a , which demonstrates that Gradient
Descent (GD) converges _asymptotically_ ($t\rightarrow\infty$) to
$\bm{W}_{\text{mm}}$ only under an appropriate _local_ initialization
direction. The only previously known global convergence result in Tarzanagh et
al., 2023b ; Tarzanagh et al., 2023a  requires Assumption 2 and $n=1$,
limiting its applicability to training with a single sample in the training
set. Our theorem addresses this limitation and additionally provides
convergence rates. For completeness, we demonstrate in Theorem 6 in the
Appendix that GD (with constant step size) also converges globally to
$\bm{W}_{\text{mm}}$ under the conditions specified in Theorem 1, without
requiring $n=1$.

##### Extension to Polyak step-size.

In Theorem 7 in the Appendix, we establish an analogue of Theorem 1
specifically for the Polyak step-size rule (Eq. (3)). This extension required
us to further establish a lower bound on the gradient norm. This is because
Polyak step-size has an additional
$\lVert\nabla_{\bm{\theta}}\widehat{L}({\bm{\theta}}_{t})\rVert$ factor in the
denominator compared to NGD. To accomplish this, we applied the variational
form of the Euclidean norm for a “good” direction $\bm{W}_{\text{mm}}$,
yielding
$\lVert\nabla_{\bm{\theta}}\widehat{L}({\bm{\theta}}_{t})\rVert\geq\operatorname{poly}\left(\Lambda^{-1}\right)\widehat{L}({\bm{\theta}}_{t})$.

As seen in Fig. 2, Polyak step-size leads to even faster convergence than NGD
across all metrics, including the train loss. In our implementation, we use a
hyperparameter $\eta_{\max}$ as a threshold on the step-size value, which has
been shown to stabilize training (Loizou et al.,, 2021). Further, since the
selection of opt tokens minimizes the loss, $\widehat{L}^{*}$ is calculated by
setting $\varphi_{i,\texttt{opt}_{i}}=1$ for all $i\in[n]$. Note that after
some time, when $\eta_{t}\geq\eta_{\max}$, the Polyak step updates used in
practice are identical to standard GD updates with step-size $\eta_{\max}$.
Hence, the rates are much faster in the beginning, but slow down as training
progresses.

##### Comparison to Implicit Bias rates for linear logistic regression.

At a high level, the result of Theorem 1 bears resemblance to corresponding
results in the theory of implicit bias for linear logistic regression.
Specifically, Ji and Telgarsky, (2021) demonstrate that NGD on logistic loss
under separable data converges globally, in direction, at rate
$\widetilde{\mathcal{O}}(1/t)$ to the max-margin separator of data points
belonging to distinct classes. Despite the resemblance of the results, it is
important to note that in our setting, the SVM program (W-SVM) separates
optimal tokens from the rest, rather than data based on class membership. More
importantly, the presence of softmax in (1) induces a more challenging non-
convex landscape. As we have seen, in addition to NGD, our convergence results
also apply to Polyak step-size, which empirically outperforms NGD (see Fig.
2). In the same figure, we also present results for NGD with momentum, which
results in a performance boost. In the simpler setting of linear logistic
regression, NGD with momentum is proven to have a faster implicit bias rate of
$\widetilde{\mathcal{O}}(1/t^{2})$ (Ji et al.,, 2021). Extending our theory to
incorporate momentum in the self-attention setting is an interesting direction
for future work.

###  3.3 Proof Sketch of Theorem 1

We first prove the following key result, which states that after a certain
number of iterations, the negative gradient of the loss at $\bm{W}_{t}$ is
more correlated with $\bm{W}_{\text{mm}}$ compared to $\bm{W}_{t}$ itself.

######  Lemma 2.

Under small initialization scale and Assumptions 1 and 2, for any
$\epsilon\in(0,1)$, there exists

$\displaystyle R_{\epsilon}$
$\displaystyle:=2\Lambda\epsilon^{-1}(\log(4n(B^{2}\Lambda)^{-1}\Upsilon\mathrm{T}^{3}\epsilon^{-1})\vee
5\log(20\mathrm{T}B^{2}\Lambda\epsilon^{-1})),$

such that for every $t$ such that $\lVert\bm{W}_{t}\rVert\geq R_{\epsilon}$,
it holds that

$\displaystyle\left\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq(1-\epsilon)\left\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}\right\rangle.$
(6)

###### Proof sketch of Lemma 2.

First, under Assumption 2, for any $\bm{W}$, we have

$\displaystyle\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\bm{W}\rangle$
$\displaystyle=-\frac{1}{n}\sum\limits_{i=1}^{n}\ell^{\prime}_{t,i}\left\langle
y_{i}\nabla_{\bm{W}}\Phi(\bm{W}_{t},{\bm{X}}_{i}),\bm{W}\right\rangle=\frac{1}{n}\sum\limits_{i=1}^{n}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})h_{t,i},$
(7) where $\displaystyle
h_{t,i}=a_{i,\texttt{opt}_{i}}-\frac{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}a_{i,\tau}}{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}},\quad\bm{a}_{i}={\bm{X}}_{i}\bm{W}{\bm{x}}_{i,1}.$

Define $\tilde{h}$ and $h^{*}$ analogously to $h$, using
$\tilde{\bm{a}}=\Lambda{\bm{X}}\overline{\bm{W}}_{t}{\bm{x}}_{1}$ and
$\bm{a}^{*}={\bm{X}}\bm{W}_{\text{mm}}{\bm{x}}_{1}$, respectively. Then, in
order to prove Eq. (6), it suffices that

$\displaystyle\tilde{h}_{t,i}\leq(1+\nu)h^{*}_{t,i},\quad\text{where
$\nu=\epsilon(1-\epsilon)^{-1}$.}$ (8)

Suppose that
$\lVert\overline{\bm{W}}_{t}-\overline{\bm{W}}_{\text{mm}}\rVert>\frac{\nu}{2B^{2}\Lambda}$,
since otherwise the claim follows easily (see App. A.2 for details).

We categorize the set of training samples into three subsets based on the
satisfaction of the constraints in (W-SVM): $\mathcal{I}_{1}$ for samples
where the constraints are violated or barely satisfied, $\mathcal{I}_{2}$ for
samples where the constraints are satisfied with a small gap, and
$\mathcal{I}_{3}$ for samples where the constraints are well satisfied.
Concretely, let

$\displaystyle\mathcal{I}_{1}:=\\{i\in[n]\,\,|\,\,\delta^{\min}_{i}\leq
0.3\nu\\},\quad\mathcal{I}_{2}:=\\{i\in[n]\,\,|\,\,0.3\nu<\delta^{\min}_{i}\leq
0.8\nu\\},\quad\mathcal{I}_{3}:=\\{i\in[n]\,\,|\,\,\delta^{\min}_{i}\geq
0.8\nu\\},$

where for any $i\in[n]$,
$\delta^{\min}_{i}:=\min_{\tau\notin\\{\texttt{opt}_{i}\\}}\tilde{a}_{i,\texttt{opt}_{i}}-\tilde{a}_{i,\tau}-1$.

Since $\lVert\overline{\bm{W}}_{t}-\overline{\bm{W}}_{\text{mm}}\rVert$ is
large, there is at least one point $i\in[n]$ that violates the constraints,
i.e. $\mathcal{I}_{1}$ is non-empty. We first show the following for examples
in $\mathcal{I}_{1}$ and $\mathcal{I}_{2}$

$\displaystyle\tilde{h}_{t,i}\leq(1+0.5\nu)h^{*}_{t,i},\quad\text{for
any}\,\,i\in\mathcal{I}_{1},$ (9)
$\displaystyle\tilde{h}_{t,i}\leq(1+\nu)h^{*}_{t,i},\quad\text{for
any}\,\,i\in\mathcal{I}_{2}.$ (10)

Note that we prove a tighter inequality for $\mathcal{I}_{1}$ than required by
(8). We do so, in order to use this residual to get the desired inequality for
examples in $\mathcal{I}_{3}$. Specifically, when $R_{\epsilon}$ is large
enough, we wish to show (following Eq. (7)) that

$\displaystyle\sum\limits_{i\in\mathcal{I}_{3}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})\leq
0.5\nu(2B^{2}\Lambda)^{-1}\sum\limits_{i\in\mathcal{I}_{1}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})$
(11)

Intuitively, for the LHS, any sample $i\in\mathcal{I}_{3}$ satisfies (W-SVM)
constraints with large gap giving $\varphi_{i,\texttt{opt}_{i}}$ close to 1
and growing with $\lVert\bm{W}_{t}\rVert$. Similarly, for the RHS, we expect
$\varphi_{i,\texttt{opt}_{i}}$ to be smaller. Thus, for a sufficiently large
$\lVert\bm{W}_{t}\rVert$ and controlling $\varphi_{i,\texttt{opt}_{i}}$ for
any $i\in\mathcal{I}_{1}$ and $t>0$ would allow us to show the above
inequality.

This is achieved by establishing a lower bound on the softmax scores for
$i\in\mathcal{I}_{1}$ that holds throughout training (Lemma 5 in the
Appendix). Finally, since $h^{*}\geq 1$ and $\tilde{h}\leq 2B^{2}\Lambda$, we
combine (9), (10), (11) using the formulation in (7) and complete the proof
for Eq. (6). ∎

######  Remark 2.

Lemma 2 above improves Lemma 10 of Tarzanagh et al., 2023a , which only holds
for $n=1$, i.e., for a single training sample (same restriction appears also
in Tarzanagh et al., 2023b ). Our key idea to extend the result to $n\geq 1$
is to divide the set of $n$ samples into three sets, based on whether the
constraints in (W-SVM) are violated ($\mathcal{I}_{1}$), satisfied with a
small margin ($\mathcal{I}_{2}$) or well satisfied ($\mathcal{I}_{3}$). We
prove a tighter version for the samples in $\mathcal{I}_{1}$ and the residual
gap allows us to prove the result for the desired sum over all samples.

The next key ingredient in the proof of Theorem 1 is showing that the iterate
norm grows at rate $\lVert\bm{W}_{t}\rVert$ is $\Theta(t)$. Specifically, we
prove in Lemma 7 in the Appendix that

$\displaystyle 2\eta
t\geq\lVert\bm{W}_{t}\rVert\geq\eta(4B^{2}\Lambda)^{-1}t,$ (12)

The upper bound follows by applying Cauchy-Schwarz. To obtain the lower bound,
we first write the iterate $\bm{W}_{t}$ using its gradient updates as

$\displaystyle\lVert\bm{W}_{t}\rVert$
$\displaystyle=\left\|\bm{W}_{0}-\sum\limits_{t^{\prime}=0}^{t-1}\eta_{t^{\prime}}\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})\right\|.$

Then, we use the dual norm characterization of the $\ell_{2}$-norm, and select
the $\bm{W}_{\text{mm}}$ direction to get

$\displaystyle\lVert\bm{W}_{t}\rVert\geq\left\langle\bm{W}_{0},\overline{\bm{W}}_{\text{mm}}\right\rangle+\sum\limits_{t^{\prime}=0}^{t-1}\eta\left\langle-\frac{\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})}{\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})\rVert},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle,$

where we use the key property that for any $\bm{W}$ the gradient correlates
positively with $\bm{W}_{mm}$, i.e.

$\displaystyle\left\langle-\frac{\nabla_{\bm{W}}\widehat{L}(\bm{W})}{\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W})\rVert},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq(2B^{2}\Lambda)^{-1}>0,$

which is shown in Lemma 6 in the Appendix.

Next, in order to use Lemma 2 to obtain the convergence rate, we first use Eq.
(12) to show that when $t\geq t_{0}$, we can select
$\epsilon_{t}=(t_{0}t)^{-3/4}$ (and get corresponding $R_{\epsilon}$ in Lemma
2). Next, we use Eq. (4) to substitute the gradient update in Eq. (6) and get

$\displaystyle\left\langle\bm{W}_{t+1}-\bm{W}_{t},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq(1-\epsilon_{t})\left\langle\bm{W}_{t+1}-\bm{W}_{t},\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}\right\rangle$
$\displaystyle=\frac{1}{2\lVert\bm{W}_{t}\rVert}(\lVert\bm{W}_{t+1}\rVert^{2}-\lVert\bm{W}_{t}\rVert^{2}-\lVert\bm{W}_{t+1}-\bm{W}_{t}\rVert^{2})-\epsilon_{t}\left\langle\bm{W}_{t+1}-\bm{W}_{t},\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}\right\rangle.$

Further, using Eq. (4), we show that

$\displaystyle\left\langle\bm{W}_{t+1}-\bm{W}_{t},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle$
$\displaystyle\geq\lVert\bm{W}_{t+1}\rVert-\lVert\bm{W}_{t}\rVert-\frac{\eta^{2}}{2\lVert\bm{W}_{t}\rVert}-\epsilon_{t}\eta$
$\displaystyle\geq\lVert\bm{W}_{t+1}\rVert-\lVert\bm{W}_{t}\rVert-\eta(t_{0}t^{-1})^{3/4}-2\eta
B^{2}\Lambda(t+1)^{-1},$

where we substitute $\epsilon_{t}$ and use Eq. (12) in the second inequality.
Then, telescoping and dividing by $\lVert\bm{W}_{t}\rVert$ throughout, we get
a $\mathcal{O}(t^{1/4}\lVert\bm{W}_{t}\rVert^{-1})$ rate. Using Eq. (12) then
finishes the proof.

##  4 Training Dynamics for Joint Optimization

This section studies training dynamics when jointly optimizing prediction head
$\bm{u}$ and attention weights $\bm{W}$. Compared to Section 3, the key
challenge is that the token scores $\bm{\gamma}^{t}$ evolve with time, driven
by the changing nature of $\bm{u}_{t}$. Additionally, the objective function
becomes non-smooth in this context, given the dynamic changes in $\bm{u}_{t}$.
Addressing these challenges necessitates additional technical considerations
compared to Section 3, which we also discuss in this section.

###  4.1 Setup

We consider the following updates for $\bm{u}_{t}$ and $\bm{W}_{t}$.

$\displaystyle\bm{u}_{t+1}$
$\displaystyle=\bm{u}_{t}-\eta_{t}^{u}\nabla_{\bm{u}}\widehat{L}({\bm{\theta}}_{t}),$
$\displaystyle\vspace{-18mm}\text{ where
}\eta_{t}^{u}=\frac{\eta(t+1)^{-2/3}}{\lVert\nabla_{\bm{u}}\widehat{L}({\bm{\theta}}_{t})\rVert},$
$\displaystyle\bm{W}_{t+1}$
$\displaystyle=\bm{W}_{t}-\eta_{t}^{W}\nabla_{\bm{W}}\widehat{L}({\bm{\theta}}_{t}),$
$\displaystyle\vspace{-18mm}\text{ where
}\eta_{t}^{W}=\frac{\eta(t+1)^{-1}}{\lVert\nabla_{\bm{W}}\widehat{L}({\bm{\theta}}_{t})\rVert}.$
(13)

Here, we focus on exponential loss, for which
$|\ell^{\prime}(\cdot)|=\ell(\cdot)$. Note that for $\eta_{t}^{u}$, the factor
$(t+1)^{-1}$ can be replaced with $(t+1)^{-p}$, where $p\in[2/3,1)$. In our
results, we present findings using $p=2/3$, as it yields the fastest rate. We
also remark that common adaptive learning rates like Adam (Kingma and Ba,,
2014), AdaGrad (Duchi et al.,, 2011), RMSProp, etc. similarly set per-
parameter learning rates that change dynamically based on the past and current
gradients.

Similar to Def. 1, define token scores with trainable $\bm{u}_{t}$ at any
iteration $t$ as follows. Given a prediction head $\bm{u}_{t}$, the token
score vector for a sample $(y,{\bm{X}})$ is given by
$\bm{\gamma}^{t}=y{\bm{X}}\bm{u}_{t}$.

##### Data Model.

In this section, we consider the following data model. Within each example
${\bm{X}}$, a single $\texttt{opt}\sim\operatorname{Unif}([\mathrm{T}])$ token
and an additional $T-1$ token are sampled i.i.d. Gaussian as follows:

$\displaystyle\quad{\bm{x}}_{\texttt{opt}}\sim\,\mathcal{N}({\bm{0}},\alpha^{2}\rho^{2}{\bm{I}}_{d}),\quad$
$\displaystyle{\bm{x}}_{\tau}\sim\mathcal{N}({\bm{0}},\rho^{2}{\bm{I}}_{d}),\,\,\forall\,\tau\in[\mathrm{T}]\setminus\\{\texttt{opt}\\},$
(DM) $\displaystyle y=$
$\displaystyle\,\texttt{sign}(\bm{u}_{*}^{\top}{\bm{x}}_{\texttt{opt}}),$

where $\bm{u}_{*}\in\mathbb{R}^{d}$ is some fixed vector used to generate the
label $y\in\\{+1,-1\\}$. Here, $\rho$ controls the norm of the tokens, and
$\alpha>1$ allows us to separate the opt token from the other tokens. We
consider the following mild conditions on the initialization, token norms,
overparameterization and step-sizes.

######  Condition 1.

Let $B=\alpha\rho\sqrt{1.5d}\geq 1$, and $C_{0}>1$, $C_{1},C_{2}>0$ be some
absolute constants. We consider the following conditions for any
$\delta\in(0,1)$,

  * •

Zero initialization: $\lVert\bm{u}_{0}\rVert=0,\lVert\bm{W}_{0}\rVert=0$,

  * •

opt token has larger norm: $\alpha\geq C_{0}$,

  * •

Sufficient overparameterization: $d\geq
C_{1}\alpha^{4}n^{4}\log\left(\frac{10n^{2}}{\delta}\right)$,

  * •

Small step-size: $\eta\leq(C_{2}\sqrt{n}B^{4}(B\vee d))^{-1}$.

The condition on the initialization is for ease of exposition; our analysis
and results extend to random initialization with small scale (similar to the
conditions of Lemma 4 when optimizing only $\bm{W}$). We also remark that
under Condition 1 and data model DM, the tokens are nearly orthogonal, similar
to Assumption 1 in Section 3. We formalize this in Lemma 11 in the Appendix.

###  4.2 Main Results

We now present our main results for joint training. For ease of exposition, we
consider $\mathrm{T}=2$, but note that our results directly extend to
$\mathrm{T}>2$ assuming token scores $\bm{\gamma}^{t}$ satisfy Assumption 2 at
every iteration $t$. We note that Assumption 2 is directly satisfied when
$\mathrm{T}=2$. Note that it is also possible to extend our results to the
case where Assumption 2 is not exactly satisfied, but the token scores for the
non-opt tokens show small deviation.

Our first result shows that the train loss converges to $0$ at an
$\mathcal{O}(\exp(-t^{1/3}))$ rate. See Section 4.3 for a proof sketch.

######  Theorem 2 (Train loss convergence).

Under Condition 1, data model DM, using the updates in Eq. (13), it holds for
any $t>0$ that

$\displaystyle\widehat{L}({\bm{\theta}}_{t+1})\leq\mathcal{O}\left(\exp\left(-C_{0}\frac{\eta
B}{\sqrt{n}}(t+1)^{1/3}\right)\right),$

where $B=\alpha\rho\sqrt{1.5d}$ and $C_{0}>0$ is an absolute constant.

The next theorem illustrates that $\bm{W}_{t}$ converges to
$\bm{W}_{\text{mm}}$ in direction at an $\mathcal{O}(1/\log t)$ rate.

######  Theorem 3 (IB Rate under Joint Training).

Under Condition 1 and the data model DM, using the updates in Eq. (13), for
any $t\geq
t_{\epsilon}=\exp(\operatorname{poly}\left(\eta,B,\Lambda,n,d,\log\delta^{-1},\epsilon^{-1}\right))$,

$\displaystyle\left\langle\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq
1-\epsilon-\frac{C(\eta,B,\Lambda,\epsilon)}{\log t},$

where $B=\alpha\rho\sqrt{1.5d}$ and
$C(\eta,B,\Lambda,\epsilon)=\operatorname{poly}\left(\eta,B,\Lambda,\epsilon\right)$.

Assuming $\rho=\mathcal{O}(\sqrt{\eta}(n\Lambda)^{-1}d^{-3/4})$, for any
$t\geq t_{\epsilon}=\Omega(\exp(\epsilon^{-4/3}))$,

$\displaystyle\left\|\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}-\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\|\leq\mathcal{O}\left(\epsilon\,\vee\,(\log
t)^{-1}\right).$

Intuitively, since the updates for $\bm{W}_{t}$ have an additional
$(t+1)^{-1}$ factor due to the smaller step-size $\eta_{t}$, the iterate norm
$\lVert\bm{W}_{t}\rVert$ grows as $\Theta(\log t)$ in contrast to the
$\Theta(t)$ rate in the previous case, where we were only optimizing
$\bm{W}_{t}$ (Section 3). Consequently, the convergence to implicit bias is
much slower — $\mathcal{O}(1/\log t)$ instead of $\mathcal{O}(t^{-3/4})$. We
note that the additional $(t+1)^{-1}$ factor in the updates for $\bm{W}_{t}$
helps to account for the non-smoothness of the objective, which will become
apparent in Section 4.3.

The proof of Theorem 3 is quite similar to the proof of Theorem 1, so we defer
the details to the Appendix. The key challenge is that the token scores are
not fixed and evolve as training progress. As discussed in Section 4.3, we
show a key property that the gap between the opt token score and the non-opt
token grows with time, for every $i\in[n]$.

We briefly discuss the implicit bias and convergence of $\bm{u}_{t}$ below.
From previous results (Soudry et al.,, 2018; Ji and Telgarsky,, 2019), one can
expect $\bm{u}_{t}$ to converge to the max-margin predictor separating the set
of samples
$\\{(\varphi_{i,\texttt{opt}_{t}}^{t}{\bm{x}}_{i,\texttt{opt}_{i}}+(1-\varphi_{i,\texttt{opt}_{i}}^{t}){\bm{x}}_{i,\tau},y_{i})_{i=1}^{n}\\}$.
As shown above, when $t\rightarrow\infty$, $\varphi_{\texttt{opt}}\rightarrow
1$. Consequently, we consider the following hard-margin SVM problem,

$\displaystyle\bm{u}_{\text{mm}}=\underset{\bm{u}}{\operatorname*{argmin}}\lVert\bm{u}\rVert\quad\text{subject
to}\quad y_{i}{\bm{x}}_{i,\texttt{opt}_{i}}^{\top}\bm{u}\geq 1\,\,\forall
i\in[n].$ (u-SVM)

The following result confirms that the learned model attains margin that is
asymptotically a constant factor ($1/4$) of the maximum margin
$\overline{\gamma}\\!:=\\!\underset{\bm{u}:\lVert\bm{u}\rVert\leq
1}{\max}\min_{i}y_{i}\bm{u}^{\top}{\bm{x}}_{i,\texttt{opt}_{i}}$.

######  Theorem 4 (IB Rate of $\bm{u}$).

Under Condition 1 and the data model DM, using the updates in Eq. (13), for
any $t\geq
t_{\epsilon}\,\vee\,\exp(C(\eta,B,\Lambda,\epsilon)(\epsilon^{-1}\,\vee\,(8B^{2}\Lambda)^{4}))$,

$\displaystyle\min_{i}y_{i}\overline{\bm{u}_{t}}^{\top}{\bm{x}}_{i,\texttt{opt}_{i}}\geq\frac{\overline{\gamma}}{4}-\frac{1}{1+\exp(\eta(8B^{2}\Lambda^{2})^{-1}\log
t)}.$

######  Remark 3.

(Comparison to linear logistic regression) The loss convergence rate in Thm. 2
is analogous to the $\mathcal{O}(\exp(-t))$ loss rate for linear predictors on
separable data with
$\eta_{t}=\eta/\lVert\nabla\widehat{L}(\bm{\theta}_{t})\rVert$ (Ji and
Telgarsky,, 2019). The difference $\mathcal{O}(\exp(-t^{1/3}))$ vs.
$\mathcal{O}(\exp(-t))$ shows up as we use smaller step-size. Nacson et al.,
(2019) similarly show a slower rate of $\mathcal{O}(\exp(-t^{1/2}))$ for
linear predictors with
$\eta_{t}=\frac{\eta(t+1)^{-1/2}}{\lVert\nabla\widehat{L}(\bm{\theta}_{t})\rVert}$.
Contrary to these works, our results hold in a non-convex setting unlike that
of linear logistic regression. Further, the implicit bias rate in Theorem 3 is
comparable to linear models trained with standard GD on separable data (Soudry
et al.,, 2018).

###  4.3 Proof Sketch of Theorem 2

The first challenge in establishing rate of loss convergence is the non-
smoothness of the objective function. To overcome this challenges, we show
three key properties that hold throughout the training trajectory: i) constant
loss ratio for any two samples, ii) growing token score gap between opt and
non-opt tokens for every sample, and iii) a non-decreasing opt-token softmax
score. These are formalized in Lem. 13 in the App. Property (iii) is crucial
in obtaining a PL-like inequality which is the key challenging step in proving
the loss convergence rate in Thm. 2. Further, property (ii) is crucial in
proving the implicit bias rate in Thm. 3.

The following result shows some properties of the loss function which are
useful to prove the train loss convergence rate, as we will see later.

######  Lemma 3.

Under Condition 1 and the data model DM, using the updates in Eq. (13), the
exponential loss has the following properties:

  * •

$\lVert\nabla_{\bm{u}}\widehat{L}({\bm{\theta}}_{t})\rVert\geq\omega_{1}\widehat{L}({\bm{\theta}}_{t})$,

  * •

$\underset{{\bm{\theta}}^{\prime}\in[{\bm{\theta}}_{t},{\bm{\theta}}_{t+1}]}{\max}\widehat{L}({\bm{\theta}}^{\prime})\leq
8\widehat{L}({\bm{\theta}}_{t})$,

  * •

$\underset{{\bm{\theta}}^{\prime}\in[{\bm{\theta}}_{t},{\bm{\theta}}_{t+1}]}{\max}\lVert\nabla^{2}_{{\bm{\theta}}}\widehat{L}({\bm{\theta}}^{\prime})\rVert\leq
8(\omega({\bm{\theta}}_{t})\vee\omega({\bm{\theta}}_{t+1}))\widehat{L}({\bm{\theta}}_{t})$,

where $\omega_{1}=\omega_{1}(B,n)=\frac{B}{8\sqrt{1.5n}}$,
$\omega({\bm{\theta}}_{t}):=13(B\vee 1)^{5}(B\vee
d)(\lVert\bm{u}_{t}\rVert\vee 1)^{2}$, $B=\alpha\rho\sqrt{1.5d}$.

###### Proof sketch of Lemma 3.

To prove the first bullet point, we use the dual norm characterization of the
$\ell_{2}$-norm,

$\displaystyle\lVert\nabla_{\bm{u}}\widehat{L}({\bm{\theta}}_{t})\rVert$
$\displaystyle=\sup_{{\bm{v}}:\lVert{\bm{v}}\rVert=1}\left\langle\frac{1}{n}\sum\limits_{i=1}^{n}|\ell^{\prime}_{i,t}|y_{i}\nabla_{\bm{u}}\Phi(\bm{u}_{t},{\bm{X}}_{i}),{\bm{v}}\right\rangle$
$\displaystyle\geq\widehat{L}({\bm{\theta}}_{t})\,\sup_{{\bm{v}}:\lVert{\bm{v}}\rVert=1}\min_{i}y_{i}(\varphi_{i,\texttt{opt}_{i}}^{t}{\bm{x}}_{i,\texttt{opt}_{i}}+(1-\varphi_{i,\texttt{opt}_{i}}^{t}){\bm{x}}_{i,\tau})^{\top}{\bm{v}}.$

To proceed, we lower bound the softmax scores for the opt tokens for every
$i\in[n]$ and $t>0$ as

$\displaystyle\varphi_{i,\texttt{opt}_{i}}^{t}\geq\frac{1}{2}.$ (14)

Using this, we show that the set of samples
$\\{(0.5({\bm{x}}_{i,\texttt{opt}_{i}}+{\bm{x}}_{i,\tau}),y_{i})\\}_{i=1}^{n}$,
where $\tau\neq\texttt{opt}_{i}$, are separable using
$\tilde{\bm{u}}=\sum\limits_{i\in[n]}y_{i}{\bm{x}}_{i,\texttt{opt}_{i}}$
(formalized in Lemma 14 in the Appendix). Then, selecting ${\bm{v}}$ as
$\tilde{\bm{u}}$ finishes the proof for the first bullet point.

We now briefly discuss the process to prove Eq. (14). We can show this by
induction. Since $\bm{W}_{0}=\mathbf{0}$, then for any $i\in[n]$,
$\varphi_{i,\texttt{opt}_{i}}^{0}=1/2$ at initialization. To show the
inductive step, note that for any $i\in[n]$, if

$\displaystyle({\bm{x}}_{i,\texttt{opt}_{i}}-{\bm{x}}_{i,\tau})^{\top}(-\nabla_{\bm{W}}\widehat{L}({\bm{\theta}}_{t})){\bm{x}}_{i,1}>0,$
(15)
$\displaystyle\implies\frac{\varphi_{j,\texttt{opt}_{j}}^{t+1}/\varphi_{j,\tau}^{t+1}}{\varphi_{j,\texttt{opt}_{j}}^{t}/\varphi_{j,\tau}^{t}}$
$\displaystyle=\exp\left(({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}(\bm{W}_{t+1}-\bm{W}_{t}){\bm{x}}_{j,1}\right)\geq
1.$

We now consider two cases. If $\varphi_{i,\texttt{opt}_{i}}^{t}\geq 3/4$, Eq.
(15) follows as $\eta$ is small. Specifically, for any $j\in[n]$,
$\tau\neq\texttt{opt}_{j}$, we have

$\displaystyle\varphi_{j,\texttt{opt}_{j}}^{t}=\frac{1}{1+\exp(({\bm{x}}_{i,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t}{\bm{x}}_{j,1})}>3/4$
$\displaystyle\implies\quad$
$\displaystyle({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t-1}{\bm{x}}_{j,1}\leq-\log(3).$

After the gradient step, we have

$\displaystyle({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t+1}{\bm{x}}_{j,1}$
$\displaystyle=({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t}{\bm{x}}_{j,1}-\eta_{t-1}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}(-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})){\bm{x}}_{j,1}$
$\displaystyle\leq-\log(3)+2\eta B^{2}\leq 0,$

since $\eta\leq(2B^{2})^{-1}\log(3)$.

If $\varphi_{i,\texttt{opt}_{i}}^{t}\leq 3/4$, we proceed as follows. Consider
the expansion of the LHS of Eq. (15),

$\displaystyle({\bm{x}}_{i,\texttt{opt}_{i}}-{\bm{x}}_{i,\tau})^{\top}($
$\displaystyle-\nabla_{\bm{W}}\widehat{L}({\bm{\theta}}_{t})){\bm{x}}_{i,1}=\frac{\ell_{t,i}}{n}(\gamma^{t}_{i,\texttt{opt}_{i}}-\gamma^{t}_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})\lVert{\bm{x}}_{i,1}\rVert^{2}\lVert{\bm{x}}_{i,\texttt{opt}_{i}}-{\bm{x}}_{i,\tau}\rVert^{2}$
$\displaystyle-\frac{1}{n}\sum\limits_{j\neq
i}\ell_{t,j}(\gamma_{j,\texttt{opt}_{j}}^{t}-\gamma_{j}^{t})\varphi_{j,\texttt{opt}_{j}}^{t}(1-\varphi_{j,\texttt{opt}_{j}}^{t}){\bm{x}}_{i,1}^{\top}{\bm{x}}_{j,1}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}({\bm{x}}_{i,\texttt{opt}_{i}}-{\bm{x}}_{i,\tau^{\prime}}).$

In order to use this to prove Eq. (15), we first show that the token scores
$\bm{\gamma}^{t}$ satisfy Defn. 1
($\gamma_{\texttt{opt}_{i}}^{t}>\gamma_{i}^{t}$ for any $i\in[n]$), and that
the loss ratio $\max_{i,j\in[n]}\frac{\ell_{t,i}}{\ell_{t,j}}$ is bounded by a
constant throughout training, i.e., for any $i\in[n]$,

$\displaystyle\gamma^{t}_{i,\texttt{opt}_{i}}-\gamma^{t}_{i}\geq\Omega(t^{1/3})\quad\text{and}\quad\max_{i,j\in[n]}\frac{\ell_{t,i}}{\ell_{t,j}}\leq
C,$ (16)

where $C>0$ is a universal constant. This is formalized in Lemma 13 in the
Appendix. We then prove Eqs. (14) and (16) jointly by induction.

For the second bullet point, we first show that for any
${\bm{\theta}},{\bm{\theta}}^{\prime}$,

$\displaystyle|y\Phi({\bm{\theta}},{\bm{X}})-y\Phi({\bm{\theta}}^{\prime},{\bm{X}})|$
$\displaystyle=|\bm{u}^{\top}{\bm{X}}^{\top}\bm{\varphi}({\bm{X}}\bm{W}^{\top}{\bm{x}}_{1})-\bm{u}^{\prime\top}{\bm{X}}^{\top}\bm{\varphi}({\bm{X}}\bm{W}^{\prime\top}{\bm{x}}_{1})|$
$\displaystyle\leq\lVert{\bm{X}}\rVert_{2,\infty}\lVert\bm{u}-\bm{u}^{\prime}\rVert+2\lVert{\bm{X}}\rVert^{3}_{2,\infty}\lVert\bm{u}^{\prime}\rVert\lVert\bm{W}-\bm{W}^{\prime}\rVert.$

Since we are using exponential loss, we use this to get,

$\displaystyle\max_{\lambda\in[0,1]}$
$\displaystyle\frac{\ell(y\Phi({\bm{\theta}}_{t}+\lambda({\bm{\theta}}_{t+1}-{\bm{\theta}}_{t}),{\bm{X}}))}{\ell(y\Phi({\bm{\theta}}_{t},{\bm{X}}))}\leq\max_{\lambda\in[0,1]}\exp\left(2\lambda\lVert{\bm{X}}\rVert^{3}_{2,\infty}\lVert\bm{u}_{t}\rVert\lVert\bm{W}_{t+1}-\bm{W}_{t}\rVert+\lambda\lVert{\bm{X}}\rVert_{2,\infty}\lVert\bm{u}_{t+1}-\bm{u}_{t}\rVert\right).$

Further, we show that the iterate norm $\lVert\bm{u}_{t}\rVert$ grows as
$\mathcal{O}(t^{1/3})$ (formalized in Lemma 12 in the Appendix). Then, using
the updates in Eq. (13), we get

$\displaystyle\max_{i\in[n]}\max_{\lambda\in[0,1]}$
$\displaystyle\frac{\ell(y_{i}\Phi({\bm{\theta}}_{t}+\lambda({\bm{\theta}}_{t+1}-{\bm{\theta}}_{t}),{\bm{X}}_{i}))}{\ell(y_{i}\Phi({\bm{\theta}}_{t},{\bm{X}}_{i}))}\leq\exp\left(6\eta^{2}B^{3}+\eta
B\right).$

Then, using Condition 1 for $\eta$ finishes the proof.

For the third bullet point, we first upper bound
$\lVert\nabla^{2}_{\bm{\theta}}\Phi\rVert$ as

$\displaystyle\lVert\nabla_{\bm{\theta}}^{2}\Phi({\bm{\theta}},{\bm{X}})\rVert$
$\displaystyle\leq 6dB^{5}\lVert\bm{u}\rVert+2\sqrt{d}B^{3}.$

Using this and the upper bound on $\lVert\nabla_{\bm{\theta}}\Phi\rVert$, we
show that for any
${\bm{\theta}}^{\prime}\in[{\bm{\theta}}_{t},{\bm{\theta}}_{t+1}]$,

$\displaystyle\lVert\nabla_{\bm{\theta}}^{2}\widehat{L}({\bm{\theta}}^{\prime})\rVert\leq\max_{i\in[n]}(\lVert\nabla_{\bm{\theta}}\Phi({\bm{\theta}}_{t},{\bm{X}}_{i})\rVert^{2}+\lVert\nabla_{\bm{\theta}}^{2}\Phi({\bm{\theta}}_{t},{\bm{X}}_{i})\rVert)\widehat{L}({\bm{\theta}}^{\prime})\leq\omega({\bm{\theta}}^{\prime})\widehat{L}({\bm{\theta}}^{\prime}).$

Then, we leverage the second bullet point to finish the proof. ∎

Next, to prove Theorem 2, we start with the second-order Taylor expansion of
$\widehat{L}({\bm{\theta}}_{t+1})$. For some
${\bm{\theta}}^{\prime}\in[{\bm{\theta}}_{t},{\bm{\theta}}_{t+1}]$, we have

$\displaystyle\widehat{L}({\bm{\theta}}_{t+1})$
$\displaystyle=\widehat{L}({\bm{\theta}}_{t})+\langle\nabla_{{\bm{\theta}}}\widehat{L}({\bm{\theta}}_{t}),{\bm{\theta}}_{t+1}-{\bm{\theta}}_{t}\rangle+\frac{1}{2}({\bm{\theta}}_{t+1}-{\bm{\theta}}_{t})^{\top}\nabla_{{\bm{\theta}}}^{2}\widehat{L}({\bm{\theta}}^{\prime})({\bm{\theta}}_{t+1}-{\bm{\theta}}_{t})$
$\displaystyle\leq\widehat{L}({\bm{\theta}}_{t})-\eta_{t}^{u}\lVert\nabla_{\bm{u}}\widehat{L}({\bm{\theta}}_{t})\rVert^{2}-\eta_{t}^{W}\lVert\nabla_{\bm{W}}\widehat{L}({\bm{\theta}}_{t})\rVert^{2}+\frac{1}{2}\lVert{\bm{\theta}}_{t+1}-{\bm{\theta}}_{t}\rVert^{2}\max_{{\bm{\theta}}^{\prime}\in[{\bm{\theta}}_{t},{\bm{\theta}}_{t+1}]}\lVert\nabla_{{\bm{\theta}}}^{2}\widehat{L}({\bm{\theta}}^{\prime})\rVert.$

Using the updates in Eq. (13) and the properties of the loss function from
Lemma 3, we show that

$\displaystyle\widehat{L}({\bm{\theta}}_{t+1})\leq\widehat{L}({\bm{\theta}}_{t})-\frac{\eta\omega_{1}}{(t+1)^{2/3}}\widehat{L}({\bm{\theta}}_{t})+\frac{8\eta^{2}}{(t+1)^{4/3}}(\omega({\bm{\theta}}_{t})\vee\omega({\bm{\theta}}_{t+1}))\widehat{L}({\bm{\theta}}_{t}).$

Further, since $\lVert\bm{u}_{t}\rVert\leq\mathcal{O}(t^{1/3})$, we can show
that $\omega({\bm{\theta}}_{t})\leq\mathcal{O}(t^{2/3})$. Then, using a small
enough step-size $\eta$ as per Condition 1, we get,

$\displaystyle\widehat{L}({\bm{\theta}}_{t+1})\leq\left(1-\frac{\eta\omega_{1}}{2(t+1)^{2/3}}\right)\widehat{L}({\bm{\theta}}_{t})\leq\exp\left(-\frac{\eta\omega_{1}}{2}(t+1)^{-2/3}\right)\widehat{L}({\bm{\theta}}_{t}).$

Then, telescoping over $t>0$ gives the final result.

######  Remark 4.

(Lemma 3) We note that the first property is the most challenging to show in
the analysis and it yields a PL-inequality like form, e.g. for neural networks
(Frei and Gu,, 2021; Nguyen and Mondelli,, 2020; Liu et al.,, 2022). Analogous
properties but for simpler settings, e.g. to obtain loss convergence results
for linear predictors using NGD on separable data have been shown in Nacson et
al., (2019). The second point is an analog of the log-Lipschitzness property
shown in Taheri and Thrampoulidis, 2023a  for the analysis of two-layer neural
networks trained with NGD and controls the loss on a line between current
iterate $\bm{\theta}_{t}$ and the next one $\bm{\theta}_{t+1}$. The third
property of second-order self-boundedness has been seen previously in the
convergence analysis of multi-head self-attention (Deora et al.,, 2023) and
MLPs (Taheri and Thrampoulidis, 2023b, ) trained with GD.

##  5 Experimental Results

In this section, we present some experiments on synthetic and real-world
datasets to show that (S)NGD-based training leads to faster convergence for
various metrics compared to standard (S)GD training.

###  5.1 Experiments with Synthetic Data

We first consider a single-head self-attention model, as mentioned in Eq. (1)
and analyze the experimental results using data generated as per data model
DM.

Figure 3: Training dynamics of a single-head self-attention model (Eq. (1))
under joint optimization with data generated using the data model DM.

In Fig. 3, we plot the train loss, iterate norms $\lVert\bm{u}_{t}\rVert$ and
$\lVert\bm{W}_{t}\rVert$, average softmax score for the opt tokens, and
alignment of the iterates $\bm{u}_{t}$ and $\bm{W}_{t}$ to the respective max-
margin solutions $\bm{W}_{\text{mm}}$ and $\bm{u}_{\text{mm}}$, respectively.
We compare standard GD updates, the NGD updates in Eq. (13) (without the
additional time-dependent $t^{-2/3},t^{-1}$ factors in the step-size), updates
with Polyak step-size, and the joint NGD updates (Eq. (2)). We observe that
all other algorithms are faster than GD, with Polyak step being the fastest.
We also see that the train loss converges at a similar rate for NGD and joint-
NGD, while for the other metrics, NGD is slightly faster compared to joint-
NGD.

###  5.2 Experiments with Language Data

Figure 4: Comparison of train and test dynamics of various optimizers—SGD,
SNGD, SPS, and Adam—while fine-tuning a pre-trained BERT model on the
CivilComments dataset; see App. C for details.

Figs. 1 and 4 compare the train and test dynamics of various optimizers while
fine-tuning a pre-trained BERT model on MNLI and CivilComments datasets,
respectively. As evident, both gradient normalized step-sizes SNGD and SPS
train significantly faster than SGD.

Further, observations on vision datasets (see App. C) suggest that SNGD trains
slower than SGD, which could be explored further in future work. Additionally,
we see no significant gap between Adam and SGD as observed in Xie et al.,
(2023).

##  6 Related Work

##### Transformers theory.

Several studies, such as those by Baldi and Vershynin, (2022); Dong et al.,
(2021); Yun et al., 2020a ; Yun et al., 2020b ; Sanford et al., (2023); Bietti
et al., (2023) have delved into exploring the expressivity of attention, while
Baldi and Vershynin, (2022); Dong et al., (2021); Yun et al., 2020a ; Yun et
al., 2020b ; Mahdavi et al., (2023) have initiated an investigation of its
memory capacity. To gain insights into the optimization aspects of training
attention models, Sahiner et al., (2022); Ergen et al., (2022) have explored
convex relaxations of attention. Furthermore, a subdomain experiencing growing
attention involves the theoretical exploration of in-context learning, as
evidenced by recent studies (von Oswald et al.,, 2022; Akyürek et al.,, 2023;
Zhang et al.,, 2023; Li et al., 2023c, ).

In this context, we discuss studies that seek to understand the optimization
and generalization dynamics of transformers. Jelassi et al., (2022) show that
Vision Transformers (ViTs) learn spatially localized patterns in a binary
classification task using gradient-based methods. For a three-layer ViT
starting from some structured initialization to mimic a pre-trained network,
Li et al., 2023b  show sample complexity bounds to achieve zero generalization
and show attention maps sparsify as SGD training proceeds. As a step towards
understanding the training dynamics of the closely related prompt-attention,
Oymak et al., (2023) study the initial trajectory of GD for one-layer
attention. Further, there has also been work to obtain optimization and
generalization guarantees of GD in multi-head attention models (Deora et al.,,
2023). Additionally, recent work by Tian et al., 2023a  attempts to understand
SGD-dynamics for the task of next-token prediction for one-layer transformer
with a linear decoder showing a similar sparsifying effect in the attention
map. More recently, Tian et al., 2023b  extend this by analyzing the joint
training dynamics of multi-layer transformer with an MLP. As detailed in the
Introduction, our work is most closely related and improves upon Tarzanagh et
al., 2023b ; Tarzanagh et al., 2023a , which adopts an implicit-bias view to
attention training.

##### Implicit bias of NNs.

Since the first few works characterizing the implicit bias of linear
predictors on separable data (see Introduction), there has been an abundance
of work studying the implicit bias of gradient based methods for both linear
predictors and NNs. Nacson et al., (2019); Ji and Telgarsky, (2021); Ji et
al., (2021) show fast convergence of GD-based methods to the max-margin
predictor. For MLPs, early works study the implicit bias of GD/gradient flow
using exponentially-tailed classification losses towards the KKT points of the
corresponding max-margin problem in finite (Ji and Telgarsky,, 2020; Lyu and
Li,, 2020) and infinite width (Chizat and Bach,, 2020). Additionally, works by
Phuong and Lampert, (2021); Frei et al., 2022b ; Kou et al., 2023a  study the
implicit bias of GD trained ReLU/Leaky-ReLU networks on orthogonal data. Other
works also study the implicit bias towards rank minimization with square loss
in regression settings (Vardi and Shamir,, 2021; Arora et al.,, 2019; Li et
al.,, 2021). We encourage the reader to go through these works and a recent
survey (Vardi,, 2022) for a thorough understanding.

##  7 Concluding Remarks

We studied the implicit bias of GD for self-attention in binary
classification. For both fixed and trainable decoder settings, we identified
data conditions under which global convergence holds and characterized, for
the first time, the convergence rates. Our results improve upon recent works
Tarzanagh et al., 2023a ; Tarzanagh et al., 2023b , which in addition to the
equal token-score (Assumption 2), further needed to assume a single sample
$n=1$ and lacked convergence rates. Future work aims to further relax
Assumption 2 and identify other conditions (e.g. overparameterization) for
global convergence. Additionally, extending our theory to incorporate momentum
and explain its fast convergence is also an interesting future direction.
Finally, motivated by our experiments, we want to further investigate why NGD
trains faster on language datasets. This might require extending our results
to the next-token prediction setting which is also of interest.

## References

  * Akyürek et al., (2023) Akyürek, E., Schuurmans, D., Andreas, J., Ma, T., and Zhou, D. (2023).  What learning algorithm is in-context learning? Investigations with linear models.  In Int. Conference on Learning Representations.
  * Arora et al., (2019) Arora, S., Cohen, N., Hu, W., and Luo, Y. (2019).  Implicit regularization in deep matrix factorization.  In Advances in Neural Information Processing Systems, volume 32.
  * Baldi and Vershynin, (2022) Baldi, P. and Vershynin, R. (2022).  The quarks of attention.  arXiv preprint arXiv:2202.08371.
  * Bietti et al., (2023) Bietti, A., Cabannes, V., Bouchacourt, D., Jegou, H., and Bottou, L. (2023).  Birth of a transformer: A memory viewpoint.  arXiv preprint arXiv:2306.00802.
  * Borkan et al., (2019) Borkan, D., Dixon, L., Sorensen, J., Thain, N., and Vasserman, L. (2019).  Nuanced metrics for measuring unintended bias with real data for text classification.  In Companion Proceedings of The 2019 World Wide Web Conference.
  * Brown et al., (2020) Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., Agarwal, S., Herbert-Voss, A., Krueger, G., Henighan, T., Child, R., Ramesh, A., Ziegler, D. M., Wu, J., Winter, C., Hesse, C., Chen, M., Sigler, E., Litwin, M., Gray, S., Chess, B., Clark, J., Berner, C., McCandlish, S., Radford, A., Sutskever, I., and Amodei, D. (2020).  Language models are few-shot learners.
  * Cao et al., (2022) Cao, Y., Chen, Z., Belkin, M., and Gu, Q. (2022).  Benign overfitting in two-layer convolutional neural networks.  arXiv preprint arXiv:2202.06526.
  * Cao and Gu, (2019) Cao, Y. and Gu, Q. (2019).  Generalization bounds of stochastic gradient descent for wide and deep neural networks.  Advances in neural information processing systems, 32.
  * Chatterji and Long, (2021) Chatterji, N. S. and Long, P. M. (2021).  Finite-sample analysis of interpolating linear classifiers in the overparameterized regime.  Journal of Machine Learning Research, 22(129):1–30.
  * Chizat and Bach, (2020) Chizat, L. and Bach, F. (2020).  Implicit bias of gradient descent for wide two-layer neural networks trained with the logistic loss.  In Conference on Learning Theory, pages 1305–1338. PMLR.
  * Deora et al., (2023) Deora, P., Ghaderi, R., Taheri, H., and Thrampoulidis, C. (2023).  On the optimization and generalization of multi-head attention.
  * Devlin et al., (2019) Devlin, J., Chang, M.-W., Lee, K., and Toutanova, K. (2019).  BERT: Pre-training of deep bidirectional transformers for language understanding.  In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long and Short Papers), pages 4171–4186, Minneapolis, Minnesota. Association for Computational Linguistics.
  * Dong et al., (2021) Dong, Y., Cordonnier, J.-B., and Loukas, A. (2021).  Attention is not all you need: Pure attention loses rank doubly exponentially with depth.  In International Conference on Machine Learning, pages 2793–2803. PMLR.
  * Dosovitskiy et al., (2021) Dosovitskiy, A., Beyer, L., Kolesnikov, A., Weissenborn, D., Zhai, X., Unterthiner, T., Dehghani, M., Minderer, M., Heigold, G., Gelly, S., Uszkoreit, J., and Houlsby, N. (2021).  An image is worth 16x16 words: Transformers for image recognition at scale.
  * Duchi et al., (2011) Duchi, J., Hazan, E., and Singer, Y. (2011).  Adaptive subgradient methods for online learning and stochastic optimization.  Journal of Machine Learning Research, 12(61):2121–2159.
  * Ergen et al., (2022) Ergen, T., Neyshabur, B., and Mehta, H. (2022).  Convexifying transformers: Improving optimization and understanding of transformer networks.  arXiv preprint arXiv:2211.11052.
  * (17) Frei, S., Chatterji, N. S., and Bartlett, P. (2022a).  Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data.  In Conference on Learning Theory, pages 2668–2703. PMLR.
  * Frei and Gu, (2021) Frei, S. and Gu, Q. (2021).  Proxy convexity: A unified framework for the analysis of neural networks trained by gradient descent.  In Beygelzimer, A., Dauphin, Y., Liang, P., and Vaughan, J. W., editors, Advances in Neural Information Processing Systems.
  * Frei et al., (2023) Frei, S., Vardi, G., Bartlett, P., and Srebro, N. (2023).  Benign overfitting in linear classifiers and leaky relu networks from kkt conditions for margin maximization.  In Neu, G. and Rosasco, L., editors, Proceedings of Thirty Sixth Conference on Learning Theory, volume 195 of Proceedings of Machine Learning Research, pages 3173–3228. PMLR.
  * (20) Frei, S., Vardi, G., Bartlett, P. L., Srebro, N., and Hu, W. (2022b).  Implicit bias in leaky relu networks trained on high-dimensional data.  arXiv preprint arXiv:2210.07082.
  * Hazan et al., (2015) Hazan, E., Levy, K. Y., and Shalev-Shwartz, S. (2015).  Beyond convexity: Stochastic quasi-convex optimization.  In Proceedings of the 28th International Conference on Neural Information Processing Systems - Volume 1, NIPS’15, page 1594–1602, Cambridge, MA, USA. MIT Press.
  * Jelassi et al., (2022) Jelassi, S., Sander, M. E., and Li, Y. (2022).  Vision transformers provably learn spatial structure.  In Oh, A. H., Agarwal, A., Belgrave, D., and Cho, K., editors, Advances in Neural Information Processing Systems.
  * Ji et al., (2021) Ji, Z., Srebro, N., and Telgarsky, M. (2021).  Fast margin maximization via dual acceleration.  In International Conference on Machine Learning, pages 4860–4869. PMLR.
  * Ji and Telgarsky, (2018) Ji, Z. and Telgarsky, M. (2018).  Risk and parameter convergence of logistic regression.  arXiv preprint arXiv:1803.07300.
  * Ji and Telgarsky, (2019) Ji, Z. and Telgarsky, M. (2019).  Characterizing the implicit bias via a primal-dual analysis.  In International Conference on Algorithmic Learning Theory.
  * Ji and Telgarsky, (2020) Ji, Z. and Telgarsky, M. (2020).  Directional convergence and alignment in deep learning.  Advances in Neural Information Processing Systems, 33:17176–17186.
  * Ji and Telgarsky, (2021) Ji, Z. and Telgarsky, M. (2021).  Characterizing the implicit bias via a primal-dual analysis.  In Algorithmic Learning Theory, pages 772–804. PMLR.
  * Kingma and Ba, (2014) Kingma, D. P. and Ba, J. (2014).  Adam: A method for stochastic optimization.  arXiv preprint arXiv:1412.6980.
  * Koh et al., (2021) Koh, P. W., Sagawa, S., Marklund, H., Xie, S. M., Zhang, M., Balsubramani, A., Hu, W., Yasunaga, M., Phillips, R. L., Gao, I., et al. (2021).  Wilds: A benchmark of in-the-wild distribution shifts.  In International Conference on Machine Learning, pages 5637–5664. PMLR.
  * (30) Kou, Y., Chen, Z., and Gu, Q. (2023a).  Implicit bias of gradient descent for two-layer relu and leaky relu networks on nearly-orthogonal data.
  * (31) Kou, Y., Chen, Z.-Y., Chen, Y., and Gu, Q. (2023b).  Benign overfitting in two-layer relu convolutional neural networks.  In International Conference on Machine Learning.
  * Krizhevsky, (2009) Krizhevsky, A. (2009).  Learning multiple layers of features from tiny images.  pages 32–33.
  * Kunstner et al., (2023) Kunstner, F., Chen, J., Lavington, J. W., and Schmidt, M. (2023).  Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be.  In The Eleventh International Conference on Learning Representations.
  * LeCun and Cortes, (2005) LeCun, Y. and Cortes, C. (2005).  The mnist database of handwritten digits.
  * Lee et al., (2021) Lee, S. H., Lee, S., and Song, B. C. (2021).  Vision transformer for small-size datasets.
  * (36) Li, H., Wang, M., Liu, S., and Chen, P.-Y. (2023a).  A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity.  ArXiv, abs/2302.06015.
  * (37) Li, H., Weng, M., Liu, S., and Chen, P.-Y. (2023b).  A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity.  In International Conference on Learning Representations.
  * (38) Li, Y., Ildiz, M. E., Papailiopoulos, D., and Oymak, S. (2023c).  Transformers as algorithms: Generalization and stability in in-context learning.
  * Li et al., (2021) Li, Z., Luo, Y., and Lyu, K. (2021).  Towards resolving the implicit bias of gradient descent for matrix factorization: Greedy low-rank learning.
  * Liu et al., (2022) Liu, C., Zhu, L., and Belkin, M. (2022).  Loss landscapes and optimization in over-parameterized non-linear systems and neural networks.  Applied and Computational Harmonic Analysis, 59:85–116.
  * Loizou et al., (2021) Loizou, N., Vaswani, S., Hadj Laradji, I., and Lacoste-Julien, S. (2021).  Stochastic polyak step-size for sgd: An adaptive learning rate for fast convergence.  In Banerjee, A. and Fukumizu, K., editors, Proceedings of The 24th International Conference on Artificial Intelligence and Statistics, volume 130 of Proceedings of Machine Learning Research, pages 1306–1314. PMLR.
  * Lyu and Li, (2020) Lyu, K. and Li, J. (2020).  Gradient descent maximizes the margin of homogeneous neural networks.  In International Conference on Learning Representations.
  * Mahdavi et al., (2023) Mahdavi, S., Liao, R., and Thrampoulidis, C. (2023).  Memorization capacity of multi-head attention in transformers.  arXiv preprint arXiv:2306.02010.
  * Merrill et al., (2020) Merrill, W., Ramanujan, V., Goldberg, Y., Schwartz, R., and Smith, N. A. (2020).  Effects of parameter norm growth during transformer training: Inductive bias from gradient descent.  In Conference on Empirical Methods in Natural Language Processing.
  * Muthukumar et al., (2019) Muthukumar, V., Vodrahalli, K., and Sahai, A. (2019).  Harmless interpolation of noisy data in regression.  arXiv preprint arXiv:1903.09139.
  * Nacson et al., (2019) Nacson, M. S., Lee, J., Gunasekar, S., Savarese, P. H. P., Srebro, N., and Soudry, D. (2019).  Convergence of gradient descent on separable data.  In The 22nd International Conference on Artificial Intelligence and Statistics, pages 3420–3428. PMLR.
  * Nguyen and Mondelli, (2020) Nguyen, Q. N. and Mondelli, M. (2020).  Global convergence of deep networks with one wide layer followed by pyramidal topology.  Advances in Neural Information Processing Systems, 33:11961–11972.
  * OpenAI, (2022) OpenAI (2022).  Openai: Introducing chatgpt.
  * Oymak et al., (2023) Oymak, S., Rawat, A. S., Soltanolkotabi, M., and Thrampoulidis, C. (2023).  On the role of attention in prompt-tuning.  In ICLR 2023 Workshop on Mathematical and Empirical Understanding of Foundation Models.
  * Phuong and Lampert, (2021) Phuong, M. and Lampert, C. H. (2021).  The inductive bias of re{lu} networks on orthogonally separable data.  In International Conference on Learning Representations.
  * Radford et al., (2021) Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., Krueger, G., and Sutskever, I. (2021).  Learning transferable visual models from natural language supervision.  In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 8748–8763. PMLR.
  * Raffel et al., (2020) Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. J. (2020).  Exploring the limits of transfer learning with a unified text-to-text transformer.  21(1).
  * Sahiner et al., (2022) Sahiner, A., Ergen, T., Ozturkler, B., Pauly, J., Mardani, M., and Pilanci, M. (2022).  Unraveling attention via convex duality: Analysis and interpretations of vision transformers.  International Conference on Machine Learning.
  * Sanford et al., (2023) Sanford, C., Hsu, D., and Telgarsky, M. (2023).  Representational strengths and limitations of transformers.  arXiv preprint arXiv:2306.02896.
  * Soudry et al., (2018) Soudry, D., Hoffer, E., Nacson, M. S., Gunasekar, S., and Srebro, N. (2018).  The implicit bias of gradient descent on separable data.  The Journal of Machine Learning Research, 19(1):2822–2878.
  * (56) Taheri, H. and Thrampoulidis, C. (2023a).  Fast convergence in learning two-layer neural networks with separable data.  Proceedings of the AAAI Conference on Artificial Intelligence, 37(8):9944–9952.
  * (57) Taheri, H. and Thrampoulidis, C. (2023b).  Generalization and stability of interpolating neural networks with minimal width.  arXiv preprint arXiv:2302.09235.
  * (58) Tarzanagh, D. A., Li, Y., Thrampoulidis, C., and Oymak, S. (2023a).  Transformers as support vector machines.
  * (59) Tarzanagh, D. A., Li, Y., Zhang, X., and Oymak, S. (2023b).  Max-margin token selection in attention mechanism.
  * (60) Tian, Y., Wang, Y., Chen, B., and Du, S. (2023a).  Scan and snap: Understanding training dynamics and token composition in 1-layer transformer.
  * (61) Tian, Y., Wang, Y., Zhang, Z., Chen, B., and Du, S. (2023b).  Joma: Demystifying multilayer transformers via joint dynamics of mlp and attention.
  * Touvron et al., (2021) Touvron, H., Cord, M., Douze, M., Massa, F., Sablayrolles, A., and Jegou, H. (2021).  Training data-efficient image transformers & distillation through attention.  In Meila, M. and Zhang, T., editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 10347–10357. PMLR.
  * Vardi, (2022) Vardi, G. (2022).  On the implicit bias in deep-learning algorithms.
  * Vardi and Shamir, (2021) Vardi, G. and Shamir, O. (2021).  Implicit regularization in relu networks with the square loss.  In Conference on Learning Theory, pages 4224–4258. PMLR.
  * Vaswani et al., (2017) Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. (2017).  Attention is all you need.  Advances in neural information processing systems, 30.
  * Vershynin, (2018) Vershynin, R. (2018).  High-Dimensional Probability: An Introduction with Applications in Data Science.  Cambridge Series in Statistical and Probabilistic Mathematics. Cambridge University Press.
  * von Oswald et al., (2022) von Oswald, J., Niklasson, E., Randazzo, E., Sacramento, J., Mordvintsev, A., Zhmoginov, A., and Vladymyrov, M. (2022).  Transformers learn in-context by gradient descent.  arXiv preprint arXiv:2212.07677.
  * Wang et al., (2021) Wang, K., Muthukumar, V., and Thrampoulidis, C. (2021).  Benign overfitting in multiclass classification: All roads lead to interpolation.  Advances in Neural Information Processing Systems, 34:24164–24179.
  * Wang and Thrampoulidis, (2022) Wang, K. and Thrampoulidis, C. (2022).  Binary classification of gaussian mixtures: Abundance of support vectors, benign overfitting, and regularization.  SIAM Journal on Mathematics of Data Science, 4(1):260–284.
  * Williams et al., (2018) Williams, A., Nangia, N., and Bowman, S. (2018).  A broad-coverage challenge corpus for sentence understanding through inference.  In Proceedings of the 2018 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume 1 (Long Papers), pages 1112–1122. Association for Computational Linguistics.
  * Xie et al., (2023) Xie, X., Zhou, P., Li, H., Lin, Z., and YAN, S. (2023).  Adan: Adaptive nesterov momentum algorithm for faster optimizing deep models.
  * (72) Yun, C., Bhojanapalli, S., Rawat, A. S., Reddi, S. J., and Kumar, S. (2020a).  Are transformers universal approximators of sequence-to-sequence functions?
  * (73) Yun, C., Chang, Y.-W., Bhojanapalli, S., Rawat, A. S., Reddi, S. J., and Kumar, S. (2020b).  $o(n)$ connections are expressive enough: Universal approximability of sparse transformers.
  * Zhang et al., (2023) Zhang, R., Frei, S., and Bartlett, P. L. (2023).  Trained transformers learn linear models in-context.

## Appendix

[appendix]

[appendix]1

##  Appendix A Proofs for Section 3

###  A.1 Preliminaries

We list the useful notations used in this section in Table 1 for convenience.

${\bm{x}}_{i,\tau}$ |  $\tau^{\text{th}}$ token in the $i^{\text{th}}$ sample
---|---
$\bm{a}(\bm{W})$ | ${\bm{X}}\bm{W}{\bm{x}}_{1}$
$\bm{\gamma}$ | $y{\bm{X}}\bm{u}_{*}$
$\bm{\varphi}(\cdot)$ | softmax vector
$\varphi_{i,\tau}^{t}$ |  $\bm{\varphi}(\cdot)|_{\tau}$, $\tau\in[\mathrm{T}]$ for sample $i$ at time $t$
$\Lambda$ | $\lVert\bm{W}_{\text{mm}}\rVert$
$B$ | $\max_{i}\lVert{\bm{X}}_{i}\rVert_{2,\infty}$
$\kappa$ | $\frac{\underset{i}{\max}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})}{\underset{i}{\min}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})}$
$\Upsilon$ | $\kappa\frac{\underset{i}{\min}\exp(\gamma_{i,\texttt{opt}_{i}})}{\underset{i}{\max}\exp(\gamma_{i})}$
$\zeta$ | $\frac{\underset{i}{\max}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})}{\underset{i}{\max}\exp(\gamma_{i})}$
Table 1: Useful Notation.

We note that under Assumption 2, for any $\bm{W}$, we have

$\displaystyle\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\bm{W}\rangle$
$\displaystyle=-\frac{1}{n}\sum\limits_{i=1}^{n}\ell^{\prime}_{t,i}\left\langle
y_{i}\nabla_{\bm{W}}\Phi(\bm{W}_{t},{\bm{X}}_{i}),\bm{W}\right\rangle$
$\displaystyle=\frac{1}{n}\sum\limits_{i=1}^{n}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})\left[a_{i,\texttt{opt}_{i}}-\frac{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}a_{i,\tau}}{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}\right],$
(17)

where $\bm{a}_{i}={\bm{X}}_{i}\bm{W}{\bm{x}}_{i,1}$.

###  A.2 Key Lemmas

We first prove some key Lemmas that are useful in the proof of Theorem 1.

######  Lemma 4.

Let $w_{u,v}^{0}\sim\mathcal{N}(0,\sigma_{0}^{2})$ for all $u,v\in[d]$, with
the initialization scale

$\displaystyle\sigma_{0}\leq
B^{-2}d^{-1/2}\left(C^{\prime-1}\log\left(\frac{k\mathrm{T}-1}{\mathrm{T}-1}\right)\,\wedge\,\eta
d^{-1/2}(8\Lambda)^{-1}\right),$

for some $k>1$, and absolute constant $C^{\prime}>0$. Further, if
$d\geq\log(4\delta)$ for any $\delta\in(0,1)$, then with probability at least
$1-\delta$, we have
$\varphi_{i,\texttt{opt}_{i}}^{0}\geq\frac{1}{k\mathrm{T}}$ for every
$i\in[n]$.

###### Proof.

Let $w_{u,v}^{0}\sim\mathcal{N}(0,\sigma_{0}^{2})$, for any $u,v\in[d]$. Then
for any $i\in[n]$ and $\tau\in[T]$, for any $\delta\in(0,1)$, with probability
at least $1-\frac{\delta}{2}$ we have

$\displaystyle({\bm{x}}_{i,\tau}-{\bm{x}}_{i,\texttt{opt}_{i}})^{\top}\bm{W}_{0}{\bm{x}}_{i,1}$
$\displaystyle\leq\left\lvert({\bm{x}}_{i,\tau}-{\bm{x}}_{i,\texttt{opt}_{i}})^{\top}\bm{W}{\bm{x}}_{i,1}\right\rvert$
$\displaystyle\leq\lVert{\bm{x}}_{i,\tau}-{\bm{x}}_{i,\texttt{opt}_{i}}\rVert\lVert\bm{W}_{0}\rVert_{2}\lVert{\bm{x}}_{i,1}\rVert$
$\displaystyle\leq 2B^{2}(\sigma_{0}C(2\sqrt{d}+\sqrt{\log(4/\delta)}))$
$\displaystyle\leq
6CB^{2}\sigma_{0}\sqrt{d}=:B^{2}\sigma_{0}\sqrt{d}C^{\prime}.$

where the second-last inequality follows by using a high probability upper
bound on the spectral norm of $\bm{W}_{0}$ (See Thm. 4.4.5 in Vershynin,
(2018)). Using this, we have

$\displaystyle\varphi_{i,\texttt{opt}_{i}}^{0}$
$\displaystyle=\frac{1}{1+\sum\limits_{\tau\neq\texttt{opt}_{i}}\exp(({\bm{x}}_{i,\tau}-{\bm{x}}_{i,\texttt{opt}_{i}})^{\top}\bm{W}_{0}{\bm{x}}_{i,1})}$
$\displaystyle\geq\frac{1}{1+(T-1)\exp(B^{2}\sigma_{0}\sqrt{d}C^{\prime})}$
$\displaystyle\geq\frac{1}{kT}$ (using $\sigma_{0}\leq
C^{\prime-1}B^{-2}d^{-1/2}\log\left(\frac{k\mathrm{T}-1}{\mathrm{T}-1}\right)$)

Also, $\lVert\bm{W}_{0}\rVert\leq 2\sigma_{0}d\leq\eta(4B^{2}\Lambda)^{-1}$.
The second inequality follows from the concentration of norm with independent
sub-Gaussian entries (see Thm. 3.1.1 in Vershynin, (2018)), and using
$d\geq\sqrt{\log(4/\delta)}$. The last inequality follows by using the
condition on $\sigma_{0}$. Finally, we finish the proof by applying a union
bound. ∎

######  Lemma 5.

Under the conditions of Lemma 4, Assumptions 1 and 2, using the updates in Eq.
(4) with $\eta\leq(2B^{2})^{-1}\log(2\mathrm{T}-1)$, at any time $t\geq 0$,
for any sample $j\in[n]$, the softmax weight for the $\texttt{opt}_{j}$ token
satisfies

$\displaystyle\varphi_{j,\texttt{opt}_{j}}^{t}\geq\varphi_{j,\texttt{opt}_{j}}^{0}\,\wedge\,\frac{1}{\mathrm{T}}.$
(18)

###### Proof.

We prove this by induction. Clearly, Eq. (18) is true for $t=0$. Next, we
assume that it is true at time $t-1$. Combining this with the conditions of
Lemma 4, we have
$\varphi_{j,\texttt{opt}_{j}}^{t-1}\geq\frac{1}{2\mathrm{T}}$.

Further, for any $j\in[n]$ we can have one of the following two scenarios:
Scenario 1: At time $t-1$,
$\varphi_{j,\texttt{opt}_{j}}^{t-1}>1-\frac{1}{2\mathrm{T}}$. This implies

$\displaystyle\varphi_{j,\texttt{opt}_{j}}^{t-1}=\frac{1}{1+\sum_{\tau\neq\texttt{opt}_{i}}\exp(({\bm{x}}_{i,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t-1}{\bm{x}}_{j,1})}>1-\frac{1}{2\mathrm{T}}$
$\displaystyle\implies\quad$
$\displaystyle\sum_{\tau\neq\texttt{opt}_{j}}\exp(({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t-1}{\bm{x}}_{j,1})<\frac{1}{2\mathrm{T}-1}$
$\displaystyle\implies\quad$
$\displaystyle({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t-1}{\bm{x}}_{j,1}\leq-\log(2\mathrm{T}-1),\,\,\quad\forall\tau\neq\texttt{opt}_{j}.$

After the gradient step at time $t-1$, for any $\tau\neq\texttt{opt}_{j}$, we
have,

$\displaystyle({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t}{\bm{x}}_{j,1}$
$\displaystyle=({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t-1}{\bm{x}}_{j,1}-\eta_{t-1}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}(-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t-1})){\bm{x}}_{j,1}$
$\displaystyle\leq-\log(2\mathrm{T}-1)+\eta\lVert{\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}}\rVert\lVert{\bm{x}}_{j,1}\rVert$
$\displaystyle\leq-\log(2\mathrm{T}-1)+2\eta B^{2}$ $\displaystyle\leq 0,$

where the last step follows since $\eta\leq(2B^{2})^{-1}\log(2\mathrm{T}-1)$.
This gives

$\displaystyle\sum_{\tau\neq\texttt{opt}_{j}}\exp(({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t}{\bm{x}}_{j,1})$
$\displaystyle\leq\mathrm{T}-1$
$\displaystyle\implies\varphi_{j,\texttt{opt}_{j}}^{t}$
$\displaystyle\geq\frac{1}{\mathrm{T}}.$

Scenario 2: At time $t-1$, $\varphi_{j,\texttt{opt}_{j}}^{t-1}\leq
1-\frac{1}{2\mathrm{T}}$. This gives

$\displaystyle\varphi_{j,\texttt{opt}_{j}}^{t-1}(1-\varphi_{j,\texttt{opt}_{j}}^{t-1})\geq\frac{2\mathrm{T}-1}{4\mathrm{T}^{2}}\geq\frac{1}{4\mathrm{T}}.$
(19)

We will use this to show that

$\displaystyle({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}(-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})){\bm{x}}_{j,1}\geq
0.$ (20)

Let
$\bm{a}^{\prime}=\bm{a}(({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau}){\bm{x}}_{j,1}^{\top})={\bm{X}}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau}){\bm{x}}_{j,1}^{\top}{\bm{x}}_{1}$.
Using similar calculations as Eq. (17), we get

$\displaystyle({\bm{x}}_{j,\texttt{opt}_{j}}$
$\displaystyle-{\bm{x}}_{j,\tau})^{\top}(-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t-1})){\bm{x}}_{j,1}=\frac{1}{n}\sum\limits_{i=1}^{n}\ell_{t-1,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t-1}(1-\varphi_{i,\texttt{opt}_{i}}^{t-1})\left[{a}^{\prime}_{i,\texttt{opt}_{i}}-\frac{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t-1}{a}^{\prime}_{i,\tau}}{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t-1}}\right].$
(21)

There are two cases:
Case 1: $\texttt{opt}_{j}\neq 1$. In this case, by splitting the sum over
$i\in[n]$ in Eq. (21) into $i=j$ and $i\neq j$, and using non-negativity of
the softmax weights, we get

$\displaystyle({\bm{x}}_{j,\texttt{opt}_{j}}$
$\displaystyle-{\bm{x}}_{j,\tau})^{\top}(-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t-1})){\bm{x}}_{j,1}$
$\displaystyle\geq\frac{\ell_{t-1,j}}{n}(\gamma_{j,\texttt{opt}_{j}}-\gamma_{j})\varphi_{j,\texttt{opt}_{j}}^{t-1}(1-\varphi_{j,\texttt{opt}_{j}}^{t-1})\lVert{\bm{x}}_{j,1}\rVert^{2}\min_{\tau^{\prime}\neq\texttt{opt}_{j}}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau^{\prime}})$
$\displaystyle-\max_{i\neq
j}\ell_{t-1,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t-1}(1-\varphi_{i,\texttt{opt}_{i}}^{t-1})|{\bm{x}}_{i,1}^{\top}{\bm{x}}_{j,1}|\max_{\tau^{\prime}\neq\texttt{opt}_{i}}|({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}({\bm{x}}_{i,\texttt{opt}_{i}}-{\bm{x}}_{i,\tau^{\prime}})|.$
(22)

Next, we use Assumption 1 and that $n\Upsilon\mathrm{T}\geq 1$ to get

$\displaystyle\max_{i\neq
j}\max_{\tau^{\prime}\neq\texttt{opt}_{i}}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}({\bm{x}}_{i,\texttt{opt}_{i}}-{\bm{x}}_{i,\tau^{\prime}})$
$\displaystyle\leq\max_{i\neq
j}|{\bm{x}}_{j,\texttt{opt}_{j}}^{\top}{\bm{x}}_{i,\texttt{opt}_{i}}|+3\,\max_{\mathclap{\begin{subarray}{c}i,j:i\neq
j,\tau\in[\mathrm{T}],\\\
\tau^{\prime}\in[\mathrm{T}]\setminus\\{\texttt{opt}_{j}\\}\end{subarray}}}\,|{\bm{x}}_{i,\tau}^{\top}{\bm{x}}_{j,\tau^{\prime}}|$
$\displaystyle\leq
4\frac{\lVert{\bm{x}}_{j,\texttt{opt}_{j}}\rVert^{2}}{4}=\lVert{\bm{x}}_{j,\texttt{opt}_{j}}\rVert^{2}.$
(23)

Similarly,

$\displaystyle\min_{\tau^{\prime}\neq\texttt{opt}_{j}}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau^{\prime}})$
$\displaystyle\geq\lVert{\bm{x}}_{j,\texttt{opt}_{j}}\rVert^{2}-3\,\max_{\mathclap{\begin{subarray}{c}\tau\in[\mathrm{T}],\\\
\tau^{\prime}\in[\mathrm{T}]\setminus\\{\texttt{opt}_{j}\\}\end{subarray}}}\,|{\bm{x}}_{j,\tau}^{\top}{\bm{x}}_{j,\tau^{\prime}}|$
$\displaystyle\geq\lVert{\bm{x}}_{j,\texttt{opt}_{j}}\rVert^{2}-3\frac{\lVert{\bm{x}}_{j,\texttt{opt}_{j}}\rVert^{2}}{4}=\frac{1}{4}\lVert{\bm{x}}_{j,\texttt{opt}_{j}}\rVert^{2}.$
(24)

Combining Eqs. (19), (23) and (24), we have

$\displaystyle\zeta$ $\displaystyle:=\frac{n\max_{i\neq
j}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t-1}(1-\varphi_{i,\texttt{opt}_{i}}^{t-1})}{\ell_{t,j}(\gamma_{j,\texttt{opt}_{j}}-\gamma_{j})\varphi_{j,\texttt{opt}_{j}}^{t-1}(1-\varphi_{j,\texttt{opt}_{j}}^{t-1})}\frac{\max_{i,\tau^{\prime}\neq\texttt{opt}_{i}}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}({\bm{x}}_{i,\texttt{opt}_{i}}-{\bm{x}}_{i,\tau^{\prime}})}{\min_{\tau^{\prime}\neq\texttt{opt}_{j}}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau^{\prime}})}$
$\displaystyle\leq
n\frac{\max_{i}\ell_{t-1,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})}{\min_{i}\ell_{t-1,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})}\frac{1/4}{1/(4\mathrm{T})}\frac{\lVert{\bm{x}}_{j,\texttt{opt}_{j}}\rVert^{2}}{\lVert{\bm{x}}_{j,\texttt{opt}_{j}}\rVert^{2}/4}$
$\displaystyle\leq
4n\mathrm{T}\kappa\frac{\max_{i}\ell_{t-1,i}}{\min_{i}\ell_{t-1,i}}$
$\displaystyle\leq
4n\mathrm{T}\kappa\frac{\max_{i}\exp(-(\varphi_{i,\texttt{opt}_{i}}^{t-1}\gamma_{i,\texttt{opt}_{i}}+(1-\varphi_{i,\texttt{opt}_{i}}^{t-1})\gamma_{i}))}{\min_{i}\exp(-(\varphi_{i,\texttt{opt}_{i}}^{t-1}\gamma_{i,\texttt{opt}_{i}}+(1-\varphi_{i,\texttt{opt}_{i}}^{t-1})\gamma_{i}))}$
$\displaystyle\leq
4n\mathrm{T}\kappa\frac{\max_{i}\exp(-\gamma_{i})}{\min_{i}\exp(-\gamma_{i,\texttt{opt}_{i}})}=4n\Upsilon\mathrm{T}.$
(25)

Using Assumption 1 and Eq. (25), we have

$\displaystyle\lVert{\bm{x}}_{j,1}\rVert^{2}\geq\zeta\max_{i\neq
j}\left\lvert\langle{\bm{x}}_{i,1},{\bm{x}}_{j,1}\rangle\right\rvert,$

which implies that Eq. (22) is non-negative.
Case 2: $\texttt{opt}_{j}=1$. In this case, let
$\mathcal{I}:=\\{i\in[n]:y_{i}=y_{j}\wedge\texttt{opt}_{i}=1\\}$. Using Eq.
(21), we get

$\displaystyle({\bm{x}}_{j,\texttt{opt}_{j}}$
$\displaystyle-{\bm{x}}_{j,\tau})^{\top}(-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t-1})){\bm{x}}_{j,1}$
$\displaystyle\geq\frac{\ell_{t-1,j}}{n}(\gamma_{j,\texttt{opt}_{j}}-\gamma_{j})\varphi_{j,\texttt{opt}_{j}}^{t-1}(1-\varphi_{j,\texttt{opt}_{j}}^{t-1})\lVert{\bm{x}}_{j,1}\rVert^{2}\min_{\tau^{\prime}\neq\texttt{opt}_{j}}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau^{\prime}})$
$\displaystyle+\min_{i\neq
j,i\in\mathcal{I}}\ell_{t-1,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t-1}(1-\varphi_{i,\texttt{opt}_{i}}^{t-1}){\bm{x}}_{i,1}^{\top}{\bm{x}}_{j,1}\min_{\tau^{\prime}\neq\texttt{opt}_{i}}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}({\bm{x}}_{i,\texttt{opt}_{i}}-{\bm{x}}_{i,\tau^{\prime}})$
$\displaystyle-\max_{i\neq
j,i\notin\mathcal{I}}\ell_{t-1,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t-1}(1-\varphi_{i,\texttt{opt}_{i}}^{t-1})|{\bm{x}}_{i,1}^{\top}{\bm{x}}_{j,1}|\max_{\tau^{\prime}\neq\texttt{opt}_{i}}|({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}({\bm{x}}_{i,\texttt{opt}_{i}}-{\bm{x}}_{i,\tau^{\prime}})|.$
(26)

Using Assumption 1, the second term is non-negative. For the remaining two
terms, we can proceed in a similar way as the previous case to show that using
Assumption 1, these terms are also non-negative.

Using Eq. (20), we get

$\displaystyle({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}(\bm{W}_{t}-\bm{W}_{t-1}){\bm{x}}_{j,1}$
$\displaystyle=\eta\frac{({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}(-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t-1})){\bm{x}}_{j,1}}{\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t-1})\rVert}\geq
0.$

Then, for any $\tau\neq\texttt{opt}_{j}$, we have

$\displaystyle\frac{\varphi_{j,\texttt{opt}_{j}}^{t}/\varphi_{j,\tau}^{t}}{\varphi_{j,\texttt{opt}_{j}}^{t-1}/\varphi_{j,\tau}^{t-1}}$
$\displaystyle=\exp\left(({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}(\bm{W}_{t}-\bm{W}_{t-1}){\bm{x}}_{j,1}\right)\geq
1.$

By telescoping, we get
$\frac{\varphi_{j,\texttt{opt}_{j}}^{t}}{\varphi_{j,\tau}^{t}}\geq\frac{\varphi_{j,\texttt{opt}_{j}}^{0}}{\varphi_{j,\tau}^{0}}$,
which implies that
$\varphi_{j,\texttt{opt}_{j}}^{t}\geq\varphi_{j,\texttt{opt}_{j}}^{0}$.

∎

######  Lemma 6.

Under Assumption 2, for $\bm{W}_{\text{mm}}$ as defined in (W-SVM) and any
$\bm{W}$,

$\displaystyle\left\langle-\frac{\nabla_{\bm{W}}\widehat{L}(\bm{W})}{\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W})\rVert_{F}},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq(2B^{2}\Lambda)^{-1}>0.$

###### Proof.

First, using Eq. (17) for $\bm{W}_{\text{mm}}$, we have

$\displaystyle\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}),\bm{W}_{\text{mm}}\rangle$
$\displaystyle\geq\min_{i}(a^{*}_{i,\texttt{opt}_{i}}-\max_{\tau\neq\texttt{opt}_{i}}a^{*}_{i,\tau})\frac{1}{n}\sum\limits_{i=1}^{n}\ell_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}(1-\varphi_{i,\texttt{opt}_{i}})$
$\displaystyle\geq\frac{1}{n}\sum\limits_{i=1}^{n}\ell_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}(1-\varphi_{i,\texttt{opt}_{i}}),$
(27)

where $\bm{a}^{*}_{i}={\bm{X}}_{i}\bm{W}_{\text{mm}}{\bm{x}}_{i,1}$, and the
second inequality follows by the definition of $\bm{W}_{\text{mm}}$.

Let
$\hat{\bm{W}}:=-\frac{\nabla_{\bm{W}}\widehat{L}(\bm{W})}{\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W})\rVert_{F}}$,
and $\hat{\bm{a}}_{i}={\bm{X}}_{i}\hat{\bm{W}}{\bm{x}}_{i,1}$, then by similar
calculations as Eq. (17), we have

$\displaystyle\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W})\rVert_{F}$
$\displaystyle=\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}),\hat{\bm{W}}\rangle$
$\displaystyle=\frac{1}{n}\sum\limits_{i=1}^{n}\ell_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}(1-\varphi_{i,\texttt{opt}_{i}})\left[\hat{a}_{i,\texttt{opt}_{i}}-\frac{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}\hat{a}_{i,\tau}}{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}}\right]$
$\displaystyle\leq\max_{i}(\hat{a}_{i,\texttt{opt}_{i}}-\min_{\tau\neq\texttt{opt}_{i}}\hat{a}_{i,\tau})\frac{1}{n}\sum\limits_{i=1}^{n}\ell_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}(1-\varphi_{i,\texttt{opt}_{i}})$
$\displaystyle\leq
2B^{2}\frac{1}{n}\sum\limits_{i=1}^{n}\ell_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}(1-\varphi_{i,\texttt{opt}_{i}}),$
(28)

where for the last step we use

$\displaystyle\max_{i}(\hat{a}_{i,\texttt{opt}_{i}}-\min_{\tau\neq\texttt{opt}_{i}}\hat{a}_{i,\tau})$
$\displaystyle\leq
2\max_{i,\tau}|\hat{a}_{i,\tau}|=2\max_{i,\tau}|{\bm{x}}_{i,\tau}^{\top}\hat{\bm{W}}{\bm{x}}_{i,1}|$
$\displaystyle\leq
2\max_{i}\lVert{\bm{X}}_{i}\rVert_{2,\infty}\lVert{\bm{x}}_{i,1}\rVert\lVert\hat{\bm{W}}\rVert\leq
2B^{2}.$

Using Eqs. (27) and (28), we get

$\displaystyle\left\langle-\frac{\nabla_{\bm{W}}\widehat{L}(\bm{W})}{\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W})\rVert_{F}},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq(2B^{2}\Lambda)^{-1}.$

∎

######  Lemma 7 (Iterate Norm).

Using the updates in Eq. (4), under the conditions of Lemma 4 and Assumption
2, at any time $t>0$, we have

$\displaystyle
2\eta((4B^{2}\Lambda)^{-1}\,\vee\,1)t\geq\lVert\bm{W}_{t}\rVert_{F}\geq\eta(4B^{2}\Lambda)^{-1}t.$
(29)

###### Proof.

Using Eq. (4) and Lemma 6 with $\bm{W}=\bm{W}_{t}$, we have

$\displaystyle\lVert\bm{W}_{t}\rVert$
$\displaystyle=\left\|\bm{W}_{0}-\sum\limits_{t^{\prime}=0}^{t-1}\eta_{t^{\prime}}\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})\right\|$
$\displaystyle\geq\left\langle\bm{W}_{0},\overline{\bm{W}}_{\text{mm}}\right\rangle+\sum\limits_{t^{\prime}=0}^{t-1}\eta\left\langle-\frac{\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})}{\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})\rVert_{F}},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle$
$\displaystyle\geq\eta(2B^{2}\Lambda)^{-1}t-\left\lvert\left\langle\bm{W}_{0},\overline{\bm{W}}_{\text{mm}}\right\rangle\right\rvert.$

Under the conditions of Lemma 4,

$\displaystyle\left\lvert\left\langle\bm{W}_{0},\overline{\bm{W}}_{\text{mm}}\right\rangle\right\rvert\leq\lVert\bm{W}_{0}\rVert_{F}\leq\eta(4B^{2}\Lambda)^{-1},$

which gives the required result for $t>0$.

Similarly, for the upper bound we have

$\displaystyle\lVert\bm{W}_{t}\rVert_{F}$
$\displaystyle=\left\|\bm{W}_{0}-\sum\limits_{t^{\prime}=0}^{t-1}\eta_{t^{\prime}}\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})\right\|_{F}$
$\displaystyle\leq\lVert\bm{W}_{0}\rVert+\sum_{t^{\prime}=0}^{t-1}\eta$
$\displaystyle\leq\eta(4B^{2}\Lambda)^{-1}+\eta t.$

∎

######  Lemma 8 (Data Example 1).

Using $\bm{u}_{*}=\frac{1}{U^{2}}(\bm{\mu}_{+}-\bm{\mu}_{-})$ and
$d\gtrsim\Omega(n^{2}\mathrm{T}^{2}\log\left(\frac{4n}{\delta}\right))$, data
generated as per Example 1 satisfies Ass. 1 and 2 with probability at least
$1-\delta$, for any $\delta\in(0,1)$.

###### Proof.

Using $\bm{u}_{*}=\frac{1}{U^{2}}(\bm{\mu}_{+}-\bm{\mu}_{-})$, for all
$i\in[n]$ and $\tau\neq\texttt{opt}_{i}$, we have

$\displaystyle\gamma_{i,\texttt{opt}_{i}}=y\bm{u}_{*}^{\top}{\bm{x}}_{i,\texttt{opt}_{i}}=1,\quad\gamma_{i,\tau}=y\bm{u}_{*}^{\top}{\bm{x}}_{i,\tau}=0.$

Using this, we have

$\displaystyle\Upsilon$
$\displaystyle=\frac{\underset{i}{\max}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})}{\underset{i}{\min}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})}\frac{\underset{i}{\min}\exp(\gamma_{i,\texttt{opt}_{i}})}{\underset{i}{\max}\exp(\gamma_{i})}=e.$

Next, by using Bernstein inequality, we know with probability at least
$1-\frac{\delta}{2n}$, for every $i\in[n]$ and $\tau\neq\texttt{opt}_{i}$ we
have

$\displaystyle\left|\lVert{\bm{x}}_{i,\texttt{opt}_{i}}\rVert^{2}-\sigma^{2}d-U^{2}\right|\leq
c\sigma^{2}\sqrt{{d\log\left(\frac{4n}{\delta}\right)}}\leq\frac{\sigma^{2}d}{2},$
$\displaystyle\left|\lVert{\bm{x}}_{i,\tau}\rVert^{2}-\rho^{2}d\right|\leq
c\rho^{2}\sqrt{{d\log\left(\frac{4n}{\delta}\right)}}\leq\frac{\rho^{2}d}{2}.$

where $c>0$ is an absolute constant and we use $d\geq
4c^{2}\log\left(\frac{4n}{\delta}\right)$.

Further, for any $i,j\in[n],i\neq j$ and any
$\tau,\tau^{\prime}\in[\mathrm{T}]$. By applying Bernstein’s inequality, each
of the following is true with probability at least $1-\frac{\delta}{5n^{2}}$,

$\displaystyle|\langle{\bm{x}}_{i,\texttt{opt}_{i}},{\bm{x}}_{j,\texttt{opt}_{j}}\rangle-\langle\bm{\mu}_{y_{i}},\bm{\mu}_{y_{j}}\rangle|\leq
2\sigma^{2}\sqrt{d\log\left(\frac{10n^{2}}{\delta}\right)},$
$\displaystyle\quad|\langle{\bm{x}}_{i,\tau^{\prime}},{\bm{x}}_{j,\tau}\rangle|\leq
2\rho^{2}\sqrt{d\log\left(\frac{10n^{2}}{\delta}\right)},$
$\displaystyle|\langle{\bm{x}}_{i,\texttt{opt}_{i}},{\bm{x}}_{j,\tau}\rangle|$
$\displaystyle\leq
2\sigma\rho\sqrt{d\log\left(\frac{10n^{2}}{\delta}\right)},$

where $\tau^{\prime}\neq\texttt{opt}_{i},\tau\neq\texttt{opt}_{j}$.

In order to satisfy Assumption 1, ignoring the log factors, we require

$\displaystyle(U^{2}+\sigma^{2}d)\,\wedge\,\rho^{2}d\,\,\gtrsim\,\,nT(\rho^{2}\sqrt{d}\,\vee\,\sigma\rho\sqrt{d})$
$\displaystyle
U^{2}+\sigma^{2}\sqrt{d}\,\,\gtrsim\,\,nT(\sigma^{2}\sqrt{d}\,\vee\,\rho^{2}\sqrt{d}\,\vee\,\sigma\rho\sqrt{d})$

The above inequalities hold if we have
$d=\widetilde{\Omega}(n^{2}\mathrm{T}^{2})$ and
$U^{2}\geq\widetilde{\Omega}((\sigma\,\vee\,\rho)^{2}n\mathrm{T}\sqrt{d})$.
The proof finishes by the application of a union bound. ∎

####  A.2.1 Proof of Lemma 2

We first restate Lemma 2 below for convenience.

######  Lemma 9.

Under the conditions of Lemma 4 and Assumptions 1 and 2, for any
$\epsilon\in(0,1)$, there exists

$\displaystyle R_{\epsilon}$
$\displaystyle:=2\Lambda\epsilon^{-1}(\log(4n(B^{2}\Lambda)^{-1}\Upsilon\mathrm{T}^{3}\epsilon^{-1})\vee
5\log(20\mathrm{T}B^{2}\Lambda\epsilon^{-1})),$

such that for every $t$ where $\lVert\bm{W}_{t}\rVert\geq R_{\epsilon}$,

$\displaystyle\left\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq(1-\epsilon)\left\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}\right\rangle.$
(30)

###### Proof.

Let $\tilde{\bm{a}}_{i}:={\bm{X}}_{i}\widetilde{\bm{W}}_{t}{\bm{x}}_{i,1}$,
$\bm{a}^{*}_{i}:={\bm{X}}_{i}\bm{W}_{\text{mm}}{\bm{x}}_{i,1}$. We consider
two scenarios based on whether
$\lVert\widetilde{\bm{W}}_{t}-\bm{W}_{\text{mm}}\rVert\leq\frac{\nu}{2B^{2}}$
(Scenario 1) or not (Scenario 2), where

$\displaystyle\widetilde{\bm{W}}_{t}:=\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}\lVert\bm{W}_{\text{mm}}\rVert,\,\,\,\nu:=\frac{\epsilon}{(1-\epsilon)}.$
(31)

Scenario 1 of the proof is the same as Tarzanagh et al., 2023a . Below, we
consider Scenario 2.
Using the definition of $\nu$ and similar calculations to Eq. (17), Eq. (30)
translates to

$\displaystyle\frac{1}{n}\sum\limits_{i=1}^{n}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})\tilde{h}_{t,i}\leq(1+\nu)\frac{1}{n}\sum\limits_{i=1}^{n}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})h^{*}_{t,i},$
(32)

where
$\tilde{h}_{t,i}:=\tilde{a}_{i,\texttt{opt}_{i}}-\frac{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}\tilde{a}_{i,\tau}}{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}$
and $h^{*}_{t,i}$ is defined similarly.
For any sample $i\in[n]$, let

$\displaystyle\delta^{\max}_{i}(\epsilon,\omega)$
$\displaystyle:=\max_{\tau\notin\\{\texttt{opt}_{i}\\}}\tilde{a}_{i,\texttt{opt}_{i}}-\tilde{a}_{i,\tau}-1,$
$\displaystyle\delta^{\min}_{i}(\epsilon,\omega)$
$\displaystyle:=\min_{\tau\notin\\{\texttt{opt}_{i}\\}}\tilde{a}_{i,\texttt{opt}_{i}}-\tilde{a}_{i,\tau}-1.$

Define the following sets

$\displaystyle\mathcal{I}_{1}$
$\displaystyle:=\\{i\in[n]\,\,|\,\,\delta^{\min}_{i}\leq 0.3\nu\\},$
$\displaystyle\mathcal{I}_{2}$
$\displaystyle:=\\{i\in[n]\,\,|\,\,0.3\nu<\delta^{\min}_{i}\leq 0.8\nu\\},$
$\displaystyle\mathcal{I}_{3}$
$\displaystyle:=\\{i\in[n]\,\,|\,\,\delta^{\min}_{i}\geq 0.8\nu\\}.$

Intuitively, $\mathcal{I}_{1}$ is the set of samples for which some (W-SVM)
constraints are either violated or all constraints are barely satisfied.
Similarly, $\mathcal{I}_{3}$ contains samples where all the (W-SVM)
constraints are satisfied. Finally, $\mathcal{I}_{2}$ makes up the rest of the
samples.
First, we will show that for any $i\in\mathcal{I}_{1}$,

$\displaystyle\left[\tilde{a}_{i,\texttt{opt}_{i}}-\frac{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}\tilde{a}_{i,\tau}}{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}\right]\leq(1+0.5\nu)\left[a^{*}_{i,\texttt{opt}_{i}}-\frac{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}a^{*}_{i,\tau}}{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}\right].$
(33)

Next, we will show that for any $i\in\mathcal{I}_{2}$, we have
$\tilde{h}_{t,i}\leq(1+\nu)h^{*}_{t,i}$, i.e.

$\displaystyle\left[\tilde{a}_{i,\texttt{opt}_{i}}-\frac{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}\tilde{a}_{i,\tau}}{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}\right]\leq(1+\nu)\left[a^{*}_{i,\texttt{opt}_{i}}-\frac{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}a^{*}_{i,\tau}}{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}\right].$
(34)

Finally, notice that in order to complete the proof for Eq. (32), using Eqs.
(33) and (34) it suffices to show that

$\displaystyle\sum\limits_{i\in\mathcal{I}_{3}}$
$\displaystyle\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})(\tilde{h}_{t,i}-(1+\nu)h^{*}_{t,i})$
$\displaystyle\leq
0.5\nu\sum\limits_{i\in\mathcal{I}_{1}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})h^{*}_{t,i}.$
(35)

Using $h^{*}_{t,i}\geq 1$, Eq. (A.2.1) follows when

$\displaystyle\sum\limits_{i\in\mathcal{I}_{3}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})(\tilde{h}_{t,i}-1-\nu)\leq
0.5\nu\sum\limits_{i\in\mathcal{I}_{1}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t}).$
(36)

Case 1: Sample $i\in\mathcal{I}_{1}$, i.e. $\delta^{\min}_{i}\leq 0.3\nu$.
If $\delta^{\max}_{i}<0.5\nu$, (33) follows directly. For the rest of the
samples $i$ where $\delta^{\max}_{i}\geq 0.5\nu$, let

$\displaystyle\mathcal{N}:=\\{\tau\in[\mathrm{T}]:a_{i,\texttt{opt}_{i}}-a_{i,\tau}\leq
1+0.4\nu\\}.$

Let
$R^{\prime}:=\frac{\lVert\bm{W}_{t}\rVert}{\lVert\bm{W}_{\text{mm}}\rVert}$,
then by definition of $\mathcal{N},\delta_{\min}$, we have

$\displaystyle\frac{\sum\limits_{\tau\neq\texttt{opt}_{i},\tau\notin\mathcal{N}}\varphi_{i,\tau}^{t}}{\sum\limits_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}$
$\displaystyle\leq\frac{\mathrm{T}\max_{\tau\neq\texttt{opt}_{i},\tau\notin\mathcal{N}}\varphi_{i,\tau}^{t}}{\varphi_{i,\tau}^{t}}\leq\frac{\mathrm{T}\exp(-R^{\prime}(1+0.4\nu))}{\exp(-R^{\prime}(1+\delta^{\min}_{i}))}$
$\displaystyle\leq\mathrm{T}\exp(-0.1R^{\prime}\nu).$

Combining these, we have

$\displaystyle\frac{\sum\limits_{\tau\neq\texttt{opt}_{i}}(\tilde{a}_{i,\texttt{opt}_{i}}-\tilde{a}_{i,\tau})\varphi_{i,\tau}^{t}}{\sum\limits_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}$
$\displaystyle\leq(1+0.4\nu)\frac{\sum\limits_{\tau\in\mathcal{N}}\varphi_{i,\tau}^{t}}{\sum\limits_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}+(1+\delta^{\max}_{i})\frac{\sum\limits_{\tau\neq\texttt{opt}_{i},\tau\notin\mathcal{N}}\varphi_{i,\tau}^{t}}{\sum\limits_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}$
$\displaystyle\leq(1+0.4\nu)\frac{\sum\limits_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}-\sum\limits_{\tau\neq\texttt{opt}_{i}\notin\mathcal{N}}\varphi_{i,\tau}^{t}}{\sum\limits_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}+(1+\delta^{\max}_{i})\frac{\sum\limits_{\tau\neq\texttt{opt}_{i},\tau\notin\mathcal{N}}\varphi_{i,\tau}^{t}}{\sum\limits_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}$
$\displaystyle\leq
1+0.4\nu+(\delta^{\max}_{i}-0.4\nu)\mathrm{T}\exp(-0.1R^{\prime}\nu).$

To satisfy (33), we need
$(\delta^{\max}_{i}-0.4\nu)\mathrm{T}\exp(-0.1R^{\prime}\nu)\leq 0.1\nu$,
which is true when

$\displaystyle
R^{\prime}\geq(0.1\nu)^{-1}\log(\mathrm{T}(\delta^{\max}_{i}/(0.1\nu)-4)).$

This is true since $R^{\prime}\geq
10\epsilon^{-1}\log(10\mathrm{T}\max_{i\in[n]}\delta^{\max}_{i}/\epsilon)$,
where we use the definition of $\nu$ from Eq. (31).

Case 2: Sample $i\in\mathcal{I}_{2}$, i.e. $0.3\nu<\delta^{\min}\leq 0.8\nu$.
If $\delta^{\max}_{i}<\nu$, (34) follows directly. For the rest of the samples
where $\delta^{\max}_{i}\geq\nu$, we can proceed in a similar way as Case 1.
We use a threshold of $0.9\nu$ to split the tokens, and obtain

$\displaystyle\frac{\sum\limits_{\tau\neq\texttt{opt}_{i}}(a_{i,\texttt{opt}_{i}}-a_{i,\tau})\varphi_{t,i,\tau}}{\sum\limits_{\tau\neq\texttt{opt}_{i}}\varphi_{t,i,\tau}}\leq
1+\nu,$

when $R^{\prime}\geq
10\epsilon^{-1}\log(10\mathrm{T}\max_{i\in[n]}\delta^{\max}_{i}/\epsilon)$.

Case 3: Sample $i\in\mathcal{I}_{3}$, i.e. $\delta^{\min}_{i}\geq 0.8\nu$.

As
$\lVert\widetilde{\bm{W}}_{t}-\bm{W}_{\text{mm}}\rVert\geq\frac{\nu}{2B^{2}}$,
at least one sample $i\in[n]$ violates the (W-SVM) constraints, i.e.
$|\mathcal{I}_{1}|\geq 1$. We have

$\displaystyle\sum\limits_{i\in\mathcal{I}_{1}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})$
$\displaystyle\geq\min_{i\in\mathcal{I}_{1}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\min_{t}(\varphi_{i,\texttt{opt}_{i}}^{t})^{2}\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}/\varphi_{i,\texttt{opt}_{i}}^{t}$
$\displaystyle\geq\min_{i\in\mathcal{I}_{1}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})(4\mathrm{T})^{-2}\exp(-(1+0.3\nu)R^{\prime}),$
(37)

where the last inequality follows by using Lemmas 4 and 5, and that
$\delta_{i}^{\min}\leq 0.3\nu$ for any $i\in\mathcal{I}_{1}$. Also,

$\displaystyle\sum\limits_{i\in\mathcal{I}_{3}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})$
$\displaystyle\leq(n-1)\max_{i\in\mathcal{I}_{3}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\max_{i\in\mathcal{I}_{3}}\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}/\varphi_{i,\texttt{opt}_{i}}^{t}$
$\displaystyle\leq(n-1)\max_{i\in\mathcal{I}_{3}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})(\mathrm{T}-1)\exp(-(1+0.8\nu)R^{\prime}),$
(38)

where the last inequality follows because
$0.8\nu\leq\min_{i\in\mathcal{I}_{3}}\delta_{i}^{\min}$. Combining Eqs.
(A.2.1) and (A.2.1) gives us (36) when

$\displaystyle\exp(0.5\nu R^{\prime})$ $\displaystyle\geq
0.5\nu(2B^{2}\Lambda)^{-1}(n-1)(\mathrm{T}-1)(4\mathrm{T})^{2}\frac{\max_{i\in\mathcal{I}_{3}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})}{\min_{i\in\mathcal{I}_{1}}\ell_{t,i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})},$
(39)

which is true when

$\displaystyle\exp(0.5\nu R^{\prime})\geq
4n(B^{2}\Lambda)^{-1}\Upsilon\mathrm{T}^{3}\epsilon^{-1},$

which is true when $R^{\prime}\geq
2\epsilon^{-1}\log(4n(B^{2}\Lambda)^{-1}\Upsilon\mathrm{T}^{3}\epsilon^{-1})$.
Combining all conditions on $R^{\prime}$, and using $\delta_{i}^{\max}\leq
2B^{2}\Lambda$ for all $i\in[n]$, we get the desired result since

$\displaystyle
R_{\epsilon}\geq\Lambda^{-1}=2\Lambda\epsilon^{-1}(\log(4n(B^{2}\Lambda)^{-1}\Upsilon\mathrm{T}^{3}\epsilon^{-1})\vee
5\log(20\mathrm{T}B^{2}\Lambda\epsilon^{-1})).$

∎

###  A.3 Proof of Theorem 1

We first restate Theorem 1, this time with the exact constants.

######  Theorem 5 (IB Rate).

Under the conditions of Lemma 4 and Assumptions 1 and 2, using the updates in
Eq. (4) with $\eta\leq(2B^{2})^{-1}\log(2\mathrm{T}-1)$, for any $t\geq
t_{0}:=20B^{2}\Lambda^{2}\mathrm{T}(n\Upsilon(2B^{2}\Lambda)^{-1}\vee
5^{4}(2\mathrm{T}^{2})^{-1}B^{2}\Lambda)^{1/3}$,

$\displaystyle\left\langle\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq
1-\frac{C(\eta,B,\Lambda,t_{0})}{t^{3/4}},$

where
$C(\eta,B,\Lambda,t_{0})=4\eta^{-1}B^{2}\Lambda\left(\lVert\bm{W}_{t_{0}}\rVert-\left\langle\bm{W}_{t_{0}},\overline{\bm{W}}_{\text{mm}}\right\rangle+5\eta
t_{0}^{3/4}+3.2\eta B^{2}\Lambda\right)$.

###### Proof.

First, we use Lemma 7 to show that Lemma 2 is true for any $t\geq t_{0}$, by
selecting $\epsilon=\epsilon_{t}:=(t_{0}t^{-1})^{3/4}$. Clearly,
$\epsilon_{t}\leq 1$ for $t\geq t_{0}$. We also have

$\displaystyle
2\Lambda\epsilon^{-1}(\log(4n\Upsilon(B^{2}\Lambda)^{-1}\mathrm{T}^{3}\epsilon^{-1})\vee
5\log(20\mathrm{T}B^{2}\Lambda\epsilon^{-1}))$ $\displaystyle\leq
5\Lambda\mathrm{T}(n\Upsilon(2B^{2}\Lambda)^{-1}\vee
5^{4}(2\mathrm{T}^{2})^{-1}B^{2}\Lambda)^{1/3}(\epsilon^{-1})^{1/3}\epsilon^{-1}$
$\displaystyle=5\Lambda\mathrm{T}(n\Upsilon(2B^{2}\Lambda)^{-1}\vee
5^{4}(2\mathrm{T}^{2})^{-1}B^{2}\Lambda)^{1/3}t_{0}^{-1}t$
$\displaystyle=\eta(4B^{2}\Lambda)^{-1}t.$

Thus, $\lVert\bm{W}_{t}\rVert\geq R$ for $t\geq t_{0}$ and we can use Lemma 2
with $\epsilon=\epsilon_{t}$, i.e.,

$\displaystyle\left\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle$
$\displaystyle\geq(1-\epsilon_{t})\left\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}\right\rangle.$

Using the update Eq. (4), we have

$\displaystyle\left\langle\bm{W}_{t+1}-\bm{W}_{t},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq(1-\epsilon_{t})\left\langle\bm{W}_{t+1}-\bm{W}_{t},\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}\right\rangle$
$\displaystyle=\frac{1}{2\lVert\bm{W}_{t}\rVert}(\lVert\bm{W}_{t+1}\rVert^{2}-\lVert\bm{W}_{t}\rVert^{2}-\lVert\bm{W}_{t+1}-\bm{W}_{t}\rVert^{2})-\epsilon_{t}\left\langle\bm{W}_{t+1}-\bm{W}_{t},\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}\right\rangle$
$\displaystyle=\frac{1}{2\lVert\bm{W}_{t}\rVert}(\lVert\bm{W}_{t+1}\rVert^{2}-\lVert\bm{W}_{t}\rVert^{2})-\frac{\eta^{2}}{2\lVert\bm{W}_{t}\rVert}+\epsilon_{t}\eta\left\langle\overline{\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})},\overline{\bm{W}_{t}}\right\rangle$
(using Eq. (4))
$\displaystyle\geq\lVert\bm{W}_{t+1}\rVert-\lVert\bm{W}_{t}\rVert-\frac{\eta^{2}}{2\lVert\bm{W}_{t}\rVert}-\epsilon_{t}\eta$
(since $\frac{a^{2}-b^{2}}{2b}\geq a-b$ $\forall$ $a,b>0$)
$\displaystyle\geq\lVert\bm{W}_{t+1}\rVert-\lVert\bm{W}_{t}\rVert-\eta(t_{0}t^{-1})^{3/4}-2\eta
B^{2}\Lambda(t+1)^{-1}$ (substituting $\epsilon_{t}$, using Lemma 7).

Summing over $t\geq t_{0}$, we get

$\displaystyle\left\langle\bm{W}_{t},\overline{\bm{W}}_{\text{mm}}\right\rangle\geq\lVert\bm{W}_{t}\rVert-\lVert\bm{W}_{t_{0}}\rVert+\left\langle\bm{W}_{t_{0}},\overline{\bm{W}}_{\text{mm}}\right\rangle-\eta
t_{0}^{3/4}\sum\limits_{t^{\prime}=t_{0}}^{t}(t^{\prime})^{-3/4}-2\eta
B^{2}\Lambda\sum\limits_{t^{\prime}=t_{0}}^{t}(t^{\prime}+1)^{-1}.$ (40)

We bound the last two terms as follows

$\displaystyle\sum\limits_{t^{\prime}=t_{0}}^{t}(t^{\prime})^{-3/4}$
$\displaystyle\leq 4(t+1)^{1/4}\leq 4\,2^{1/4}t^{1/4}\leq 5t^{1/4},$
$\displaystyle\sum\limits_{t^{\prime}=t_{0}}^{t}(t^{\prime}+1)^{-1}$
$\displaystyle\leq\log t\leq 1.6t^{1/4}.$

Using these in Eq. (40) and dividing by $\lVert\bm{W}_{t}\rVert$ throughout,
we get

$\displaystyle\left\langle\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle$
$\displaystyle\geq
1-\frac{1}{\lVert\bm{W}_{t}\rVert}\left(\lVert\bm{W}_{t_{0}}\rVert-\left\langle\bm{W}_{t_{0}},\overline{\bm{W}}_{\text{mm}}\right\rangle+(5\eta
t_{0}^{3/4}+3.2\eta B^{2}\Lambda)t^{1/4}\right)$ $\displaystyle\geq
1-\frac{C(\eta,B,t_{0})}{t^{3/4}},$

where the last step follows by Lemma 7. ∎

###  A.4 Proofs for Additional Results in Section 3.2

######  Lemma 10 (Softmax score rate).

Under the conditions of Lemma 4 and Assumptions 1 and 2, using the updates in
Eq. (4), for any $i\in[n]$ and any

$\displaystyle t\geq
2^{16}\left(18B^{10}\Lambda^{4}(\Lambda\,\vee\,(48\mathrm{T})^{-1})t_{0}\right)^{4/3},$
$\displaystyle\text{ where }t_{0}:=20B^{2}\Lambda^{2}\mathrm{T}(n\Upsilon\vee
5^{4}(2\mathrm{T}^{2})^{-1}B^{2}\Lambda)^{1/3},$
$\displaystyle\varphi_{i,\texttt{opt}_{i}}^{t}\geq$
$\displaystyle\,\,\frac{1}{1+(T-1)\exp(-\eta(8B^{2}\Lambda^{2})^{-1}t)}.$

###### Proof.

Using Theorem 1 and Lemma 7, we have

$\displaystyle({\bm{x}}_{i,\tau}-{\bm{x}}_{i,\texttt{opt}_{i}})^{\top}(\bm{W}_{t}-\overline{\bm{W}}_{\text{mm}}\lVert\bm{W}_{t}\rVert+\overline{\bm{W}}_{\text{mm}}\lVert\bm{W}_{t}\rVert){\bm{x}}_{i,1}$
$\displaystyle\leq
2B^{2}\lVert\bm{W}_{t}\rVert\lVert\overline{\bm{W}_{t}}-\overline{\bm{W}}_{\text{mm}}\rVert-\frac{1}{\lVert\bm{W}_{\text{mm}}\rVert}\lVert\bm{W}_{t}\rVert$
$\displaystyle\leq
2\sqrt{2}B^{2}\frac{\sqrt{C(\eta,B,\Lambda,t_{0})}}{t^{3/8}}(2\eta
t)-\eta(4B^{2}\Lambda^{2})^{-1}t$ $\displaystyle=4\eta
B^{2}\sqrt{2C(\eta,B,\Lambda,t_{0})}t^{5/8}-\eta(4B^{2}\Lambda^{2})^{-1}t$
$\displaystyle\leq-\eta(8B^{2}\Lambda^{2})^{-1}t,$ (41)

since $C(\eta,B,\Lambda,t_{0})\leq
36B^{2}t_{0}(\Lambda\,\vee\,(48\mathrm{T})^{-1})$ and $t\geq
2^{16}\left(18B^{10}\Lambda^{4}(\Lambda\,\vee\,(48\mathrm{T})^{-1})t_{0}\right)^{4/3}$.
We use this to find the softmax rate as follows

$\displaystyle\varphi_{i,\texttt{opt}_{i}}^{t}$
$\displaystyle=\frac{1}{1+\sum_{\tau\neq\texttt{opt}_{i}}\exp(({\bm{x}}_{i,\tau}-{\bm{x}}_{i,\texttt{opt}_{i}})^{\top}\bm{W}_{t}{\bm{x}}_{i,1})}$
$\displaystyle=\frac{1}{1+\sum_{\tau\neq\texttt{opt}_{i}}\exp(({\bm{x}}_{i,\tau}-{\bm{x}}_{i,\texttt{opt}_{i}})^{\top}(\bm{W}_{t}-\overline{\bm{W}}_{\text{mm}}\lVert\bm{W}_{t}\rVert+\overline{\bm{W}}_{\text{mm}}\lVert\bm{W}_{t}\rVert){\bm{x}}_{i,1})}$
$\displaystyle\geq\frac{1}{1+(T-1)\exp(-\eta(8B^{2}\Lambda^{2})^{-1}t)}.$

∎

######  Theorem 6 (Asymptotic Convergence of GD).

Under the conditions of Lemma 4 and Assumptions 1 and 2, using the standard GD
updates with $\eta\leq\frac{\log(2\mathrm{T}-1)}{B^{4}\zeta}$, for any $t\geq
t_{\epsilon}$ such that $\lVert\bm{W}_{t}\rVert\geq R_{\epsilon}\,\vee\,1/2$,

$\displaystyle\left\langle\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq
1-\epsilon-\frac{C(\eta,B,\Lambda,\epsilon)}{\lVert\bm{W}_{t}\rVert},$

where
$C(\eta,B,\Lambda,\epsilon)=\frac{2B^{2}\Lambda}{\eta}(1-\epsilon)\lVert\bm{W}_{t_{\epsilon}}\rVert\left(1-(1-\epsilon)^{-1}\left\langle\overline{\bm{W}}_{t_{\epsilon}},\overline{\bm{W}}_{\text{mm}}\right\rangle-2\eta\lVert\bm{W}_{t_{\epsilon}}\rVert^{-1}\widehat{L}({\bm{\theta}}_{t_{\epsilon}})\right)$.

###### Proof.

First, we show that Lemma 5 holds for standard GD updates, i.e., for any
$t\geq 0$ and any $j\in[n]$,

$\displaystyle\varphi_{j,\texttt{opt}_{j}}^{t}\geq\varphi_{j,\texttt{opt}_{j}}^{0}\,\wedge\,\frac{1}{\mathrm{T}}.$

Following similar steps as the proof of Lemma 5, we consider two scenarios,
Scenario 1: At time $t-1$,
$\varphi_{j,\texttt{opt}_{j}}^{t-1}>1-\frac{1}{2\mathrm{T}}$. This implies

$\displaystyle\varphi_{j,\texttt{opt}_{j}}^{t-1}\leq-\log(2\mathrm{T}-1),\,\,\quad\forall\tau\neq\texttt{opt}_{j}.$

After the gradient step at time $t-1$, for any $\tau\neq\texttt{opt}_{j}$, we
have,

$\displaystyle({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t}{\bm{x}}_{j,1}$
$\displaystyle=({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t-1}{\bm{x}}_{j,1}-\eta({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}(-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t-1})){\bm{x}}_{j,1}$
$\displaystyle\leq-\log(2\mathrm{T}-1)+\eta\lVert{\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}}\rVert\lVert{\bm{x}}_{j,1}\rVert\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t-1})\rVert$
$\displaystyle\leq-\log(2\mathrm{T}-1)+\eta(2B^{2})(B^{2}\widehat{L}(\bm{W}_{t-1})\max_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t}))$
$\displaystyle\leq-\log(2\mathrm{T}-1)+\eta
B^{4}\max_{i}\exp(-\gamma_{i})\max_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})$
$\displaystyle\leq 0,$

where the last step follows since
$\eta\leq\frac{\log(2\mathrm{T}-1)}{B^{4}\zeta}$. This gives
$\varphi_{j,\texttt{opt}_{j}}^{t}\geq\frac{1}{\mathrm{T}}$.

Scenario 2: At time $t-1$, $\varphi_{j,\texttt{opt}_{j}}^{t-1}\leq
1-\frac{1}{2\mathrm{T}}$. The proof for this part is the same as Lemma 5.

Using this, we can easily show that Lemma 2 is true for GD updates, i.e., for
any $\epsilon\in(0,1)$, there exists

$\displaystyle R_{\epsilon}$
$\displaystyle:=2\Lambda\epsilon^{-1}(\log(4n(B^{2}\Lambda)^{-1}\Upsilon\mathrm{T}^{3}\epsilon^{-1})\vee
5\log(20\mathrm{T}B^{2}\Lambda\epsilon^{-1})),$

such that for every $t$ where $\lVert\bm{W}_{t}\rVert\geq R_{\epsilon}$,

$\displaystyle\left\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq(1-\epsilon)\left\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}\right\rangle.$

Once we have this, the remaining proof is the same as Theorem 4 in Tarzanagh
et al., 2023a . ∎

######  Theorem 7 (IB rate for Polyak-step).

Under the conditions of Lemma 4 and Assumptions 1 and 2, using GD updates with
Polyak-step, with $\eta\leq(2B^{2})^{-1}\omega_{1}\log(2\mathrm{T}-1)$, where
$\omega_{1}=\beta(1-\beta)\Lambda^{-1}\min_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})$,
for any $T_{0}\geq t\geq
t_{0}:=20B^{2}\Lambda^{2}\omega_{2}\upsilon\mathrm{T}(n\Upsilon(2B^{2}\Lambda)^{-1}\vee
5^{4}(2\mathrm{T}^{2})^{-1}B^{2}\Lambda)^{1/3}$, where
$\omega_{2}=0.5B^{2}\max_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})$, such
that for any $i\in[n]$, $\varphi_{i,\texttt{opt}_{i}}^{t}\leq\beta$,

$\displaystyle\left\langle\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert},\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq
1-\frac{C(\eta,B,\Lambda,t_{0})}{t^{3/4}},$

where
$C(\eta,B,\Lambda,t_{0})=4\eta^{-1}B^{2}\Lambda\left(\lVert\bm{W}_{t_{0}}\rVert-\left\langle\bm{W}_{t_{0}},\overline{\bm{W}}_{\text{mm}}\right\rangle+5\eta\omega_{1}^{-1}t_{0}^{3/4}+3.2\eta\upsilon
B^{2}\Lambda\omega_{1}^{-2}\omega_{2}\right)$.

###### Proof.

To extend our analysis and results for NGD updates to the Polyak-step, we
mainly use the following lower bound on the gradient norm,

$\displaystyle\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})\rVert$
$\displaystyle\geq\widehat{L}(\bm{W}_{t})\min_{i}\left\langle
y_{i}\nabla_{\bm{W}}\Phi({\bm{\theta}}_{t},{\bm{X}}_{i}),\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle$
$\displaystyle=\Lambda^{-1}\widehat{L}(\bm{W}_{t})\min_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\varphi_{i,\texttt{opt}_{i}}^{t}(1-\varphi_{i,\texttt{opt}_{i}}^{t})\left[a^{*}_{i,\texttt{opt}_{i}}-\frac{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}a^{*}_{i,\tau}}{\sum_{\tau\neq\texttt{opt}_{i}}\varphi_{i,\tau}^{t}}\right]$
$\displaystyle\geq\Lambda^{-1}\widehat{L}(\bm{W}_{t})\min_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\min_{i,t^{\prime}\leq
t}\varphi_{i,\texttt{opt}_{i}}^{t^{\prime}}(1-\varphi_{i,\texttt{opt}_{i}}^{t^{\prime}})$
$\displaystyle\geq\beta(1-\beta)\Lambda^{-1}\min_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})\widehat{L}(\bm{W}_{t})=\omega_{1}\widehat{L}(\bm{W}_{t}).$

Note that similar lower bounds have appeared when deriving loss convergence
rates of linear predictors trained with NGD on separable data (Nacson et al.,,
2019).

We also have the following upper bound,

$\displaystyle\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})\rVert$
$\displaystyle\leq\widehat{L}(\bm{W}_{t})(1/4)\max_{i}(\gamma_{i,\texttt{opt}_{i}}-\gamma_{i})(2B^{2})=\omega_{2}\widehat{L}(\bm{W}_{t}).$

First, we show that Lemma 5 holds for Polyak-step updates, i.e., for any
$t\geq 0$ and any $j\in[n]$,

$\displaystyle\varphi_{j,\texttt{opt}_{j}}^{t}\geq\varphi_{j,\texttt{opt}_{j}}^{0}\,\wedge\,\frac{1}{\mathrm{T}}.$

Following similar steps as the proof of Lemma 5, we consider two scenarios,
Scenario 1: At time $t-1$,
$\varphi_{j,\texttt{opt}_{j}}^{t-1}>1-\frac{1}{2\mathrm{T}}$. This implies

$\displaystyle\varphi_{j,\texttt{opt}_{j}}^{t-1}\leq-\log(2\mathrm{T}-1),\,\,\quad\forall\tau\neq\texttt{opt}_{j}.$

After the gradient step at time $t-1$, for any $\tau\neq\texttt{opt}_{j}$, we
have,

$\displaystyle({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t}{\bm{x}}_{j,1}$
$\displaystyle=({\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}})^{\top}\bm{W}_{t-1}{\bm{x}}_{j,1}-\eta_{t-1}({\bm{x}}_{j,\texttt{opt}_{j}}-{\bm{x}}_{j,\tau})^{\top}(-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t-1})){\bm{x}}_{j,1}$
$\displaystyle\leq-\log(2\mathrm{T}-1)+\eta\frac{\widehat{L}(\bm{W}_{t-1})-\widehat{L}^{*}}{\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t-1})\rVert^{2}}\lVert{\bm{x}}_{j,\tau}-{\bm{x}}_{j,\texttt{opt}_{j}}\rVert\lVert{\bm{x}}_{j,1}\rVert\lVert\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t-1})\rVert$
$\displaystyle\leq-\log(2\mathrm{T}-1)+\eta\frac{\widehat{L}(\bm{W}_{t-1})}{\omega_{1}\widehat{L}(\bm{W}_{t-1})}(2B^{2})$
$\displaystyle=-\log(2\mathrm{T}-1)+\eta\frac{2B^{2}}{\omega_{1}}$
$\displaystyle\leq 0,$

where the last step follows since
$\eta\leq\frac{\omega_{1}\log(2\mathrm{T}-1)}{2B^{2}}$. This gives
$\varphi_{j,\texttt{opt}_{j}}^{t}\geq\frac{1}{\mathrm{T}}$.

Scenario 2: At time $t-1$, $\varphi_{j,\texttt{opt}_{j}}^{t-1}\leq
1-\frac{1}{2\mathrm{T}}$. The proof for this part is the same as Lemma 5.

Using this, we can easily show that Lemma 2 is true for Polyak-step updates,
i.e., for any $\epsilon\in(0,1)$, there exists

$\displaystyle R_{\epsilon}$
$\displaystyle:=2\Lambda\epsilon^{-1}(\log(4n(B^{2}\Lambda)^{-1}\Upsilon\mathrm{T}^{3}\epsilon^{-1})\vee
5\log(20\mathrm{T}B^{2}\Lambda\epsilon^{-1})),$

such that for every $t$ where $\lVert\bm{W}_{t}\rVert\geq R_{\epsilon}$,

$\displaystyle\left\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\frac{\bm{W}_{\text{mm}}}{\lVert\bm{W}_{\text{mm}}\rVert}\right\rangle\geq(1-\epsilon)\left\langle-\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t}),\frac{\bm{W}_{t}}{\lVert\bm{W}_{t}\rVert}\right\rangle.$
(42)

Next, following similar steps as the proof of Lemma 7, the iterate norm growth
can be characterized as follows. Using Lemma 6 with $\bm{W}=\bm{W}_{t}$, for
any $t\leq T_{0}$,

$\displaystyle\lVert\bm{W}_{t}\rVert$
$\displaystyle=\left\|\bm{W}_{0}-\sum\limits_{t^{\prime}=0}^{t-1}\eta_{t^{\prime}}\nabla_{\bm{W}}\widehat{L}(\bm{W}_{t})\right\|$