# Mechanics of Next Token Prediction with Self-Attention

Yingcong Li 1 Yixiao Huang∗ 1 M. Emrullah Ildiz1
Ankit Singh Rawat2 Samet Oymak1 Equal contribution. 1 University of Michigan,
Ann Arbor
<EMAIL_ADDRESS><EMAIL_ADDRESS>
###### Abstract

Transformer-based language models are trained on large datasets to predict the
next token given an input sequence. Despite this simple training objective,
they have led to revolutionary advances in natural language processing.
Underlying this success is the self-attention mechanism. In this work, we ask:
_What does a single self-attention layer learn from next-token prediction?_ We
show that training self-attention with gradient descent learns an automaton
which generates the next token in two distinct steps: (1) Hard retrieval:
Given input sequence, self-attention precisely selects the _high-priority
input tokens_ associated with the last input token. (2) Soft composition: It
then creates a convex combination of the high-priority tokens from which the
next token can be sampled. Under suitable conditions, we rigorously
characterize these mechanics through a directed graph over tokens extracted
from the training data. We prove that gradient descent implicitly discovers
the strongly-connected components (SCC) of this graph and self-attention
learns to retrieve the tokens that belong to the highest-priority SCC
available in the context window. Our theory relies on decomposing the model
weights into a directional component and a finite component that correspond to
hard retrieval and soft composition steps respectively. This also formalizes a
related implicit bias formula conjectured in [Tarzanagh et al. 2023]. We hope
that these findings shed light on how self-attention processes sequential data
and pave the path toward demystifying more complex architectures.

Figure 1: Overview of our result on next-token prediction. We study the
implicit bias of gradient descent where a 1-layer self-attention model is
trained until convergence. We prove that, during test-time, this model
implements a hard retrieval to precisely select the high-priority tokens and
then outputs a convex combination of these as the output from which the next
token can be sampled. The notion of _high-priority_ is formalized through the
strongly-connected components of a directed graph associated to the last input
token.

##  1 Introduction

Language modeling as enabled by Transformer architecture (Vaswani et al.,
2017) and seemingly simple training objectives such as next-token prediction
(Radford et al., 2018, 2019) have not only led to breakthroughs in the field
of natural language processing (NLP) (Brown et al., 2020; Chowdhery et al.,
2022; OpenAI, 2023; Touvron et al., 2023), but rather straightforward
adaptations of this symbiosis between Transformers and next-token prediction
tasks have also realized remarkable performance in other domains, including
vision (Chen et al., 2020), speech (Chung & Glass, 2020), reinforcement
learning (Chen et al., 2021), and even protein design (Ferruz et al., 2022;
Nijkamp et al., 2022). This widespread empirical success is often attributed
to the (self-)attention mechanism of Transformers that produces high-quality
contextual representations needed to realize excellent prediction performance
in a wide range of domains. However, a rigorous understanding of how
Transformers can learn such high-quality representations by solving next-token
prediction task via natural algorithms such as gradient descent is largely
missing from the literature.

This work aims to bridge this gap between the empirical success and principled
understanding of Transformer-based language modeling by shedding light on the
optimization landscape and key implicit biases faced by the self-attention
mechanism in solving the next-token prediction task. In particular, focusing
on a single-layer self-attention model with linear classification head, and
solving the next-token prediction task, we consider the following questions:

  * •

What relationships in the training data are captured by the single-layer self-
attention model?

  * •

How exactly do these relationships dictate the optimization geometry of
natural algorithms such as gradient descent?

We show that the answers to both of these questions are intertwined which we
achieve by significantly expanding the recently proposed framework that
connects learning with Transformers to the celebrated support vector machines
(SVMs) (Tarzanagh et al., 2023b, a).

As illustrated in Figure 1, given training data as a collection of (input
sequence, next token) pairs, self-attention model learns to (1) retrieve the
high-priority tokens (highlighted with red color) to the last input token; and
then (2) build a convex combination of these high-priority tokens. The notion
of _high-priority_ is dictated by a directed graph learned from training data.
SGD training accomplishes this by learning hard and soft components of the
attention weights $\bm{W}$ to execute (1) and (2) respectively. Concretely,
the following theorem dictates the evolution of attention weights during
gradient descent.

######  Theorem 1 (_informal_)

Consider training a single-layer self-attention model with gradient descent.
The combined attention weights $\bm{W}:=\bm{W}_{K}\bm{W}_{Q}^{\top}$ evolve as

$\bm{W}_{\text{GD}}\approx C\cdot\bm{W}_{\text{hard}}+\bm{W}_{\text{soft}},$

where $C\cdot\bm{W}_{\text{hard}}$ is the hard retrieval component selecting
the high-priority tokens when $C\to\infty$; and $\bm{W}_{\text{soft}}$ is the
soft composition component allocating nonzero softmax probabilities over
selected tokens.

To capture the priority order among different tokens as observed in Figure 1,
we construct directed graphs among the tokens in the vocabulary, namely token-
priority graphs (TPGs). An illustration is provided in Figure 2, where a
strongly connected component (SCC, highlighted as dashed black rectangles) in
a TPG corresponds to the tokens that are reachable from each other, indicating
the absence of a strict priority among those tokens. The hard retrieval
component $\bm{W}_{\text{hard}}$ captures the topological order of different
SCCs (orange arrows) whereas the soft composition component
$\bm{W}_{\text{soft}}$ captures the relationships of different tokens within
each SCC (black arrows). These TPGs will geometrically capture the learning
dynamics of self-attention. Specifically, we propose the SVM problem (Graph-
SVM), solution of which describes the direction gradient descent converges to.
This way, SGD asymptotically enforces the topological order between SCCs i.e.
the $C\cdot\bm{W}_{\text{hard}}$ term in Theorem 1 as $C\rightarrow\infty$. In
practice, this implies that self-attention model favors suppressing lower
priority tokens in favour of sampling higher priority tokens.

A conjecture on the decomposition in Theorem 1 was first proposed in
(Tarzanagh et al., 2023a)***Their conjecture aims to characterize the impact
of the MLP layer that follow self-attention in a binary classification
setting. However, the high-level claim is same.. This decomposition is also
related to the implicit bias of logistic regression on non-separable data (Ji
& Telgarsky, 2019a). Our theory fully formalizes this decomposition under the
next-token prediction setting and reveals fundamental connections to graphical
structure in data (e.g. through SCCs, TPGs).

Figure 2: A token-priority graph (TPG) is a directed graph derived from
training data (see Sec 2.1 for definition). The edges in TPG capture the
input-output relationships between different tokens. A TPG can be partitioned
into several SCCs depicted as dashed black squares. In light of Theorem 1,
black intra-SCC edges within each SCC induce the soft-composition component of
the attention weights whereas the orange edges induce the hard-retrieval
component enforcing the priority orders among various SCCs.

Overall, we carefully study the gradient descent and regularization path
algorithms for attention-based next-token prediction and make the following
contributions:

  1. 1.

We study the optimization landscape of self-attention with log-loss and show
that the problem is convex under suitable assumptions. We then establish a
global convergence result to fully formalize Theorem 1 in terms of a
directional component (Graph-SVM) and a finite component (see Sec 3). Notably,
results apply to arbitrary datasets as we don’t require distributional
assumptions.

  2. 2.

Our theory reveals insightful connections between continuous and discrete
optimization, namely: _Self-attention implicitly discovers the strongly-
connected components of the TPGs during training_.

  3. 3.

Under a general setting, we establish the implicit bias of the solution
obtained by vanishing regularization (Rosset et al., 2003; Suggala et al.,
2018; Ji et al., 2020). This yields a result similar to the gradient descent
theory (see Sec 4). We also show that, in general, gradient descent can
exhibit local directional convergence rather than global. We characterize
these local directions through the SVM solutions of _pseudo TPGs_ (see Sec 5).

##  2 Problem Setup

_Notation._ Let $[n]$ denote the set $\\{1,\cdots,n\\}$. For a space
$\mathcal{S}$, let $\mathcal{S}^{\perp}$ denote the orthogonal complement of
$\mathcal{S}$ and $\bm{\Pi}_{\mathcal{S}}$ denote the projection operator on
$\mathcal{S}$ with respect to Euclidean distance.

Next-token prediction problem. Let $K$ be the vocabulary size with
${\bm{E}}=\begin{bmatrix}\bm{e}_{1}~{}\dots~{}\bm{e}_{K}\end{bmatrix}^{\top}\in\mathbb{R}^{K\times
d}$ denoting the embedding matrix consisting of $d$-dimensional token
embeddings for the $K$ tokens in the vocabulary. The next-token prediction is
a multi-class classification problem and the goal is to predict the ID
$y\in[K]$ of the next token given an input sequence
${\bm{X}}=[\bm{x}_{1}~{}\dots~{}\bm{x}_{T}]^{\top}\in\mathbb{R}^{T\times d}$,
where $\bm{x}_{t}\in{\bm{E}}$ for all $t\in[T]$.

Suppose that we have a training dataset
$\texttt{DSET}=\big{\\{}({\bm{X}}_{i},y_{i})\in\mathbb{R}^{T_{i}\times
d}\times[K]\big{\\}}_{i=1}^{n}$ consisting of $n$ sequences where we allow the
sequences to have different lengths $T_{i}$, $i\in[n]$. Throughout this paper,
we use $x_{it}\in[K]$ to denote the scalar token ID corresponding to the
$t$-th token $\bm{x}_{it}\in\mathbb{R}^{d}$ of the input sequence
${\bm{X}}_{i}$, i.e., $\bm{x}_{it}=\bm{e}_{x_{it}}$.

Self-attention model. We consider a single-layer self-attention model when
making a prediction on a given input sequence ${\bm{X}}\in\mathbb{R}^{T\times
d}$. Following the previous work (Tarzanagh et al., 2023a), we denote the
combined key-query weights by a trainable $\bm{W}\in\mathbb{R}^{d\times d}$
matrix, and assume identity value matrix. Let $\bm{\bar{x}}:=\bm{x}_{T}$ be
the last token of the input sequence ${\bm{X}}$. Then, the single-layer self-
attention outputs the following embedding to predict the next-token ID $y$:

$\displaystyle
f_{\bm{W}}({\bm{X}})={\bm{X}}^{\top}\mathbb{S}({\bm{X}}\bm{W}\bm{\bar{x}}),$
(1)

where $\mathbb{S}(\cdot)$ denotes the softmax operation which facilitates
weighing tokens of ${\bm{X}}$ based on the data-dependent probabilities
$\mathbb{S}({\bm{X}}\bm{W}\bm{\bar{x}})$. Note that the output embedding
$f_{\bm{W}}({\bm{X}})\in\mathbb{R}^{d}$ in (1) is a weighted linear
combination of the input token embeddings in ${\bm{X}}$.

Empirical risk minimization (ERM) problem. Let $\ell:\mathbb{R}\to\mathbb{R}$
be the loss function. Given training dataset DSET, we consider the ERM problem
with the following objective:

$\displaystyle{\cal{L}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\ell(\bm{c}_{y_{i}}^{\top}{\bm{X}}_{i}^{\top}\mathbb{S}({\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i})).$
(ERM)

Throughout this paper, we fix the linear classification head
$\bm{c}_{k}$†††Specifically, we assume well-pretrained head
$\bm{c}_{1},\cdots,\bm{c}_{K}$ such that $\ell(\bm{c}_{y}^{\top}\bm{e}_{k})$
returns the minimal risk when $k=y$.  and assume $\|{\bm{c}_{k}}\|$ is bounded
for all $k\in[K]$. Note that even though the classification head is linear,
the problem of learning attention parameters $\bm{W}$ via ERM is not
necessarily convex due to the softmax operator. In this work, we focus on this
exact problem and consider the following two algorithms to optimze for
$\bm{W}$:

1\. Gradient descent: Given starting point $\bm{W}(0)\in\mathbb{R}^{d\times
d}$ and step size $\eta>0$, for $\tau\geq 0$,
$\displaystyle\bm{W}(\tau+1)=\bm{W}(\tau)-\eta\nabla{\cal{L}}(\bm{W}(\tau)).$
(Algo-GD) 2\. Regularization path: Given $R>0$, $\bm{W}\in\mathbb{R}^{d\times
d}$, $\displaystyle\bar{\bm{W}}_{R}=\arg\min_{\left\|{\bm{W}}\right\|_{F}\leq
R}{\cal{L}}(\bm{W}).$ (Algo-RP)

The next-token prediction task aims to capture various patterns present in the
underlying dataset. Towards this, we introduce token-priority graph (TPG) in
the following section that summarizes the sequential priority orders presented
in the training data. As we will see later, TPGs play a crucial role in
characterizing the optimization geometry for both (Algo-GD) and (Algo-RP)
algorithms.

Figure 3: Illustration of token-priority graph (TPG). Given the input
sequences and labels (next tokens), we construct the TPGs
$\\{{\cal{G}}^{(k)}\\}_{k=1}^{K}$ according to the last token. Two TPGs
${\cal{G}}^{(1)}$ (left) and ${\cal{G}}^{(2)}$ (right) are constructed using
the samples with $\bm{e}_{1}$ and $\bm{e}_{2}$ as the last tokens,
respectively. In each graph, directed edges (label token$\to$input token) are
added between tokens/nodes. Based on these directed edges, each graph can be
partitioned into its strongly-connected components (SCCs, highlighted as
dashed grey rectangles). Each SCC is a set of tokens where each token is
reachable from every other token within that SCC. Further details are deferred
to Section 2.1.

###  2.1 Token-priority Graph of the Dataset

A token-priority graph (TPG) is a directed graph with at most $K$ nodes
corresponding to the elements in the vocabulary. We associate the dataset
$\texttt{DSET}=\\{({\bm{X}}_{i},y_{i})\\}_{i=1}^{n}$ with multiple TPGs
$\\{{\cal{G}}^{(k)}\\}_{k=1}^{K}$, with each TPG focusing on a subset of the
dataset comprising of those input sequences that agree on the last token
$\bm{\bar{x}}$. Concretely, we construct ${\cal{G}}^{(k)}$’s as follows:

  1. 1.

Split DSET into $K$ subsets $\\{\texttt{DSET}^{(k)}\\}_{k=1}^{K}$ with
$\texttt{DSET}^{(k)}$ containing all input sequences that end with the same
last token $\bm{\bar{x}}=\bm{e}_{k}$.

  2. 2.

For each $({\bm{X}},y)\in\texttt{DSET}^{(k)}$ and for all $(x,y)$ pairs in
$({\bm{X}},y)$ where $x$ is the corresponding token ID of $\bm{x}\in{\bm{X}}$,
add a directed edge $(y\rightarrow x)$ to ${\cal{G}}^{(k)}$.

An illustration is provided in Fig. 3, where we construct two TPGs
(${\cal{G}}^{(1)}$ and ${\cal{G}}^{(2)}$) based on the last tokens (depicted
in yellow), and the directed edges are presented as arrows starting from
labels (orange) to input tokens (blue/yellow) within each sequence. Note that
nodes of each ${\cal{G}}^{(k)}$ constitute a subset of the indices $[K]$. The
edges in ${\cal{G}}^{(k)}$ capture the priorities across the tokens in an
extended data sequence, conditioned on the last token of the input being
$\bm{\bar{x}}=\bm{e}_{k}$. We will see that if there is a cycle, i.e.,
$y\rightarrow x$ and $x\rightarrow y$ are both directionally reachable in the
graph, then the self-attention learnt via next-token prediction task can
assign comparable priorities to the tokens $x$ and $y$. In contrast, if $y$
always _dominates_ $x$, i.e., $y\to x$ is reachable but $x\not\to y$, then,
when $x$ and $y$ are both present in an input sequence, self-attention will be
learnt to suppress $x$ and select $y$ through an SVM mechanism along the line
of Tarzanagh et al. (2023a).

Strongly-connected components in TPGs.To formalize the aforementioned SVM
mechanism, we need the notion of _strongly-connected components_ (SCCs). A
directed graph is strongly connected if every node in the graph is reachable
from every other node. SCCs of a directed graph form a partition into
subgraphs that are themselves strongly connected. Given the TPGs
$\\{{\cal{G}}^{(k)}\\}_{k=1}^{K}$ associated with the dataset DSET, we can
split the directed graph ${\cal{G}}^{(k)}$ into its SCCs, denoted by
$\\{\mathcal{C}^{(k)}_{i}\\}_{i=1}^{N_{k}}$. Note that the number of SCCs in
${\cal{G}}^{(k)}$, as denoted by $N_{k}$, is at most the number of nodes in
${\cal{G}}^{(k)}$, which is upper bounded by the vocabulary size $K$.
Furthermore, by definition, different SCCs within a graph consist of distinct
nodes, i.e., $\mathcal{C}^{(k)}_{i}\bigcap\mathcal{C}^{(k)}_{j}=\emptyset$,
for $i\neq j$. Now, returning to Fig. 3, each of the dashed grey rectangle
represents an SCC. ${\cal{G}}^{(1)}$(left) contains four SCCs and therefore,
all tokens within the graph have strict priority orders. In contrast,
${\cal{G}}^{(2)}$(right) consists of two SCCs, with one containing three
nodes. Following the arrows, we can see that all the tokens/nodes within this
specific SCC are directional reachable.

Before formally connecting TPGs and their SCCs to the SVM mechanism that
enables next-token prediction, we introduce some necessary graph-related
notation. Given a directed graph ${\cal{G}}$, for $i,j\in[K]$ such that $i\neq
j$:

  * •

$i\in{\cal{G}}$ denotes that the node $i$ belongs to ${\cal{G}}$.

  * •

$(i\Rightarrow j)\in{\cal{G}}$ denotes that the directed edge $(i\rightarrow
j)$ is present in ${\cal{G}}$ but $j\rightarrow i$ is not reachable.

  * •

$(i\asymp j)\in{\cal{G}}$ means that both two nodes $(i,j)$ are in the same
SCC of ${\cal{G}}$.

From the construction, for any two distinct nodes $i,j$ in the same TPG, they
either satisfy $(i\Rightarrow j)$/$(j\Rightarrow i)\text{ or }(i\asymp j)$.

###  2.2 SVM Bias of Self-attention Learning

The main contribution of this paper is to establish the SVM equivalence that
captures the optimization geometry of the next-token prediction problem. We
will show that the self-attention model learnt via either (Algo-GD) or (Algo-
RP) converges to the solution of an SVM defined by the TPGs of the underlying
dataset DSET. In particular, given $({\cal{G}}^{(k)}_{k})_{k=1}^{K}$, we
introduce the following SVM formulation:

$\displaystyle\bm{W}^{\text{svm}}=\arg\min_{\bm{W}}\left\|{\bm{W}}\right\|_{F}\quad$
(Graph-SVM)
$\displaystyle\text{s.t.}~{}~{}(\bm{e}_{i}-\bm{e}_{j})^{\top}\bm{W}\bm{e}_{k}\begin{cases}=0\quad\forall(i\asymp
j)\in{\cal{G}}^{(k)}\\\ \geq 1\quad\forall(i\Rightarrow
j)\in{\cal{G}}^{(k)}\end{cases}\text{for all}\quad k\in[K].$

Fix last token $\bm{e}_{k}$, and consider any token IDs $i,j\in[K]$, $i\neq
j$. When $(i\Rightarrow j)$, token ID $i$ has a higher _priority_ than $j$ and
hence, the SVM problem (Graph-SVM) aims to find a $\bm{W}$ such that
$\bm{W}\bm{e}_{k}$ achieves strictly higher correlation to token embedding
$\bm{e}_{i}$ than $\bm{e}_{j}$, that is,
$\bm{e}_{i}^{\top}\bm{W}\bm{e}_{k}\geq\bm{e}_{j}^{\top}\bm{W}\bm{e}_{k}+1$,
and then softmax operation will assign higher probability to the token $i$.
While if $(i\asymp j)$, there is not strict priority order between $i$ and
$j$, and hence we set the correlation difference equal to zero to prevent the
SVM solution $\bm{W}$ from distinguishing them. The existence of the solution
$\bm{W}^{\text{svm}}$ ensures the separability of tokens $i$’s from the $j$’s
for all pairs $(i\Rightarrow j)\in{\cal{G}}^{(k)}$. Additionally, if for all
$k\in[K]$, the number of SCCs‡‡‡Note that $N_{k}=0$ implies that within DSET,
there is not training sample whose input sequence has $\bm{e}_{k}$ as its last
token; or equivalently, $\texttt{DSET}^{(k)}=\emptyset$. $N_{k}\leq 1$, then
$\bm{W}^{\text{svm}}=0$.

###### Lemma 1

Suppose that the embedding matrix ${\bm{E}}$ is full row rank. Then, (Graph-
SVM) is feasible.

Next focusing on the nodes $i$, $j$, with $(i\asymp j)$, we introduce the
following subspace definition.

######  Definition 1 (Cyclic subspace)

Define cyclic subspace $\mathcal{S}_{\text{fin}}$ as the span of all matrices
$(\bm{e}_{i}-\bm{e}_{j})\bm{e}_{k}^{\top}$ for all $(i\asymp
j)\in{\cal{G}}^{(k)}$ and $k\in[K]$.

Note that since $\bm{W}^{\text{svm}}$ satisfies all the “$=0$” constraints in
(Graph-SVM), if (Graph-SVM) is feasible and $\bm{W}^{\text{svm}}\neq 0$,
$\bm{W}^{\text{svm}}\perp\mathcal{S}_{\text{fin}}$.

###  2.3 Technical Assumptions

In what follows, we work with a few assumptions that will make the
optimization landscape of the underlying learning problem more benign, and we
introduce these assumptions along with their justifications.

###### Assumption 1

For $\forall y,k\in[K]$, $k\neq y$, $\bm{c}_{y}^{\top}\bm{e}_{y}=1$ and
$\bm{c}_{y}^{\top}\bm{e}_{k}=0$.

This assumption essentially enforces that the rows of the prediction head
${\bm{C}}$ are aligned with the corresponding vocabulary embeddings in
${\bm{E}}$. This is a variation of the _weight tying_ strategy which is
commonly employed in language models (Press & Wolf, 2017; Vaswani et al.,
2017). It should be noted that Assumption 1 implies $K\leq d$, which further
establishes the feasibility of (Graph-SVM) as demonstrated by Lemma 1. Given
objective (ERM) and a decreasing loss function $\ell$, our ideal goal is for
the attention (1) to output $\bm{e}_{y}$ which minimizes the training risk.
Recall that single-layer self-attention outputs a convex combination of the
input tokens (cf. (1)). If tokens are linearly independent, the only way model
can output the embedding $\bm{e}_{y}$ corresponding to the target label $y$
would be if $\bm{e}_{y}$ was among the input sequence. This motivates the
following realizability assumption.

###### Assumption 2

For any $({\bm{X}},y)\in\texttt{DSET}$, the token $\bm{e}_{y}$ is contained in
the input sequence ${\bm{X}}$.

In the scenario where $({\bm{X}},y)$ is not realizable, self-attention would
select $\bm{e}_{\hat{y}}\neq\bm{e}_{y}$ and the SVM formula would be
established via separating $\hat{y}$ from the other tokens in the sequence
instead of the true label $y$. Additionally, when Assumption 1 holds, the
model can only make a random prediction over the output labels (since any
output of the self-attention model would result in the same training risk);
consequently, such non-realizable examples will not play roles in optimizing
$\bm{W}$, i.e.
$\nabla_{\bm{W}}\ell(\bm{c}_{y_{i}}^{\top}{\bm{X}}_{i}^{\top}\mathbb{S}({\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i}))=0$.

##  3 Global Convergence of Gradient Descent

In this section, we assume the log-loss function, i.e., $\ell(u)=-\log(u)$,
and establish the gradient descent convergence of attention weight $\bm{W}$
via the convexity of ${\cal{L}}(\bm{W})$. Note that although loss function
$\ell$ is convex and the classification head is linear, due to the non-
convexity of softmax, the convexity of ${\cal{L}}(\bm{W})$ is not immediately
clear. Towards this, we introduce the following lemma:

###### Lemma 2

Suppose Assumptions 1 and 2 hold and consider the log-loss $\ell(u)=-\log(u)$,
then ${\cal{L}}(\bm{W})$ is convex. Furthermore, ${\cal{L}}(\bm{W})$ is
strictly convex on $\mathcal{S}_{\text{fin}}$.

Let $({\bm{X}},y)\in\texttt{DSET}$ be any sample and set
$\gamma_{t}=\bm{c}_{y}^{\top}\bm{x}_{t}$. Assumption 1 guarantees that
$\gamma_{t}=1$ when $x_{t}=y$, otherwise $\gamma_{t}=0$. Consider the
attention output
$\bm{c}_{y}^{\top}{\bm{X}}^{\top}\mathbb{S}({\bm{X}}\bm{W}\bm{\bar{x}})$ (cf.
(ERM)) and let softmax probabilities be
$s_{t}=\mathbb{S}({\bm{X}}\bm{W}\bm{\bar{x}})_{t}$, where $\sum_{t}s_{t}=1$.
Then, the loss of this single sample ${\bm{X}}$ is $\ell(\bar{\gamma})$ where
$\bar{\gamma}=\sum_{t}\gamma_{t}s_{t}=\sum_{x_{t}=y}s_{t}$. Given log-loss,
note that when $\bar{\gamma}\to 0^{+}$, $-\log(\bar{\gamma})$ results in the
infinite loss, which suggests that, once attention weight $\bm{W}$ diverges to
saturate the softmax probability, the finite training risk is achievable only
when $s_{t}\not\to 0$ for all $t$ satisfying $x_{t}=y$. Given different label
$y$ for different input and recalling the SCC definition in Section 2.1,
attention selects all $\bm{x}_{t}$’s within the same SCC as $y$. The following
result characterizes the global directional convergence of the GD iterates to
the solution of (Graph-SVM).

###### Theorem 2

Consider a dataset DSET and suppose Assumptions 1 and 2 hold. Set loss
function as $\ell(u)=-\log(u)$. Let
$\bm{W}^{\text{svm}}{\in\mathcal{S}_{\text{fin}}^{\perp}}$ be the solution of
(Graph-SVM). Starting from any $\bm{W}(0)$ with constant step size $\eta$, the
algorithm Algo-GD satisfies
$\lim_{\tau\to\infty}\left\|{\bm{W}(\tau)}\right\|_{F}=\infty$ and

$\displaystyle\qquad\lim_{\tau\rightarrow\infty}\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}(\tau))=\bm{W}^{\text{fin}}.$
(2)

Here $\bm{W}^{\text{fin}}$ is the unique finite minima of the loss
$\tilde{{\cal{L}}}(\bm{W}):=\lim_{R\rightarrow\infty}{\cal{L}}(\bm{W}+R\cdot\bm{W}^{\text{svm}})$
over $\mathcal{S}_{\text{fin}}$. Additionally, if $\bm{W}^{\text{svm}}\neq 0$,

$\lim_{\tau\rightarrow\infty}\frac{\bm{W}(\tau)}{\left\|{\bm{W}(\tau)}\right\|_{F}}=\frac{\bm{W}^{\text{svm}}}{\left\|{\bm{W}^{\text{svm}}}\right\|_{F}}.$

Otherwise, $\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\bm{W}(\tau))$ remains
unchanged throughout the optimization.

This theorem demonstrates the directional convergence of attention weight
$\bm{W}$, and the limits imply the decomposition $\bm{W}(\tau)\approx
C(\tau)\cdot\bm{W}^{\text{svm}}+\bm{W}^{\text{fin}}$ for an appropriate
$C(\tau)>0$ with $\lim_{\tau\to\infty}C(\tau)=\infty$. Note that, for
directional convergence to happen, we need $\bm{W}^{\text{svm}}\neq 0$ which
happens if and only if (Graph-SVM) has $``\geq 1^{\prime\prime}$ constraints.
That is, the token graph contains a strict priority order. This is consistent
with our Theorem 1 where $\bm{W}^{\text{svm}}$ corresponds to the hard
retrieval component ($\bm{W}_{\text{hard}}$) that selects the high-priority
tokens and $\bm{W}^{\text{fin}}$ is the soft composition component
($\bm{W}_{\text{soft}}$) that determines the softmax probability assignments
among these selected tokens. Importantly, this is a global convergence result
thanks to the convexity of the optimization problem, which is enabled by the
log likelihood optimization combined with Assumption 1. In Section 5, we will
demonstrate that global convergence of gradient descent does not hold in
general.

Iterations ($\tau$)Correlation coefficient (a) Evolution of
$\frac{\bm{W}(\tau)}{\left\|{\bm{W}(\tau)}\right\|_{F}}\to\frac{\bm{W}^{\text{svm}}}{\left\|{\bm{W}^{\text{svm}}}\right\|_{F}}$

Iterations
$(\tau)$$\left\|{\widehat{\bm{W}}^{\text{fin}}-\bm{W}^{\text{fin}}}\right\|_{F}$
(b) Evolution of
$\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}(\tau))\to\bm{W}^{\text{fin}}$

Figure 4: GD convergence of attention weight $\bm{W}$ when training with
general dataset. (a) shows the directional convergence of $\bm{W}(\tau)$;
while (b) presents the convergence of
$\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}(\tau))$.

Iterations ($\tau$)Correlation coefficient Figure 5: GD convergence of
attention weight $\bm{W}$ when training with acyclic dataset (Def. 2).
Correlation coefficient between $\bm{W}(\tau)$ and $\bm{W}^{\text{svm}}$ are
presented.

To illustrate Theorem 2, we conduct experiments with results presented in
Figure 5. We create embedding tables with $K=6,~{}d=8$ and randomly generate
dataset with $n=6,~{}T=4$. Here we choose step size $\eta=0.01$ and perform
the normalized gradient descent method to accelerate the increase in the norm
of attention weight, so that softmax can easily saturate. Specifically, we
update attention weight $\bm{W}$ via
$\bm{W}(\tau+1)=\bm{W}(\tau)-\eta\frac{\nabla{\cal{L}}(\bm{W}(\tau))}{\left\|{\nabla{\cal{L}}(\bm{W}(\tau))}\right\|_{F}}.$
At each iteration $\tau$, correlation coefficient is computed by
$\left<\bm{W}(\tau),\bm{W}^{\text{svm}}\right>/(\left\|{\bm{W}(\tau)}\right\|_{F}\left\|{\bm{W}^{\text{svm}}}\right\|_{F})$,
and results averaged over $100$ random instances are displayed in Fig. 4(a),
which end in correlation $\approx 0.987$ after training with $4000$
iterations. In addition to the directional convergence of $\bm{W}(\tau)$, we
also verify the convergence of finite component by tracking the matrix
distance
$\left\|{\widehat{\bm{W}}^{\text{fin}}-\bm{W}^{\text{fin}}}\right\|_{F}$ where
$\widehat{\bm{W}}^{\text{fin}}=\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}(\tau))$,
and results are displayed in Fig. 4(b), which obtains $<0.01$ final distance.
Both results validate our Theorem 2.

##  4 Implicit Bias of Self-attention

In Section 3, we have discussed that gradient descent with log loss guarantees
the global convergence. To proceed, in this section, we discuss the implicit
bias of attention via analysis of regularization path (RP) as employed in
Algo-RP and identify the implicit bias of self-attention on more general next-
token prediction problems.

###### Theorem 3

Consider any dataset DSET and suppose Assumptions 1 and 2 hold. Additionally,
assume loss $\ell:\mathbb{R}\rightarrow\mathbb{R}$ is strictly decreasing and
$|\ell^{\prime}|$ is bounded. Let
$\bm{W}^{\text{svm}}{\in\mathcal{S}_{\text{fin}}^{\perp}}$ be the solution of
(Graph-SVM) and suppose $\bm{W}^{\text{svm}}\neq 0$. Then the solution of
regularization path Algo-RP obeys

$\displaystyle\lim_{R\rightarrow\infty}\frac{\bar{\bm{W}}_{R}}{R}=\frac{\bm{W}^{\text{svm}}}{\left\|{\bm{W}^{\text{svm}}}\right\|_{F}}~{}~{}\text{and}~{}~{}\lim_{R\rightarrow\infty}\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bar{\bm{W}}_{R})\in{\cal{W}}^{\text{fin}}.$

Here
${\cal{W}}^{\text{fin}}=\arg\min_{\bm{W}\in\mathcal{S}_{\text{fin}}}\lim_{R\rightarrow\infty}{\cal{L}}(\bm{W}+R\cdot\bm{W}^{\text{svm}})$
and we assume that ${\cal{W}}^{\text{fin}}$ is a bounded set.

Here, we allow more general loss function $\ell$ which is different from the
log-loss employed in Section 3, and the ERM problem (cf. (ERM)) is not
guaranteed to be convex.

###  4.1 Acyclic Dataset

Below, in contrast, we introduce the concept of acyclic dataset which implies
that the next-token prediction task always encounters a strict priority order
among tokens within each TPG. This corresponds to the setting where all SCCs
$((\mathcal{C}^{(k)}_{i})_{i=1}^{N_{k}})_{k=1}^{K}$ are all singletons; or
equivalently, $\overline{\texttt{DSET}}=\emptyset$.

######  Definition 2 (Acyclic dataset)

We call DSET acyclic if all of its TPGs are directed acyclic graphs.

For an acyclic dataset, there are not $i$, $j\in[K]$ satisfying $(i\asymp
j)\in{\cal{G}}^{(k)}$, for all $k\in[K]$. Thus, SVM formulation (Graph-SVM)
reduces to the following simpler form:

$\displaystyle\bm{W}^{\text{svm}}=\arg\min_{\bm{W}}\left\|{\bm{W}}\right\|_{F}\quad$
(Acyc-SVM)
$\displaystyle\text{s.t.}\quad(\bm{e}_{i}-\bm{e}_{j})^{\top}\bm{W}\bm{e}_{k}\geq
1\quad\forall(i\Rightarrow j)\in{\cal{G}}^{(k)},~{}k\in[K].$

We next make the following assumption on the linear head. Notably, Assumption
1 is a special case of Assumption 3.

###### Assumption 3

For $\forall y\in[K]$, $\arg\max_{k\in[K]}\bm{c}_{y}^{\top}\bm{e}_{k}=y$.

###### Lemma 3

Consider acyclic dataset DSET per Def. 2 and suppose Assumptions 2 and 3 hold.
Additionally, assume loss function $\ell$ is strictly decreasing and
$|\ell^{\prime}|$ is bounded. Then for any finite
$\bm{W}\in\mathbb{R}^{d\times d}$, training risk obeys
${\cal{L}}(\bm{W})>{\cal{L}}_{\star}:=\frac{1}{n}\sum_{i=1}^{n}\ell(\bm{c}_{y_{i}}^{\top}\bm{e}_{y_{i}})$.
Additionally, if (Acyc-SVM) is feasible, then for any $\bm{W}$ that satisfies
the constraints in (Acyc-SVM),
$\lim_{R\to\infty}{\cal{L}}(R\cdot\bm{W})={\cal{L}}_{\star}$.

Assumption 3 and Lemma 3 ensure that the best way for attention to make a
correct prediction on class $k$ is to output the vector $\bm{e}_{k}$, i.e.,
$f_{\bm{W}}({\bm{X}})=\bm{e}_{k}$. The next theorem states the directional
bias of self-attention on the acyclic dataset towards the solution of (Acyc-
SVM).

###### Theorem 4

Suppose DSET is acyclic per Definition 2 and Assumptions 2 and 3 hold.
Additionally, assume loss $\ell:\mathbb{R}\rightarrow\mathbb{R}$ is strictly
decreasing and $|\ell^{\prime}|$ is bounded. Suppose (Acyc-SVM) is feasible
with $\bm{W}^{\text{svm}}$ denoting its solution. Then, Algorithm Algo-RP
satisfies

$\lim_{R\to\infty}\frac{\bar{\bm{W}}_{R}}{R}=\frac{\bm{W}^{\text{svm}}}{\left\|{\bm{W}^{\text{svm}}}\right\|_{F}}.$

Recall that a dataset being acyclic implies that there is strict priority
order among all tokens in each TPG. Theorem 4 establishes the implicit bias of
self-attention model for next-token prediction problem in the presence of such
strict priority order. It demonstrates that once the SVM problem (Acyc-SVM) is
feasible, the regularized path of optimizing (ERM) converges directionally
toward it solution $\bm{W}^{\text{svm}}$.

Following the same implementation setting as in Section 3, in Fig. 5, we again
conduct $100$ trials but with randomly generated acyclic dataset DSET under
the setting of $K=d=8,n=4$ and $T=6$. The results averaged over $100$ random
instances are presented in Fig. 5 with correlation coefficient exceeds $0.99$
after training with $4000$ iterations. Note that in these experiments, log-
loss is employed as loss function and Assumption 1 is satisfied which
guarantee the convexity of ${\cal{L}}(\bm{W})$ following Lemma 2 and hence,
the connection between Algo-GD and Algo-RP is built by Ji et al. (2020).

Iterations ($\tau$)Correlation coefficientGlobalLocal (a)
$\frac{\bm{W}(\tau)}{\left\|{\bm{W}(\tau)}\right\|_{F}}\to\frac{\widetilde{\bm{W}}^{\text{svm}}}{\left\|{\widetilde{\bm{W}}^{\text{svm}}}\right\|_{F}}$

Iterations
($\tau$)$\left\|{\widehat{\bm{W}}^{\text{fin}}-\widetilde{\bm{W}}^{\text{fin}}}\right\|_{F}$GlobalLocal
(b)
$\bm{\Pi}_{\widetilde{\mathcal{S}}_{\text{fin}}}(\bm{W}(\tau))\to\widetilde{\bm{W}}^{\text{fin}}$

Figure 6: Squared loss with general classifier

Iterations ($\tau$)Correlation coefficientGlobalLocal (c)
$\frac{\bm{W}(\tau)}{\left\|{\bm{W}(\tau)}\right\|_{F}}\to\frac{\widetilde{\bm{W}}^{\text{svm}}}{\left\|{\widetilde{\bm{W}}^{\text{svm}}}\right\|_{F}}$

Iterations
($\tau$)$\left\|{\widehat{\bm{W}}^{\text{fin}}-\widetilde{\bm{W}}^{\text{fin}}}\right\|_{F}$GlobalLocal
(d)
$\bm{\Pi}_{\widetilde{\mathcal{S}}_{\text{fin}}}(\bm{W}(\tau))\to\widetilde{\bm{W}}^{\text{fin}}$

Figure 7: Cross-entropy loss with general classifier

##  5 Further Investigation on Local Convergence

So far, we have proved the global GD convergence of attention weight when one
employs the log-loss (Section 3) and studied the implicit bias of self-
attention over next-token prediction problem using RP analysis (Section 4). In
this section, we investigate further on the convergence performance of GD and
ask:

_When does the GD exhibit local convergence rather than global?
Can we characterize its implicit bias? _

Convergence performance of learning 1-layer attention has been analyzed in the
previous work Tarzanagh et al. (2023b, a), and they have observed the local
convergence phenomenon, and also provided the theoretical explanation and
empirical evidence. Inspired by their work, we define the _pseudo_ TPGs for
obtaining _locally-optimal_ SVM equivalence $\widetilde{\bm{W}}^{\text{svm}}$
and cyclic component $\widetilde{\bm{W}}^{\text{fin}}$ as follows:

  1. 1.

Given any dataset DSET, consider GD solution $\bm{W}^{\text{GD}}$. For each
training example $({\bm{X}},y)\in\texttt{DSET}$, let
$\bm{s}=\mathbb{S}({\bm{X}}\bm{W}^{\text{GD}}\bm{\bar{x}})$.

  2. 2.

Construct TPGs based on $\bm{s}$ by adding directed edge $(x_{t_{1}}\to
x_{t_{2}})$ to ${\cal{G}}^{(k)}$ if $\bm{s}_{t_{1}}>0$, where $k,x_{t}$ are
the token IDs of last token and $\bm{x}_{t}$, respectively.

Different from the TPGs defined in Section 2.1 which is uniquely determined by
the dataset and the ground truth labels, pseudo-TPGs build edges based on the
tokens selected by GD solution $\bm{W}^{\text{GD}}$. To further investigate
under which scenarios local convergence phenomenon exists, we consider the
following cases and provide experimental evidence.

General loss function $\ell$. In Section 3, we analyze the convergence
performance of gradient descent when employing log-loss. As we have discussed,
such loss guarantees the convexity of the problem and therefore, GD of
attention weight (directional) converges to its global minima. Here, we
investigate the performance of more general loss function, i.e., squared loss,
and find empirical evidence of local convergence (Figures 7 and 7).

Extended linear head. In this work, GD experiments are conducted under
Assumptions 1 and 2, which implies the convex training loss
${\cal{L}}(\bm{W})$ and global convergence performance. Now consider a more
general linear head (i.e., Assumption 3). As have also been observed and
discussed in Tarzanagh et al. (2023b, a), Algo-GD can converge to a locally-
optimal solution.

Figures 7 and 7 display our local convergence results where Fig. 7 employs
squared loss, i.e. $\ell(u)=(1-u)^{2}$ and Fig. 7 utilizes cross-entropy loss.
Both apply general head following Assumption 3. Similar to Fig. 5, we present
(directional) convergence performance of $\bm{W}$ towards
$\bm{W}^{\text{svm}}$ and $\bm{W}^{\text{fin}}$. Results indicate that instead
of converging to the global solution (blue curves), attention weights trained
via GD align more closely with the locally-optimal SVM solution defined via
the pseudo TPGs constructed by $\bm{W}^{\text{GD}}$ (orange curves). In Fig.
6(b), the norm difference to $\widetilde{\bm{W}}^{\text{fin}}$ remains zero,
indicating that all SCCs in the pseudo TPGs are singleton and GD optimizes
attention weights towards selecting one token per sequence. While in Fig. 7,
multiple tokens can be selected by $\bm{W}^{\text{GD}}$. Note that in Fig.
6(d), the norm of difference does not end with zero value on average. The
potential explanations can be: Due to the non-convexity of training loss,
training $\bm{W}$ with GD may not fully capture its RP solution
$\widetilde{\bm{W}}^{\text{fin}}$ over the cyclic subspace, and general
classification head induces correlation among tokens, leading the attention
mechanism to generate more intricate composed tokens. Nevertheless, our
empirical results indicate that $\bm{W}$ more closely aligns with the local
$\widetilde{\bm{W}}^{\text{fin}}$ within its cyclic subspace. We defer a
rigorous definition of local $\widetilde{\bm{W}}^{\text{fin}}$ and guarantees
related to gradient descent for future exploration. Experimental details are
deferred to the appendix.

##  6 Related Work

Inspired by the increasing popularity of Transformer-based models, a large
number of research efforts have focused on developing theoretical
understanding of various aspects of such models. Yun et al. (2020a); Fu et al.
(2023); Bombari & Mondelli (2024) studied the expressive power of Transformers
and showed that they are universal approximators for sequence-to-sequence
functions. A similar result for efficient variants of Transformers based on
sparse attention was presented in Yun et al. (2020b). Edelman et al. (2022)
studied bias of single attention layer towards representing sparse functions
of input sequence with favourable generalization behaviour. Interestingly,
Baldi & Vershynin (2023) explored key building blocks of attention mechanism
beyond modern neural networks and studied the functional capacity of the
resulting attention-based models. Other lines of theoretical efforts have
focused on explaining various properties of Transformer-based models,
including rank collapse (Dong et al., 2021) and realization of in-context
learning (Xie et al., 2022; Garg et al., 2022; Akyürek et al., 2023; Von
Oswald et al., 2023; Li et al., 2023c; Huang et al., 2023; Li et al., 2023b;
Collins et al., 2024; Jeon et al., 2024; Chen et al., 2024; Li et al., 2024).

Unlike these prior work, we focus on optimization-theoretic analysis of
attention-based models for the next-token prediction objective. Our work sheds
light on the implicit bias of underlying optimization problem towards SVM
formulations, which builds on the recent research efforts (Tarzanagh et al.,
2023b, a). However, different from these prior efforts that deal with
traditional (supervised) classification tasks, we focus on next-token
prediction task – the main workhorse of Transformer-based language modeling.
The recent work Thrampoulidis (2024) also explores the next-token prediction
problem under a classification-like setting, employing a related SVM
formulation. Since we study transformers, the main messages are fairly
different, e.g., our theory relies on graph-theoretic concepts such as SCCs
and token-priority graphs to capture the generative process learned by SGD.
Notably, several recent efforts (Jelassi et al., 2022; Li et al., 2023a, d;
Oymak et al., 2023; Deora et al., 2023; Chen & Li, 2024) have also analyzed
optimization and generalization dynamics of attention-based models. However,
these works again only focus on traditional classification tasks and consider
simplifications of the attention mechanism (Jelassi et al., 2022) or work with
strict statistical data assumptions (Jelassi et al., 2022; Li et al., 2023a;
Oymak et al., 2023). In contrast, we provide a detailed optimization-theoretic
treatment of the original (non-linear input dependent) attention mechanism
without any statistical assumption on the underlying data. Related work by
Tian et al. (2023) studies the training dynamics of next-token prediction.
Compared to us, their analysis is restricted to a specific statistical data
model, including the requirement of long input sequences ($T\to\infty$). Ildiz
et al. (2024); Makkuva et al. (2024) build connections between self-attention
and Markov chains. In contrast, we characterize the implicit bias of self-
attention learning to novel SVM formulations without any such assumptions on
the data model or sequence lengths.

We would also like to note the rich literature on studying implicit bias of
gradient-based optimization methods (see, e.g., Soudry et al. (2018);
Gunasekar et al. (2018); Ji et al. (2020); Ji & Telgarsky (2021); Kini et al.
(2021); Li et al. (2019); Blanc et al. (2020); Qian & Qian (2019); Wang et al.
(2021) and references therein). However, this prior work does not focus on the
optimization landscape of learning Transformer-based models and thus, does not
provide specific insights into their inner-workings, which is the main
objective of our work.

##  7 Discussion

In this work we set out to demystify Transformer-based language modeling via
next-token prediction task. We established that single-layer self-attention
learning has implicit bias towards the solution of a support vector machine
(SVM) formulation based on token-priority graphs which encode the priority
order among the tokens as per the training data. Our analysis shows that a
self-attention model learned via next-token prediction objective implements a
selection mechanism to suppress the lower priority tokens in order to predict
the higher priority tokens as the next-token for an input sequence. At the
same time, such an attention model would distribute its softmax probabilities
among all equal priority tokens as modeled by the strongly-connected
components of the next-token graph. Ultimately, our results comprehensively
capture the automaton implemented by a 1-layer self-attention under realistic
assumptions.

A natural future direction is relaxing our assumptions in SGD analysis and
providing a comprehensive characterization of the training dynamics,
accounting for non-convexities. It would also be interesting to extend our
analysis to multi-layer multi-head self-attention models or explore how feed-
forward layers (a.k.a. MLP layers) in Transformers affect the optimization
dynamics and aid in the aforementioned token selection and composition
mechanisms during next-token prediction.

## Acknowledgements

This work was supported in part by the NSF grants CCF-2046816, CCF-2212426,
CNS-1932254, UMich’s MIDAS PODS program, a Google Research Scholar award, and
an Adobe Data Science Research award.

## References

  * Akyürek et al. (2023) Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou.  What learning algorithm is in-context learning? investigations with linear models.  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=0g0X4H8yN4I.
  * Baldi & Vershynin (2023) Pierre Baldi and Roman Vershynin.  The quarks of attention: Structure and capacity of neural attention building blocks.  _Artificial Intelligence_ , 319:103901, 2023.  ISSN 0004-3702.  doi: https://doi.org/10.1016/j.artint.2023.103901.  URL https://www.sciencedirect.com/science/article/pii/S0004370223000474.
  * Blanc et al. (2020) Guy Blanc, Neha Gupta, Gregory Valiant, and Paul Valiant.  Implicit regularization for deep neural networks driven by an ornstein-uhlenbeck like process.  In _Conference on learning theory_ , pp. 483–513. PMLR, 2020.
  * Bombari & Mondelli (2024) Simone Bombari and Marco Mondelli.  Towards understanding the word sensitivity of attention layers: A study via random features.  _arXiv preprint arXiv:2402.02969_ , 2024.
  * Brown et al. (2020) Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel Ziegler, Jeffrey Wu, Clemens Winter, Chris Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 1877–1901. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper_files/paper/2020/file/1457c0d6bfcb4967418bfb8ac142f64a-Paper.pdf.
  * Chen et al. (2021) Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch.  Decision transformer: Reinforcement learning via sequence modeling.  In M. Ranzato, A. Beygelzimer, Y. Dauphin, P.S. Liang, and J. Wortman Vaughan (eds.), _Advances in Neural Information Processing Systems_ , volume 34, pp. 15084–15097. Curran Associates, Inc., 2021.  URL https://proceedings.neurips.cc/paper_files/paper/2021/file/7f489f642a0ddb10272b5c31057f0663-Paper.pdf.
  * Chen et al. (2020) Mark Chen, Alec Radford, Rewon Child, Jeffrey Wu, Heewoo Jun, David Luan, and Ilya Sutskever.  Generative pretraining from pixels.  In Hal Daumé III and Aarti Singh (eds.), _Proceedings of the 37th International Conference on Machine Learning_ , volume 119 of _Proceedings of Machine Learning Research_ , pp. 1691–1703. PMLR, 13–18 Jul 2020.  URL https://proceedings.mlr.press/v119/chen20s.html.
  * Chen & Li (2024) Sitan Chen and Yuanzhi Li.  Provably learning a multi-head attention layer.  _arXiv preprint arXiv:2402.04084_ , 2024.
  * Chen et al. (2024) Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang.  Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality.  _arXiv preprint arXiv:2402.19442_ , 2024.
  * Chowdhery et al. (2022) Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung, Charles Sutton, Sebastian Gehrmann, et al.  Palm: Scaling language modeling with pathways.  _arXiv preprint arXiv:2204.02311_ , 2022.
  * Chung & Glass (2020) Yu-An Chung and James Glass.  Generative pre-training for speech with autoregressive predictive coding.  In _ICASSP 2020-2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP)_ , pp. 3497–3501. IEEE, 2020.
  * Collins et al. (2024) Liam Collins, Advait Parulekar, Aryan Mokhtari, Sujay Sanghavi, and Sanjay Shakkottai.  In-context learning with transformers: Softmax attention adapts to function lipschitzness.  _arXiv preprint arXiv:2402.11639_ , 2024.
  * Deora et al. (2023) Puneesh Deora, Rouzbeh Ghaderi, Hossein Taheri, and Christos Thrampoulidis.  On the optimization and generalization of multi-head attention.  _arXiv preprint arXiv:2310.12680_ , 2023.
  * Dong et al. (2021) Yihe Dong, Jean-Baptiste Cordonnier, and Andreas Loukas.  Attention is not all you need: pure attention loses rank doubly exponentially with depth.  In Marina Meila and Tong Zhang (eds.), _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 2793–2803. PMLR, 18–24 Jul 2021.  URL https://proceedings.mlr.press/v139/dong21a.html.
  * Edelman et al. (2022) Benjamin L Edelman, Surbhi Goel, Sham Kakade, and Cyril Zhang.  Inductive biases and variable creation in self-attention mechanisms.  In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato (eds.), _Proceedings of the 39th International Conference on Machine Learning_ , volume 162 of _Proceedings of Machine Learning Research_ , pp. 5793–5831. PMLR, 17–23 Jul 2022.  URL https://proceedings.mlr.press/v162/edelman22a.html.
  * Ferruz et al. (2022) Noelia Ferruz, Steffen Schmidt, and Birte Höcker.  Protgpt2 is a deep unsupervised language model for protein design.  _Nature communications_ , 13(1):4348, 2022.
  * Fu et al. (2023) Hengyu Fu, Tianyu Guo, Yu Bai, and Song Mei.  What can a single attention layer learn? a study through the random features lens.  _arXiv preprint arXiv:2307.11353_ , 2023.
  * Garg et al. (2022) Shivam Garg, Dimitris Tsipras, Percy S Liang, and Gregory Valiant.  What can transformers learn in-context? a case study of simple function classes.  In S. Koyejo, S. Mohamed, A. Agarwal, D. Belgrave, K. Cho, and A. Oh (eds.), _Advances in Neural Information Processing Systems_ , volume 35, pp. 30583–30598. Curran Associates, Inc., 2022.  URL https://proceedings.neurips.cc/paper_files/paper/2022/file/c529dba08a146ea8d6cf715ae8930cbe-Paper-Conference.pdf.
  * Gunasekar et al. (2018) Suriya Gunasekar, Jason Lee, Daniel Soudry, and Nathan Srebro.  Characterizing implicit bias in terms of optimization geometry.  In Jennifer Dy and Andreas Krause (eds.), _Proceedings of the 35th International Conference on Machine Learning_ , volume 80 of _Proceedings of Machine Learning Research_ , pp. 1832–1841. PMLR, 10–15 Jul 2018.  URL https://proceedings.mlr.press/v80/gunasekar18a.html.
  * Huang et al. (2023) Yu Huang, Yuan Cheng, and Yingbin Liang.  In-context convergence of transformers.  _arXiv preprint arXiv:2310.05249_ , 2023.
  * Ildiz et al. (2024) M Emrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak.  From self-attention to markov models: Unveiling the dynamics of generative transformers.  _arXiv preprint arXiv:2402.13512_ , 2024.
  * Jelassi et al. (2022) Samy Jelassi, Michael Eli Sander, and Yuanzhi Li.  Vision transformers provably learn spatial structure.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), _Advances in Neural Information Processing Systems_ , 2022.  URL https://openreview.net/forum?id=eMW9AkXaREI.
  * Jeon et al. (2024) Hong Jun Jeon, Jason D Lee, Qi Lei, and Benjamin Van Roy.  An information-theoretic analysis of in-context learning.  _arXiv preprint arXiv:2401.15530_ , 2024.
  * Ji & Telgarsky (2019a) Ziwei Ji and Matus Telgarsky.  The implicit bias of gradient descent on nonseparable data.  In _Conference on learning theory_ , pp. 1772–1798. PMLR, 2019a.
  * Ji & Telgarsky (2019b) Ziwei Ji and Matus Telgarsky.  Risk and parameter convergence of logistic regression.  2019b.
  * Ji & Telgarsky (2021) Ziwei Ji and Matus Telgarsky.  Characterizing the implicit bias via a primal-dual analysis.  In Vitaly Feldman, Katrina Ligett, and Sivan Sabato (eds.), _Proceedings of the 32nd International Conference on Algorithmic Learning Theory_ , volume 132 of _Proceedings of Machine Learning Research_ , pp. 772–804. PMLR, 16–19 Mar 2021.  URL https://proceedings.mlr.press/v132/ji21a.html.
  * Ji et al. (2020) Ziwei Ji, Miroslav Dudík, Robert E Schapire, and Matus Telgarsky.  Gradient descent follows the regularization path for general losses.  In _Conference on Learning Theory_ , pp. 2109–2136. PMLR, 2020\.
  * Kini et al. (2021) Ganesh Ramachandra Kini, Orestis Paraskevas, Samet Oymak, and Christos Thrampoulidis.  Label-imbalanced and group-sensitive classification under overparameterization.  _Advances in Neural Information Processing Systems_ , 34:18970–18983, 2021.
  * Li et al. (2023a) Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen.  A theoretical understanding of shallow vision transformers: Learning, generalization, and sample complexity.  In _The Eleventh International Conference on Learning Representations_ , 2023a.  URL https://openreview.net/forum?id=jClGv3Qjhb.
  * Li et al. (2023b) Hongkang Li, Meng Wang, Songtao Lu, Hui Wan, Xiaodong Cui, and Pin-Yu Chen.  Transformers as multi-task feature selectors: Generalization analysis of in-context learning.  In _NeurIPS 2023 Workshop on Mathematics of Modern Machine Learning_ , 2023b.
  * Li et al. (2024) Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin-Yu Chen.  Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis.  _arXiv preprint arXiv:2402.15607_ , 2024.
  * Li et al. (2023c) Yingcong Li, Muhammed Emrullah Ildiz, Dimitris Papailiopoulos, and Samet Oymak.  Transformers as algorithms: Generalization and stability in in-context learning.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 19565–19594. PMLR, 23–29 Jul 2023c.  URL https://proceedings.mlr.press/v202/li23l.html.
  * Li et al. (2019) Yuanzhi Li, Colin Wei, and Tengyu Ma.  Towards explaining the regularization effect of initial large learning rate in training neural networks.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Li et al. (2023d) Yuchen Li, Yuanzhi Li, and Andrej Risteski.  How do transformers learn topic structure: Towards a mechanistic understanding.  In _International Conference on Machine Learning_ , pp. 19689–19729. PMLR, 2023d.
  * Makkuva et al. (2024) Ashok Vardhan Makkuva, Marco Bondaschi, Adway Girish, Alliot Nagle, Martin Jaggi, Hyeji Kim, and Michael Gastpar.  Attention with markov: A framework for principled analysis of transformers via markov chains.  _arXiv preprint arXiv:2402.04161_ , 2024.
  * Nijkamp et al. (2022) Erik Nijkamp, Jeffrey Ruffolo, Eli N Weinstein, Nikhil Naik, and Ali Madani.  Progen2: exploring the boundaries of protein language models.  _arXiv preprint arXiv:2206.13517_ , 2022.
  * OpenAI (2023) OpenAI.  Gpt-4 technical report.  _arXiv preprintarXiv:2303.08774_ , 2023.
  * Oymak et al. (2023) Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis.  On the role of attention in prompt-tuning.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 26724–26768. PMLR, 23–29 Jul 2023.  URL https://proceedings.mlr.press/v202/oymak23a.html.
  * Press & Wolf (2017) Ofir Press and Lior Wolf.  Using the output embedding to improve language models.  In _Proceedings of the 15th Conference of the European Chapter of the Association for Computational Linguistics: Volume 2, Short Papers_. Association for Computational Linguistics, 2017.
  * Qian & Qian (2019) Qian Qian and Xiaoyuan Qian.  The implicit bias of adagrad on separable data.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Radford et al. (2018) Alec Radford, Karthik Narasimhan, Tim Salimans, Ilya Sutskever, et al.  Improving language understanding by generative pre-training.  2018\.
  * Radford et al. (2019) Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al.  Language models are unsupervised multitask learners.  _OpenAI blog_ , 1(8):9, 2019.
  * Rosset et al. (2003) Saharon Rosset, Ji Zhu, and Trevor Hastie.  Margin maximizing loss functions.  _Advances in neural information processing systems_ , 16, 2003.
  * Soudry et al. (2018) Daniel Soudry, Elad Hoffer, Mor Shpigel Nacson, Suriya Gunasekar, and Nathan Srebro.  The implicit bias of gradient descent on separable data.  _Journal of Machine Learning Research_ , 19(70):1–57, 2018.  URL http://jmlr.org/papers/v19/18-188.html.
  * Suggala et al. (2018) Arun Suggala, Adarsh Prasad, and Pradeep K Ravikumar.  Connecting optimization and regularization paths.  _Advances in Neural Information Processing Systems_ , 31, 2018.
  * Tarjan (1972) Robert Tarjan.  Depth-first search and linear graph algorithms.  _SIAM Journal on Computing_ , 1(2):146–160, 1972\.  doi: 10.1137/0201010.  URL https://doi.org/10.1137/0201010.
  * Tarzanagh et al. (2023a) Davoud Ataee Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak.  Transformers as support vector machines.  _arXiv preprint arXiv:2308.16898_ , 2023a.
  * Tarzanagh et al. (2023b) Davoud Ataee Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak.  Margin maximization in attention mechanism.  _arXiv preprint arXiv:2306.13596_ , 2023b.
  * Thrampoulidis (2024) Christos Thrampoulidis.  Implicit bias of next-token prediction.  _arXiv preprint arXiv:2402.18551_ , 2024.
  * Tian et al. (2023) Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du.  Scan and snap: Understanding training dynamics and token composition in 1-layer transformer.  _arXiv preprint arXiv:2305.16380_ , 2023.
  * Touvron et al. (2023) Hugo Touvron, Thibaut Lavril, Gautier Izacard, Xavier Martinet, Marie-Anne Lachaux, Timothée Lacroix, Baptiste Rozière, Naman Goyal, Eric Hambro, Faisal Azhar, et al.  Llama: Open and efficient foundation language models.  _arXiv preprint arXiv:2302.13971_ , 2023.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), _Advances in Neural Information Processing Systems_ , volume 30. Curran Associates, Inc., 2017.  URL https://proceedings.neurips.cc/paper_files/paper/2017/file/3f5ee243547dee91fbd053c1c4a845aa-Paper.pdf.
  * Von Oswald et al. (2023) Johannes Von Oswald, Eyvind Niklasson, Ettore Randazzo, Joao Sacramento, Alexander Mordvintsev, Andrey Zhmoginov, and Max Vladymyrov.  Transformers learn in-context by gradient descent.  In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), _Proceedings of the 40th International Conference on Machine Learning_ , volume 202 of _Proceedings of Machine Learning Research_ , pp. 35151–35174. PMLR, 23–29 Jul 2023.  URL https://proceedings.mlr.press/v202/von-oswald23a.html.
  * Wang et al. (2021) Bohan Wang, Qi Meng, Wei Chen, and Tie-Yan Liu.  The implicit bias for adaptive optimization algorithms on homogeneous neural networks.  In Marina Meila and Tong Zhang (eds.), _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 10849–10858. PMLR, 18–24 Jul 2021.  URL https://proceedings.mlr.press/v139/wang21q.html.
  * Xie et al. (2022) Sang Michael Xie, Aditi Raghunathan, Percy Liang, and Tengyu Ma.  An explanation of in-context learning as implicit bayesian inference.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=RdJVFCHjUMI.
  * Yun et al. (2020a) Chulhee Yun, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar.  Are transformers universal approximators of sequence-to-sequence functions?  In _International Conference on Learning Representations_ , 2020a.  URL https://openreview.net/forum?id=ByxRM0Ntvr.
  * Yun et al. (2020b) Chulhee Yun, Yin-Wen Chang, Srinadh Bhojanapalli, Ankit Singh Rawat, Sashank Reddi, and Sanjiv Kumar.  O (n) connections are expressive enough: Universal approximability of sparse transformers.  _Advances in Neural Information Processing Systems_ , 33:13783–13794, 2020b.

1Contents

Figure 8: Illustration of token-priority graph (TPG). Given the input
sequences and labels (next tokens), we construct the TPGs
$\\{{\cal{G}}^{(k)}\\}_{k=1}^{K}$ according to the last token. Left hand side:
Two TPGs ${\cal{G}}^{(1)}$ and ${\cal{G}}^{(2)}$ are constructed using the
samples with $\bm{e}_{1}$ and $\bm{e}_{2}$ as the last tokens, respectively.
In each graph, directed edges (label$\to$input token) are added, and based on
the token relations, each graph can be partitioned into different strongly-
connected components (SCCs, highlighted as dashed grey rectangles). Right hand
side: We consider a cyclic subdataset $\overline{\texttt{DSET}}$ by removing
all the singleton SCCs, as well as their corresponding edges (see Definition
3). Then, $\overline{\texttt{DSET}}$ contains non-singleton SCCs only as shown
on the right.

##  Appendix A Auxiliary Results

###  A.1 Soft-composition Component

In Section 3, we have theoretically shown that when training a single-layer
self-attention model with gradient descent and log-loss function, once
$\bm{W}^{\text{svm}}\neq 0$, the composed attention weight $\bm{W}(\tau)$ will
diverge in Frobenius norm and $\bm{W}(\tau)$ converges towards direction
$\bm{W}^{\text{svm}}/\left\|{\bm{W}^{\text{svm}}}\right\|_{F}$; while in the
subspace of $\mathcal{S}_{\text{fin}}$,
$\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}(\tau))$ converges to a finite
$\bm{W}^{\text{fin}}$ which is the unique solution minimizing the training
loss over subspace $\mathcal{S}_{\text{fin}}$ as described in Theorem 2. Here,
$\bm{W}^{\text{svm}}$ follows the solution of (Graph-SVM) and plays a role in
separating tokens from different SCCs within the same TPG. Specifically, nodes
satisfy $(i\Rightarrow j)\in{\cal{G}}^{(k)}$.

As for the nodes contained within the same SCC (e.g., $(i\asymp j)$), to
ensure that $i$ and $j$ will not suppress each other, (Graph-SVM) solves the
SVM problem with the constraint
$(\bm{e}_{i}-\bm{e}_{j})^{\top}\bm{W}\bm{e}_{k}=0$. This essentially
disregards the influence of distinct tokens within the same SCC. Consequently,
$\bm{W}^{\text{svm}}$ does not truly capture the essence of the ERM solution.
In the following, we introduce cyclic subdataset and the so-called _cyclic-
component_ , and an equivalence between the cyclic term and
$\bm{W}^{\text{fin}}$ can be established under mild assumptions.

######  Definition 3 (Cyclic subdataset)

Given any training sample $({\bm{X}},y)\in\texttt{DSET}$, we obtain the
corresponding sample $({\bm{X}}^{\prime},y)\in\overline{\texttt{DSET}}$ by
removing all tokens in ${\bm{X}}$ that satisfy $(y\Rightarrow x)$ in the
corresponding TPG.

In short, cyclic subdataset focuses on the input tokens that are part of the
same SCC as the label token. Fig. 8(Right) presents the cyclic subdataset
$\overline{\texttt{DSET}}$ of DSET given in Fig. 8(Left), which is the same as
Fig. 3. In ${\cal{G}}^{(1)}$, all nodes are separated into different SCCs, and
therefore, none of them is present in $\overline{\texttt{DSET}}$; while in
${\cal{G}}^{(2)}$, token $\bm{e}_{1},\bm{e}_{2}$ and $\bm{e}_{3}$ are
reachable from each other, and then are utilized to construct
$\overline{\texttt{DSET}}$ while $\bm{e}_{4}$ is removed from the dataset.
Note that $\overline{\texttt{DSET}}$ provides a self-contained sub-problem
that solely focuses on intra-SCC edges.

######  Definition 4 (Cyclic component)

${\cal{W}}^{\text{fin}}$ is obtained as the solution set of the ERM problem
over the cyclic subdataset $\overline{\texttt{DSET}}$ per Definition 3.
Concretely,

$\displaystyle{\cal{W}}^{\text{fin}}=\arg\min_{\bm{W}\in\mathcal{S}_{\text{fin}}}\bar{\cal{L}}(\bm{W})\quad$
$\displaystyle\text{where}\quad\bar{\cal{L}}(\bm{W})=\frac{1}{n}\sum_{({\bm{X}},y)\in\overline{\texttt{DSET}}}\ell(\bm{c}_{y}^{\top}{\bm{X}}^{\top}\mathbb{S}({\bm{X}}\bm{W}\bm{\bar{x}})).$

###### Lemma 4

Consider a dataset DSET and let $\bm{W}^{\text{svm}}$ be the corresponding SVM
solution of (Graph-SVM) with $\bm{W}^{\text{svm}}\neq 0$. Then we have
$\bm{W}^{\text{svm}}\perp\mathcal{S}_{\text{fin}}$, and for any
${\bm{\bar{W}}}^{\text{fin}}\neq 0\in{\cal{W}}^{\text{fin}}$,
${\bm{\bar{W}}}^{\text{fin}}$ and $\bm{W}^{\text{svm}}$ are orthogonal.

###### Lemma 5

Let $\bm{W}\in\mathbb{R}^{d\times d}$ be an arbitrary matrix, then we have
${\bar{\cal{L}}}(\bm{W})={\bar{\cal{L}}}(\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}))$.

###### Lemma 6

Suppose Assumptions 1 and 2 hold, and loss function $\ell(u)=-\log(u)$, then
for any finite $\bm{W}$, $\tilde{\cal{L}}(\bm{W})={\bar{\cal{L}}}(\bm{W})$
where $\tilde{\cal{L}}(\bm{W})$ and ${\bar{\cal{L}}}(\bm{W})$ are defined in
Theorem 2 and Definition 4, respectively.

###  A.2 Useful Notations

In this section, we introduce additional notations used in the subsequent
proofs.

$\bullet$ Token index sets ${\cal{O}}_{i}$, $\bar{{\cal{O}}}_{i}$,
$\mathcal{R}_{i}$, $\bar{\mathcal{R}}_{i}$, $i\in[n]$. Consider dataset DSET.
Throughout, for any sample $({\bm{X}}_{i},y_{i})\in\texttt{DSET}$, $i\in[n]$,
we define

$\displaystyle{\cal{O}}_{i}:=\left\\{t~{}\Big{|}~{}x_{it}=y_{i},\forall
t\in[T_{i}]\right\\}\quad\text{and}\quad\bar{{\cal{O}}}_{i}=[T_{i}]-{\cal{O}}_{i},$
(3a)
$\displaystyle\mathcal{R}_{i}:={\cal{O}}_{i}\bigcup\left\\{t~{}\Big{|}~{}(x_{it}\asymp
y_{i})\in{\cal{G}}^{(\bar{x}_{i})},\forall
t\in[T_{i}]\right\\}\quad\text{and}\quad\bar{\mathcal{R}}_{i}=[T_{i}]-\mathcal{R}_{i}$
(3b)

where $x_{it}$ is the token ID of $\bm{x}_{it}$, $T_{i}$ is the number of
tokens in the input sequence ${\bm{X}}_{i}$ and ${\cal{G}}^{(\bar{x}_{i})}$ is
the corresponding token-priority graph (TPG) associated with the last/query
token of ${\bm{X}}_{i}$. Concretely, ${\cal{O}}_{i}$ returns the token indices
of $i$-th input that have the same token ID as label $y_{i}$, while
$\mathcal{R}_{i}$ returns the token indices of $i$-th input that are included
in the same strongly-connected component (SCC) as label $y_{i}$ in the
corresponding TPG. Then for any $t\in\bar{\mathcal{R}}_{i}$, we have
$(y_{i}\Rightarrow x_{it})\in{\cal{G}}^{(\bar{x}_{i})}$. Take the last input
sequence in Figure 8(left) as an example, where ${\cal{O}}=\\{3\\}$ and
$\mathcal{R}=\\{1,3\\}$.

$\bullet$ Datasets DSET, $\overline{\texttt{DSET}}$ and sample index set
${\mathcal{I}}$, $\bar{\mathcal{I}}$. Recap the training dataset
$\texttt{DSET}=({\bm{X}}_{i},y_{i})_{i=1}^{n}$. Based on the relationships
between input tokens and label token, following instructions in Section 2.1 we
can construct the TPGs of dataset DSET. Then, let ${\mathcal{I}}\subseteq[n]$
be the sample index set such that for any $i\in{\mathcal{I}}$, ${\bm{X}}_{i}$
contains distinct tokens from the same SCC as label $y_{i}$ in their
corresponding TPG. Or equivalently,

$\displaystyle{\mathcal{I}}=\left\\{i~{}\Big{|}~{}\mathcal{R}_{i}-{\cal{O}}_{i}\neq\emptyset,i\in[n]\right\\}\quad\text{and}\quad\bar{\mathcal{I}}=[n]-{\mathcal{I}}.$
(4)

Then the cyclic subset defined in Definition 3 can be written by

$\displaystyle\overline{\texttt{DSET}}=(\bar{\bm{X}}_{i},y_{i})_{i\in{\mathcal{I}}},$
(5)

where $\bar{\bm{X}}_{i}$ is obtained by removing all input tokens of
${\bm{X}}_{i}$ that are in the different SCCs from the label token $y_{i}$, or
equivalently, removing $\bm{x}_{it}$, $t\in\bar{\mathcal{R}}_{i}$. Hence, for
all $i\in\bar{\mathcal{I}}$, ${\bm{X}}_{i}$ only contains input tokens
(ignoring the ones with the same token ID as label) that have strictly lower
priority than its label token, i.e., $(y_{i}\Rightarrow
x_{it})\in{\cal{G}}^{(\bar{x}_{i})}$ for $t\in\bar{{\cal{O}}}_{i}$. In Figure
8(left), we have ${\mathcal{I}}=\\{4,5,6,7\\}$ and
$\bar{\mathcal{I}}=\\{1,2,3\\}$.

$\bullet$ Token scores ${\bm{\gamma}}_{i},i\in[n]$ and loss
${\cal{L}}(\bm{W})$ under Assumption 1. Let
${\bm{\gamma}}_{i}={\bm{X}}_{i}\bm{c}_{y_{i}}$ be the token score vectors.
Then under Assumption 1, we have

$\displaystyle\gamma_{it}=\begin{cases}1,\quad t\in{\cal{O}}_{i}\\\ 0,\quad
t\in\bar{{\cal{O}}}_{i}\end{cases}\quad\text{for all }i\in[n].$ (6)

Additionally, letting
$\bm{s}_{i}^{\bm{W}}=\mathbb{S}({\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i})$, we can
rewrite the training risk as follows:

$\displaystyle{\cal{L}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\ell\left(\sum_{t\in{\cal{O}}_{i}}s_{it}^{\bm{W}}\right).$
(7)

###  A.3 Proof of Lemma 1

Proof. Recap the constraints in (Graph-SVM) problem:

$\displaystyle(\bm{e}_{i}-\bm{e}_{j})^{\top}\bm{W}\bm{e}_{k}\begin{cases}=0\quad\forall(i\asymp
j)\in{\cal{G}}^{(k)}\\\ \geq 1\quad\forall(i\Rightarrow
j)\in{\cal{G}}^{(k)}\end{cases}\text{for all}\quad k\in[K].$ (8)

Since
${\bm{E}}=[\bm{e}_{1}~{}\bm{e}_{2}~{}\cdots~{}\bm{e}_{K}]^{\top}\in\mathbb{R}^{K\times
d}$ is full row rank, then $K\leq d$ and $\bm{e}_{k},k\in[K]$ are linearly
independent. Let $\bar{\bm{E}}\in\mathbb{R}^{K\times d}$ satisfying
$\bar{\bm{E}}{\bm{E}}^{\top}={\bm{I}}$. Then for any
$\bar{\bm{W}}\in\mathbb{R}^{K\times K}$, we get

$\displaystyle(\bm{e}_{i}-\bm{e}_{j})^{\top}\bar{\bm{E}}^{\top}\bar{\bm{W}}\bar{\bm{E}}\bm{e}_{k}\begin{cases}=0\quad\forall(i\asymp
j)\in{\cal{G}}^{(k)}\\\ \geq 1\quad\forall(i\Rightarrow
j)\in{\cal{G}}^{(k)}\end{cases}\text{for all}\quad k\in[K]$ (9)

and feasibility of (9) implies $\bm{W}\in\mathbb{R}^{d\times d}$ in (8) is
feasible. Since we can set
$\bm{W}=\bar{\bm{E}}^{\top}\bar{\bm{W}}\bar{\bm{E}}$. Next let
${\bm{u}}_{i}=\bar{\bm{E}}\bm{e}_{i}$ $i\in[K]$ be $K$-dimensional one-hot
vectors. Then it remains to show that there exists
$\bar{\bm{W}}\in\mathbb{R}^{K\times K}$ such that

$\displaystyle({\bm{u}}_{i}-{\bm{u}}_{j})^{\top}\bar{\bm{W}}{\bm{u}}_{k}\begin{cases}=0\quad\forall(i\asymp
j)\in{\cal{G}}^{(k)}\\\ \geq 1\quad\forall(i\Rightarrow
j)\in{\cal{G}}^{(k)}\end{cases}\text{for all}\quad k\in[K]$ (10)

is feasible. Additionally, it is equivalent with showing that for any
$k\in[K]$, there exists $\bm{w}\in\mathbb{R}^{K}$, such that

$\displaystyle({\bm{u}}_{i}-{\bm{u}}_{j})^{\top}\bm{w}\begin{cases}=0\quad\forall(i\asymp
j)\in{\cal{G}}^{(k)}\\\ \geq 1\quad\forall(i\Rightarrow
j)\in{\cal{G}}^{(k)}.\end{cases}$ (11)

To start with, we first derive the priority order of each graph (referring to
the topological sorting of directed graph). Let $M_{i}$ be the order of
${\bm{u}}_{i}$ where $M_{i}$’s $i\in[K]$ are positive integers. Then if
$(i\asymp j)$, $M_{i}=M_{j}$; if $(i\Rightarrow j)$, $M_{i}>M_{j}$. Then let
$\bm{w}=\sum_{i\in[K]}M_{i}{\bm{u}}_{i}$. We obtain that for any $k\in[K]$,

$\displaystyle\forall(i\asymp
j)\in{\cal{G}}^{(k)},~{}({\bm{u}}_{i}-{\bm{u}}_{j})^{\top}\bm{w}=({\bm{u}}_{i}-{\bm{u}}_{j})^{\top}(M_{i}{\bm{u}}_{i}+M_{j}{\bm{u}}_{j})=M_{i}-M_{j}=0$
$\displaystyle\forall(i\Rightarrow
j)\in{\cal{G}}^{(k)},~{}({\bm{u}}_{i}-{\bm{u}}_{j})^{\top}\bm{w}=({\bm{u}}_{i}-{\bm{u}}_{j})^{\top}(M_{i}{\bm{u}}_{i}+M_{j}{\bm{u}}_{j})=M_{i}-M_{j}\geq
1$

which indicates that (11) is feasible for any $k\in[K]$ and it completes the
proof.

###  A.4 Proof of Lemma 4

Proof. Recall from Definition 1 that $\mathcal{S}_{\text{fin}}$ is the span of
all matrices $(\bm{e}_{i}-\bm{e}_{j})\bm{e}_{k}^{\top}$ for all $(i\asymp
j)\in{\cal{G}}^{(k)}$ and $k\in[K]$. Then for any matrix
$\bm{W}\in\mathcal{S}_{\text{fin}}$, there exist $a_{ijk}$’s satisfying

$\bm{W}=\sum_{i,j,k}a_{ijk}(\bm{e}_{i}-\bm{e}_{j})\bm{e}_{k}^{\top}$

where $(i\asymp j)\in{\cal{G}}^{(k)}$ and $k\in[K]$. Since
$\bm{W}^{\text{svm}}$ is the solution of (Graph-SVM) that satisfies the all
“$=0$” constraints, for any matrix $\bm{W}\in\mathcal{S}_{\text{fin}}$, we
have

$\left<\bm{W}^{\text{svm}},\bm{W}\right>=\sum_{i,j,k}a_{ijk}\left<\bm{W}^{\text{svm}},(\bm{e}_{i}-\bm{e}_{j})\bm{e}_{k}^{\top}\right>=0.$

Therefore, $\bm{W}^{\text{svm}}\perp\mathcal{S}_{\text{fin}}$.

###  A.5 Proof of Lemma 5

Proof. Recap the definition of ${\bar{\cal{L}}}(\bm{W})$ from Def. 4 and
${\mathcal{I}}$, $\bar{\bm{X}}_{i}$ from (4), (5). Then

${\bar{\cal{L}}}(\bm{W})=\frac{1}{n}\sum_{i\in{\mathcal{I}}}\ell(\bm{c}_{y_{i}}^{\top}\bar{\bm{X}}_{i}^{\top}\mathbb{S}(\bar{\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i})).$

Let $\bm{W}^{\perp}=\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\bm{W})$ and
$\bm{W}^{\parallel}=\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W})$. Then it
remains to show that for any $(\bar{\bm{X}},y)\in\overline{\texttt{DSET}}$,
$\mathbb{S}(\bar{\bm{X}}\bm{W}\bm{\bar{x}})=\mathbb{S}(\bar{\bm{X}}\bm{W}^{\parallel}\bm{\bar{x}})$.

For simplification, let $\bm{\bar{x}}=\bm{e}_{k}$, and following the
definition of TPG, SCC and $\overline{\texttt{DSET}}$, we have that all tokens
$\bm{x}\in\bar{\bm{X}}$ are in the same SCC and denote the token set as
$\mathcal{C}^{(k)}$. Then $\mathcal{S}_{\text{fin}}$ spans the matrices
$(\bm{e}_{i}-\bm{e}_{j})\bm{e}_{k}^{\top}$ for $i,j\in\mathcal{C}^{(k)}$. For
any $i\in\mathcal{C}^{(k)}$, we get

$\bm{e}_{i}^{\top}\bm{W}\bm{e}_{k}=\bm{e}_{i}^{\top}\bm{W}^{\parallel}\bm{e}_{k}+\bm{e}_{i}^{\top}\bm{W}^{\perp}\bm{e}_{k}.$

Next, let $a_{ik}=\bm{e}_{i}^{\top}\bm{W}^{\perp}\bm{e}_{k}$,
$i\in\mathcal{C}^{(k)}$. Since $\bm{W}^{\perp}\perp\mathcal{S}_{\text{fin}}$,
and $(\bm{e}_{i}-\bm{e}_{j})\bm{e}_{k}^{\top}\in\mathcal{S}_{\text{fin}}$, we
obtain

$\displaystyle(\bm{e}_{i}-\bm{e}_{j})^{\top}\bm{W}^{\perp}\bm{e}_{k}=0$ (12)
$\displaystyle\Longrightarrow~{}$
$\displaystyle\bm{e}_{i}^{\top}\bm{W}^{\perp}\bm{e}_{k}-\bm{e}_{j}^{\top}\bm{W}^{\perp}\bm{e}_{k}=0$
$\displaystyle\Longrightarrow~{}$ $\displaystyle a_{ik}-a_{jk}=0$
$\displaystyle\Longrightarrow~{}$ $\displaystyle a_{ik}=a_{jk}=:\bar{a}_{k}.$

Then we have that for any $\bm{x}\in\bar{\bm{X}}$,
$\bm{x}^{\top}\bm{W}^{\perp}\bm{\bar{x}}=\bar{a}_{k}$ where $\bar{a}_{k}$ is
associated with the last/query token $\bm{\bar{x}}$ and hence

$\displaystyle\bar{{\bm{X}}}\bm{W}\bm{\bar{x}}=\bar{{\bm{X}}}\bm{W}^{\parallel}\bm{\bar{x}}+\bar{{\bm{X}}}\bm{W}^{\perp}\bm{\bar{x}}=\bar{{\bm{X}}}\bm{W}^{\parallel}\bm{\bar{x}}+\bar{a}_{k}\mathbf{1}$
$\displaystyle\mathbb{S}(\bar{{\bm{X}}}\bm{W}\bm{\bar{x}})=\mathbb{S}(\bar{{\bm{X}}}\bm{W}^{\parallel}\bm{\bar{x}}+\bar{a}_{k}\mathbf{1})=\mathbb{S}(\bar{{\bm{X}}}\bm{W}^{\parallel}\bm{\bar{x}}),$

which completes the proof.

###  A.6 Proof of Lemma 6

In the following, we present an additional lemma that incorporates Lemma 6.

###### Lemma 7

Suppose Assumptions 1 and 2 hold and loss function
$\ell:\mathbb{R}\to\mathbb{R}$ is strictly decreasing. For any finite
$\bm{W}$, once $\overline{\texttt{DSET}}\neq\texttt{DSET}$, we have that

$\displaystyle{\cal{L}}(\bm{W})>\frac{|\bar{\mathcal{I}}|}{n}\ell(1)+{\bar{\cal{L}}}(\bm{W}).$
(13)

Additionally, we have

$\displaystyle\min_{\bm{W}^{\prime}\in\mathcal{S}_{\text{fin}}^{\perp}}{\cal{L}}(\bm{W}^{\prime}+\bm{W})=\tilde{\cal{L}}(\bm{W})=\frac{|\bar{\mathcal{I}}|}{n}\ell(1)+{\bar{\cal{L}}}(\bm{W}),$
(14a)
$\displaystyle\min_{\bm{W}}{\cal{L}}(\bm{W})=\min_{\bm{W}}\tilde{\cal{L}}(\bm{W})=\frac{|\bar{\mathcal{I}}|}{n}\ell(1)+\min_{\bm{W}}{\bar{\cal{L}}}(\bm{W}).$
(14b)

Proof. We start with proving that for any finite $\bm{W}$,
${\cal{L}}(\bm{W})\geq\frac{|\bar{\mathcal{I}}|}{n}\ell(1)+{\bar{\cal{L}}}(\bm{W})$.
Let $\bm{W}^{\perp}=\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\bm{W})$,
$\bm{W}^{\parallel}=\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W})$ where we have
$\bm{W}=\bm{W}^{\perp}+\bm{W}^{\parallel}$. Let
$\bm{a}_{i}={\bm{X}}_{i}\bm{W}^{\perp}\bm{\bar{x}}_{i}$,
$\bm{b}_{i}={\bm{X}}_{i}\bm{W}^{\parallel}\bm{\bar{x}}_{i}$ and
$\bm{s}_{i}=\mathbb{S}({\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i})=\mathbb{S}(\bm{a}_{i}+\bm{b}_{i})$.
Following (7) and the definition of ${\bar{\cal{L}}}(\bm{W})$, training losses
obey

${\cal{L}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\ell\left(\sum_{t\in{\cal{O}}_{i}}s_{it}\right)\quad\text{and}\quad{\bar{\cal{L}}}(\bm{W})=\frac{1}{n}\sum_{i\in{\mathcal{I}}}\ell\left(\frac{\sum_{t\in{\cal{O}}_{i}}s_{it}}{\sum_{t\in\mathcal{R}_{i}}s_{it}}\right).$

From proof of Lemma 5 (more specifically (12)), we have that for any
$t\in\mathcal{R}_{i}$,

$\bm{x}_{it}^{\top}\bm{W}^{\perp}\bm{\bar{x}}_{i}=\bar{a}_{i}~{}\Longrightarrow
a_{it}=\bar{a}_{i},~{}\forall t\in\mathcal{R}_{i}$

where $\bar{a}_{i}$ is some constant associated with $\bm{W}^{\perp}$. Then
for any $i\in[n]$, we get

$\displaystyle\sum_{t\in{\cal{O}}_{i}}s_{it}=\frac{\sum_{t\in{\cal{O}}_{i}}e^{\bar{a}_{i}+b_{it}}}{\sum_{t\in\mathcal{R}_{i}}e^{\bar{a}_{i}+b_{it}}+\sum_{t\in\bar{\mathcal{R}}_{i}}e^{a_{it}+b_{it}}}=\frac{\sum_{t\in{\cal{O}}_{i}}e^{b_{it}}}{\sum_{t\in\mathcal{R}_{i}}e^{b_{it}}+\sum_{t\in\bar{\mathcal{R}}_{i}}e^{b_{it}+a_{it}-\bar{a}_{i}}}\leq\frac{\sum_{t\in{\cal{O}}_{i}}e^{b_{it}}}{\sum_{t\in\mathcal{R}_{i}}e^{b_{it}}},$
$\displaystyle\frac{\sum_{t\in{\cal{O}}_{i}}s_{it}}{\sum_{t\in\mathcal{R}_{i}}s_{it}}=\frac{\sum_{t\in{\cal{O}}_{i}}e^{\bar{a}_{i}+b_{it}}}{\sum_{t\in\mathcal{R}_{i}}e^{\bar{a}_{i}+b_{it}}}=\frac{\sum_{t\in{\cal{O}}_{i}}e^{b_{it}}}{\sum_{t\in\mathcal{R}_{i}}e^{b_{it}}}.$

Next following (4) we have that for $i\in\bar{\mathcal{I}}$,
$\mathcal{R}_{i}={\cal{O}}_{i}$ and therefore

$\frac{\sum_{t\in{\cal{O}}_{i}}s_{it}}{\sum_{t\in\mathcal{R}_{i}}s_{it}}=1\quad\text{for
all}\quad i\in\bar{\mathcal{I}}.$

Note that since $a_{it},b_{it}$ are finite and
$\overline{\texttt{DSET}}\neq\texttt{DSET}$, there exists $i\in[n]$ such that
$\sum_{t\in{\cal{O}}_{i}}s_{it}<\frac{\sum_{t\in{\cal{O}}_{i}}s_{it}}{\sum_{t\in\mathcal{R}_{i}}s_{it}}$.
Given strictly decreasing loss function $\ell$ and any finite $\bm{W}$, the
training risks obey

$\displaystyle{\cal{L}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\ell\left(\sum_{i\in{\cal{O}}_{i}}s_{it}\right)>\frac{1}{n}\sum_{i\in\bar{\mathcal{I}}}\ell(1)+\frac{1}{n}\sum_{i\in{\mathcal{I}}}\ell\left(\frac{\sum_{t\in{\cal{O}}_{i}}s_{it}}{\sum_{t\in\mathcal{R}_{i}}s_{it}}\right)=\frac{|\bar{\mathcal{I}}|}{n}\ell(1)+{\bar{\cal{L}}}(\bm{W}).$

It completes the proof of (13).

We next show that

$\min_{\bm{W}^{\prime}\in\mathcal{S}_{\text{fin}}^{\perp}}{\cal{L}}(\bm{W}^{\prime}+\bm{W})=\tilde{\cal{L}}(\bm{W})=\frac{|\bar{\mathcal{I}}|}{n}\ell(1)+{\bar{\cal{L}}}(\bm{W}).$

Recap from Theorem 2 that
$\tilde{\cal{L}}(\bm{W})=\lim_{R\to\infty}{\cal{L}}(\bm{W}+R\cdot\bm{W}^{\text{svm}})$.
Let $\bm{W}^{\perp}=\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\bm{W})$,
$\bm{W}^{\parallel}=\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W})$ where we have
$\bm{W}=\bm{W}^{\perp}+\bm{W}^{\parallel}$. Let
$\bm{a}_{i}={\bm{X}}_{i}\bm{W}^{\perp}\bm{\bar{x}}_{i}$,
$\bm{b}_{i}={\bm{X}}_{i}\bm{W}^{\parallel}\bm{\bar{x}}_{i}$,
$\bm{c}_{i}={\bm{X}}_{i}\bm{W}^{\text{svm}}\bm{\bar{x}}_{i}$ and
$\bm{s}^{R}_{i}=\mathbb{S}({\bm{X}}_{i}(R\cdot\bm{W}^{\text{svm}}+\bm{W})\bm{\bar{x}}_{i})=\mathbb{S}(\bm{a}_{i}+\bm{b}_{i}+R\cdot\bm{c}_{i})$.
Similarly, for any $t\in\mathcal{R}_{i}$,

$\bm{x}_{it}^{\top}\bm{W}^{\perp}\bm{\bar{x}}_{i}=\bar{a}_{i}~{}\Longrightarrow
a_{it}=\bar{a}_{i},~{}\forall t\in\mathcal{R}_{i}$

where $\bar{a}_{i}$ is some constant associated with $\bm{W}^{\perp}$.
Additionally, since $\bm{W}^{\text{svm}}$ follows (Graph-SVM), we have

$\displaystyle(\bm{x}_{i\tau}-\bm{x}_{it})^{\top}\bm{W}^{\text{svm}}\bm{\bar{x}}_{i}\begin{cases}=0&\forall
t\in\mathcal{R}_{i}\\\ \geq 1&\forall
t\in\bar{\mathcal{R}}_{i}\end{cases},\quad\text{for all
}\tau\in\mathcal{R}_{i}\quad\Longrightarrow\quad\begin{cases}c_{it}=\bar{c}_{i},&~{}\forall
t\in\mathcal{R}_{i},\\\ c_{it}\leq\bar{c}_{i}-1,&~{}\forall
t\in\bar{\mathcal{R}}_{i}.\end{cases}$ (15)

Then for any $i\in[n]$, we get

$\sum_{t\in{\cal{O}}_{i}}s_{it}^{R}=\frac{\sum_{t\in{\cal{O}}_{i}}e^{\bar{a}_{i}+b_{it}+R\bar{c}_{i}}}{\sum_{t\in\mathcal{R}_{i}}e^{\bar{a}_{i}+b_{it}+R\bar{c}_{i}}+\sum_{t\in\bar{\mathcal{R}}_{i}}e^{a_{it}+b_{it}+Rc_{it}}}=\frac{\sum_{t\in{\cal{O}}_{i}}e^{b_{it}}}{\sum_{t\in\mathcal{R}_{i}}e^{b_{it}}+\sum_{t\in\bar{\mathcal{R}}_{i}}e^{b_{it}+a_{it}-\bar{a}_{i}+R(c_{it}-\bar{c}_{i})}}\leq\frac{\sum_{t\in{\cal{O}}_{i}}e^{b_{it}}}{\sum_{t\in\mathcal{R}_{i}}e^{b_{it}}}.$

Case 1: $\bm{W}^{\text{svm}}=0$. Then for all $i\in[n]$,
$\bar{\mathcal{R}}_{i}=\emptyset$ and the equality holds for all $i\in[n]$.

Case 2: $\bm{W}^{\text{svm}}\neq 0$. Since $a_{it},b_{it}$ are finite and
$c_{it}-\bar{c}_{i}\leq-1$ for $t\in\bar{\mathcal{R}}_{i}$ following (15), the
equality holds when $R\to\infty$, and therefore we have for any $i\in[n]$,

$\lim_{R\to\infty}\sum_{t\in{\cal{O}}_{i}}s_{it}^{R}=\frac{\sum_{t\in{\cal{O}}_{i}}e^{b_{it}}}{\sum_{t\in\mathcal{R}_{i}}e^{b_{it}}}$

and

$\displaystyle\tilde{\cal{L}}(\bm{W})=\lim_{R\to\infty}{\cal{L}}(R\cdot\bm{W}^{\text{svm}}+\bm{W})$
$\displaystyle=\lim_{R\to\infty}\frac{1}{n}\sum_{i=1}^{n}\ell\left(\sum_{i\in{\cal{O}}_{i}}s_{it}^{R}\right)$
$\displaystyle=\frac{|\bar{\mathcal{I}}|}{n}\ell(1)+\frac{1}{n}\sum_{i\in{\mathcal{I}}}\ell\left(\frac{\sum_{t\in{\cal{O}}_{i}}e^{b_{it}}}{\sum_{t\in\mathcal{R}_{i}}e^{b_{it}}}\right)$
$\displaystyle=\frac{|\bar{\mathcal{I}}|}{n}\ell(1)+{\bar{\cal{L}}}(\bm{W}).$
(16)

Additionally, we have for any
$\bm{W}^{\prime}\in\mathcal{S}_{\text{fin}}^{\perp}$,

${\cal{L}}(\bm{W}^{\prime}+\bm{W})\geq\frac{|\bar{\mathcal{I}}|}{n}\ell(1)+{\bar{\cal{L}}}(\bm{W}^{\prime}+\bm{W})=\frac{|\bar{\mathcal{I}}|}{n}\ell(1)+{\bar{\cal{L}}}(\bm{W})$

where the inequality uses (13) and the equality comes from Lemma 5. Since
bound is achievable (by choosing
$\bm{W}^{\prime}=\lim_{R\to\infty}R\cdot\bm{W}^{\text{svm}}$ as in (16)), then
combining it with (16) completes the proof of (14a). (14b) is directly
obtained from (14a).

##  Appendix B Global Convergence of Gradient Descent

###  B.1 Supporting Results under the Setting of Theorem 2

In this section, we introduce results useful for the main proof. Recap the
setting of Theorem 2 where $\ell(u)=-\log(u)$. Therefore loss defined in (7)
is

$\displaystyle{\cal{L}}(\bm{W})=-\frac{1}{n}\sum_{i=1}^{n}\log\left(\sum_{t\in{\cal{O}}_{i}}s_{it}^{\bm{W}}\right)$
(17)

where $\bm{s}_{i}^{\bm{W}}=\mathbb{S}({\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i})$ and
${\cal{O}}_{i}$’s follow (3).

$\bullet$ $\nabla{\cal{L}}(\bm{W})$ under the setting of Theorem 2. For any
$\bm{W}\in\mathbb{R}^{d\times d}$, let
$\bm{h}_{i}={\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i},\bm{s}_{i}=\mathbb{S}(\bm{h}_{i}),{\bm{\gamma}}_{i}={\bm{X}}_{i}\bm{c}_{y_{i}}$.

$\begin{split}\nabla{\cal{L}}(\bm{W})&=\frac{1}{n}\sum_{i=1}^{n}\ell^{\prime}({\bm{\gamma}}_{i}^{\top}\bm{s}_{i}){\bm{X}}_{i}^{\top}\mathbb{S}^{\prime}(\bm{h}_{i}){\bm{\gamma}}_{i}\bm{\bar{x}}_{i}^{\top}\\\
&=\frac{1}{n}\sum_{i=1}^{n}-\frac{1}{{\bm{\gamma}}_{i}^{\top}\bm{s}_{i}}{\bm{X}}_{i}^{\top}(\text{diag}(\bm{s}_{i})-\bm{s}_{i}\bm{s}_{i}^{\top}){\bm{\gamma}}_{i}\bm{\bar{x}}_{i}^{\top}\\\
&=\frac{1}{n}\sum_{i=1}^{n}\sum_{t\in\bar{{\cal{O}}}_{i}}s_{it}(\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top}\end{split}$
(18)

where the last equation uses the fact that for any example
$({\bm{X}}_{i},y_{i})\in\texttt{DSET}$, $i\in[n]$,

$\begin{split}\frac{{\bm{X}}_{i}^{\top}(\text{diag}(\bm{s}_{i})-\bm{s}_{i}\bm{s}_{i}^{\top}){\bm{\gamma}}_{i}}{{\bm{\gamma}}_{i}^{\top}\bm{s}_{i}}&=\frac{{\bm{X}}_{i}^{\top}\text{diag}(\bm{s}_{i}){\bm{\gamma}}_{i}}{{\bm{\gamma}}_{i}^{\top}\bm{s}_{i}}-{\bm{X}}_{i}^{\top}\bm{s}_{i}\\\
&=\frac{\sum_{t\in{\cal{O}}_{i}}s_{it}\bm{e}_{y_{i}}}{\sum_{t\in{\cal{O}}_{i}}s_{it}}-{\bm{X}}_{i}^{\top}\bm{s}_{i}\\\
&=\bm{e}_{y_{i}}-{\bm{X}}_{i}^{\top}\bm{s}_{i}\\\
&=\sum_{t\in\bar{{\cal{O}}}_{i}}s_{it}(\bm{e}_{y_{i}}-\bm{x}_{it}).\end{split}$

Here, the second equality comes from (6).

$\bullet$ Lipschitzness of $\nabla{\cal{L}}(\bm{W})$ in (18). For any
$\bm{W},\dot{\bm{W}}\in\mathbb{R}^{d\times d}$, let
$\bm{s}_{i}=\mathbb{S}({\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i})$ and
$\dot{\bm{s}}_{i}=\mathbb{S}({\bm{X}}_{i}\dot{\bm{W}}\bm{\bar{x}}_{i})$.
Consider bounded tokens and let $M:=\max_{k\in[K]}\|\bm{e}_{k}\|$. Following
(18), we have:

$\begin{split}\left\|{\nabla\mathcal{L}(\bm{W})-\nabla\mathcal{L}(\dot{\bm{W}})}\right\|_{F}&=\left\|{\frac{1}{n}\sum_{i=1}^{n}\sum_{t\in\bar{{\cal{O}}}_{i}}(s_{it}-\dot{s}_{it})(\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top}}\right\|_{F}\\\
&\leq\frac{1}{n}\sum_{i=1}^{n}\sum_{t\in\bar{{\cal{O}}}_{i}}|s_{it}-\dot{s}_{it}|~{}\left\|{(\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top}}\right\|_{F}\\\
&\leq\frac{2M^{2}}{n}\sum_{i=1}^{n}\sum_{t\in\bar{{\cal{O}}}_{i}}|s_{it}-\dot{s}_{it}|\\\
&\leq\frac{2M^{2}}{n}\sum_{i=1}^{n}\|\bm{s}_{i}-\dot{\bm{s}}_{i}\|_{1}\\\
&\leq\frac{2M^{2}}{n}\sum_{i=1}^{n}\sqrt{T_{i}}\cdot\|\bm{s}_{i}-\dot{\bm{s}}_{i}\|.\end{split}$
(19)

Next for any $\bm{s},\dot{\bm{s}}$, we get

$\begin{split}\|\bm{s}-\dot{\bm{s}}\|&=\|\mathbb{S}({\bm{X}}\bm{W}\bm{\bar{x}})-\mathbb{S}({\bm{X}}\dot{\bm{W}}\bm{\bar{x}})\|\\\
&\leq\|{\bm{X}}\bm{W}\bm{\bar{x}}-{\bm{X}}\dot{\bm{W}}\bm{\bar{x}}\|\\\ &\leq
M^{2}\left\|{\bm{W}-\dot{\bm{W}}}\right\|_{F}.\end{split}$ (20)

Combining results in that

$\begin{split}\left\|{\nabla{\cal{L}}(\bm{W})-\nabla{\cal{L}}(\dot{\bm{W}})}\right\|_{F}\leq
2M^{4}\sqrt{T_{\text{max}}}\cdot\left\|{\bm{W}-\dot{\bm{W}}}\right\|_{F}\end{split}$
(21)

where $T_{\text{max}}:=\max_{i\in[n]}T_{i}$. Then let

$\displaystyle L:=2M^{4}\sqrt{T_{\text{max}}}$ (22)

and $\nabla{\cal{L}}(\bm{W})$ is $L$-Lipschitz continuous.

###  B.2 Proof of Lemma 2

In this subsection, we provide and prove a general version of Lemma 2. To to
that, we first introduce the following new subspaces.

###### Definition 5

Define the subspace $\mathcal{S}_{\text{active}}$ as the span of all matrices
$(\bm{e}_{i}-\bm{e}_{j})\bm{e}_{k}^{\top}$ for all $(i\rightarrow
j)\in{\cal{G}}^{(k)}$ and $k\in[K]$.

######  Definition 6 (Cyclic subspace (Restated))

Define cyclic subspace $\mathcal{S}_{\text{fin}}$ as the span of all matrices
$(\bm{e}_{i}-\bm{e}_{j})\bm{e}_{k}^{\top}$ for all $(i\asymp
j)\in{\cal{G}}^{(k)}$ and $k\in[K]$.

Observe that $\mathcal{S}_{\text{fin}}$ is a subspace of
$\mathcal{S}_{\text{active}}$. This is because if two nodes
$i,j\in{\cal{G}}^{(k)}$ and $i\asymp j$, this also implies that $i\rightarrow
j$ and $j\rightarrow i$.

######  Definition 7 (SVM subspace)

Define svm subspace ${\mathcal{S}_{\text{svm}}}$ as the orthogonal complement
of the subspace $\mathcal{S}_{\text{fin}}$ inside the subspace
$\mathcal{S}_{\text{active}}$.

###### Lemma 8

We have the following:

  1. (i)

Recall the definition of $\bm{W}^{\text{svm}}$ in (Graph-SVM).
$\bm{W}^{\text{svm}}\in{\mathcal{S}_{\text{svm}}}$.

  2. (ii)

Let $\mathcal{S}_{\text{active}}^{\perp}$ be the orthogonal complement of
$\mathcal{S}_{\text{active}}$ inside $\mathbb{R}^{d\times d}$. Then,

$\displaystyle\left\|{\bm{\Pi}_{\mathcal{S}_{\text{active}}^{\perp}}\left(\nabla{\cal{L}}(\bm{W})\right)}\right\|_{F}=0,\quad\forall\bm{W}\in\mathbb{R}^{d\times
d}.$

Proof.

$\bullet$ (i): Recall the definition of $\bm{W}^{\text{svm}}$:

$\displaystyle\bm{W}^{\text{svm}}=\arg\min_{\bm{W}}\left\|{\bm{W}}\right\|_{F}\quad$
$\displaystyle\text{s.t.}~{}~{}(\bm{e}_{i}-\bm{e}_{j})^{\top}\bm{W}\bm{e}_{k}\begin{cases}=0\quad\forall(i\asymp
j)\in{\cal{G}}^{(k)}\\\ \geq 1\quad\forall(i\Rightarrow
j)\in{\cal{G}}^{(k)}\end{cases}\text{for all}\quad k\in[K].$

Assume that the statement is not correct. Then, either
$\left\|{\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}^{\text{svm}})}\right\|_{F}>0$
or
$\left\|{\bm{\Pi}_{\mathcal{S}_{\text{active}}^{\perp}}(\bm{W}^{\text{svm}})}\right\|_{F}>0$.

By definition,
$\left\|{\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}^{\text{svm}})}\right\|_{F}=0$
since for all $(i\asymp j)\in{\cal{G}}^{(k)}$,
$(\bm{e}_{i}-\bm{e}_{j})^{\top}\bm{W}^{\text{svm}}\bm{e}_{k}=0$.

On the other hand, if
$\left\|{\bm{\Pi}_{\mathcal{S}_{\text{active}}^{\perp}}(\bm{W}^{\text{svm}})}\right\|_{F}>0$,
then
$\bm{W}^{\text{svm}}-\bm{\Pi}_{\mathcal{S}_{\text{active}}^{\perp}}(\bm{W}^{\text{svm}})$
also satisfies all of the constraints of (Graph-SVM), and

$\displaystyle\left\|{\bm{W}^{\text{svm}}-\bm{\Pi}_{\mathcal{S}_{\text{active}}^{\perp}}(\bm{W}^{\text{svm}})}\right\|_{F}<\left\|{\bm{W}^{\text{svm}}}\right\|_{F},$

which is a contradiction. Therefore,
$\bm{W}^{\text{svm}}\in{\mathcal{S}_{\text{svm}}}$.

$\bullet$ (ii): From (18), we know that

$\displaystyle\nabla{\cal{L}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\sum_{t\in\bar{{\cal{O}}}_{i}}s_{it}(\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top}$

where $\bar{{\cal{O}}}_{i}$ is given by (3). By definition of
$\mathcal{S}_{\text{active}}$,
$\left\|{\bm{\Pi}_{\mathcal{S}_{\text{active}}^{\perp}}(\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top}}\right\|_{F}=0$
for any $i\in[n]$ and $t\in\bar{{\cal{O}}}_{i}$. As $\nabla{\cal{L}}(\bm{W})$
is the summation of these terms, the advertised result is proved.

Now, we are ready to prove a stronger version of Lemma 2.

######  Lemma 9 (Stronger version of Lemma 2)

Suppose Assumptions 1 and 2 hold and consider the log-loss $\ell(u)=-\log(u)$,
then ${\cal{L}}(\bm{W})$ is convex. Furthermore, ${\cal{L}}(\bm{W})$ is
strictly convex on $\mathcal{S}_{\text{active}}$.

Proof. Let $\mathcal{S}_{K}$ be the span of all $\bm{e}_{i}\bm{e}_{j}^{\top}$
where $i,j\in[K]$.

$\bullet$ First Case: $\bm{W}\in\mathcal{S}_{K}$. Let
$g:\mathcal{S}_{K}\xrightarrow[]{}\mathbb{R}^{K\times K}$ such that
$g(\bm{W})={\bm{E}}\bm{W}{\bm{E}}^{\top}$. By definition, this function is
linear. In addition to that, this function $g$ is invertible by Assumption 1
and the domain of the function is $\mathcal{S}_{K}$. Note that Assumption 1
ensures $\text{rank}({\bm{E}})=K$.

Let ${\bm{E}}^{\prime}={\bm{C}}^{\prime}={\bm{I}}_{k}$,
$({\bm{X}}^{\prime}_{i},y_{i}^{\prime})$ be a DSET such that
$y_{i}^{\prime}=y_{i}$ and
${\bm{X}}_{i}^{\prime}={\bm{X}}_{i}{\bm{E}}^{\dagger}$. Then, for any
$\bm{W}^{\prime}\in\mathbb{R}^{K\times K}$, we have the following:

$\displaystyle{\cal{L}}\circ
g^{-1}(\bm{W}^{\prime})=\frac{1}{n}\sum_{i=1}^{n}\ell\left((\bm{c}^{\prime}_{y_{i}})^{\top}({\bm{X}}_{i}^{\prime})^{\top}\mathbb{S}({\bm{X}}_{i}^{\prime}\bm{W}^{\prime}\bar{\bm{x}_{i}}^{\prime})\right)$

Using Lemma 10 and 11, we know that ${\cal{L}}\circ g^{-1}(\bm{W}^{\prime})$
is convex on $\mathbb{R}^{K\times K}$ and strictly convex on
$g(\mathcal{S}_{\text{active}})$. Using these two facts and Lemma 10, we have
${\cal{L}}(\bm{W})$ is convex on $\mathcal{S}_{K}$ and strictly convex on
$\mathcal{S}_{\text{active}}\cap\mathcal{S}_{K}=\mathcal{S}_{\text{active}}$.

$\bullet$ Second Case: $\bm{W}\not\in\mathcal{S}_{K}$. By definition of loss
function in (ERM), we have

$\displaystyle{\cal{L}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\ell(\bm{c}_{y_{i}}^{\top}{\bm{X}}_{i}^{\top}\mathbb{S}({\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i}))=\frac{1}{n}\sum_{i=1}^{n}\ell(\bm{c}_{y_{i}}^{\top}{\bm{X}}_{i}^{\top}\mathbb{S}({\bm{X}}_{i}\bm{\Pi}_{\mathcal{S}_{K}}(\bm{W})\bm{\bar{x}}_{i}))={\cal{L}}(\bm{\Pi}_{\mathcal{S}_{K}}(\bm{W}))$
(23)

Let $\bm{W}_{1},\bm{W}_{2}\in\mathbb{R}^{d\times d}$ be arbitrary variables.
For any $0<\lambda<1$, we have the following:

$\displaystyle{\cal{L}}(\lambda\bm{W}_{1}+(1-\lambda)\bm{W}_{2})={\cal{L}}(\lambda\bm{\Pi}_{\mathcal{S}_{K}}(\bm{W}_{1})+\lambda\bm{\Pi}_{\mathcal{S}_{K}}(\bm{W}_{2}))$
(24)

Then, using (23) and (24), we have the following:

$\displaystyle\lambda{\cal{L}}(\bm{W}_{1})+(1-\lambda){\cal{L}}(\bm{W}_{2})$
$\displaystyle=\lambda{\cal{L}}(\bm{\Pi}_{\mathcal{S}_{K}}(\bm{W}_{1}))+(1-\lambda){\cal{L}}(\bm{\Pi}_{\mathcal{S}_{K}}(\bm{W}_{2}))$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{\geq}}{\cal{L}}(\lambda\bm{\Pi}_{\mathcal{S}_{K}}(\bm{W}_{1})+\lambda\bm{\Pi}_{\mathcal{S}_{K}}(\bm{W}_{2}))={\cal{L}}(\lambda\bm{W}_{1}+(1-\lambda)\bm{W}_{2})$

where (a) follows from the convexity of ${\cal{L}}(\bm{W})$ inside
$\mathcal{S}_{K}$. This implies that ${\cal{L}}(\bm{W})$ is convex when
$\bm{W}\not\in\mathcal{S}_{K}$. Note that
$\mathcal{S}_{\text{active}}\subset\mathcal{S}_{K}$, therefore we do not look
at the strict convexity in this case.

###### Lemma 10

Let $T:\mathcal{X}\xrightarrow[]{}\mathcal{Y}$ be an invertible linear map. If
a function $f:\mathcal{Y}\xrightarrow[]{}\mathbb{R}$ is convex/strictly convex
on $\mathcal{Y}$, then $f\circ T(x)$ is a convex/strictly convex function on
$\mathcal{X}$.

Proof. Let $x_{1}\neq x_{2}\in\mathcal{X}$ be arbitrary variables. Let
$y_{1}=T(x_{1})$ and $y_{2}=T(x_{2})$. Since $T$ is an invertible map,
$y_{1}\neq y_{2}$. Since $T$ is a linear map, $T(\lambda
x_{1}+(1-\lambda)x_{2})=\lambda y_{1}+(1-\lambda)y_{2}$ for $0<\lambda<1$.
Then, we obtain the following

$\displaystyle\lambda(f\circ T(x_{1}))+(1-\lambda)(f\circ T(x_{2}))$
$\displaystyle=\lambda f(y_{1})+(1-\lambda)f(y_{2})$
$\displaystyle\stackrel{{\scriptstyle(a)}}{{>}}f(\lambda
y_{1}+(1-\lambda)y_{2})$ $\displaystyle=f\circ T(\lambda
x_{1}+(1-\lambda)x_{2})$

where (a) follows from the strict convexity of the function $f$. This implies
that $f\circ T(x)$ is a strictly convex function on $\mathcal{X}$. Note that
if $y_{1}=y_{2}$, then we cannot achieve (a). Additionally, if $f$ is convex
instead of strictly convex, then $>$ in (a) is changed to $\geq$, and $f\circ
T(x)$ is convex.

###### Lemma 11

Suppose that Assumption 2 holds and ${\bm{E}}=\bm{I}_{d}$. Let
$f:\mathbb{R}^{d\times d}\xrightarrow[]{}\mathbb{R}^{d^{2}}$ be a linear
transformation defined as $f(\bm{W})=\bm{v}$ where $v_{i\times
d+j}={\bm{e}_{i}^{\top}\bm{W}\bm{e}_{j}}$. Then, ${\cal{L}}\circ
f^{-1}(\bm{v})$ is convex. Furthermore, ${\cal{L}}\circ f^{-1}(\bm{v})$ is
strictly convex on $f(\mathcal{S}_{\text{active}})$, where
$\mathcal{S}_{\text{active}}$ is defined in Definition 5.

Proof. $\bullet$ We first prove that ${\cal{L}}\circ f^{-1}(\bm{v})$ is
convex.  Let $\bar{\ell}:\mathbb{R}^{d^{2}}\times\mathbb{R}^{T\times
d}\times\mathbb{R}\xrightarrow[]{}\mathbb{R}$ be defined as follows:

$\displaystyle\bar{\ell}(\bm{v},{\bm{X}},y):=\ell\left(\bm{c}_{y}^{\top}{\bm{X}}^{\top}\mathbb{S}({\bm{X}}\big{(}f^{-1}(\bm{v})\big{)}\bar{\bm{x}})\right).$

Then, using (ERM), we have the following:

$\displaystyle{\cal{L}}\circ
f^{-1}(\bm{v})=\frac{1}{n}\sum_{i=1}^{n}\ell\left(\bm{c}_{y_{i}}^{\top}{\bm{X}}_{i}^{\top}\mathbb{S}({\bm{X}}_{i}\big{(}f^{-1}(\bm{v})\big{)}\bar{\bm{x}_{i}})\right)=\frac{1}{n}\sum_{i=1}^{n}\bar{\ell}(\bm{v},{\bm{X}}_{i},y_{i}).$
(25)

Note that the summation of convex functions is convex. Therefore, it is
sufficient to prove the convexity of ${\cal{L}}\circ f^{-1}(\bm{v})$ by
proving the convexity of $\bar{\ell}(\bm{v},{\bm{X}},y)$ for an arbitrary pair
of input sequence and label $({\bm{X}},y)$. For the simplicity of notation, we
use $\bar{\ell}(\bm{v})$ instead of $\bar{\ell}(\bm{v},{\bm{X}},y)$. Let
$m_{j}$ be the number of token ID $j$ inside input sequence ${\bm{X}}$ for
$j\in[K]$. Let $k$ be the last token of ${\bm{X}}$. By Assumption 1 and log-
loss, we know that

$\displaystyle\bar{\ell}(\bm{v}):=\bar{\ell}(\bm{v},{\bm{X}},y)=-\log\left(\frac{m_{y}\cdot
e^{\bm{v}_{y\times d+k}}}{\sum_{j\in[K]}m_{j}\cdot e^{\bm{v}_{j\times
d+k}}}\right)=\log\left(\sum_{j\in[K]}m_{j}\cdot e^{\bm{v}_{j\times
d+k}}\right)-\log(m_{y}\cdot e^{v_{y\times d+k}}).$

Let ${\bm{z}}\in\mathbb{R}^{d^{2}}$ be a vector such that the $(j\times
d+k)^{\text{th}}$ element of ${\bm{z}}$ is $z_{j\times d+k}=m_{j}\cdot
e^{v_{j\times d+k}}$ for $k\in[K]$, otherwise $z_{i}=0$. Then, the Hessian
matrix of $\bar{\ell}(\bm{v})$ is

$\displaystyle\nabla^{2}\bar{\ell}(\bm{v})=\frac{1}{(\bm{1}^{\top}{\bm{z}})^{2}}\left((\bm{1}^{\top}{\bm{z}})\text{diag}({\bm{z}})-{\bm{z}}{\bm{z}}^{\top}\right)$

For any ${\bm{u}}\in\mathbb{R}^{d^{2}}$, we obtain that

$\displaystyle{\bm{u}}^{\top}\nabla^{2}\bar{\ell}(\bm{v}){\bm{u}}=\frac{1}{(\bm{1}^{\top}{\bm{z}})^{2}}\left(\left(\sum_{j=1}^{d^{2}}z_{j}\right)\left(\sum_{j=1}^{d^{2}}u_{j}^{2}z_{j}\right)-\left(\sum_{j=1}^{d^{2}}u_{j}z_{j}\right)^{2}\right)\geq
0.$ (26)

Since ${\bm{z}}_{i}\geq 0$, $i\in[d^{2}]$, (26) follows from the Cauchy-
Schwarz inequality
$({\bm{\alpha}}^{\top}{\bm{\alpha}})({\bm{\beta}}^{\top}{\bm{\beta}})\geq({\bm{\alpha}}^{\top}{\bm{\beta}})^{2}$
applied to the vectors with $\alpha_{i}=u_{i}\sqrt{z_{i}}$ and
$\beta_{i}=\sqrt{z_{i}}$. The equality condition holds
$k{\bm{\alpha}}={\bm{\beta}}$ for $k\neq 0$. This means that
$\bar{\ell}(\bm{v})$ is convex.

$\bullet$ Next, we will show that ${\cal{L}}\circ f^{-1}(\bm{v})$ is strictly
convex on $f(\mathcal{S}_{\text{active}})$. Assume that ${\cal{L}}\circ
f^{-1}(\bm{v})$ is not strictly convex on $f(\mathcal{S}_{\text{active}})$.
Using the convexity of ${\cal{L}}\circ f^{-1}(\bm{v})$, this implies that
there exist ${\bm{u}},\bm{v}\in f(\mathcal{S}_{\text{active}})$,
$\|{{\bm{u}}}\|_{2}>0$ such that

$\displaystyle{\bm{u}}^{\top}\left(\nabla^{2}{\cal{L}}\circ
f^{-1}(\bm{v})\right){\bm{u}}=0$

Using the convexity of $\bar{\ell}(\bm{v})$ and (25), we have the following:

$\displaystyle{\bm{u}}^{\top}\left(\nabla^{2}\bar{\ell}(\bm{v},{\bm{X}}_{i},\bm{y}_{i})\right){\bm{u}}=0\qquad\forall
i\in[n]$ (27)

Now, we are going to prove that $\|{\bm{u}}\|_{2}=0$ if (27) holds. As
${\bm{u}}\in f(\mathcal{S}_{\text{active}})$, there exists
$\bm{W}\in\mathcal{S}_{\text{active}}$ such that $f(\bm{W})={\bm{u}}$. As the
function $f$ preserves the norm, $\left\|{\bm{W}}\right\|_{F}>0$. By
definition of $\mathcal{S}_{\text{active}}$, there exist
$\bar{i},\bar{j},\bar{k}\in[K]$ and
$({\bm{X}}_{\bar{n}},y_{\bar{n}})\in\texttt{DSET}$ such that
$\langle(\bm{e}_{\bar{i}}-\bm{e}_{\bar{j}})\bm{e}_{\bar{k}}^{\top},\bm{W}\rangle>0$,
${\bm{X}}_{\bar{n}}$ includes the $\bar{j}^{\text{th}}$ token, the last token
of ${\bm{X}}_{\bar{n}}$ is the $\bar{k}^{\text{th}}$ token, and
$y_{\bar{n}}=\bar{i}$. On the other hand, by Assumption 2, $z_{\bar{i}\times
d+\bar{k}}$ and $z_{\bar{j}\times d+\bar{k}}$ in (26) are non-zero for this
input sequence ${\bm{X}}_{\bar{n}}$. Using the equality condition of Cauchy-
Schwartz Inequality in (26), we obtain that $u_{\bar{i}\times
d+\bar{k}}-u_{\bar{j}\times d+\bar{k}}=0$. This implies that

$\displaystyle 0$ $\displaystyle=u_{\bar{i}\times d+\bar{k}}-u_{\bar{j}\times
d+\bar{k}}$
$\displaystyle=\bm{e}_{\bar{i}}^{\top}\bm{W}\bm{e}_{\bar{k}}-\bm{e}_{\bar{j}}^{\top}\bm{W}\bm{e}_{\bar{k}}$
$\displaystyle=(\bm{e}_{\bar{i}}-\bm{e}_{\bar{j}})^{\top}\bm{W}\bm{e}_{\bar{k}}=\langle(\bm{e}_{\bar{i}}-\bm{e}_{\bar{j}})\bm{e}_{\bar{k}}^{\top},\bm{W}\rangle$

which contradicts with the fact that $\|{\bm{u}}\|_{2}>0$. This completes the
proof.

###  B.3 Divergence of
$\left\|{\bf{\Pi}_{{\mathcal{S}_{\text{svm}}}}({\bf{W}}(\tau))}\right\|_{F}$

We first introduce the following lemmas establishing the descent property of
gradient descent for ${\cal{L}}(\bm{W})$ (Lemma 12) and the correlation
between $\nabla{\cal{L}}(\bm{W})$ and the solution of (Graph-SVM)
$\bm{W}^{\text{svm}}$ (Lemma 13) under the setting of Theorem 2. The proofs in
this section follow Appendix B.1 of Tarzanagh et al. (2023a).

######  Lemma 12 (Descent Lemma)

Consider the loss in (17) and choose step size $\eta\leq 1/L$ where $L$ is the
Lipschitzness of $\nabla{\cal{L}}(\bm{W})$ defined in (22). Then from any
initialization $\bm{W}(0)$, Algorithm Algo-GD satisfies:

${\cal{L}}(\bm{W}(\tau+1))-{\cal{L}}(\bm{W}(\tau))\leq-\frac{\eta}{2}\left\|{\nabla{\cal{L}}(\bm{W}(\tau))}\right\|_{F}^{2}$

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

Proof. From (Algo-GD), for $\tau\geq 0$, we have that
$\bm{W}(\tau+1)=\bm{W}(\tau)-\eta\nabla{\cal{L}}(\bm{W}(\tau))$. Since
$\mathcal{L}(\bm{W})$ is $L$-smooth with $L$ defined in (22), we get

$\begin{split}{\cal{L}}(\bm{W}(\tau+1))&\leq{\cal{L}}(\bm{W}(\tau))+\langle{\nabla{\cal{L}}(\bm{W}(\tau)),\bm{W}(\tau+1)-\bm{W}(\tau)}\rangle+\frac{L}{2}\left\|{\bm{W}(\tau+1)-\bm{W}(\tau)}\right\|_{F}^{2}\\\
&={\cal{L}}(\bm{W}(\tau))-\eta\cdot\left\|{\nabla{\cal{L}}(\bm{W}(\tau))}\right\|_{F}^{2}+\frac{L\eta^{2}}{2}\left\|{\nabla{\cal{L}}(\bm{W}(\tau))}\right\|_{F}^{2}\\\
&={\cal{L}}(\bm{W}(\tau))-\eta\left(1-\frac{L\eta}{2}\right)\left\|{\nabla{\cal{L}}(\bm{W}(\tau))}\right\|_{F}^{2}\\\
&\leq{\cal{L}}(\bm{W}(\tau))-\frac{\eta}{2}\left\|{\nabla{\cal{L}}(\bm{W}(\tau))}\right\|_{F}^{2}.\end{split}$

The inequality above also indicates that

$\sum_{\tau=0}^{\infty}\left\|{\nabla{\cal{L}}(\bm{W}(\tau))}\right\|_{F}^{2}\leq\frac{2}{\eta}(\mathcal{L}(\bm{W}(0))-\mathcal{L}^{*})<\infty,~{}~{}~{}\text{and}~{}~{}~{}\lim_{\tau\to\infty}\left\|{\nabla{\cal{L}}(\bm{W}(\tau))}\right\|_{F}^{2}=0.$

###### Lemma 13

Let $\bm{W}^{\text{svm}}$ be the SVM solution of (Graph-SVM) and suppose
$\bm{W}^{\text{svm}}\neq 0$. For any $\bm{W}$ with
$\left\|{\bm{W}}\right\|_{F}<\infty$, the training loss ${\cal{L}}(\bm{W})$ in
(17) obeys $\langle{\nabla\mathcal{L}(\bm{W}),\bm{W}^{\text{svm}}}\rangle<0$.
Equivalently,
$\langle{\bm{\Pi}_{{\mathcal{S}_{\text{svm}}}}(\nabla\mathcal{L}(\bm{W})),\bm{W}^{\text{svm}}}\rangle<0$.

Proof. Recap
${\cal{O}}_{i},\bar{{\cal{O}}}_{i},\mathcal{R}_{i},\bar{\mathcal{R}}_{i},i\in[n]$
in (3). From (18), for any $\bm{W}\in\mathbb{R}^{d\times d}$, we obtain the
gradient

$\nabla{\cal{L}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\sum_{t\in\bar{{\cal{O}}}_{i}}s_{it}(\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top}.$

Then

$\displaystyle\left<\nabla{\cal{L}}(\bm{W}),\bm{W}^{\text{svm}}\right>$
$\displaystyle=\frac{1}{n}\sum_{i\in[n]}\sum_{t\in\bar{{\cal{O}}}_{i}}\left<s_{it}(\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top},\bm{W}^{\text{svm}}\right>$
$\displaystyle=\frac{1}{n}\sum_{i\in[n]}\sum_{t\in\bar{{\cal{O}}}_{i}}s_{it}\cdot\text{trace}\left((\bm{W}^{\text{svm}})^{\top}(\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top}\right)$
$\displaystyle=\frac{1}{n}\sum_{i\in[n]}\sum_{t\in\bar{{\cal{O}}}_{i}}s_{it}\cdot(\bm{x}_{it}-\bm{e}_{y_{i}})^{\top}\bm{W}^{\text{svm}}\bm{\bar{x}}_{i}.$

From the (Graph-SVM) formulation, we have that
$(\bm{x}_{it}-\bm{e}_{y_{i}})^{\top}\bm{W}^{\text{svm}}\bm{\bar{x}}_{i}=0$ for
$t\in\mathcal{R}_{i}$ and
$(\bm{x}_{it}-\bm{e}_{y_{i}})^{\top}\bm{W}^{\text{svm}}\bm{\bar{x}}_{i}\leq-1$
for $t\in\bar{\mathcal{R}}_{i}$. Then $\bm{W}^{\text{svm}}\neq 0$ ensures that
there exists $i\in[n]$ such that $\bar{\mathcal{R}}_{i}\neq\emptyset$, which
implies that

$\left<\nabla{\cal{L}}(\bm{W}),\bm{W}^{\text{svm}}\right><0.$

Using the fact that $\bm{W}^{\text{svm}}\in{\mathcal{S}_{\text{svm}}}$ (Lemma
4) completes the proof.

The next theorem proves the divergence of norm of the iterates $\bm{W}(\tau)$.

###### Theorem 5

Consider the same setting as in Theorem 2, then there is no finite
$\bm{W}\in\mathbb{R}^{d\times d}$ satisfying ${\nabla{\cal{L}}(\bm{W})}=0$.
Furthermore, Algorithm Algo-GD with the step size $\eta\leq 1/L$ where $L$ is
the Lipschitzness of $\nabla{\cal{L}}(\bm{W})$ defined in (22) and any
starting point $\bm{W}(0)$ satisfies
$\lim_{\tau\to\infty}\left\|{\bm{\Pi}_{{\mathcal{S}_{\text{svm}}}}(\bm{W}(\tau))}\right\|_{F}=\infty$.

Proof. Following Lemma 12, when using log-loss $\ell(u)=-\log(u)$, for any
starting point $\bm{W}(0)$, the Algorithm Algo-GD satisfies
$\lim_{\tau\to\infty}\left\|{\nabla{\cal{L}}(\bm{W}(\tau))}\right\|_{F}^{2}=0$.
Moreover, assume that the first claim is wrong and that there is a finite
critical point $\bm{W}$ that satisfies ${\nabla{\cal{L}}(\bm{W})}=0$. We then
have
$\langle\bm{\Pi}_{{\mathcal{S}_{\text{svm}}}}({\nabla{\cal{L}}(\bm{W})}),\bm{W}^{\text{svm}}\rangle=0$.
This leads to a contradiction with Lemma 13 which says that for any finite
$\bm{W}$,
$\langle\bm{\Pi}_{{\mathcal{S}_{\text{svm}}}}(\nabla\mathcal{L}(\bm{W})),\bm{W}^{\text{svm}}\rangle<0$.
This implies that
$\left\|{\bm{\Pi}_{{\mathcal{S}_{\text{svm}}}}(\bm{W}(\tau))}\right\|_{F}\to\infty$.

###  B.4 Uniqueness and Finiteness of ${{{\bf{W^{\text{fin}}}}}}$

###### Lemma 14

Consider the setting of Theorem 2. $\bm{W}^{\text{fin}}$ defined in Theorem 2
is unique and finite.

Proof. Following Lemma 6, it is equivalent to show ${\cal{W}}^{\text{fin}}$
defined in Def. 3 has unique element ${\bm{\bar{W}}}^{\text{fin}}$ and
${\bm{\bar{W}}}^{\text{fin}}=\bm{W}^{\text{fin}}$ is unique and finite. To
start with, recap the definition of $\overline{\texttt{DSET}}$ (Definition 3).
Denote ${\mathcal{I}}\subset[n]$ as in (5), and let
$\bm{s}_{i}=\mathbb{S}(\bar{\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i})$ where
$(\bar{\bm{X}}_{i},y_{i})\in\overline{\texttt{DSET}}$. What’s more, recap the
ERM loss from (17) and loss function $\ell(u)=-\log(u)$. Then we have

$\displaystyle{\bm{\bar{W}}}^{\text{fin}}=\arg\min_{\bm{W}\in\mathcal{S}_{\text{fin}}}{\bar{\cal{L}}}(\bm{W})\quad\text{where}\quad{\bar{\cal{L}}}(\bm{W})=\frac{1}{n}\sum_{i\in{\mathcal{I}}}-\log\left(\sum_{t\in{\cal{O}}_{i}}s_{it}\right).$
(28)

Different from (3), ${\cal{O}}_{i}$ for dataset $\overline{\texttt{DSET}}$ is
defined as follows:

${\cal{O}}_{i}:=\left\\{t~{}|~{}x_{it}={y_{i}},~{}\bm{x}_{it}\in\bar{\bm{X}}_{i},~{}t\in[\bar{T}_{i}]\right\\},$

where $\bar{T}_{i}$ is the number of tokens – tokens that are in the same SCC
as label token ${y_{i}}$ within their corresponding TPG – in
$\bar{\bm{X}}_{i}$ and if recap the notation of $\mathcal{R}_{i}$ in (3), here
we have $|\mathcal{R}_{i}|=\bar{T}_{i}$.

We will first prove that ${\bm{\bar{W}}}^{\text{fin}}$ is finite by
contradiction. Specifically, we will show that for any
$\bm{W}\in\mathcal{S}_{\text{fin}}$ with $\left\|{\bm{W}}\right\|_{F}\neq 0$,
$\lim_{R\to\infty}{\bar{\cal{L}}}(R\cdot\bm{W})=\infty$ which implies that the
optimal solution ${\bm{\bar{W}}}^{\text{fin}}$ has to be finite.

Let $\bm{W}\in\mathcal{S}_{\text{fin}}$ be arbitrary attention weight.
Following the definition of $\mathcal{S}_{\text{fin}}$ as in Def. 1, we have
that $(\bm{e}_{i}-\bm{e}_{j})\bm{e}_{k}^{\top}\in\mathcal{S}_{\text{fin}}$ for
all $(i\asymp j)\in{\cal{G}}^{(k)}$ and $k\in[K]$. For any $\bar{\bm{X}}$, let
${\cal{O}},\bar{{\cal{O}}}$ correspond to the token index sets. Then we have

$\sum_{t\in{\cal{O}}}s_{t}=\frac{|{\cal{O}}|e^{\bm{e}_{y}^{\top}(R\cdot\bm{W})\bm{e}_{k}}}{|{\cal{O}}|e^{\bm{e}_{y}^{\top}(R\cdot\bm{W})\bm{e}_{k}}+\sum_{t\in\bar{{\cal{O}}}}e^{\bm{e}_{t}^{\top}(R\cdot\bm{W})\bm{e}_{k}}}=\frac{1}{1+\sum_{t\in\bar{{\cal{O}}}}e^{(\bm{e}_{t}-\bm{e}_{y})^{\top}(R\cdot\bm{W})\bm{e}_{k}}/|{\cal{O}}|}.$

Given the sample loss $\ell=-\log\left(\sum_{t\in{\cal{O}}}s_{t}\right)$ and
to prevent it from divergence as $R\to\infty$, that is,
$\sum_{t\in{\cal{O}}}s_{t}\not\to 0$, we have

$\sum_{t\in\bar{{\cal{O}}}}e^{(\bm{e}_{t}-\bm{e}_{y})^{\top}(R\cdot\bm{W})\bm{e}_{k}}\not\to\infty\quad\Longrightarrow\quad(\bm{e}_{y}-\bm{e}_{t})^{\top}\bm{W}\bm{e}_{k}\geq
0\text{ for all }t\in\bar{{\cal{O}}}$

where $\bm{e}_{y}$ is the label token and $\bm{e}_{t}$ is any other token in
$\bar{{\cal{O}}}$. Recap from the construction of TPG in Section 2.1, the
directed edge $y\to t$ exists in the graph ${\cal{G}}^{(k)}$. Since SCC is
bidirectionally reachable, which means there exists route from $t$ to $y$,
e.g., $t\to p_{1}\to p_{2}\to\cdots p_{m}\to y$, similarly we have

$(\bm{e}_{t}-\bm{e}_{p_{1}})^{\top}\bm{W}\bm{e}_{k},~{}(\bm{e}_{p_{1}}-\bm{e}_{p_{2}})^{\top}\bm{W}\bm{e}_{k},~{}\cdots,(\bm{e}_{p_{m}}-\bm{e}_{y})^{\top}\bm{W}\bm{e}_{k}\geq
0\Longrightarrow(\bm{e}_{t}-\bm{e}_{y})^{\top}\bm{W}\bm{e}_{k}\geq 0.$

Combining results in that $(\bm{e}_{t}-\bm{e}_{y})^{\top}\bm{W}\bm{e}_{k}=0$.
This implies that for all $(i\asymp j)\in{\cal{G}}^{(k)}$ and $R\to\infty$, to
ensure the training loss ${\bar{\cal{L}}}(R\cdot\bm{W})$ finite,
$(\bm{e}_{i}-\bm{e}_{j})^{\top}\bm{W}\bm{e}_{k}=0$, which contradicts the
facts that $\bm{W}\in\mathcal{S}_{\text{fin}}$ and $\bm{W}\neq 0$.

Next, we prove that there is at most one local minimum for
${\bar{\cal{L}}}(\bm{W})$ based on Lemma 2. Suppose to the contrary that we
have two optimal solutions satisfying
$\min_{\bm{W}}{\bar{\cal{L}}}(\bm{W})={\bar{\cal{L}}}(\bm{W}^{\text{fin}}_{1})={\bar{\cal{L}}}(\bm{W}^{\text{fin}}_{2}),\bm{W}^{\text{fin}}_{1}\neq\bm{W}^{\text{fin}}_{2}$.
From Lemma 2, since ${\cal{L}}(\bm{W})$ is strictly convex over subspace
$\mathcal{S}_{\text{fin}}$, for any
$\bm{W}_{1},\bm{W}_{2}\in\mathcal{S}_{\text{fin}},\lambda\in(0,1)$, if
$\bm{W}_{1}\neq\bm{W}_{2}$, we have

${\bar{\cal{L}}}((1-\lambda)\bm{W}_{1}+\lambda\bm{W}_{2})<(1-\lambda){\bar{\cal{L}}}(\bm{W}_{1})+\lambda{\bar{\cal{L}}}(\bm{W}_{2})$
(29)

Substitute
$\bm{W}_{1}=\bm{W}^{\text{fin}}_{1},\bm{W}_{2}=\bm{W}^{\text{fin}}_{2}$, we
get

${\bar{\cal{L}}}((1-\lambda)\bm{W}^{\text{fin}}_{1}+\lambda\bm{W}^{\text{fin}}_{2})<(1-\lambda){\bar{\cal{L}}}(\bm{W}^{\text{fin}}_{1})+\lambda{\bar{\cal{L}}}(\bm{W}^{\text{fin}}_{2})=\min_{\bm{W}}{\bar{\cal{L}}}(\bm{W})$
(30)

which leads to a contradiction to the assumption that
$\bm{W}^{\text{fin}}_{1}\text{ and }\bm{W}^{\text{fin}}_{2}$ are both optimal
solutions. Combining this with the fact that $\bm{W}^{\text{fin}}$ is not
attained at infinity, there exists one unique and finite solution with
${\bm{\bar{W}}}^{\text{fin}}=\arg\min_{\bm{W}\in\mathcal{S}_{\text{fin}}}{\bar{\cal{L}}}(\bm{W})$.

###  B.5 Proof of Theorem 2

###### Lemma 15

Consider the same setting of Theorem 2. Given any $\pi>0$, there exists
$R_{\pi}>0$ such that for any $\bm{W}$ with
$\left\|{\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W})}\right\|_{F}<\infty$ and
$\left\|{\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\bm{W})}\right\|_{F}>R_{\pi}$,

${\cal{L}}(\bm{W})\geq{\cal{L}}\left((1+\pi)\frac{\left\|{\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\bm{W})}\right\|_{F}}{\left\|{\bm{W}^{\text{svm}}}\right\|_{F}}\bm{W}^{\text{svm}}+\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W})\right).$

Proof. Recap
${\cal{O}}_{i},\bar{{\cal{O}}}_{i},\mathcal{R}_{i},\bar{\mathcal{R}}_{i}$,
$i\in[n]$ from (3) and recap from (17), we get that for any $\bm{W}$,

${\cal{L}}(\bm{W})=-\frac{1}{n}\sum_{i=1}^{n}\log\left(\sum_{t\in{\cal{O}}_{i}}s_{it}\right)$

where $\bm{s}_{i}=\mathbb{S}({\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i})$.

To obtain the result, we establish a refined softmax probability control by
studying the distance to ${\bar{\cal{L}}}(\bm{W})$ as defined in Definition 4.
Let $\bm{W}^{\parallel}=\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W})$,
$\bm{W}^{\perp}=\bm{W}-\bm{W}^{\parallel}$,
$\left\|{\bm{W}^{\perp}}\right\|_{F}=R$, and
$\Theta=1/\left\|{\bm{W}^{\text{svm}}}\right\|_{F}$. Let
$\bm{b}_{i}={\bm{X}}_{i}\bm{W}^{\parallel}\bm{\bar{x}}_{i}$,
$\bm{a}_{i}^{\star}={\bm{X}}_{i}((1+\pi)R\Theta\cdot\bm{W}^{\text{svm}})\bm{\bar{x}}_{i}$,
and
$\bm{s}_{i}^{\star}=\mathbb{S}({\bm{X}}_{i}((1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}+\bm{W}^{\parallel})\bm{\bar{x}}_{i})=\mathbb{S}(\bm{a}_{i}^{\star}+\bm{b}_{i})$.
Additionally, let $\bm{a}_{i}={\bm{X}}_{i}\bm{W}^{\perp}\bm{\bar{x}}_{i}$,
$\bm{s}_{i}=\mathbb{S}({\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i})=\mathbb{S}(\bm{a}_{i}+\bm{b}_{i})$,
$\gamma_{i}^{\star}:=\bm{c}_{y_{i}}^{\top}{\bm{X}}_{i}^{\top}\bm{s}_{i}^{\star}$,
and $\gamma_{i}:=\bm{c}_{y_{i}}^{\top}{\bm{X}}_{i}^{\top}\bm{s}_{i}$.

From proof of Lemma 5 (more specifically (12)), we get for all
$t,t^{\prime}\in\mathcal{R}_{i}$

$(\bm{x}_{it}-\bm{x}_{it^{\prime}})^{\top}{\bm{V}}\bm{\bar{x}}_{i}=0\quad\text{for
any }{\bm{V}}\perp\mathcal{S}_{\text{fin}}\Longrightarrow
a^{\star}_{it}-a^{\star}_{it^{\prime}}=a_{it}-a_{it^{\prime}}=0.$

Additionally, since
$\frac{\bm{W}}{\left\|{\bm{W}}\right\|_{F}}\neq\Theta\bm{W}^{\text{svm}}$,
there exist $i\in[n],t\in{\cal{O}}_{i},t^{\prime}\in\bar{\mathcal{R}}_{i}$
such that
$(\bm{x}_{it}-\bm{x}_{it^{\prime}})^{\top}\bm{W}\bm{\bar{x}}_{i}<R\Theta$.
Then,

$\displaystyle\sum_{t\in{\cal{O}}_{i}}s_{it}$
$\displaystyle=\frac{\sum_{t\in{\cal{O}}_{i}}e^{a_{it}+b_{it}}}{\sum_{t\in[T_{i}]}e^{a_{it}+b_{it}}}\leq\frac{\sum_{t\in{\cal{O}}_{i}}e^{b_{it}}}{\sum_{t\in\mathcal{R}_{i}}e^{b_{it}}+e^{-R\Theta+b_{it^{\prime}}}}\leq\frac{c_{i}}{d_{i}+e^{-R\Theta-\bar{b}}},\quad\exists
i\in[n]$ $\displaystyle\sum_{t\in{\cal{O}}_{i}}s_{it}^{\star}$
$\displaystyle=\frac{\sum_{t\in{\cal{O}}_{i}}e^{a^{\star}_{it}+b_{it}}}{\sum_{t\in[T_{i}]}e^{a^{\star}_{it}+b_{it}}}\geq\frac{\sum_{t\in{\cal{O}}_{i}}e^{b_{it}}}{\sum_{t\in\mathcal{R}_{i}}e^{b_{it}}+\sum_{t\in\bar{\mathcal{R}}_{i}}e^{-(1+\pi)R\Theta+b_{it}}}\geq\frac{c_{i}}{d_{i}+Te^{-(1+\pi)R\Theta+\bar{b}}},\quad\forall
i\in[n],$

where $c_{i}=\sum_{t\in{\cal{O}}_{i}}e^{b_{it}}$,
$d_{i}=\sum_{t\in\mathcal{R}_{i}}e^{b_{it}}$, and
$\bar{b}:=\max_{t\in\bar{\mathcal{R}}_{i},i\in[n]}|b_{it}|$, and we have

${\bar{\cal{L}}}(\bm{W})={\bar{\cal{L}}}(\bm{W}^{\parallel})=-\frac{1}{n}\sum_{i\in{\mathcal{I}}}\log\left(\frac{c_{i}}{d_{i}}\right).$

Then we obtain

$\displaystyle{\cal{L}}(\bm{W})-{\bar{\cal{L}}}(\bm{W})$
$\displaystyle\geq-\frac{1}{n}\left(\log\left(\frac{c_{i}}{d_{i}+e^{-R\Theta-\bar{b}}}\right)-\log\left(\frac{c_{i}}{d_{i}}\right)\right)$
$\displaystyle\geq\frac{1}{n}\log\left(1+e^{-R\Theta-\bar{b}}/d_{i}\right)$

and let
$j:=\arg\max_{i\in[n]}\left(-\log\left(\sum_{t\in{\cal{O}}_{i}}s_{it}^{\star}\right)+\log\left(\frac{c_{i}}{d_{i}}\right)\right)$.
We can upper-bound the loss difference for
$(1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}+\bm{W}^{\parallel}$ as follows:

$\displaystyle{\cal{L}}((1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}+\bm{W}^{\parallel})-{\bar{\cal{L}}}(\bm{W})$
$\displaystyle\leq\max_{i\in[n]}\left(-\log\left(\sum_{t\in{\cal{O}}_{i}}s_{it}^{\star}\right)+\log\left(\frac{c_{i}}{d_{i}}\right)\right)$
$\displaystyle=\log\left(1+Te^{-(1+\pi)R\Theta+\bar{b}}/d_{j}\right).$

Combining them together results in that,
${\cal{L}}(\bm{W})>{\cal{L}}((1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}+\bm{W}^{\parallel})$
whenever

$\displaystyle\frac{1}{n}\log\left(1+e^{-R\Theta-\bar{b}}/d_{i}\right)\geq\log\left(1+Te^{-(1+\pi)R\Theta+\bar{b}}/d_{j}\right).$

Given that $x/2\leq\log(1+x)$ for any $0\leq x\leq 1$, we get
${\cal{L}}(\bm{W})>{\cal{L}}((1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}+\bm{W}^{\parallel})$
whenever

$\displaystyle\frac{e^{-R\Theta-\bar{b}}}{2nd_{i}}\geq\frac{Te^{-(1+\pi)R\Theta+\bar{b}}}{d_{j}}\quad\text{when}\quad
R\geq-\frac{\log(d_{j})+\bar{b}}{\Theta}$ $\displaystyle\Longrightarrow$
$\displaystyle
R>R_{\pi}:=\max\left\\{\frac{1}{\pi\Theta}\log\left(\frac{2nTd_{i}}{d_{j}}\right)+\frac{2\bar{b}}{\pi\Theta},-\frac{\log(d_{j})+\bar{b}}{\Theta}\right\\}.$
(31)

Here, since $\left\|{\bm{W}^{\parallel}}\right\|_{F}<\infty$,
$d_{i},d_{j},\bar{b}<\infty$.

Proof of Theorem 2. Now gathering all the results so far, we are ready to
prove the gradient descent convergence. The divergence of
$\left\|{\bm{W}(\tau)}\right\|_{F}$ as $\tau\to\infty$ has been proven by
Theorem 5.

$\bullet$ We first show that
$\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}(\tau))\to\bm{W}^{\text{fin}}$.
Lemma 6 has established the equivalence between ${\bm{\bar{W}}}^{\text{fin}}$
and $\bm{W}^{\text{fin}}$. Since ${\cal{L}}(\bm{W})$ is convex following Lemma
2, we have that
${\cal{L}}(\bm{W}(\tau))\to{\cal{L}}_{\star}:=\min_{\bm{W}}{\cal{L}}(\bm{W})$.
Additionally, Lemma 2 shows that ${\cal{L}}(\bm{W})$ is strictly convex on
$\mathcal{S}_{\text{fin}}$ and Lemma 14 shows that $\bm{W}^{\text{fin}}$ is
the unique finite solution. Suppose
$\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}(\tau))\not\to\bm{W}^{\text{fin}}$.
Let $\bm{W}$ be any matrix with
$\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W})\neq\bm{W}^{\text{fin}}$. Then
Lemma 7 and Lemma 5 give that
${\cal{L}}(\bm{W})\geq{\bar{\cal{L}}}(\bm{W})={\bar{\cal{L}}}(\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}))>{\cal{L}}_{\star}$
where
${\cal{L}}_{\star}=\min_{\bm{W}}{\cal{L}}(\bm{W})=\min_{\bm{W}}{\bar{\cal{L}}}(\bm{W})$.
Given that ${\cal{L}}_{\star}$ is achievable, the proof is done by
contradiction and we have that
$\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}(\tau))\to\bm{W}^{\text{fin}}$.

$\bullet$ We next prove that when $\bm{W}^{\text{svm}}=0$,
$\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\bm{W}(\tau))=\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\bm{W}(0))$.
Recap the gradient in (18) where

$\nabla{\cal{L}}(\bm{W})=\frac{1}{n}\sum_{i=1}^{n}\sum_{t\in\bar{{\cal{O}}}_{i}}s_{it}(\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top}.$

Since $\bm{W}^{\text{svm}}=0$ implies that for all $i\in[n]$,
$\bar{\mathcal{R}}_{i}=\emptyset$, then
$\bar{{\cal{O}}}_{i}\subseteq\mathcal{R}_{i}$. Additionally, since following
definition of $\mathcal{S}_{\text{fin}}$ from Def. 1, for any
$t\in\mathcal{R}_{i}$,
$(\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top}\in\mathcal{S}_{\text{fin}}$.
Then we obtain

$\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\nabla{\cal{L}}(\bm{W}))=\frac{1}{n}\sum_{i=1}^{n}\sum_{t\in\bar{{\cal{O}}}_{i}}s_{it}\bm{\Pi}_{\mathcal{S}_{\text{fin}}}\left((\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top}\right)=\frac{1}{n}\sum_{i=1}^{n}\sum_{t\in\bar{{\cal{O}}}_{i}}s_{it}(\bm{x}_{it}-\bm{e}_{y_{i}})\bm{\bar{x}}_{i}^{\top}=\nabla{\cal{L}}(\bm{W}).$

Therefore,
$\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\nabla{\cal{L}}(\bm{W}))=0$ for
any $\bm{W}$, which completes the proof.

$\bullet$ Last, we show that when $\bm{W}^{\text{svm}}\neq 0$,
$\bm{W}(\tau)/\left\|{\bm{W}(\tau)}\right\|_{F}\to\bm{W}^{\text{svm}}/\left\|{\bm{W}^{\text{svm}}}\right\|_{F}$.

Consider any $\bm{W}\in\mathbb{R}^{d\times d}$, and let
$\bm{W}^{\perp}=\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\bm{W})$,
$\bm{W}^{\parallel}=\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W})$,
$R=\left\|{\bm{W}^{\perp}}\right\|_{F}$, and
$\Theta=1/\left\|{\bm{W}^{\text{svm}}}\right\|_{F}$. Next, from Lemma 12, for
any $\tau\geq 0$, ${\cal{L}}(\bm{W}(\tau+1))\leq{\cal{L}}(\bm{W}(\tau))$. Let
$\bm{W}^{\perp}(\tau)=\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\bm{W}(\tau))$
and
$\bm{W}^{\parallel}(\tau)=\bm{\Pi}_{\mathcal{S}_{\text{fin}}}(\bm{W}(\tau))$.
Following Lemma 7, since loss satisfies $\ell(1)=-\log(1)=0$, we obtain

$\displaystyle{\cal{L}}(\bm{W}(\tau))\geq{\bar{\cal{L}}}(\bm{W}^{\parallel}(\tau)).$

Since following Lemma 14, the training risk
${\bar{\cal{L}}}(\bm{W}^{\parallel}(\tau))$ is infinite if
$\left\|{\bm{W}^{\parallel}(\tau)}\right\|_{F}\to\infty$, which implies
$\left\|{\bm{W}^{\parallel}(\tau)}\right\|_{F}<\infty$ for any $\tau\geq 0$.
Additionally, Theorem 5 proves the divergence of $\bm{W}(\tau)$ as
$\tau\to\infty$, hence we have
$\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}\to\infty$.

Applying Lemma 15, as well as the fact that
$\left\|{\bm{W}^{\parallel}(\tau)}\right\|_{F}<\infty$ and
$\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}\to\infty$, there exists
sufficiently large $R_{\pi}$ as defined in (31) such that once
$\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}=R>R_{\pi}$,
${\cal{L}}(\bm{W}(\tau))-{\cal{L}}((1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}+\bm{W}^{\parallel}(\tau))\geq
0$. Since ${\cal{L}}(\bm{W})$ is convex following Lemma 2, we have that

$\begin{split}{\cal{L}}(\bm{W})&\leq{\cal{L}}((1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}+\bm{W}^{\parallel})+\left<\nabla{\cal{L}}(\bm{W}),\bm{W}-\left((1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}+\bm{W}^{\parallel}\right)\right>\\\
&={\cal{L}}((1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}+\bm{W}^{\parallel})+\left<\nabla{\cal{L}}(\bm{W}),\left(\bm{W}^{\perp}-(1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}\right)\right>\\\
&={\cal{L}}((1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}+\bm{W}^{\parallel})+\left<\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\nabla{\cal{L}}(\bm{W})),\left(\bm{W}^{\perp}-(1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}\right)\right>.\end{split}$
(32)

Here, the first inequality uses the convexity of ${\cal{L}}(\bm{W})$ and last
equation is obtained from the fact that
$\bm{W}^{\text{svm}}\perp\mathcal{S}_{\text{fin}}$. It implies that once
$\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}=R>R_{\pi}$,

$\left<\bm{\Pi}_{\mathcal{S}_{\text{fin}}^{\perp}}(\nabla{\cal{L}}(\bm{W})),\left(\bm{W}^{\perp}-(1+\pi)R\Theta\cdot\bm{W}^{\text{svm}}\right)\right>\geq
0.$

Now we choose $\tau_{0}$ such that for all $\tau\geq\tau_{0}$,
$\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}>R_{\pi}$. Then for $\tau>\tau_{0}$,
we get

$\displaystyle\left<\bm{W}^{\perp}(\tau+1)-\bm{W}^{\perp}(\tau),\frac{\bm{W}^{\text{svm}}}{\left\|{\bm{W}^{\text{svm}}}\right\|_{F}}\right>\geq\frac{1}{1+\pi}\left<\bm{W}^{\perp}(\tau+1)-\bm{W}^{\perp}(\tau),\frac{\bm{W}^{\perp}(\tau)}{\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}}\right>$
(33)

where

$\displaystyle\left<\bm{W}^{\perp}(\tau+1)-\bm{W}^{\perp}(\tau),\frac{\bm{W}^{\perp}(\tau)}{\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}}\right>$
$\displaystyle=\frac{1}{2\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}}\left(\left\|{\bm{W}^{\perp}(\tau+1)}\right\|_{F}^{2}-\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}^{2}-\left\|{\bm{W}^{\perp}(\tau+1)-\bm{W}^{\perp}(\tau)}\right\|_{F}^{2}\right)$
$\displaystyle\geq\frac{\left\|{\bm{W}^{\perp}(\tau+1)}\right\|_{F}^{2}-\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}^{2}}{2\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}}-\left\|{\bm{W}^{\perp}(\tau+1)-\bm{W}^{\perp}(\tau)}\right\|_{F}^{2}$
(34)
$\displaystyle\geq\left\|{\bm{W}^{\perp}(\tau+1)}\right\|_{F}-\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}-\left\|{\bm{W}^{\perp}(\tau+1)-\bm{W}^{\perp}(\tau)}\right\|_{F}^{2}$
(35)
$\displaystyle\geq\left\|{\bm{W}^{\perp}(\tau+1)}\right\|_{F}-\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}-\left\|{\bm{W}(\tau+1)-\bm{W}(\tau)}\right\|_{F}^{2}$
(36)
$\displaystyle\geq\left\|{\bm{W}^{\perp}(\tau+1)}\right\|_{F}-\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}+2\eta\left({\cal{L}}(\bm{W}(\tau+1))-{\cal{L}}(\bm{W}(\tau))\right).$
(37)

Here, (33) is obtained from (32) and holds for all $\tau>\tau_{0}$; (34) comes
from the fact that $\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}>0.5$; (35)
follows that for any $a,b>0$, $(a^{2}-b^{2})/2b>a-b$; (36) follows the
projection property that
$\left\|{\bm{W}(\tau+1)-\bm{W}(\tau)}\right\|_{F}^{2}=\left\|{\bm{W}^{\perp}(\tau+1)-\bm{W}^{\perp}(\tau)}\right\|_{F}^{2}+\left\|{\bm{W}^{\parallel}(\tau+1)-\bm{W}^{\parallel}(\tau)}\right\|_{F}^{2}$;
and (37) is obtained via Lemma 12.

Summing the above inequality over $\tau\geq\tau_{0}$ obtains

$\displaystyle\left<\bm{W}^{\perp}(\tau)-\bm{W}^{\perp}(\tau_{0}),\frac{\bm{W}^{\text{svm}}}{\left\|{\bm{W}^{\text{svm}}}\right\|_{F}}\right>\geq\frac{1}{1+\pi}\left(\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}-\left\|{\bm{W}^{\perp}(\tau_{0})}\right\|_{F}+2\eta\left({\cal{L}}(\bm{W}(\tau))-{\cal{L}}(\bm{W}(\tau_{0}))\right)\right)$
$\displaystyle\Longrightarrow\left<\frac{\bm{W}^{\perp}(\tau)}{\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}},\frac{\bm{W}^{\text{svm}}}{\left\|{\bm{W}^{\text{svm}}}\right\|_{F}}\right>\geq\frac{1}{1+\pi}\left(1+\frac{C}{\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}}\right)$

where

$C:=\left<\bm{W}^{\perp}(\tau_{0}),\frac{\bm{W}^{\text{svm}}}{\left\|{\bm{W}^{\text{svm}}}\right\|_{F}}\right>-\left\|{\bm{W}^{\perp}(\tau_{0})}\right\|_{F}+2\eta\left({\cal{L}}(\bm{W}(\tau))-{\cal{L}}(\bm{W}(\tau_{0}))\right).$

Since $\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}\to\infty$ and
$0<{\cal{L}}(\bm{W}(\tau))\leq{\cal{L}}(\bm{W}(0))<\infty$, we get

$\displaystyle\lim_{\tau\to\infty}\left<\frac{\bm{W}^{\perp}(\tau)}{\left\|{\bm{W}^{\perp}(\tau)}\right\|_{F}},\frac{\bm{W}^{\text{svm}}}{\left\|{\bm{W}^{\text{svm}}}\right\|_{F}}\right>\geq\frac{1}{1+\pi}.$
(38)

Choosing $\pi\to 0$ and combining (38) with the fact that
$\lim_{\tau\to\infty}\left\|{\bm{W}^{\parallel}(\tau)}\right\|_{F}<\infty$
completes the proof.

##  Appendix C Global Convergence of Regularization Path

###  C.1 Proof of Theorem 3

###### Lemma 16

Suppose Assumptions 1 and 2 hold. Additionally, assume loss
$\ell:\mathbb{R}\rightarrow\mathbb{R}$ is strictly decreasing and
$|\ell^{\prime}|$ is bounded. Define
$\bar{\bm{W}}_{R}^{\perp}:=\bar{\bm{W}}_{R}^{\perp}(\bm{W}^{\parallel})\in\mathcal{S}_{\text{fin}}^{\perp}$
by

$\displaystyle\bar{\bm{W}}_{R}^{\perp}:=\arg\min_{\bm{W}\in\mathcal{S}_{\text{fin}}^{\perp},\left\|{\bm{W}}\right\|_{F}\leq
R}{\cal{L}}(\bm{W}+\bm{W}^{\parallel}).$ (39)

Let $\bm{W}^{\text{svm}}\neq 0$ denote the solution of (Graph-SVM). Then we
have that for any $\bm{W}^{\parallel}\in\mathcal{S}_{\text{fin}}$ with
$\left\|{\bm{W}^{\parallel}}\right\|_{F}<\infty$

$\lim_{R\to\infty}\frac{\bar{\bm{W}}_{R}^{\perp}}{R}=\frac{\bm{W}^{\text{svm}}}{\left\|{\bm{W}^{\text{svm}}}\right\|_{F}}.$

Proof. Recap
${\cal{O}}_{i},\bar{{\cal{O}}}_{i},\mathcal{R}_{i},\bar{\mathcal{R}}_{i}$,
$i\in[n]$ from (3). Let ${\bm{\gamma}}_{i}={\bm{X}}_{i}\bm{c}_{y_{i}}$ where
following Assumption 1 we have

$\gamma_{it}=\bm{x}_{it}^{\top}\bm{c}_{y_{i}}=\begin{cases}1,\quad
t\in{\cal{O}}_{i}\\\ 0,\quad t\in\bar{{\cal{O}}}_{i}.\end{cases}$

Recap ${\mathcal{I}},\bar{\mathcal{I}}$ from (4). To proceed, define
${\gamma}^{\text{max}}_{i}$ as follows:

  * •

Consider $i\in\bar{\mathcal{I}}$. Given
$\min_{\bm{W}}\ell(\bm{c}_{y_{i}}^{\top}{\bm{X}}_{i}^{\top}\mathbb{S}({\bm{X}}_{i}\bm{W}\bm{\bar{x}}_{i}))\geq\ell(\bm{c}_{y_{i}}^{\top}\bm{e}_{y_{i}})$,
we define the maximal score
${\gamma}^{\text{max}}_{i}:=\bm{c}_{y_{i}}^{\top}\bm{e}_{y_{i}}=1$.

  * •

Consider $i\in{\mathcal{I}}$.

    1. 1.

Assumption 1 ensures that all tokens, excluding the ones with token ID
$x_{it}=y_{i}$, return zero score, that is,
$\bm{c}_{y_{i}}^{\top}\bm{e}_{k}=0$ for $k\neq y_{i}$.

    2. 2.

From proof of Lemma 4, for any
$\bm{W}^{\perp}\in\mathcal{S}_{\text{fin}}^{\perp}$ and $t\in\mathcal{R}_{i}$,
$\bm{x}_{it}^{\top}(\bm{W}^{\perp}+\bm{W}^{\parallel})\bm{\bar{x}}_{i}=\bm{x}_{it}^{\top}\bm{W}^{\parallel}\bm{\bar{x}}_{i}+\bar{a}_{i}$,
where $\bar{a}_{i}$ is some constant associated with $\bm{W}^{\perp}$ and
remains the same value within the same SCC. Let
$\bm{b}_{i}={\bm{X}}_{i}\bm{W}^{\parallel}\bm{\bar{x}}_{i}$. Then the
probabilities for $t\in\mathcal{R}_{i}$ (if denoted by $s_{it}$) obey

$\displaystyle\frac{s_{it}}{\sum_{t^{\prime}\in\mathcal{R}_{i}}s_{it^{\prime}}}=\frac{e^{b_{it}+\bar{a}_{i}}}{\sum_{t^{\prime}\in\mathcal{R}_{i}}e^{b_{it^{\prime}}+\bar{a}_{i}}}=\frac{e^{b_{it}}}{\sum_{t^{\prime}\in\mathcal{R}_{i}}e^{b_{it^{\prime}}}},$
(40)

which means that the probability distribution over set $\mathcal{R}_{i}$
remains the same with varying $\bm{W}^{\perp}$.

Combining both, we define the maximal score as follows:

${\gamma}^{\text{max}}_{i}:=|{\cal{O}}_{i}|\cdot\bar{s}_{i}\quad\text{where}\quad\bar{s}_{i}=\frac{e^{\bm{e}_{y_{i}}^{\top}\bm{W}^{\parallel}\bm{\bar{x}}_{i}}}{\sum_{t^{\prime}\in\mathcal{R}_{i}}e^{b_{it^{\prime}}}}.$

Note that if consider the cyclic subdataset $\overline{\texttt{DSET}}$ as in
(5). Let $(\bar{{\bm{X}}}_{i},y_{i})\in\overline{\texttt{DSET}}$ where
$\bar{{\bm{X}}}_{i}$ is the corresponding sequence by removing the tokens in
$\bar{\mathcal{R}}_{i}$. Then we have
${\gamma}^{\text{max}}_{i}=\bm{c}_{y_{i}}^{\top}\bar{\bm{X}}_{i}^{\top}\mathbb{S}(\bar{\bm{X}}_{i}\bm{W}^{\parallel}\bm{\bar{x}}_{i})$.

Hence, given $\bm{b}_{i}={\bm{X}}_{i}\bm{W}^{\parallel}\bm{\bar{x}}_{i}$, we
obtain

${\gamma}^{\text{max}}_{i}=\begin{cases}1,&i\in\bar{\mathcal{I}}\\\
\frac{|{\cal{O}}_{i}|e^{\bm{e}_{y_{i}}^{\top}\bm{W}^{\parallel}\bm{\bar{x}}_{i}}}{\sum_{t^{\prime}\in\mathcal{R}_{i}}e^{b_{it^{\prime}}}},&i\in{\mathcal{I}}.\end{cases}$

Then we define the optimal risk of (39) and its corresponding softmax
probabilities $\bm{s}^{\max}_{i}$, $i\in[n]$ as follows:

${\cal{L}}_{\star}^{\bm{W}^{\parallel}}:=\frac{1}{n}\sum_{i=1}^{n}\ell({\gamma}^{\text{max}}_{i}),\quad\text{and}\quad
s^{\max}_{it}=\begin{cases}0,&t\in\bar{\mathcal{R}}_{i}\\\
\frac{e^{b_{it}}}{\sum_{t^{\prime}\in\mathcal{R}_{i}}e^{b_{it^{\prime}}}},&t\in\mathcal{R}_{i}\end{cases}\quad\text{for
all }i\in[n].$

Note that we also have