On the Optimization and Generalization of Two-layer Transformers with Sign Gradient Descent

## Abstract

Abstract The Adam optimizer is widely used for transformer optimization in practice, which makes understanding the underlying optimization mechanisms an important problem.
However, due to the Adam’s complexity, theoretical analysis of how it optimizes transformers remains a challenging task.
Fortunately, Sign Gradient Descent (SignGD) serves as an effective surrogate for Adam.
Despite its simplicity, theoretical understanding of how SignGD optimizes transformers still lags behind.
In this work, we study how SignGD optimizes a two-layer transformer – consisting of a softmax attention layer with trainable query-key parameterization followed by a linear layer – on
a linearly separable noisy dataset.
We identify four stages in the training dynamics, each exhibiting intriguing behaviors.
Based on the training dynamics, we prove the fast convergence but poor generalization of the learned transformer on the noisy dataset.
We also show that Adam behaves similarly to SignGD in terms of both optimization and generalization in this setting.
Additionally, we find that the poor generalization of SignGD is not solely due to data noise,
suggesting that both SignGD and Adam requires high-quality data for real-world tasks.
Finally, experiments on synthetic and real-world datasets empirically support our theoretical results.

## 1 Introduction

The transformer architecture (Vaswani et al., ) has become ubiquitous across various domains, achieving state-of-the-art results in areas such as language modeling (Devlin et al., ; Brown et al., ), computer vision (Dosovitskiy et al., ; Peebles & Xie, ), and reinforcement learning (Chen et al., ).
Regardless of the specific task or data modality, the Adam optimizer (Kingma & Ba, ) is typically employed to train large transformer models, making it the de facto choice in practice.
This widespread use highlights that understanding the inner mechanism on how Adam optimizes transformers is an important problem.
However,
the complexity of Adam’s formulation presents significant challenges for rigorous analysis.
Many of the underlying mechanisms of how Adam optimizes transformers are still poorly understood.

Recent theoretical works (Jelassi et al., ; Tarzanagh et al., ; Tian et al., ) study the training dynamics of transformers across various datasets and objectives.
Training dynamics analysis allows us to trace the evolution of model parameters throughout the training process. In doing so, it enables a precise description of the optimization process, which can ultimately lead to new insights on convergence and generalization results.
However, analyzing the training dynamics of transformers presents many challenges.
The transformer architecture is inherently more complex than simpler models like MLPs (Wang & Ma, ; Xu & Du, ) and CNNs (Cao et al., ; Kou et al., ), making a detailed analysis of its training dynamics more challenging.
To facilitate such analyses, researchers often introduced relaxed assumptions, such as using linear attention (Zhang et al., ) or unrealistic initialization (Li et al., ).
A more commonly employed assumption in theoretical works is reparameterizing the query and key matrices into a single joint attention matrix, as seen in many studies (e.g., Tian et al. ()). While this assumption simplifies the analysis, it remains unrealistic in practice.
Moreover, existing analyses focus primarily on Gradient Descent (GD) or Stochastic Gradient Descent (SGD), with little attention paid to optimizers like Adam. The analysis of transformer training dynamics remains an active area of research.

Our work addresses a crucial gap by analyzing the training dynamics of Sign Gradient Descent (SignGD), which is an effective surrogate for understanding Adam.
SignGD is a simple gradient-based algorithm that updates parameters using only the sign of the gradient, discarding the gradient’s magnitude.
Over the years, SignGD has been extensively studied  (Balles & Hennig, ; Balles et al., ; Bernstein et al., ; ), and has inspired the development of optimizers like Adam (Kingma & Ba, ) and Lion (Chen et al., ).
More importantly, SignGD shares many similarities with Adam, making it an effective proxy for gaining insights into Adam’s optimization behavior (Balles & Hennig, ; Bernstein et al., ; Kunstner et al., ; ; Wu et al., ; Zou et al., ).
For example, Kunstner et al. () has shown that while a performance gap between Adam and GD persists in the full-batch setting on transformers, SignGD can effectively bridge this gap, achieving performance closer to Adam.
Despite its simplicity, however, theoretical understanding of how SignGD optimizes transformers remains an open problem.

Our contributions. In this work, we study the problem of how SignGD optimizes transformers in a binary classification task with linearly separable datasets with signal and noise.

- •
We provide a theoretical characterization of the entire training dynamics of SignGD.
Specifically, we identify four different stages in the training dynamics, each exhibiting unique behaviors, as summarized in Tab. [1](#S1.T1).
This detailed four-stage analysis captures the complex yet systematic dynamics within the attention layer,
and offers a precise description of how SignGD optimizes transformers in our setting.
- •
Based on the training dynamics, we prove the convergence and generalization results. On our noisy dataset, SignGD demonstrates fast convergence but poor generalization, achieving a linear convergence rate in training loss but maintaining a high constant test loss, leading to a sparse attention matrix through noise memorization.
Additionally, we provide evidence that Adam exhibits similar behaviors to SignGD in terms of training dynamics, convergence, and generalization, suggesting that SignGD is a strong proxy for understanding Adam.
We also find that the poor generalization of SignGD is not solely due to data noise, but is also related to its inherent algorithmic properties, indicating that SignGD and Adam require higher data quality in practice compared to GD.
Our results and findings are further validated through experiments on both synthetic and real-world datasets.

**Table 1: Overview of the four-stage dynamics: corresponding behaviors and theoretical results.**
| Stage I | The mean value noise shifts early, then stabilizes. | Lemma [4.1](#S4.Thmtheorem1) |
| --- | --- | --- |
| Stage II | The query & key noise align their sign to each other. | Lemma [4.2](#S4.Thmtheorem2), [4.3](#S4.Thmtheorem3) |
| Stage III | Majority voting determines the sign of query & key signals. | Lemma [4.4](#S4.Thmtheorem4) |
| Stage IV | The noise-signal softmax outputs decay fast exponentially, | Lemma [4.5](#S4.Thmtheorem5) |
| then the query & key noise align their sign to signals. | Lemma [4.6](#S4.Thmtheorem6), [4.7](#S4.Thmtheorem7) |  |

Technical novelties.
We use the feature learning framework (Allen-Zhu & Li, ; Cao et al., ) for our theoretical analysis. Our technical novelties include: Firstly, we analyze an softmax attention layer with trainable query-key parameterization, which is not carefully studied in the literature. Secondly, we perform a multi-stage analysis for transformers by breaking down the complex dynamics into simple sub-stages. In each sub-stage, only one or two key behaviors dominate. Finally, we cleverly combined SignGD and the sparse data model, greatly simplifying the analysis.

In summary, our work investigates the training dynamics of transformers using SignGD. To the best of our knowledge, this is the first provable result characterizing the training dynamics of transformers with SignGD. Our findings offer valuable insights into the inner workings of both SignGD and Adam, advancing our theoretical understanding of transformers and their optimization.

## 2 Preliminaries

##### Notations.

We use lower case letters, lower case bold face letters, and upper case bold face letters to denote
scalars, vectors, and matrices respectively.
For a vector $\boldsymbol{\mathbf{v}}=[v_{1},\dots,v_{d}]^{\top}$, we denote the $\ell_{2}$ and $\ell_{1}$ norm by $\left\lVert\boldsymbol{\mathbf{v}}\right\rVert$ and $\left\lVert\boldsymbol{\mathbf{v}}\right\rVert_{1}$, respectively.
For two fixed non-negative sequences $\left\{x_{n}\right\}$ and $\left\{y_{n}\right\}$, we denote $x_{n}=O(y_{n})$ if there exist some
absolute constant $C>0$ and $N>0$ such that $\left|x_{n}\right|\leq C\left|y_{n}\right|$ for all $n\geq N$.
We say $x_{n}=\Omega(y_{n})$ if $y_{n}=O(x_{n})$,
and say $x_{n}=\Theta(y_{n})$ if $x_{n}=O(y_{n})$ and $x_{n}=\Omega(y_{n})$.
We use $\tilde{O}(\cdot)$, $\tilde{\Omega}(\cdot)$ and $\tilde{\Theta}(\cdot)$ to hide logarithmic factors in these notations, respectively.
Moreover, we denote $x_{n}=\operatorname{poly}(y_{n})$ if $x_{n}=O(y^{D}_{n})$ for some constant $D>0$,
and $x_{n}=\operatorname{polylog}(y_{n})$ if $x_{n}=\operatorname{poly}(\log(y_{n}))$.
We use $[d]$ to denote the set $\left\{1,2,\dots,d\right\}$.
We use $\operatorname{sgn}(x)=x/|x|$ when $x\neq 0$ and $\operatorname{sgn}(0)=0$.
We denote a $n$-dim all-ones vector by $\textbf{1}_{n}$.

##### Data model.

We consider a binary classification task where each data point contains signal vector and sparse noise vector. The data model is formally defined in Definition [2.1](#S2.Thmtheorem1).

###### Definition 2.1 .

Let $\boldsymbol{\mathbf{\mu}}\in\mathbb{R}^{d}$ be a fixed vector representing the signal contained in each data point.
We assume $\boldsymbol{\mathbf{\mu}}=[1,0,\dots,0]^{\top}$.
For each data point $(\boldsymbol{\mathbf{X}},y)$, the predictor $\boldsymbol{\mathbf{X}}=[\boldsymbol{\mathbf{x}}^{(1)},\boldsymbol{\mathbf{x}}^{(2)}]\in\mathbb{R}^{d\times 2}$ consists of two patches (or tokens, vectors),
where $\boldsymbol{\mathbf{x}}^{(1)},\boldsymbol{\mathbf{x}}^{(2)}\in\mathbb{R}^{d}$, and the label $y$ is binary, i.e., $y\in\{\pm 1\}$.
The data is generated from a distribution $\mathcal{D}$, which we specify as follows:

- 1
The label $y$ is generated as a Rademacher random variable.
- 2
Randomly select $s$ coordinates from $[d]\backslash\left\{1\right\}$ uniformly, denoted as a vector $\boldsymbol{\mathbf{s}}\in\left\{0,1\right\}^{d}$. Generate each coordinate in $\boldsymbol{\mathbf{\xi}}$ from distribution $N(0,\sigma_{p}^{2})$, and then mask off the first coordinate and other $d-s-1$ coordinates, i.e., $\boldsymbol{\mathbf{\xi}}=\boldsymbol{\mathbf{\xi}}\odot\boldsymbol{\mathbf{s}}$.
- 3
One of $\boldsymbol{\mathbf{x}}^{(1)},\boldsymbol{\mathbf{x}}^{(2)}$ is randomly selected and then assigned as $y\boldsymbol{\mathbf{\mu}}$, representing the signal, while the other is designated as $\boldsymbol{\mathbf{\xi}}$, representing noise.

The design of the signal patch $y\boldsymbol{\mathbf{\mu}}$ and the noise patch $\boldsymbol{\mathbf{\xi}}$ can be viewed as a simplification of real-world image classification problems where only certain patches contain useful features that are correlated with the label, e.g., the wheel of a car, while many other patches contain uninformative features or consist solely of noise, e.g., the background of the image.
Specifically, $y\boldsymbol{\mathbf{\mu}}$ represents useful, label-correlated features (referred to as signal), whereas $\boldsymbol{\mathbf{\xi}}$ represents non-informative features or irrelevant noise (referred to as noise).

Remarks on data assumptions.
We make several assumptions regarding sparsity, orthogonality, and context length for the data.
Specifically, we assume $\boldsymbol{\mathbf{\mu}}$ is 1-sparse (i.e., it has only one non-zero entry), $\boldsymbol{\mathbf{\xi}}$ is $s$-sparse, and that $\boldsymbol{\mathbf{\mu}}$ and $\boldsymbol{\mathbf{\xi}}$ are orthogonal to each other.
The sparsity assumption is essential for analysing optimizers that are not invariant under orthogonal transformations, such as Adam and SignGD.
Our results can be easily extended to any $C$-sparse signal vector $\boldsymbol{\mathbf{\mu}}$, where $C=O(1)$ is a constant, with non-zero entries in arbitrary positions and constant magnitude.
The orthogonality assumption holds with high probability under the sparsity assumption, which confirms its validity (see Lemma [C.2](#A3.Thmtheorem2) for details).
We also assume a context length of $L=2$, which is a technical assumption to simplify the analysis.
With additional appropriate assumptions, our analysis can be extended to data with longer contexts (see discussion in Appendix [F.2](#A6.SS2)).
We empirically validate our theoretical results for non-sparse, non-orthogonal, and multi-patch data in Appendix [B](#A2).
Data models comprised of signal and noise patches with similar assumptions have also been studied in recent works (Allen-Zhu & Li, ; Cao et al., ; Jelassi et al., ; Zou et al., ).

###### Definition 2.1 .

##### Two-layer transformers.

Motivated by vision transformers (Dosovitskiy et al., ),
we consider a two-layer transformer, where the first layer is a single-head softmax attention layer and the second layer is a linear head layer.
The attention layer is a sequence-to-sequence mapping, of which the parameters are $\boldsymbol{\mathbf{W}}:=\left(\boldsymbol{\mathbf{W}}_{Q},\boldsymbol{\mathbf{W}}_{K},\boldsymbol{\mathbf{W}}_{V,j}\right)$, where $\boldsymbol{\mathbf{W}}_{Q},\boldsymbol{\mathbf{W}}_{K}\in\mathbb{R}^{m_{k}\times d}$ and $\boldsymbol{\mathbf{W}}_{V,j}\in\mathbb{R}^{m_{v}\times d}$ for $j\in\left\{\pm 1\right\}$.
The parameters of the second layer are fixed as $1/m_{v}$ and $-1/m_{v}$ respectively. We also talk about learnable linear head in Appendix [F.3](#A6.SS3).
Then, the network can be written as $f(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}}):=F_{1}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})-F_{-1}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$,
where $F_{1}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ and $F_{-1}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ are defined as:

$$ $\displaystyle F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}}):=\frac{1}{m_{v}}\sum_{l=1}^{L}\textbf{1}^{\top}_{m_{v}}\boldsymbol{\mathbf{W}}_{V,j}\boldsymbol{\mathbf{X}}\text{softmax}\left(\boldsymbol{\mathbf{X}}^{\top}\boldsymbol{\mathbf{W}}_{K}^{\top}\boldsymbol{\mathbf{W}}_{Q}\boldsymbol{\mathbf{x}}^{(l)}\right).$ $$

Let $\boldsymbol{\mathbf{w}}_{Q,s}:=\boldsymbol{\mathbf{W}}_{Q,(\cdot,s)}^{\top}\in\mathbb{R}^{d}$, $\boldsymbol{\mathbf{w}}_{K,s}:=\boldsymbol{\mathbf{W}}_{K,(\cdot,s)}^{\top}\in\mathbb{R}^{d}$, $\boldsymbol{\mathbf{w}}_{V,j,r}:=\boldsymbol{\mathbf{W}}_{V,j,(\cdot,r)}^{\top}\in\mathbb{R}^{d}$ be the $s$-th or $r$-th row of the parameter $\boldsymbol{\mathbf{W}}_{Q}$, $\boldsymbol{\mathbf{W}}_{K}$, $\boldsymbol{\mathbf{W}}_{V,j}$, respectively.
Let $\bar{\boldsymbol{\mathbf{w}}}_{V,j}:=\sum_{r\in[m_{v}]}\boldsymbol{\mathbf{w}}_{V,j,r}/m_{v}$ be the mean value in $F_{j}$.
Let $\boldsymbol{\mathbf{v}}=\bar{\boldsymbol{\mathbf{w}}}_{V,1}-\bar{\boldsymbol{\mathbf{w}}}_{V,-1}$ be the
mean value.
We can write the model in a simpler form:

$$ $\displaystyle F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}\left[\left(s_{11}+s_{21}\right)\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(1)}\right\rangle+\left(s_{12}+s_{22}\right)\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\right],$ (1) $$

where $s_{la}:=\text{softmax}\left(z_{l1},\dots z_{lL}\right)_{a}$, and
$z_{la}:=\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s},\boldsymbol{\mathbf{x}}^{(l)}\rangle\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(a)}\right\rangle$.
Unless otherwise specified, we set $L=2$.

We refer to the softmax outputs with the noise vector as the query and the signal vector as the key as noise-signal softmax outputs, formally written as

$$ $\displaystyle s_{i,21}^{(t)}=\frac{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle\right)}{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle\right)+\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}.$ $$

Similarly, we refer to $s_{i,11}^{(t)}$, $s_{i,12}^{(t)}$, $s_{i,22}^{(t)}$ as signal-signal, signal-noise, and noise-noise softmax outputs, respectively.
When $L=2$, a key fact about the softmax function is that $s_{i,l1}^{(t)}+s_{i,l2}^{(t)}\equiv 1$, for $l\in[2]$.
The subscript is used only to distinguish between signal and noise, without imposing any restrictions on the permutation of patches. Although we use the symbol $s$ for sparsity, softmax outputs, and the indices of query and key neurons simultaneously, the context clearly indicates which one is being referred to.

##### Training algorithm.

We train our transformer model by minimizing the empirical cross-entropy loss function
$L_{S}(\boldsymbol{\mathbf{W}}):=\frac{1}{n}\sum_{i=1}^{n}\ell(y_{i}\cdot f(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}}_{i}))$,
where $\ell(x):=\log\left(1+\exp(-x)\right)$ is the logistic loss function, and $S:=\left\{(\boldsymbol{\mathbf{X}}_{i},y_{i})\right\}_{i=1}^{n}$ is the training dataset. We further define the test loss
$L_{\mathcal{D}}(\boldsymbol{\mathbf{W}}):=\mathbb{E}_{(\boldsymbol{\mathbf{X}},y)\sim\mathcal{D}}[\ell(y\cdot f(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}}))]$.

We study SignGD starting from Gaussian initialization, where each entry of $\boldsymbol{\mathbf{W}}_{Q}$, $\boldsymbol{\mathbf{W}}_{K}$, $\boldsymbol{\mathbf{W}}_{V,j}$ for $j\in\left\{\pm 1\right\}$ is sampled from a Gaussian distribution $N(0,\sigma_{0}^{2})$.
The SignGD update for the parameters can be written as

$$ $\displaystyle\boldsymbol{\mathbf{w}}^{(t+1)}_{V,j,r}=\boldsymbol{\mathbf{w}}^{(t)}_{V,j,r}-\eta\operatorname{sgn}(\nabla_{\boldsymbol{\mathbf{w}}_{V,j,r}}L_{S}(\boldsymbol{\mathbf{W}}^{(t)})),$ $\displaystyle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s}=\boldsymbol{\mathbf{w}}^{(t)}_{Q,s}-\eta\operatorname{sgn}(\nabla_{\boldsymbol{\mathbf{w}}_{Q,s}}L_{S}(\boldsymbol{\mathbf{W}}^{(t)})),\leavevmode\nobreak\ \leavevmode\nobreak\ \boldsymbol{\mathbf{w}}^{(t+1)}_{K,s}=\boldsymbol{\mathbf{w}}^{(t)}_{K,s}-\eta\operatorname{sgn}(\nabla_{\boldsymbol{\mathbf{w}}_{K,s}}L_{S}(\boldsymbol{\mathbf{W}}^{(t)})),$ $$

for all $s\in[m_{k}]$, $j\in\left\{\pm 1\right\}$ and $r\in[m_{v}]$.
The expanded gradient formulas and update rules can be seen in Appendix [D](#A4).

## 3 Main Results

In this section, we present our main results.
Firstly, we provide a detailed characterization on training dynamics with a four-stage analysis, each exhibiting different behaviors. Then, based on the training dynamics, we analyze the convergence and generalization at the end of the training.
We further give an evidence that Adam exhibits similar behaviors to SignGD, and provide new insights that SignGD and Adam requires higher data quality compared to GD.

Before presenting the main results, we state our main condition. All of our theoretical results are based on the following condition.

###### Condition 3.1 .

Suppose that

- 1
[Data dimension and Sparsity] Data dimension $d$ is sufficiently large with $d=\Omega(\operatorname{poly}(n))$.
Sparsity $s$ satisfies: $s=\Theta(d^{1/2}n^{-2})$.
- 2
[Noise strength] The standard variance of noise $\sigma_{p}$ satisfies: $\sigma_{p}=\Omega(d^{-1/4}n^{3})$.
- 3
[Network width and initialization] Network width of value $m_{v}$ and of query & key $m_{k}$ satisfy: $m_{k},m_{v}=\Omega(\operatorname{polylog}(d))$. Network initialization $\sigma_{0}$ satisfies: $\sigma_{0}=o(\sigma_{p}^{-1}s^{-1}m_{k}^{-1/2})$.
- 4
[Training dataset size] The training sample size $n$ satisfies: $n=\Omega(m_{k}^{4})$.
- 5
[Learning rate] The learning rate $\eta$ satisfies: $\eta=O(\operatorname{poly}(d^{-1}))$ is sufficiently small.

Remarks on Condition [3.1](#S3.Thmtheorem1).
Our Condition [3.1](#S3.Thmtheorem1) is frequently used in the literature and realistic in practice (Cao et al., ; Chatterji & Long, ; Frei et al., ).
The conditions on $d$ and $s$ make sure the different noise patches have disjoint support with high probability.
The condition on $\sigma_{p}$ implies $\sigma_{p}\sqrt{s}=\Omega(n^{2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)$, which indicates the noise in the dataset is strong.
The conditions on $m_{v}$, $m_{k}$, $n$, $\eta$ are technical and mild.
$m_{v}$ and $m_{k}$ affect the convergence of mean value noise and softmax outputs, respectively.
The size of $n$ affects the concentration of $\left\lVert\boldsymbol{\mathbf{\xi}}\right\rVert_{1}$.
Finally, the conditions on $\sigma_{0}$ ensures the network weights are small enough at initialization, which makes the learning process fall into the feature learning regime.
Note if we set $m_{k}=\Theta(\operatorname{polylog}(d))$, then our condition become $\sigma_{0}=\Theta(d^{-1/2})$, which is realistic in practice.

###### Condition 3.1 .

### 3.1 Training Dynamics Analysis: Four-stage Analysis

Based on Condition [3.1](#S3.Thmtheorem1),
we aim to explore the underlying optimization mechanism of SignGD through training dynamics.
We identify four distinct stages in the training dynamics, where the primary behaviors and theoretical results for each stage are summarized in Tab. [1](#S1.T1).
This four-stage analysis captures the complex yet systematic dynamics within the attention layer and provides a valuable tool for further analysis of convergence and generalization.
In this subsection, we informally describe the core phenomena in each stage, while the formal results are presented in Sec. [4](#S4).

Figure: Figure 1: The training dynamics of two-layer transformers with SignGD. (a) Key noise dynamics over $t=0$ to $t=2$. (b) Mean value noise dynamics over $t=0$ to $t=2$. While mean value noise stabilizes into a linear relationship with $t$ early, key noise remains close to initialization. (c) Softmax output dynamics over $t=0$ to $t=900$. The softmax outputs decay exponentially. At $t=150$, $s_{i,21}^{(t)}$ approaches zero, while $s_{i,11}^{(t)}$ remains close to $1/2$. (d) Dynamics of query noise, key noise, and query signals over $t=0$ to $t=900$: The dotted lines represent positive query and key noise at $t=100$, and the solid lines represent negative noise at the same point. By Stage III, the majority of positive noise makes the query signal positive through majority voting. In Stage IV, sign alignment of key noise starts at about $t=150$, coinciding with $s_{i,21}^{(t)}$ approaching zero, while delayed sign alignment of query noise begins around $t=300$, about twice as late as the key noise.
Refer to caption: /html/2410.04870/assets/x1.png

To better understand the four-stage training dynamics, the dynamics of key quantities at key timesteps or during the entire dynamics are illustrated in the Fig. [1](#S3.F1) and Tab. [2](#S3.T2).

Fig. [1](#S3.F1) (a), (b) shows the early shift and stabilization of mean value noise in Stage I.
The mean value noise, i.e.,
$\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ increases monotonically from the random initialization and becomes positive, then stabilizes into a linear relationship with $t$ by the stage’s end. This period is so rapid that other quantities, including value signals and query & key noise all remain close to their initialization.
This linear behavior of mean value noise mitigates the negative impact of value noise on the gradients.

Tab. [2](#S3.T2) illustrates the sign alignment between query & key noise, i.e., $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and
$\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, in Stage II.
At initialization, the signs of the query & key noise are independent. By the end of Stage II, however, the signs of the noise for each neuron align, becoming either jointly positive or jointly negative, and continue to grow. Additionally, the number of positive and negative neurons is nearly equal.

Fig. [1](#S3.F1) (d) shows how the sign of query & key signals, i.e. $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$, are determined by majority voting in Stage III.
At the start of Stage III, the sum of key noise $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ dominates the gradients of query signals, making the update direction of query signals aligned with its sign.
For a given neuron $s$, key noise is nearly uniform across all samples, thus effectively acting as majority voting.
The sign of key signals is determined in a symmetric manner by the sum of query noise.

Fig. [1](#S3.F1) (c) illustrates the fast exponential decay of noise-signal softmax outputs, i.e. $s_{i,21}^{(t)}$, in Stage IV.
The duration from Stage I to Stage III can be relatively short, with all softmax outputs
$s_{i,11}^{(t)}$, $s_{i,21}^{(t)}$ concentrated at $1/2$.
However, in Stage IV, $s_{i,21}^{(t)}$ decreases exponentially to zero,
while $s_{i,11}^{(t)}$ remains stuck at $1/2$.
This difference is due to the varying rates of increase between query signals and query noise.

Finally, Fig. [1](#S3.F1) (d) shows
the sign alignment of query & key noise with signals in Stage IV.
We focus on a single neuron with a positive query signal, as shown in Fig.[1](#S3.F1) (d), though this applies to all neurons.
Before the noise-signal softmax outputs decay to zero, the dynamics of all signals and noise remain unchanged.
At a critical point, all negative key noise begins aligning with the positive query signal, gradually decreasing in magnitude and crossing zero.
Once the negative key noise approaches zero and fully aligns, the negative query noise begins aligning, eventually becoming positive by the end of the stage.
From that point on, the sign of all signals and noise remains unchanged.

Overall, the dynamics exhibit complex and intriguing sign alignment behaviors within the attention layer, and they ultimately reach stabilization.

**Table 2: Sign alignment between query and key noise in Stage II. $S_{K+,Q+}^{(t)}$, defined as $S_{K+,Q+}^{(t)}:=\{(s,i)\in[m_{k}]\times[n]:\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0,\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0\}$, represents the number of neurons and samples having positive query noise and positive key noise. The definitions for $S_{K+,Q-}^{(t)}$, $S_{K-,Q+}^{(t)}$, $S_{K-,Q-}^{(t)}$ are similar. Each element in the middle of the table represents the size of the intersection of the corresponding row set and the corresponding column set. For example, $|S_{K+,Q+}^{(0)}\cap S_{K+,Q+}^{(t)}|=486$. The signs of query and key noise are independent at initialization but aligned at $t=10$, which can be seen as an estimate of $T_{2}^{\text{SGN}}$.**
| init($t=0$)$\backslash t=10$ | $|S_{K+,Q+}^{(t)}|$ | $|S_{K+,Q-}^{(t)}|$ | $|S_{K-,Q+}^{(t)}|$ | $|S_{K-,Q-}^{(t)}|$ | Row sum |
| --- | --- | --- | --- | --- | --- |
| $|S_{K+,Q+}^{(0)}|$ | 486 | 1 | 0 | 25 | 512 |
| $|S_{K+,Q-}^{(0)}|$ | 244 | 4 | 9 | 250 | 507 |
| $|S_{K-,Q+}^{(0)}|$ | 223 | 10 | 4 | 221 | 458 |
| $|S_{K-,Q-}^{(0)}|$ | 37 | 2 | 3 | 481 | 523 |
| Column sum | 990 | 17 | 16 | 977 | 2000 |

### 3.2 Convergence and Generalization Analysis: Fast Convergence but Poor Generalization

Beyond the training dynamics, we characterize the convergence and generalization result at the end of the training.
Additionally, we provide evidence that Adam exhibits similar behaviors to SignGD in optimization and generalization,
and suggest that SignGD and Adam requires high data quality.

###### Theorem 3.2 .

For any $\epsilon>0$, under Condition [3.1](#S3.Thmtheorem1), with probability at least $1-n^{-1/3}$,
there exists
$T=O(\log(\epsilon^{-1})\eta^{-1}\sigma_{p}^{-1}s^{-1})$ ,
and $T_{\mathrm{attn}}=\tilde{O}(\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1/2}s^{-1/2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{-1/2})$
such that

- 1
[Training loss] The training loss converges to $\epsilon$: $L_{S}(\boldsymbol{\mathbf{W}}^{(T)})\leq\epsilon$.
- 2
[Test loss] The trained transformer has a constant order test loss: $L_{\mathcal{D}}(\boldsymbol{\mathbf{W}}^{(T)})=\Theta(1)$.
- 3
[Noise memorization of value] The value matrix in attention layer memorizes noises in the training data:
For all $i\in[n]$,
$|\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\xi}}_{i}\rangle|=\Omega(1)$,
$|\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle|=\tilde{O}(\sigma_{p}^{-1}s^{-1})$.
- 4
[Noise memorization of query & key] The softmax outputs of attention layer attends all the weights to the noise patch in the training data:
For all $i\in[n]$,
$s_{i,11}^{(T_{\mathrm{attn}})}=o(1)$, $s_{i,21}^{(T_{\mathrm{attn}})}=o(1)$.

Theorem [3.2](#S3.Thmtheorem2) outlines the training and test loss at the end of training. In this setting, SignGD achieves a fast linear convergence rate in training loss, but test loss remains high, summarizing the behavior as fast convergence but poor generalization.
Theorem [3.2](#S3.Thmtheorem2) also presents new results on noise memorization in the attention layer.
The post-softmax attention matrix concentrates all its weights on the noise patch, resulting in a sparse matrix, which is consistent with previous analyses (Tian et al., ; ; Li et al., ).
In contrast to prior works (Cao et al., ; Zou et al., ) that focus on noise memorization in linear or convolutional layers, our results address the more complex issue of noise memorization in the attention layer.
The proof of Theorem [3.2](#S3.Thmtheorem2) is in Appendix [E.8](#A5.SS8).

Remark.
Theorem [3.2](#S3.Thmtheorem2) focus on the logistic test loss. We further give a result about final 0-1 test loss in Appendix [F.1](#A6.SS1) since bad logistic loss doesn’t necessarily imply bad 0-1 loss in binary classification task.
Interestingly, the size of 0-1 test loss depends on the network initialization $\sigma_{0}$.

Fast convergence but poor generalization contradcits the ‘train faster, generalize better’ argument.
Algorithmic stability (Bousquet & Elisseeff, ; Hardt et al., ) is a widely used technique in generalization analysis.
One typical argument of algorithmic stability is “train faster, generalize better” (Hardt et al., ).
Our results provide a counterexample to this argument by showing SignGD trains fast but generalizes poorly.
Notably, Teng et al. () introduced a measure explaining why SGD trains slower but generalizes better than GD, aligning with our viewpoint.

Adam exhibits similar behaviors to SignGD in optimization and generalization.
We conducted experiments with Adam on both synthetic and real-world datasets, tracing the dynamics on synthetic data and measuring test loss on noisy MNIST dataset.
On the synthetic data, Adam follows a similar four-stage dynamics as SignGD, with Stage III and Stage IV shown in Fig. [2](#S3.F2) (a),(b).
Further similarities in other stages, and the results across different $\beta_{1}$ values are given in Appendix [B.3](#A2.SS3).
In the noisy MNIST data, Adam also shows high test loss, like SignGD, especially under strong noise.
This indicates that Adam shares key similarities with SignGD in training dynamics, convergence, and generalization, further supporting the use of SignGD as a proxy for understanding Adam.
However, when $\beta_{1}$ and $\beta_{2}$ in Adam are close to 1, which is commonly used in practice, SignGD does not always behave like Adam. In our experiments, unlike SignGD, Adam did not exhibit the sign alignment of query noise with query signals in Stage IV. We suspect this difference is due to the momentum term in Adam.

SignGD and Adam require higher data quality than GD.
We compare SignGD and GD from both theoretical and empirical perspectives.
Theoretically, our results reveal that SignGD achieves a linear convergence rate, while GD typically converges more slowly with a sublinear rate in learning CNNs for the same task (Cao et al., ; Kou et al., ) or transformers for different tasks (Nichani et al., ; Huang et al., ).
Empirically, our experiments demonstrate that SignGD trains faster than GD (Fig. [2](#S3.F2) (c)), but GD generalizes better consistently across different levels of data noise (Fig. [2](#S3.F2) (d)).
Furthermore, our experiments show that Adam, like SignGD, is also sensitive to data noise, underperforming in generalization compared to GD in noisy conditions.
This indicates that GD is better at learning true useful features from noisy data, while SignGD and Adam are more sensitive to noise.
Therefore, the poor generalization of SignGD, as indicated by our theoretical results, is not solely due to data noise but is also related to its inherent algorithmic properties.
This highlights that both SignGD and Adam require higher data quality than GD.
We recommend that practitioners using SignGD or Adam as optimizers consider improving data quality to mitigate their sensitivity to noise.
Experiment details and additional results can be found in Appendix [B](#A2).

Figure: Figure 2: Comparison of SignGD with Adam and GD on synthetic and real-world datasets. (a) Dynamics of query noise, key noise, and query signals with SignGD on the synthetic dataset. The legend follows Fig. [1](#S3.F1) (d). (b) Dynamics of query noise, key noise, and query signals with Adam($\beta_{1}=0.9$) on the synthetic dataset. (c) Training loss curve (log scale) on the synthetic data for different optimizers. The training loss with SignGD decays exponentially. Note that the training losses for Adam($\beta_{1}=0.9$), Adam($\beta_{1}=0.5$), Adam($\beta_{1}=0.0$) overlap. (d) Test loss on the noisy MNIST dataset across varying noise levels. A larger scaled SNR indicates less noise in the dataset.
Refer to caption: /html/2410.04870/assets/x2.png

In summary, we prove that SignGD achieves fast convergence but poor generalization on a noisy dataset.
We provide empirical evidence that Adam exhibits similar behaviors to SignGD in terms of training dynamics, convergence, and generalization, offering new insights into understanding Adam through the lens of SignGD.
By comparing it with GD,
we find that the poor generalization of SignGD is not solely due to data noise but is related to the inherent properties of SignGD, suggesting that both SignGD and Adam require higher data quality in practice.

###### Theorem 3.2 .

## 4 Proof Sketch

In this section, we present the proof sketch of the training dynamics, which can subsequently easily lead to the convergence and generalization results.

We use the feature learning framework (Allen-Zhu & Li, ; Cao et al., ), which studies the dynamics of parameter-data inner product. This inner product shows a simpler and clearer pattern compared with parameter itself.
Specifically, these quantities are $\langle\boldsymbol{\mathbf{w}}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle$, $\langle\boldsymbol{\mathbf{w}}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, $\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle$, $\langle\boldsymbol{\mathbf{w}}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$,
$\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{\mu}}\rangle$, $\langle\boldsymbol{\mathbf{w}}_{V,j,r},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, and we name them by query signals, query noise, key signals, key noise,value signals, value noise, respectively. The mean value signals (noise) is the mean of the value signals (noise), denoted by
$\langle\boldsymbol{\mathbf{v}},\boldsymbol{\mathbf{\mu}}\rangle$ ($\langle\boldsymbol{\mathbf{v}},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$).
Conditional on these quantities, the output $f(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}}_{i})$ is independent of parameters $\boldsymbol{\mathbf{W}}$.

##### Technical Novelties.

We summarize our technical novelties in three points as below.
Firstly, we analyze an softmax attention layer with trainable query-key parameterization. In this case, the dynamics of query and key parameters are strongly correlated, and we need to carefully consider the inner interaction between query and key.
This is a quite challenging setting, and a detailed comparison with previous works is given in Appendix [A](#A1).

Secondly,
we perform a multi-stage analysis for transformers by breaking down the complex dynamics into sub-stages. In each stage, only one or two key behaviors dominate, while other patterns remain mostly unchanged or have minimal impact. This breakdown works because of the varying rates of change in key quantities, which is influenced by the parameters $\sigma_{p}$, $\sigma_{0}$, $s$, etc.

Thirdly, we cleverly combined SignGD and the sparse data model by observing that the magnitude of the update of any inner products of interest remains constant across all iterations. Formally, for all $s$, $j$, $r$, and $i$, with high probability, we have:

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{V,j,r}^{(t+1)}-\boldsymbol{\mathbf{w}}_{V,j,r}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|=\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{V,j,r}^{(t+1)}-\boldsymbol{\mathbf{w}}_{V,j,r}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|=\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$ $$

This property implies that each iteration’s update is determined solely by its sign, allowing us to simplify the analysis by focusing only on the update direction’s sign.

In the following four subsections, we present the key theoretical results for each stage, respectively.

### 4.1 Stage I. Mean Value Noise Early Shifts and Stabilizes.

In Stage I, we focus on mean value noise & signals, i.e. $\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$.
Let $\beta_{\boldsymbol{\mathbf{\xi}}}:=\max_{i,s,j,r}\{|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle|,|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle|,|\langle\boldsymbol{\mathbf{w}}_{V,j,r}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle|\}$ and
$T_{1}:=4\beta_{\boldsymbol{\mathbf{\xi}}}m_{v}^{-1/2}\eta^{-1}\sigma_{p}^{-1}s^{-1}=\tilde{\Theta}(\sigma_{0}m_{v}^{-1/2}\eta^{-1}s^{-1/2})$, we call $t\in[0,T_{1}]$ Stage I.
The following lemma is the main result in Stage I.

###### Lemma 4.1 (Stage I) .

We have
(1) (Magnitude). For all $i\in[n]$ and $t\geq T_{1}$,
$\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\Theta(t\eta\sigma_{p}s)$.
(2) (Negligibility of mean value signals). For all $i\in[n]$ and $t\geq T_{1}$,
$\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle=o(\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle).$
(3) (Query & Key noise barely move). For all $s\in[m_{k}]$, $i\in[n]$ and $t\leq T_{1}$,
$\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\cdot(1\pm o(1))$,
$\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\cdot(1\pm o(1))$.

Lemma [4.1](#S4.Thmtheorem1) (1) shows the
mean value noise quickly stabilize into a linear relationship with $t$.
Lemma [4.1](#S4.Thmtheorem1) (2) allows us to disregard the influence of the mean value signals on the query and key gradients, as it is negligible compared to the impact of the mean value noise.
Lemma [4.1](#S4.Thmtheorem1) (3) states that the query & key noise remain close to their initialization before $T_{1}$, allowing us to approximate them as still being at their initial values at $T_{1}$.
Lemma [4.1](#S4.Thmtheorem1) (1)(3) jointly show the early shift and stabilization of mean value noise.
The proof of this section is in Appendix [E.4](#A5.SS4).

###### Lemma 4.1 (Stage I) .

### 4.2 Stage II. Query and Key Noise Align Their Signs to Each Other

In Stage II, we focus on query and key noise, i.e., $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$.
Let $T_{2}:=50\sqrt{2}n\beta_{\boldsymbol{\mathbf{\xi}}}\eta^{-1}\sigma_{p}^{-1}s^{-1}=\tilde{\Theta}(\sigma_{0}n\eta^{-1}s^{-1/2})$, we call $t\in[T_{1},T_{2}]$ Stage II.

###### Lemma 4.2 (Sign Alignment Between Query and Key Noise) .

Let $T_{2}^{\text{SGN}}:=3\sqrt{2}\beta_{\boldsymbol{\mathbf{\xi}}}\eta^{-1}\sigma_{p}^{-1}s^{-1}$.
Then, with probability $1-\delta$, $\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2}^{\text{SGN}})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(T_{2}^{\text{SGN}})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)$.
Specifically:
(1)
If $\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)$,
then $\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2}^{\text{SGN}})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)$.
(2)
If $\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)=-\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)$,
then the conditional probability of the event $\{\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2}^{\text{SGN}})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)=j\}$ are at least $1/2-O(\delta)$ for $j\in\left\{\pm 1\right\}$.

Lemma [4.2](#S4.Thmtheorem2) indicates that all query and key noise can be divided into two groups.
If they initially have the same sign,
they reinforce each other and grow from initialization.
However, if not, they first both decay toward zero, evolving in opposite directions.
As they approach zero, the smaller value crosses zero, and the larger one flips direction, aligning their signs and growing in the same direction.

###### Lemma 4.3 (End of Stage II) .

Let $\beta_{\boldsymbol{\mathbf{\mu}}}:=\max_{s,j,r}\{|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle|,|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle|,|\langle\boldsymbol{\mathbf{w}}_{V,j,r}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle|\}$.
We have:
(1) (Magnitude)
At $T_{2}$, for all $s\in[m_{k}]$ and $i\in[n]$,
$|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle|,|\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle|=\Theta(t\eta\sigma_{p}s)$.
(2) (Concentration of softmax)
$s_{i,11}^{(t)}=1/2\pm o(1)$, $s_{i,21}^{(t)}=1/2\pm o(1)$ for all $i\in[n]$ and $t\leq T_{2}$.
(3) (Sum of noise)
With probability at least $1-\delta$, for all $s\in[m_{k}]$,
$|\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle|,|\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle|\geq 2n\beta_{\boldsymbol{\mathbf{\mu}}}$.
(4) For all $s\in[m_{k}]$, and $t\leq T_{2}$,
$\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\cdot(1\pm o(1))$,
$\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\cdot(1\pm o(1))$.

Lemma [4.3](#S4.Thmtheorem3) (1) shows that
query & key noise become linear with $t$ at $T_{2}$.
However, according to Lemma [4.3](#S4.Thmtheorem3) (2), the softmax outputs of attention layer remain concentrated around $1/2$ at $T_{2}$.
Lemma [4.3](#S4.Thmtheorem3) (3) estimates the magnitude of the sum of the query & key noise, i.e., $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ at $T_{2}$,
which plays a role in Stage III.
Lastly, Lemma [4.3](#S4.Thmtheorem3) (4) states that the query & key signals remain near their initialization before $T_{2}$.
The proof of this section is in Appendix [E.5](#A5.SS5).

###### Lemma 4.2 (Sign Alignment Between Query and Key Noise) .

###### Lemma 4.3 (End of Stage II) .

### 4.3 Stage III. Majority Voting Determines the Sign of Query and Key Signals.

In Stage III, we focus on query & key signals, i.e., $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$.
Recall $\beta_{\boldsymbol{\mathbf{\mu}}}:=\max_{s,j,r}\{|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle|,|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle|,|\langle\boldsymbol{\mathbf{w}}_{V,j,r}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle|\}$.
Let $T_{3}:=3\beta_{\boldsymbol{\mathbf{\mu}}}\eta^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{-1}=\tilde{\Theta}(\sigma_{0}\eta^{-1})$,
we call $t\in[T_{2},T_{3}]$ Stage III.

###### Lemma 4.4 (Stage III) .

We have:
(1) (Noise determines the sign of signals by majority voting)
For all $s\in[m_{k}]$ and $t\in[T_{2},T_{3}]$,
$\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle)=\operatorname{sgn}(\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle)$,
and
$\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle)=-\operatorname{sgn}(\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle)$.
(2) (Magnitude)
At $T_{3}$, for all $s\in[m_{k}]$,
$\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},\boldsymbol{\mathbf{\mu}}\rangle=\operatorname{sgn}(\sum_{i}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\cdot\Theta(T_{3}\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)$,
and
$\langle\boldsymbol{\mathbf{w}}_{K,s}^{(T_{3})},\boldsymbol{\mathbf{\mu}}\rangle=\operatorname{sgn}(-\sum_{i}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\cdot\Theta(T_{3}\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)$.
(3) (Concentration of softmax)
$s_{i,11}^{(t)}=1/2\pm o(1)$, $s_{i,21}^{(t)}=1/2\pm o(1)$ for all $i\in[n]$ and $t\in[T_{2},T_{3}]$.

Lemma [4.4](#S4.Thmtheorem4) (1) states that the update direction during Stage III and the final sign of the query & key signals are determined by sum of key & query noise, respectively, which remains dominant in their gradients throughout this stage.
Since query & key noise have roughly the same magnitude for all $i\in[n]$, the sign dictation can be viewed as a majority voting mechanism.
Lemma [4.4](#S4.Thmtheorem4) (2) shows that query & key signals become linear with $t$ at $T_{3}$.
Although much smaller than query & key noise, they will become dominant in Stage IV when noise-signal softmax outputs approach zero.
Lemma [4.4](#S4.Thmtheorem4) (3) also states that the softmax outputs remain around $1/2$ at this point, indicating that Stages I-III can occur within a short time in practice.
Additionally, all query & key noise continues to grow throughout Stage III. The proof of this section is in Appendix [E.6](#A5.SS6).

###### Lemma 4.4 (Stage III) .

### 4.4 Stage IV.
Query and Key Noise Align Their Sign to Signals by
by the Fast Decay of Noise-Signal Softmax Outputs

In Stage IV, we focus on the noise-signal softmax outputs $s_{i,21}^{(t)}$ and
all query & key signals and noise that could be affected by $s_{i,21}^{(t)}$.
Let $T_{4}:=C_{3}\log(C_{3}\sigma_{p}s\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{-1})\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1}=\tilde{\Theta}(m_{k}^{-1/2}\eta^{-1}\sigma_{p}^{-1}s^{-1})$,
where $C_{3}=\Theta(1)$ is a large constant,
we call $t\in[T_{3},T_{4}]$ Stage IV.

###### Lemma 4.5 (Exponentially Fast Decay of Noise-Signal Softmax Outputs) .

Let $T_{4}^{-}\geq T_{3}$ be the last time such that for all $t\in[T_{3},T_{4}^{-}]$, $s\in[m_{k}]$ and $i\in[n]$,
$\left|\sum_{i=1}^{n}s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq\frac{1}{2}n\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|$
and $s_{i,21}^{(t)}s_{i,22}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq 2s_{i,11}^{(t)}s_{i,12}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|$.
Then, we have
$T_{4}^{-}=\tilde{\Theta}(\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1})$, and for all $i\in[n]$:
(1)
$s_{i,21}^{(t)}=\exp(-O(m_{k}t^{2}\eta^{2}\sigma_{p}^{2}s^{2}))$ for $t\in[T_{3},T_{4}^{-}]$ and
$s_{i,21}^{(T_{4}^{-})}=\exp(O(\log(n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/\sigma_{p}s)))=o(1)$.
(2)
$s_{i,11}^{(t)}=1/2\pm o(1)$, for $t\in[T_{3},T_{4}^{-}]$.

Lemma [4.5](#S4.Thmtheorem5) states that the noise-signal softmax outputs decay exponentially and approach zero during $[T_{3},T_{4}^{-}]$, while other softmax outputs stay around $1/2$.
All signals and noise continue to grow as in Stage III until just before $T_{4}^{-}$.
Shortly after $T_{4}^{-}$, the final sign alignment of query and key noise begins.

###### Lemma 4.6 (Sign Alignment of Key Noise) .

There exists a small constant $\theta_{c}\in(0,1)$ such that for all $s\in[m_{k}]$ and $i\in[n]$, we have:
(1)
$\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle)=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle)$, for $t\in[T_{3},T_{4}^{-}]$.
(2)
$\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle)=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},\boldsymbol{\mathbf{\mu}}\rangle)$, for $t\geq(1+\theta_{c})T_{4}^{-}$.

Negative key noise alignment with signals.
Consider a single neuron with a positive query signal.
Lemma [4.6](#S4.Thmtheorem6) states that
after $(1+\theta_{c})T_{4}^{-}$, all negative key noise flip direction and begin aligning with query signal.
Before $T_{4}^{-}$, the dominant gradient terms for key noise are related to both query noise and noise-signal softmax outputs.
However, after
$(1+\theta_{c})T_{4}^{-}$, due to the exponential decay of $s_{i,21}^{(t)}$, the query signal term becomes dominant.

###### Lemma 4.7 (Delayed Sign Alignment of Query Noise) .

With the same $\theta_{c}$ in Lemma [4.6](#S4.Thmtheorem6), for all $s\in[m_{k}]$ and $i\in[n]$, we have:
(1)
$\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle)=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle)$, for $t\in[T_{3},(2-\theta_{c})T_{4}^{-}]$.
(2)
$\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle)=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},\boldsymbol{\mathbf{\mu}}\rangle)$, for $t\geq(2+3\theta_{c})T_{4}^{-}$.

Different alignment times for negative query and key noise.
Lemma [4.7](#S4.Thmtheorem7) indicates a time gap between the alignment of negative key noise and negative query noise.
Negative query noise continues to grow until $(2-\theta_{c})T_{4}^{-}$
since $s_{i,21}^{(t)}$ does not directly influence its gradient.
However, as key noise approaches or crosses zero,
the query signal term begins to dominate the gradient of query noise due to their correlation.
Intuitively, $(2-\theta_{c})T_{4}^{-}$ is the point where key noise still has significant magnitude,
while $(2+3\theta_{c})T_{4}^{-}$ indicates the completion of key noise alignment.

###### Lemma 4.8 (End of Stage IV) .

(1)
For all $t\geq T_{4}$, $s\in[m_{k}]$, and $i\in[n]$ $\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle)$.
(2)
Loss does not converge. $L_{S}(\boldsymbol{\mathbf{W}}^{(T_{4})})=\Theta(1)$.

Lemma [4.8](#S4.Thmtheorem8) shows that final sign alignment completes, but the training loss does not convergence at the end of Stage IV.
After Stage IV, the dynamics simplify, with all quantities growing continuously, and the training loss decreasing exponentially.
This simplified behavior makes it easier to prove convergence and generalization results.
The proof of this section is in Appendix [E.7](#A5.SS7).

###### Lemma 4.5 (Exponentially Fast Decay of Noise-Signal Softmax Outputs) .

###### Lemma 4.6 (Sign Alignment of Key Noise) .

###### Lemma 4.7 (Delayed Sign Alignment of Query Noise) .

###### Lemma 4.8 (End of Stage IV) .

## 5 Conclusion and Limitations

In conclusion, we present a theoretical analysis of the training dynamics of a two-layer transformers using SignGD for a binary classification task involving a dataset with both signal and noise.
We identify four distinct stages in the training dynamics, characterizing the complex and intriguing behaviors within the attention layer.
Our results demonstrate the fast convergence but poor generalization for SignGD.
Additionally, we provide new insights into the similarities between SignGD and Adam, suggesting that both require higher data quality compared to GD in practical applications.
We hope that our work contributes to a deeper theoretical understanding and aids in the design of more efficient optimization methods for transformers.

Limitations.
Several relaxed assumptions were made to simplify the theoretical analysis. Specifically, the datasets we considered are linearly separable, signal and noise vectors are sparse, and the context length is limited to 2 in our main theory, creating a significant gap from real-world datasets. Furthermore, our data model is motivated by image data, leaving open the question of how SignGD optimizes transformers on language data, which could be a promising direction for future research. The transformer model we analyzed includes only a single-head self-attention layer, whereas deeper transformers and multi-head attention are more commonly used in practice. We also discuss possible extensions for longer contexts and learnable linear head in Appendix [F](#A6).

## Appendix A Detailed Related Work

##### Training dynamics of transformers.

Currently, the theoretical analysis of transformers training dynamics is still in a stage of flourishing, and there is no fixed research paradigm.
Several recent works studied the training dynamics of transformers with different data, models and focus.
Jelassi et al. () showed how a one-layer single-head attention model with only position embedding learns the spatial structures in the data by GD.
Tian et al. () studied the training dynamics of a one-layer single-head attention model with a data model for next-token prediction trained with SGD, and Tian et al. () analyzed the joint dynamics of an attention and MLP layer.
Li et al. () studied the training dynamics of a single-head attention layer on the $\ell_{2}$ loss with the data modeled by topic modeling.
Li et al. () studied a three-layer vision transformer with query and key parameterization trained by SGD on the hinge loss.
Oymak et al. () studied a single-head attention layer on the prompt-tuning setting, where the query and key parameters in the attention are fixed during the training process.
Tarzanagh et al. (; ) studied the dynamics of a single-head attention layer and a tunable token, and connects the training dynamics to a certain SVM problem.
While Tarzanagh et al. (; ) only presented an
asymptotic convergence result,
Vasudeva et al. () showed the global convergence, provided a convergence rate $t^{-3/4}$ for GD, as well as removing the restriction on a fixed linear head.
Also extending Tarzanagh et al. (; ), Sheen et al. () studied this problem with a query-key parameterized transformer, which gives different implicit regularization compared with the single attention matrix parameterization.
Huang et al. () studied the GD dynamics of a single-head attention layer with a self-supervised learning objective.
Nichani et al. () studied the GD dynamics of a disentangled two-layer attention-only transformer on random sequences with causal structure
and proved that it can learn the causal structure in the first attention layer.
Wang et al. () studied the data model introduced by Sanford et al. (), showing that an one-layer transformer can efficiently learn this task via GD, while fully-connected networks cannot express the task.
Jiang et al. ()
studied the GD dynamics of a two-layer transformer with a single-head self-attention layer and a fixed linear head on a dataset with signal and noise, which has the similar data and model setting to our work. However, Jiang et al. () targets on the benign overfitting of GD, while our work focus on the optimization and generalization of SignGD.

Recently, there has been some works studying the behaviors of transformers on in-context learning tasks to understand the powerful in-context learning abilities of large language models.
We mainly focus on works with convergence guarantee via training dynamics analysis.
Ahn et al. (); Mahankali et al. () showed that a one-layer
transformer implements single-step gradient descent to minimize the pre-training loss for an in-context learning task.
Zhang et al. () studied the Gradient Flow (GF) of a single-head linear attention layer on in-context linear regression task.
Huang et al. () studied how a single-head attention layer trained by GD solves an in-context linear regression task where the the input data are orthogonal.
Chen et al. () studied the GF for an one-layer multi-head attention model for an in-context linear regression task, where the groundtruth function admits a multi-task structure.
Kim & Suzuki () studied the mean-field dynamics of a transformer with one linear attention layer and one MLP layer on an in-context regression task.
Li et al. () studied the SGD dynamics of one-layer transformer with a single-head self-attention layer and a two-layer MLP on an in-context classification task.

Comparison with other works using query-key parameterization.
Most of works mentioned above take simplifications that replaces the query and key matrices with a single attention matrix.
Only Li et al. (; ); Sheen et al. (); Jiang et al. () studied the softmax attention with trainable query-key parameterziation, but there are also some limitations. Li et al. (; ) introduced relaxed assumptions about initialization, which are too stringent and make softmax outputs not concentrated at $1/2$ at initialization anymore. Sheen et al. () started from diagonal query and key matrices and used a data-correlated "Alignment Property" assumption for the general query and key initialization, which seems hard to verify whether it holds practically.
Compared with those works, we study query and key matrix from Gaussian initialization which is commonly used in practice.
Finally, Jiang et al. () studied trainable query and key with Gaussian initialization. However, their initializations for the queries, keys, and values differ, whereas our initialization for the queries, keys, and values is the same.

Additionally, all of these studies analyzed the dynamics of (S)GD or GF, while we focus on SignGD.

##### Understanding of Adam on transformers.

While Adam may fail to converge in convex objective (Reddi et al., ), it performs so well on transformers and
is better than SGD, which means Adam converges faster and achieves lower training loss (Ahn et al., ; Jiang et al., ; Kunstner et al., ; ; Pan & Li, ; Zhang et al., ; ).

Previous works tried to give an explanation about this fact from different perspectives.
Liu et al. () observed unbalanced gradients in transformers and Adam can give uniform parameter update.
Zhang et al. () suggested that heavy-tailed stochastic gradient noise in language data on transformers compared with image data on CNN models is the main cause that adaptive methods are good.
However, Kunstner et al. () showed that the heavy-tailed stochastic gradient noise may not be the main factor. They compared the performance of deterministic Adam and GD in full-batch settings, and observed that Adam is still better than GD.
Jiang et al. () showed that Adam could bias the trajectories towards regions
where Hessian has relatively more uniform diagonals while SGD cannot.
Zhang et al. () also studied from the perspective of Hessian. They showed the distances between Hessian spectrum of different parameter blocks are large in transformers which may hamper SGD but can be handled by Adam.
Pan & Li () showed that Adam can lead to smaller directional smoothness values which may imply better optimization.
Kunstner et al. () showed that heavy-tailed class imbalance in language modeling tasks is a difficulty for GD but Adam and SignGD do not suffer from this problem.

Furthermore, many works have focused on proving the convergence rate of Adam in the framework of classical convergence analysis.  Zhang et al. (); Crawshaw et al. () sought alternative relaxed assumptions for Adam, and Li et al. (); Wang et al. () improved the analysis of Adam under those assumptions

## Appendix B Experimental Details and More Experiments

### B.1 Experimental Settings

We perform numerical experiments on the synthetic and real-world datasets to verify our main results.

Experimental setting for synthetic dataset. The synthetic dataset is generated according to our Definition [2.1](#S2.Thmtheorem1).
For data hyerparameters, they can be uniquely determined by one row in Tab. [3](#A2.T3) and the value of $d$. In the Fig. [1](#S3.F1) and Tab. [2](#S3.T2) of main text, we use (a) with $d=2000$.
We always use 500 samples for computing test loss.

For optimizers, we use following default hyperparameters.
For sign gradient descent, we use the learning rate $\eta=\text{1e-4}$.
For Adam, we use the learning rate $\eta=\text{1e-4}$, $\beta_{1}=0.9$, $\beta_{2}=0.999$, and $\epsilon=\text{1e-15}$.
For gradient descent, we use the learning rate $\eta=\text{1e-1}$.

Also, we use neuron $s=0$ in Fig. [1](#S3.F1) by default and in following experiments.
We use sign gradient descent with the learning rate $\eta=\text{1e-7}$ in 2000 iterations to simulate
sign gradient descent with the learning rate $\eta=\text{1e-4}$ in 2 iterations in Fig. [1](#S3.F1) (a),(b) and in following experiments.
We use a learning rate of 1e-4 for all optimizers in Fig. [2](#S3.F2) (c) for a fair comparison.

**Table 3: Experimental settings of data model. $n$ is the training sample size and there are always equal samples in both classes. $s$ is the noise sparsity level. $\sigma_{p}$ is the standard deviation of noise. $\sigma_{0}$ is the network initialization standard deviation. $m_{v}$ is the value dimension. $m_{k}$ is the query and key dimension. ‘orthogonal’ means whether the noise patch and signal patch are orthogonal. If they are not orthogonal, the $s$ coordinates are selcted form $[d]$ instead of $[d]\backslash\left\{1\right\}$. ‘iters’ is the total iteration/epoch number in one run. Signal patch $\boldsymbol{\mathbf{\mu}}$ is always $[1,0,\dots,0]^{\top}$.**
|  | $n$ | $s$ | $\sigma_{p}$ | orthogonal | $\sigma_{0}$ | $m_{v}$ | $m_{k}$ | iters |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
| (a) | $0.01d$ | $0.04d$ | $2.0/\sqrt{s}$ | True | $0.1/\sqrt{d}$ | $0.01d$ | $0.05d$ | $2000$ |
| (b) | $0.01d$ | $0.04d$ | $2.0/\sqrt{s}$ | False | $0.1/\sqrt{d}$ | $0.01d$ | $0.05d$ | $2000$ |
| (c) | $0.01d$ | $0.4d$ | $2.0/\sqrt{s}$ | True | $0.1/\sqrt{d}$ | $0.01d$ | $0.05d$ | $2000$ |
| (d) | $0.01d$ | $0.4d$ | $2.0/\sqrt{s}$ | False | $0.1/\sqrt{d}$ | $0.01d$ | $0.05d$ | $2000$ |

Experimental setting for real-world dataset.
We conduct real-world experiments on the MNIST dataset.
We introduce the noise to the dataset in the following way.
For each image, we first multiply each pixel in the image with a factor $\lambda$, which we call “scaled SNR”, and then add gaussian random noises with standard deviation $1-\lambda$ to the outer regions with a width of 7.
Additionally, we only the class 3 and 7 for classification, to make a binary classification task which is consistent with our theoretical settings. To input the data into the transformers, we patchify the data with a size of $7\times 7$.

We train a two-layer transformer model consistent with our theoretical setting. We set $d=49$, $m_{v}=m_{k}=10$, $\sigma_{0}=0.1/\sqrt{d}$.
We only use 128 training data points and use deterministic optimizers to train the network. The network is trained for 200 epochs in each setting.
We use learning rate $\eta=$1e-1 for GD, and $\eta=$5e-4 for SignGD, and Adam across different values of $\beta_{1}$.
In all these settings, our training setup can guarantee a training loss smaller than 0.05.
We calculate the test losses on the entire test datasets (of the class 3 and 7).
Finally, we conduct three runs for each training setup and report the mean and standard deviation.

### B.2 Comparison with Non-sparse and/or Non-orthogonal Data

In Fig. [3](#A2.F3), [4](#A2.F4), [5](#A2.F5),
we repeat the experiments in Fig. [1](#S3.F1) with non-orthogonal sparse data, orthogonal non-sparse data, and non-orthogonal non-sparse data, respectively.
The figures show that our theoretical results hold empirically in those data settings.

Figure: Figure 3: Data setting (b) in Tab. [3](#A2.T3) with $d=2000$.
Refer to caption: /html/2410.04870/assets/x3.png

Figure: Figure 4: Data setting (c) in Tab. [3](#A2.T3) with $d=2000$.
Refer to caption: /html/2410.04870/assets/x4.png

Figure: Figure 5: Data setting (d) in Tab. [3](#A2.T3) with $d=2000$.
Refer to caption: /html/2410.04870/assets/x5.png

### B.3 Comparison with Adam and Gradient Descent

In Fig. [6](#A2.F6), [7](#A2.F7), [8](#A2.F8),
we plot the dynamics of softmax outputs, query signal, query and key noise, and training and test loss with different optimizers.
The figures show that our theoretical results almost hold empirically in Adam except the sign alignment of negative query noise in the Stage IV. This is due to Adam utilizes the information of history gradients to modify the current update.
We give an explanation about this. When the key signal becomes dominant in the gradients of query noise in sign gradient descent, it is actually very small in magnitude, which implies the small gradients can be easily dominated by the momentum in Adam.
But otherwise, symbolic gradient descent is almost similar to how Adam behaves under different hyperparameters.

Also, we can observe that the convergence of softmax outputs in sign gradient descent and Adam is much faster than that in gradient descent. At around $t=600$, the softmax outputs just start to leave $1/2$ in gradient descent but is almost converged in sign gradient descent and Adam.
It is also noted that gradient descent can leads to small generalization gap in this setting.

Figure: Figure 6: Sign alignment of signal to noise by majority voting in Stage III and sign alignment of negative noise to query signal by decay of noise-signal softmax outputs in Stage IV. We use data setting (d) in Tab. [3](#A2.T3) with $d=2000$. We always use $\beta_{2}=0.99$ and $\epsilon=\text{1e-15}$ in Adam.
Refer to caption: /html/2410.04870/assets/x6.png

Figure: Figure 7: The dynamics of softmax outputs. We use data setting (d) in Tab. [3](#A2.T3) with $d=2000$. We always use $\beta_{2}=0.99$ and $\epsilon=\text{1e-15}$ in Adam.
Refer to caption: /html/2410.04870/assets/x7.png

Figure: Figure 8: Training and test loss. We use data setting (d) in Tab. [3](#A2.T3) with $d=2000$. We always use $\beta_{2}=0.99$ and $\epsilon=\text{1e-15}$ in Adam.
Refer to caption: /html/2410.04870/assets/x8.png

### B.4 Comparison with Greater Context Length

In this section, we investigate when will happen when context length is greater than 2, i.e., $L>2$. At this time, the model is defined as

$$ $\displaystyle F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}}):=\frac{1}{m_{v}}\sum_{l=1}^{L}\textbf{1}^{\top}_{m_{v}}\boldsymbol{\mathbf{W}}_{V,j}\boldsymbol{\mathbf{X}}\text{softmax}\left(\boldsymbol{\mathbf{X}}^{\top}\boldsymbol{\mathbf{W}}_{K}^{\top}\boldsymbol{\mathbf{W}}_{Q}\boldsymbol{\mathbf{x}}^{(l)}\right).$ $$

For each data point $(\boldsymbol{\mathbf{X}},y)$, predictor $\boldsymbol{\mathbf{X}}=[\boldsymbol{\mathbf{x}}^{(1)},\boldsymbol{\mathbf{x}}^{(2)},\dots,\boldsymbol{\mathbf{x}}^{(L)}]\in\mathbb{R}^{d\times L}$ have $L$ patches (or tokens),
where $\boldsymbol{\mathbf{x}}^{(1)},\boldsymbol{\mathbf{x}}^{(2)},\dots,\boldsymbol{\mathbf{x}}^{(L)}\in\mathbb{R}^{d}$, and label $y$ is binary, i.e., $y\in\left\{\pm 1\right\}$.
The data generation is similar to the $L=2$ case, except that we randomly select $L/2$ patches and assign them by $y\boldsymbol{\mathbf{\mu}}$ as signal patches, while the remaining $L/2$ patches are noise patches. The noise patches in one data sample are mutually independent.

In our experiments, we use $L=10$, and Fig. [9](#A2.F9), [10](#A2.F10), [11](#A2.F11) plot the dynamics of softmax outputs, query signal, query and key noise, and training and test loss with different optimizers.
In those figures, for all $i\in[n]$, $a_{i}$ is defined as

$$ $\displaystyle a_{i}=\arg\max_{l\in[L]}\{s_{i,1l}^{(T)}\}.$ $$

We empirically observe that for all data point, only one element in each line of $L\times L$ post-softmax attention matrix is activated, while other elements in the line are almost zero, which means that the attention attends to only one patch.
This patch for each line is also uniform in different lines and therefore uniquely corresponds to one data point, which is exactly $a_{i}$. In the $L=2$ case, we have $a_{i}=2$ for all $i\in[n]$.
But when $L>2$ and there are many noise patches in one data sample, $a_{i}$ can be varied across samples but $X^{(a_{i})}$ must be a noise patch.

The greater context length case is consistent with the $L=2$ case, and thus consistent with our theoretical results in the sense that
(1) For all $i\in[n]$, the dominated query noise $\langle\boldsymbol{\mathbf{w}}_{Q,s},y_{i}X_{i}^{(a_{i})}\rangle$ and key noise $\langle\boldsymbol{\mathbf{w}}_{K,s},y_{i}X_{i}^{(a_{i})}\rangle$ are the same in sign with the query signal $\langle\boldsymbol{\mathbf{w}}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle$.
(2) The convergence of noise-signal/noise softmax outputs is very fast and faster than signal-signal/noise softmax outputs.
For gradient descent, while the loss fast converges, but the query and key parameters are basically not learned, even with learning rate $\eta=1$.

Figure: Figure 9: The dynamics of query noise, key noise and query signal. We use data setting (d) in Tab. [3](#A2.T3) with $d=2000$. We use $\beta_{2}=0.99$ and $\epsilon=\text{1e-15}$ in Adam and we use $\eta=1.0$ for gradient descent.
Refer to caption: /html/2410.04870/assets/x9.png

Figure: Figure 10: The dynamics of softmax outputs. We use data setting (d) in Tab. [3](#A2.T3) with $d=2000$. We use $\beta_{2}=0.99$ and $\epsilon=\text{1e-15}$ in Adam and we use $\eta=1.0$ for gradient descent.
Refer to caption: /html/2410.04870/assets/x10.png

Figure: Figure 11: Training and test loss. We use data setting (d) in Tab. [3](#A2.T3) with $d=2000$. We use $\beta_{2}=0.99$ and $\epsilon=\text{1e-15}$ in Adam and we use $\eta=1.0$ for gradient descent. In this case, gradient descent converges much faster than sign gradient descent and Adam since we use a large learning rate $\eta=1.0$, which is 1e4 times of the learning rate used in sign gradient descent and Adam.
Refer to caption: /html/2410.04870/assets/x11.png

## Appendix C Preliminary Lemmas

The following lemma studies non-overlap support property in sparse data model.

###### Lemma C.1 (Non-overlapping support of noise, Lemma C.1 in Zou et al. ( 2023 ) ) .

Suppose $s=\Omega(d^{1/2}n^{-2})$.
Let $\left\{(\boldsymbol{\mathbf{X}}_{i},y_{i})\right\}_{i=1,\dots n}$
be the training dataset sampled according
to Definition [2.1](#S2.Thmtheorem1).
Moreover, let $\mathcal{B}_{i}=\operatorname{supp}(\boldsymbol{\mathbf{\xi}}_{i})$ be the support of $\boldsymbol{\mathbf{\xi}}_{i}$.
Then with probability at least $1-n^{-2}$, $\mathcal{B}_{i}\cap\mathcal{B}_{j}=\emptyset$ for all $i,j\in[n]$.

The following lemma studies the relation between sparsity assumption and orthogonality assumption.
Sparsity assumption can imply orthogonality assumption with high probability, which states the generality of orthogonal assumption under sparsity assumption.

###### Lemma C.2 (Sparsity implies orthogonality) .

Suppose $s=\Theta(d^{1/2}n^{-2})$, $d=\operatorname{poly}(n)$. Suppose that the training datasets are generated following Definition. [2.1](#S2.Thmtheorem1) except that the non-zero coordinates of noise vectors are uniformly selected from $[d]$ instead of $[d]\setminus\{1\}$.
Then, with probability at least $1-O(1/(n\sqrt{d}))$, we have $\boldsymbol{\mathbf{\mu}}$ is orthogonal to $\boldsymbol{\mathbf{\xi}}_{i}$ for all $i\in[n]$.

###### Proof.

We have

$$ $\displaystyle\mathbb{P}[\exists i\in[n],(\boldsymbol{\mathbf{\xi}}_{i})_{0}\neq 0]$ $\displaystyle=1-\left(1-\frac{s}{d}\right)^{n}$ $\displaystyle\leq 1-\exp(-2ns/d)$ $\displaystyle\leq 2ns/d$ $\displaystyle=O(1/(n\sqrt{d})).$ $$

Note that this lemma can be extended to the signal vectors with constant non-zero entries without modifying the proof idea.
∎

Let $S_{1}=\left\{i|y_{i}=1\right\}$ and $S_{-1}=\left\{i|y_{i}=1\right\}$.
We have the following lemmas characterizing their sizes.

###### Lemma C.3 .

Suppose that $\delta>0$ and $n\geq 8\log(4/\delta)$. Then with probability at least $1-\delta$, $\left|S_{1}\right|,\left|S_{-1}\in[n/4,3n/4]\right|$.

###### Proof.

Since $\left|S_{1}\right|=\sum_{i\in[n]}\mathds{1}(y_{i}=1)$,
$\left|S_{-1}\right|=\sum_{i\in[n]}\mathds{1}(y_{i}=-1)$,
we have $\mathbb{E}[\left|S_{1}\right|]=\mathbb{E}[\left|S_{-1}\right|]=n/2$.
By Hoeffding’s inequality, for arbitrary $t>0$ the following holds:

$$ $\displaystyle\mathbb{P}{\left|\left|S_{1}\right|-\mathbb{E}{\left|S_{1}\right|}\right|\geq t}\leq 2\exp\left(-\frac{2t^{2}}{n}\right),\mathbb{P}{\left|\left|S_{-1}\right|-\mathbb{E}{\left|S_{-1}\right|}\right|\geq t}\leq 2\exp\left(-\frac{2t^{2}}{n}\right).$ $$

Setting $t=\sqrt{(n/2)\log(4/\delta)}$ and taking a union bound, it follows that with probability at least $1-\delta$, we have

$$ $\displaystyle\left|\left|S_{1}\right|-\frac{n}{2}\right|\leq\sqrt{\frac{n}{2}\log\left(\frac{4}{\delta}\right)},\leavevmode\nobreak\ \left|\left|S_{-1}\right|-\frac{n}{2}\right|\leq\sqrt{\frac{n}{2}\log\left(\frac{4}{\delta}\right)}.$ $$

Therefore, as long as $n\geq 8\log(4/\delta)$, we have $\sqrt{n\log(4/\delta)/2}\leq n/4$ and hence $n/4\leq\left|S_{1}\right|,\left|S_{-1}\right|\leq 3n/4$.
∎

The following lemma estimates the norms of the noise vectors
$\boldsymbol{\mathbf{\xi}}_{i}$ for all $i\in[n]$.

###### Lemma C.4 .

Suppose $\delta>0$ and $s=\Omega(\log(4n/\delta))$.
Then with probability at least $1-\delta$,

$$ $\displaystyle\sigma_{p}^{2}s/2\leq\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{2}^{2}\leq 3\sigma_{p}^{2}s/2,$ $\displaystyle\sigma_{p}s/\sqrt{2}\leq\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\leq\sigma_{p}s,$ $\displaystyle\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}=\sqrt{\frac{2}{\pi}}\sigma_{p}s\pm O(\sqrt{log(4n/\delta)}s^{-1/2}),$ $$

for all $i\in[n]$.

###### Proof.

By Bernstein’s inequality, with probability at least $1-\delta/(2n)$, we have

$$ $\displaystyle\left|\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{2}^{2}-\sigma_{p}^{2}s\right|=O(\sigma_{p}^{2}\sqrt{s\log(4n/\delta)}).$ $$

Therefore, if we set appropriately $s=\Omega(\log(4n/\delta))$, we get

$$ $\displaystyle\sigma_{p}^{2}s/2\leq\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{2}^{2}\leq 3\sigma_{p}^{2}s/2.$ $$

Let $k\in\mathcal{B}_{i}$. Since $\boldsymbol{\mathbf{\xi}}_{i}[k]$ is Gaussian, we have $\left|\boldsymbol{\mathbf{\xi}}_{i}[k]\right|$ is sub-gaussian satisfying

$$ $\displaystyle\left\lVert\left|\boldsymbol{\mathbf{\xi}}_{i}[k]\right|-\mathbb{E}[\left|\boldsymbol{\mathbf{\xi}}_{i}[k]\right|]\right\rVert_{\psi_{2}}$ $\displaystyle\leq 2\left\lVert\left|\boldsymbol{\mathbf{\xi}}_{i}[k]\right|\right\rVert_{\psi_{2}}$ $\displaystyle=2\left\lVert\boldsymbol{\mathbf{\xi}}_{i}[k]\right\rVert_{\psi_{2}}$ $\displaystyle\leq C\sigma_{p}.$ $$

By sub-gaussian tail bounds, with probability at least $1-\delta/2n$, we have

$$ $\displaystyle\left|\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}-\sqrt{\frac{2}{\pi}}\sigma_{p}s\right|=O(\sigma_{p}\sqrt{s\log(4n/\delta)}).$ $$

Therefore, if we set appropriately $s=\Omega(\log(4n/\delta))$, we get

$$ $\displaystyle\sigma_{p}s/\sqrt{2}\leq\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\leq\sigma_{p}s.$ $$

By union bound, we complete the proof.
∎

###### Lemma C.5 .

Suppose $\delta>0$ and $s=\Omega(n^{2}\log(4n/\delta))$.
Let $S_{1},S_{2}\subset[n]$ satisfying $S_{1}\cup S_{2}=[n]$, $S_{1}$ and $S_{2}$ are disjoint, and
$\left|S_{1}\right|-\left|S_{2}\right|=c>0$ for some constant $c$, then

$$ $\displaystyle c\sigma_{p}s/\sqrt{2}\leq\sum_{i\in S_{1}}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}-\sum_{i\in S_{2}}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\leq c\sigma_{p}s.$ $$

###### Proof.

By sub-gaussian tail bounds, with probability at least $1-\delta/2n$, we have

$$ $\displaystyle\left|\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}-\sqrt{\frac{2}{\pi}}\sigma_{p}s\right|=O(\sigma_{p}\sqrt{s\log(4n/\delta)}).$ $$

Then we have

$$ $\displaystyle\sum_{i\in S_{1}}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\geq\left|S_{1}\right|\left(\sqrt{\frac{2}{\pi}}\sigma_{p}s-O(\sigma_{p}\sqrt{s\log(4n/\delta)})\right),$ $\displaystyle\sum_{i\in S_{2}}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\leq\left|S_{2}\right|\left(\sqrt{\frac{2}{\pi}}\sigma_{p}s+O(\sigma_{p}\sqrt{s\log(4n/\delta)})\right),$ $$

which imply

$$ $\displaystyle\sum_{i\in S_{1}}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}-\sum_{i\in S_{2}}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\geq c\sqrt{\frac{2}{\pi}}\sigma_{p}s-O(n\sigma_{p}\sqrt{s\log(4n/\delta)}).$ $$

If we set appropriately $s=\Omega(n^{2}\log(4n/\delta))$, we have

$$ $\displaystyle\sum_{i\in S_{1}}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}-\sum_{i\in S_{2}}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\geq c\sigma_{p}s/\sqrt{2}.$ $$

Similarly, we have

$$ $\displaystyle\sum_{i\in S_{1}}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}-\sum_{i\in S_{2}}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\leq c\sigma_{p}s.$ $$

∎

Now turning to network initialization, the following lemmas study the inner product between a randomly initialized value parameter vector $\boldsymbol{\mathbf{w}}_{V,j,r}$ for $j\in\left\{\pm 1\right\}$ and $r\in[m_{v}]$ or query/key parameter vector $\boldsymbol{\mathbf{w}}_{Q,s}$/$\boldsymbol{\mathbf{w}}_{K,s}$ for $s\in[m_{k}]$ in attention layer,
and the signal/noise vectors in the training data. The calculations characterize how the neural network at initialization randomly captures signal and noise information.

###### Lemma C.6 .

Suppose that $s=\Omega(\log(m_{v}n/\delta))$, $m_{v}=\Omega(\log(1/\delta))$. Then with probability at least $1-\delta$,

$$ $\displaystyle\sigma_{0}^{2}d/2\leq\left\lVert\boldsymbol{\mathbf{w}}_{V,j,r}^{(0)}\right\rVert_{2}^{2}\leq 3\sigma_{0}^{2}d/2,$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{V,j,r}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq\sqrt{2\log(12m_{v}/\delta)}\cdot\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert_{2},$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{V,j,r}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq 2\sqrt{\log(12m_{v}n/\delta)}\cdot\sigma_{0}\sigma_{p}\sqrt{s},$ $$

for all $r\in[m_{v}]$, $j\in\left\{\pm 1\right\}$ and $i\in[n]$.

###### Lemma C.7 .

Suppose that $s=\Omega(\log(m_{v}n/\delta))$, $m_{v}=\Omega(\log(1/\delta))$. Then with probability at least $1-\delta$,

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{v}}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq 2\sqrt{\log(12m_{v}/\delta)}\cdot m_{v}^{-1/2}\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert_{2},$ $\displaystyle\left|\langle\boldsymbol{\mathbf{v}}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq 2\sqrt{2\log(12m_{v}n/\delta)}\cdot m_{v}^{-1/2}\sigma_{0}\sigma_{p}\sqrt{s},$ $$

for all $i\in[n]$.

###### Lemma C.8 .

Suppose that $s=\Omega(\log(m_{k}n/\delta))$, $m_{k}=\Omega(\log(1/\delta))$. Then with probability at least $1-\delta$,

$$ $\displaystyle\sigma_{0}^{2}d/2\leq\left\lVert\boldsymbol{\mathbf{w}}_{Q,s}^{(0)}\right\rVert_{2}^{2}\leq 3\sigma_{0}^{2}d/2,$ $\displaystyle\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/2\leq\max_{s\in[m_{k}]}\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq\sqrt{2\log(12m_{k}/\delta)}\cdot\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert_{2},$ $\displaystyle\sigma_{0}\sigma_{p}\sqrt{s}/4\leq\max_{s\in[m_{k}]}\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq 2\sqrt{\log(12m_{k}n/\delta)}\cdot\sigma_{0}\sigma_{p}\sqrt{s},$ $\displaystyle\sigma_{0}^{2}d/2\leq\left\lVert\boldsymbol{\mathbf{w}}_{K,s}^{(0)}\right\rVert_{2}^{2}\leq 3\sigma_{0}^{2}d/2,$ $\displaystyle\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/2\leq\max_{s\in[m_{k}]}\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq\sqrt{2\log(12m_{k}/\delta)}\cdot\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert_{2},$ $\displaystyle\sigma_{0}\sigma_{p}\sqrt{s}/4\leq\max_{s\in[m_{k}]}\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq 2\sqrt{\log(12m_{k}n/\delta)}\cdot\sigma_{0}\sigma_{p}\sqrt{s},$ $$

for all $s\in[m_{k}]$, and $i\in[n]$.

###### Proof.

The proof of Lemma [C.6](#A3.Thmtheorem6), [C.7](#A3.Thmtheorem7) and [C.8](#A3.Thmtheorem8) are the same as the Lemma B.3 in Cao et al. ().
∎

Define

$$ $\displaystyle\beta_{\boldsymbol{\mathbf{\xi}}}=\max_{i,s,j,r}\left\{\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{V,j,r}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\right\},$ (2) $\displaystyle\beta_{\boldsymbol{\mathbf{\mu}}}=\max_{s,j,r}\left\{\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{V,j,r}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|\right\}.$ (3) $$

The following lemmas study the anti-concentration behaviour in a randomly initialized neural network.

###### Lemma C.9 .

Let $x\sim N(0,\sigma^{2})$. Then

$$ $\displaystyle\mathbb{P}(\left|x\right|\leq c)\leq\operatorname{erf}\left(\frac{c}{\sqrt{2}\sigma}\right)\leq\sqrt{1-\exp(-\frac{2c^{2}}{\pi\sigma^{2}})}.$ $$

###### Proof.

The probability density function for $x$ is given by

$$ $\displaystyle f(x)=\frac{1}{\sqrt{2\pi}\sigma}\exp(-\frac{x^{2}}{2\sigma^{2}}).$ $$

Then we know that

$$ $\displaystyle\mathbb{P}(|x|\leq c)=\frac{1}{\sqrt{2\pi}\sigma}\int_{-c}^{c}\exp(-\frac{x^{2}}{2\sigma^{2}})dx.$ $$

By the definition of $\mathrm{erf}$ function

$$ $\displaystyle\mathrm{erf}(c)=\frac{2}{\sqrt{\pi}}\int_{0}^{c}\exp(-{x^{2}})dx,$ $$

and variable substitution yields

$$ $\displaystyle\mathrm{erf}(\frac{c}{\sqrt{2}\sigma})=\frac{1}{\sqrt{2\pi}\sigma}\int_{0}^{c}\exp(-\frac{x^{2}}{2\sigma^{2}})dx.$ $$

Therefore, we first conclude $\mathbb{P}(|x|\leq c)=2\mathrm{erf}(\frac{c}{\sqrt{2}\sigma})$.

Next, by the inequality $\mathrm{erf}(x)\leq\sqrt{1-\exp(-4x^{2}/\pi)}$, we finally obtain

$$ $\displaystyle\mathbb{P}(|x|\leq c)\leq 2\sqrt{1-\exp(-\frac{2c^{2}}{\sigma^{2}\pi})}.$ $$

∎

###### Lemma C.10 .

Suppose the results in Lemma [C.8](#A3.Thmtheorem8) and Lemma [C.4](#A3.Thmtheorem4) hold.
Suppose that
$\sigma_{p}^{2}s\geq\frac{32\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2}\log(12m_{k}/\delta)}{-\pi\log(1-\delta^{2}/4m_{k}^{2}n^{2})}$.
Then with probability at least $1-\delta$, we have

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq 2\beta_{\boldsymbol{\mathbf{\mu}}},\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq 2\beta_{\boldsymbol{\mathbf{\mu}}},$ $$

for all $s\in[m_{k}]$ and $i\in[n]$.

###### Proof.

Given $\boldsymbol{\mathbf{\xi}}_{i}$, $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\sim N(0,\sigma_{0}^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2})$.
By Lemma [C.9](#A3.Thmtheorem9), we have

$$ $\displaystyle\mathbb{P}\left(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq c\right)\leq\sqrt{1-\exp(-\frac{2c^{2}}{\pi\sigma_{0}^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2}})}.$ $$

Let $c=2\beta_{\boldsymbol{\mathbf{\mu}}}$.
By Lemma [C.8](#A3.Thmtheorem8),
$\beta_{\boldsymbol{\mathbf{\mu}}}$ can be upper bounded by $\sqrt{2\log(12m_{k}/\delta)}\cdot\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert$.
By Lemma [C.4](#A3.Thmtheorem4),

$$ $\displaystyle\sigma_{p}^{2}s/2\leq\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2}\leq 3\sigma_{p}^{2}s/2,$ $$

for all $i\in[n]$. Therefore, if we set appropriately

$$ $\displaystyle\sigma_{p}^{2}s\geq\frac{32\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2}\log(12m_{k}/\delta)}{-\pi\log(1-\delta^{2}/4m_{k}^{2}n^{2})},$ $$

then we have

$$ $\displaystyle\mathbb{P}\left(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq 2\beta_{\boldsymbol{\mathbf{\mu}}}\right)\leq\frac{\delta}{2m_{k}n}.$ $$

The analysis for $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ is exactly the same.
By union bound, we get the final result.
∎

###### Lemma C.11 .

Suppose the results in Lemma [C.8](#A3.Thmtheorem8) and Lemma [C.4](#A3.Thmtheorem4) hold.
Suppose that
$m_{v}\geq\frac{256\log^{2}(12m_{k}n/\delta)}{\pi^{2}\log^{2}(1-\delta^{2}/4m_{k}^{2}n^{2})}$.
Then with probability at least $1-\delta$, we have

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq\beta_{\boldsymbol{\mathbf{\xi}}}m_{v}^{-1/4},\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq\beta_{\boldsymbol{\mathbf{\xi}}}m_{v}^{-1/4},$ $$

for all $s\in[m_{k}]$ and $i\in[n]$.

###### Proof.

Given $\boldsymbol{\mathbf{\xi}}_{i}$, $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\sim N(0,\sigma_{0}^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2})$.
By Lemma [C.9](#A3.Thmtheorem9), we have

$$ $\displaystyle\mathbb{P}\left(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq c\right)\leq\sqrt{1-\exp(-\frac{2c^{2}}{\pi\sigma_{0}^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2}})}.$ $$

Let $c=\beta_{\boldsymbol{\mathbf{\xi}}}m_{v}^{-1/4}$.
By Lemma [C.8](#A3.Thmtheorem8),
$\beta_{\boldsymbol{\mathbf{\xi}}}$ can be upper bounded by $2\sqrt{\log(12m_{k}n/\delta)}\cdot\sigma_{0}\sigma_{p}\sqrt{s}$.
By Lemma [C.4](#A3.Thmtheorem4),

$$ $\displaystyle\sigma_{p}^{2}s/2\leq\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2}\leq 3\sigma_{p}^{2}s/2,$ $$

for all $i\in[n]$. Therefore, if we set appropriately

$$ $\displaystyle m_{v}\geq\frac{256\log^{2}(12m_{k}n/\delta)}{\pi^{2}\log^{2}(1-\delta^{2}/4m_{k}^{2}n^{2})},$ $$

then we have

$$ $\displaystyle\mathbb{P}\left(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq\beta_{\boldsymbol{\mathbf{\xi}}}m_{v}^{-1/4}\right)\leq\frac{\delta}{2m_{k}n}.$ $$

The analysis for $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ is exactly the same.
By union bound, we get the final result.
∎

###### Lemma C.12 .

Suppose that
$s\geq n^{3}\left(\frac{2}{-\pi\log(1-\delta^{2}/4m_{k}^{2})}\right)^{3/2}$.
Then with probability at least $1-\delta$, we have

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|\geq\sigma_{0}ns^{-1/3}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ $$

for all $s\in[m_{k}]$.

###### Proof.

Note that $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\sim N(0,\sigma_{0}^{2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2})$.
By Lemma [C.9](#A3.Thmtheorem9), we have

$$ $\displaystyle\mathbb{P}\left(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq c\right)\leq\sqrt{1-\exp(-\frac{2c^{2}}{\pi\sigma_{0}^{2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2}})}.$ $$

Let $c=\sigma_{0}ns^{-1/3}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert$.
If we set appropriately

$$ $\displaystyle s\geq n^{3}\left(\frac{2}{-\pi\log(1-\delta^{2}/4m_{k}^{2})}\right)^{3/2},$ $$

then we have

$$ $\displaystyle\mathbb{P}\left(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq\sigma_{0}ns^{-1/3}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\right)\leq\frac{\delta}{2m_{k}}.$ $$

The analysis for $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle$ is exactly the same.
By union bound, we get the final result.
∎

The following lemma studies the magnitude of output of softmax operation.

###### Lemma C.13 .

Suppose that $s=\Omega(\log(m_{k}n/\delta))$, $m_{k}=\Omega(\log(1/\delta))$. Then with probability at least $1-\delta$,

$$ $\displaystyle\left|s_{i,11}^{(0)}-1/2\right|\leq 16m_{k}\log(12m_{k}n/\delta)\cdot\sigma_{0}^{2}\max\left\{\sigma_{p}^{2}s,\left\lVert\boldsymbol{\mathbf{\mu}}^{2}\right\rVert\right\},$ $\displaystyle\left|s_{i,21}^{(0)}-1/2\right|\leq 16m_{k}\log(12m_{k}n/\delta)\cdot\sigma_{0}^{2}\max\left\{\sigma_{p}^{2}s,\left\lVert\boldsymbol{\mathbf{\mu}}^{2}\right\rVert\right\},$ $$

for all $i\in[n]$. When $m_{k}\log(12m_{k}n/\delta)\cdot\sigma_{0}^{2}\max\left\{\sigma_{p}^{2}s,\left\lVert\boldsymbol{\mathbf{\mu}}^{2}\right\rVert\right\}=o(1)$, we can simplify the results into

$$ $\displaystyle s_{i,11}^{(0)},s_{i,21}^{(0)}=1/2+o(1),$ $$

for all $i\in[n]$.

###### Proof.

We directly have

$$ $\displaystyle\left|\sum_{s=1}^{m_{k}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq m_{k}\beta_{\boldsymbol{\mathbf{\mu}}}^{2},$ $\displaystyle\left|\sum_{s=1}^{m_{k}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq m_{k}\beta_{\boldsymbol{\mathbf{\xi}}}^{2},$ $\displaystyle\left|\sum_{s=1}^{m_{k}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\sum_{s=1}^{m_{k}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq m_{k}\beta_{\boldsymbol{\mathbf{\mu}}}\beta_{\boldsymbol{\mathbf{\xi}}},$ $$

for all $i\in[n]$. Then we have

$$ $\displaystyle s_{i,21}^{(0)}$ $\displaystyle\leq\frac{\exp\left(m_{k}\beta_{\boldsymbol{\mathbf{\xi}}}\beta_{\boldsymbol{\mathbf{\mu}}}\right)}{\exp\left(m_{k}\beta_{\boldsymbol{\mathbf{\xi}}}\beta_{\boldsymbol{\mathbf{\mu}}}\right)+\exp\left(-m_{k}\beta_{\boldsymbol{\mathbf{\xi}}}^{2}\right)}$ $\displaystyle\leq\frac{1}{1+\exp\left(-m_{k}\beta_{\boldsymbol{\mathbf{\xi}}}(\beta_{\boldsymbol{\mathbf{\xi}}}+\beta_{\boldsymbol{\mathbf{\mu}}})\right)}$ $\displaystyle\leq 1/2+2m_{k}\beta_{\boldsymbol{\mathbf{\xi}}}(\beta_{\boldsymbol{\mathbf{\xi}}}+\beta_{\boldsymbol{\mathbf{\mu}}}),$ $$

and

$$ $\displaystyle s_{i,21}^{(0)}$ $\displaystyle\geq\frac{\exp\left(-m_{k}\beta_{\boldsymbol{\mathbf{\xi}}}\beta_{\boldsymbol{\mathbf{\mu}}}\right)}{\exp\left(-m_{k}\beta_{\boldsymbol{\mathbf{\xi}}}\beta_{\boldsymbol{\mathbf{\mu}}}\right)+\exp\left(m_{k}\beta_{\boldsymbol{\mathbf{\xi}}}^{2}\right)}$ $\displaystyle\geq\frac{1}{1+\exp\left(m_{k}\beta_{\boldsymbol{\mathbf{\xi}}}(\beta_{\boldsymbol{\mathbf{\xi}}}+\beta_{\boldsymbol{\mathbf{\mu}}})\right)}$ $\displaystyle\geq 1/2-2m_{k}\beta_{\boldsymbol{\mathbf{\xi}}}(\beta_{\boldsymbol{\mathbf{\xi}}}+\beta_{\boldsymbol{\mathbf{\mu}}}),$ $$

for all $i\in[n]$. Similarly, we have

$$ $\displaystyle 1/2-2m_{k}\beta_{\boldsymbol{\mathbf{\mu}}}(\beta_{\boldsymbol{\mathbf{\xi}}}+\beta_{\boldsymbol{\mathbf{\mu}}})\leq s_{i,11}^{(0)}\leq 1/2+2m_{k}\beta_{\boldsymbol{\mathbf{\mu}}}(\beta_{\boldsymbol{\mathbf{\xi}}}+\beta_{\boldsymbol{\mathbf{\mu}}})$ $$

for all $i\in[n]$. By Lemma [C.8](#A3.Thmtheorem8), with probability $1-\delta$, $\beta_{\boldsymbol{\mathbf{\xi}}}$ can be bounded by $2\sqrt{\log(12m_{k}n/\delta)}\cdot\sigma_{0}\sigma_{p}\sqrt{s}$ and $\beta_{\boldsymbol{\mathbf{\mu}}}$ can be bounded by $\sqrt{2\log(12m_{k}/\delta)}\cdot\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert$, and

$$ $\displaystyle\max\left\{\beta_{\boldsymbol{\mathbf{\mu}}},\beta_{\boldsymbol{\mathbf{\xi}}}\right\}\leq 2\sqrt{\log(12m_{k}n/\delta)}\cdot\sigma_{0}\max\left\{\sigma_{p}\sqrt{s},\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\right\},$ $$

which completes the proof.
∎

###### Lemma C.14 .

Suppose that $\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert=\Omega(d^{1/5})$.
Let

$$ $\displaystyle\leavevmode\nobreak\ X_{i}:=\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle,$ $\displaystyle\leavevmode\nobreak\ Y_{i}:=-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle.$ $$

Then with probability at least $1-nd^{-1/5}$, for all $i\in[n]$, we have

$$ $\displaystyle\frac{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\cdot(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle)\geq-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle,$ $$

or

$$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\geq\frac{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\cdot(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle).$ $$

###### Proof.

We only need to show that for all $i\in[n]$ the probability of following event can be controlled.

$$ $\displaystyle\frac{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}X_{i}<Y_{i}<\frac{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}X_{i}.$ $$

Conditioning on $y_{i}\boldsymbol{\mathbf{\xi}}_{i}$, we have $X_{i}\,{\buildrel d\over{=}}\,Y_{i}\sim N(0,\sigma_{0}^{2}(\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2}+\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2}))$,
this is because $\langle\boldsymbol{\mathbf{\xi}}_{i},\boldsymbol{\mathbf{\mu}}\rangle=0$. Denote

$$ $\displaystyle\leavevmode\nobreak\ \sigma^{2}:=\sigma_{0}^{2}(\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2}+\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2}),$ $\displaystyle\leavevmode\nobreak\ D:=\left\{(x,y)\in\mathbb{R}^{2}:\frac{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}x<y<\frac{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}x\right\}.$ $$

Therefore, we have

$$ $\displaystyle\leavevmode\nobreak\ \mathbb{P}\left[\frac{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}X_{i}<Y_{i}<\frac{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}X_{i}|y_{i}\boldsymbol{\mathbf{\xi}}_{i}\right]$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \iint_{D}1\cdot dP_{X_{i}}dP_{Y_{i}}$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \iint_{D}\frac{1}{2\pi\sigma^{2}}\exp\left(-\frac{x^{2}+y^{2}}{2\sigma^{2}}\right)dxdy$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \iint_{D}\frac{1}{2\pi\sigma^{2}}r\exp\left(-\frac{r^{2}}{2\sigma^{2}}\right)drd\theta$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \frac{\pi}{2}-2\arctan\left(\frac{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\right)$ $\displaystyle\leq$ $\displaystyle\leavevmode\nobreak\ \frac{4}{3}\frac{2\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}$ $\displaystyle\leq$ $\displaystyle\leavevmode\nobreak\ O(d^{-1/5}).$ $$

where the fifth step is by $\arctan(1+x)\geq\frac{\pi}{4}+\frac{2}{3}(x-1)$ when $x=1-o(1)$,
and the last step is by the condition $\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert=\Omega(d^{1/5})$.
Then, by integrating over $y_{i}\boldsymbol{\mathbf{\xi}}_{i}$ and union bound, we have

$$ $\displaystyle\mathbb{P}\left[\exists\leavevmode\nobreak\ i\in[n],\leavevmode\nobreak\ \frac{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}X_{i}<Y_{i}<\frac{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}X_{i}\right]\leq O(nd^{-1/5}).$ $$

∎

Next, let

$$ $\displaystyle S_{s,1,K+,Q-}^{(0)}$ $\displaystyle=\left\{i\in[n]:\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle>0,\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle<0,y_{i}=1\right\},$ (4) $\displaystyle S_{i,K+,Q-}^{(0)}$ $\displaystyle=\left\{s\in[m_{k}]:\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle>0,\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle<0\right\},$ (5) $\displaystyle S_{\boldsymbol{\mathbf{\mu}},K+,Q-}^{(0)}$ $\displaystyle=\left\{s\in[m_{k}]:\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle>0,\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle<0\right\},$ (6) $\displaystyle S_{s,K+,Q-}^{(0)}$ $\displaystyle=\left\{i\in[n]:\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0,\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle<0,\right\}=S_{s,1,K+,Q-}^{(0)}\cup S_{s,-1,K-,Q+}^{(0)}.$ (7) $$

###### Lemma C.1 (Non-overlapping support of noise, Lemma C.1 in Zou et al. ( 2023 ) ) .

###### Lemma C.2 (Sparsity implies orthogonality) .

###### Proof.

###### Lemma C.3 .

###### Proof.

###### Lemma C.4 .

###### Proof.

###### Lemma C.5 .

###### Proof.

###### Lemma C.6 .

###### Lemma C.7 .

###### Lemma C.8 .

###### Proof.

###### Lemma C.9 .

###### Proof.

###### Lemma C.10 .

###### Proof.

###### Lemma C.11 .

###### Proof.

###### Lemma C.12 .

###### Proof.

###### Lemma C.13 .

###### Proof.

###### Lemma C.14 .

###### Proof.

## Appendix D Gradient

In this section, we present the gradients of the loss $L_{S}$ with respect to $\boldsymbol{\mathbf{w}}_{V,j,r}$, $\boldsymbol{\mathbf{w}}_{Q,s}$ and $\boldsymbol{\mathbf{w}}_{K,s}$ for $s\in[m_{k}]$, $j\in\left\{\pm 1\right\}$ and $r\in[m_{v}]$. Our analysis starts from a more general form:

$$ $\displaystyle F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}}):=\frac{1}{m_{v}}\sum_{l=1}^{L}\textbf{1}^{\top}_{m_{v}}\sigma(\boldsymbol{\mathbf{W}}_{V,j}\boldsymbol{\mathbf{X}}\text{softmax}\left(\boldsymbol{\mathbf{X}}^{\top}\boldsymbol{\mathbf{W}}_{K}^{\top}\boldsymbol{\mathbf{W}}_{Q}\boldsymbol{\mathbf{x}}^{(l)}\right)),$ $$

and finally focus on $L=2$ and $\sigma(x)=x$ case, which is exactly the model defined in Eq. ([1](#S2.E1)).

###### Lemma D.1 (softmax derivative) .

Let $\left(s_{1},\dots,s_{n}\right)=\text{softmax}\left(z_{1},\dots,z_{n}\right)$, then we have

$$ $\displaystyle\frac{\partial s_{i}}{\partial z_{j}}=s_{i}\left(\mathds{1}\left\{i=j\right\}-s_{j}\right).$ $$

###### Lemma D.2 (Gradient w.r.t. value parameters) .

Consider the transformer model defined in Eq. ([1](#S2.E1)).
Let $\left(\boldsymbol{\mathbf{X}},y\right)$ be a data point. The derivative of $F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ with respect to $\boldsymbol{\mathbf{w}}_{V,j,r}$ is:

$$ $\displaystyle\nabla_{\boldsymbol{\mathbf{w}}_{V,j,r}}F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})=\frac{1}{m_{v}}\left[\left(s_{1,1}+s_{2,1}\right)\boldsymbol{\mathbf{x}}^{(1)}+\left(s_{1,2}+s_{2,2}\right)\boldsymbol{\mathbf{x}}^{(2)}\right].$ $$

###### Proof.

We have

$$ $\displaystyle\nabla_{\boldsymbol{\mathbf{w}}_{V,j,r}}F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ $\displaystyle=\frac{1}{m_{v}}\sum_{l\in[L]}\sum_{a=1}^{L}s_{l,a}\sigma^{\prime}\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(a)}\right\rangle\right)\boldsymbol{\mathbf{x}}^{(a)}$ $\displaystyle=\frac{1}{m_{v}}\left[\left(s_{1,1}+s_{2,1}\right)\sigma^{\prime}\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(1)}\right\rangle\right)\boldsymbol{\mathbf{x}}^{(1)}+\left(s_{1,2}+s_{2,2}\right)\sigma^{\prime}\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\right)\boldsymbol{\mathbf{x}}^{(2)}\right]$ $\displaystyle=\frac{1}{m_{v}}\left[\left(s_{1,1}+s_{2,1}\right)\boldsymbol{\mathbf{x}}^{(1)}+\left(s_{1,2}+s_{2,2}\right)\boldsymbol{\mathbf{x}}^{(2)}\right],$ $$

where the second step is by $L=2$, the third step is by $\sigma(x)=x$.
∎

###### Lemma D.3 (Gradient w.r.t. pre-softmax quantity) .

Consider the transformer model defined in Eq. ([1](#S2.E1)).
Let $\left(\boldsymbol{\mathbf{X}},y\right)$ be a data point. The derivative of $F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ with respect to $z_{l,a}$ is:

$$ $\displaystyle\frac{\partial F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})}{\partial z_{l,a}}$ $\displaystyle=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}s_{l,1}s_{l,2}\left[\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}_{a}\right\rangle\right)-\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}_{3-a}\right\rangle\right)\right].$ $$

###### Proof.

We have

$$ $\displaystyle\frac{\partial F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{x}})}{\partial z_{l,a}}$ $\displaystyle=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}\sum_{b=1}^{L}\frac{\partial s_{l,b}}{\partial z_{l,a}}\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(b)}\right\rangle\right)$ $\displaystyle=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}\sum_{b\neq a}s_{l,a}s_{l,b}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(a)}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(b)}\right\rangle\right)\right]$ $\displaystyle=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}s_{l,a}\sum_{b\in[L]}s_{l,b}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(a)}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(b)}\right\rangle\right)\right]$ $\displaystyle=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}s_{l,1}s_{l,2}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}_{a}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}_{3-a}\right\rangle\right)\right]$ $\displaystyle=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}s_{l,1}s_{l,2}\left[\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}_{a}\right\rangle\right)-\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}_{3-a}\right\rangle\right)\right],$ $$

where the penultimate step is by $L=2$,
the final step is by $\sigma(x)=x$.
∎

###### Lemma D.4 (Gradient w.r.t. query parameters) .

Consider the transformer model defined in Eq. ([1](#S2.E1)).
Let $\left(\boldsymbol{\mathbf{X}},y\right)$ be a data point. The derivative of $F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ with respect to $\boldsymbol{\mathbf{w}}_{Q,s}$ is:

$$ $\displaystyle\nabla_{\boldsymbol{\mathbf{w}}_{Q,s}}F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ $\displaystyle=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left(s_{1,1}s_{1,2}\boldsymbol{\mathbf{x}}^{(1)}+s_{2,1}s_{2,2}\boldsymbol{\mathbf{x}}^{(2)}\right)$ $\displaystyle=\left\langle\bar{\boldsymbol{\mathbf{w}}}_{V,j},\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left(s_{1,1}s_{1,2}\boldsymbol{\mathbf{x}}^{(1)}+s_{2,1}s_{2,2}\boldsymbol{\mathbf{x}}^{(2)}\right).$ $$

###### Proof.

By Lemma [D.3](#A4.Thmtheorem3), we have

$$ $\displaystyle\nabla_{\boldsymbol{\mathbf{w}}_{Q,s}}F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ $\displaystyle=\sum_{l,a}\frac{\partial F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})}{\partial z_{l,a}}\nabla_{\boldsymbol{\mathbf{w}}_{Q,s}}z_{l,a}$ $\displaystyle=\sum_{l,a}\left[\frac{1}{m_{v}}\sum_{m\in[d_{v}]}s_{l,a}\sum_{b\in[L]}s_{l,b}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(a)}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(b)}\right\rangle\right)\right]\right]\cdot\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(a)}\right\rangle\boldsymbol{\mathbf{x}}^{(l)}$ $\displaystyle=\sum_{l,a}\frac{1}{m_{v}}\sum_{r\in[m_{v}]}s_{l,1}s_{l,2}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(a)}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(3-a)}\right\rangle\right)\right]\cdot\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(a)}\right\rangle\boldsymbol{\mathbf{x}}^{(l)}$ $\displaystyle=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}s_{1,1}s_{1,2}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(1)}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\right)\right]\cdot\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(1)}\right\rangle\boldsymbol{\mathbf{x}}^{(1)}$ $\displaystyle+s_{1,1}s_{1,2}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(1)}\right\rangle\right)\right]\cdot\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\boldsymbol{\mathbf{x}}^{(1)}$ $\displaystyle+s_{2,1}s_{2,2}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(1)}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\right)\right]\cdot\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(1)}\right\rangle\boldsymbol{\mathbf{x}}^{(2)}$ $\displaystyle+s_{2,1}s_{2,2}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(1)}\right\rangle\right)\right]\cdot\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\boldsymbol{\mathbf{x}}^{(2)}$ $\displaystyle=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(1)}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\right)\right]$ $\displaystyle\cdot\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left(s_{1,1}s_{1,2}\boldsymbol{\mathbf{x}}^{(1)}+s_{2,1}s_{2,2}\boldsymbol{\mathbf{x}}^{(2)}\right)$ $\displaystyle=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left(s_{1,1}s_{1,2}\boldsymbol{\mathbf{x}}^{(1)}+s_{2,1}s_{2,2}\boldsymbol{\mathbf{x}}^{(2)}\right)$ $\displaystyle=\left\langle\bar{\boldsymbol{\mathbf{w}}}_{V,j},\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left(s_{1,1}s_{1,2}\boldsymbol{\mathbf{x}}^{(1)}+s_{2,1}s_{2,2}\boldsymbol{\mathbf{x}}^{(2)}\right),$ $$

where the third step is by $L=2$, the last step is by $\sigma(x)=x$.
∎

###### Lemma D.5 (Gradient w.r.t. key parameters) .

Consider the transformer model defined in Eq. ([1](#S2.E1)).
Let $\left(\boldsymbol{\mathbf{X}},y\right)$ be a data point. Similarly, The derivative of $F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ with respect to $\boldsymbol{\mathbf{w}}_{K,s}$ is:

$$ $\displaystyle\nabla_{\boldsymbol{\mathbf{w}}_{K,s}}F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ $\displaystyle=\left\langle\bar{\boldsymbol{\mathbf{w}}}_{V,j},\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left\langle\boldsymbol{\mathbf{w}}_{Q,s},s_{1,1}s_{1,2}\boldsymbol{\mathbf{x}}^{(1)}+s_{2,1}s_{2,2}\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left(\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right).$ $$

###### Proof.

By Lemma [D.3](#A4.Thmtheorem3), we have

$$ $\displaystyle\nabla_{\boldsymbol{\mathbf{w}}_{K,s}}F_{j}(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ $\displaystyle=\sum_{l,a}\left[\frac{1}{m_{v}}\sum_{r\in[m_{v}]}s_{l,a}\sum_{b\in[L]}s_{l,b}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(a)}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(b)}\right\rangle\right)\right]\right]\cdot\left\langle\boldsymbol{\mathbf{w}}_{Q,s},\boldsymbol{\mathbf{x}}^{(l)}\right\rangle\boldsymbol{\mathbf{x}}^{(a)}$ $\displaystyle=\frac{1}{m_{v}}\sum_{r\in[m_{v}]}\left[\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(1)}\right\rangle\right)-\sigma\left(\left\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\right)\right]$ $\displaystyle\cdot\left\langle\boldsymbol{\mathbf{w}}_{Q,s},s_{1,1}s_{1,2}\boldsymbol{\mathbf{x}}^{(1)}+s_{2,1}s_{2,2}\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left(\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right)$ $\displaystyle=\left\langle\bar{\boldsymbol{\mathbf{w}}}_{V,j},\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left\langle\boldsymbol{\mathbf{w}}_{Q,s},s_{1,1}s_{1,2}\boldsymbol{\mathbf{x}}^{(1)}+s_{2,1}s_{2,2}\boldsymbol{\mathbf{x}}^{(2)}\right\rangle\cdot\left(\boldsymbol{\mathbf{x}}^{(1)}-\boldsymbol{\mathbf{x}}^{(2)}\right),$ $$

where the second step is by $L=2$, the third step is by $\sigma(x)=x$.
∎

###### Lemma D.6 .

Consider the transformer model defined in Eq. ([1](#S2.E1)).
Let $\boldsymbol{\mathbf{w}}^{(t)}_{V,j,r}$ for $j\in\left\{\pm 1\right\},\leavevmode\nobreak\ r\in[m_{v}]$ be the value parameters of the TF at the $t$-th iteration. If we use sign gradient descent, the update formula of $\boldsymbol{\mathbf{w}}^{(t)}_{V,j,r}$ is

$$ $\displaystyle\boldsymbol{\mathbf{w}}^{(t+1)}_{V,j,r}=\boldsymbol{\mathbf{w}}^{(t)}_{V,j,r}-\eta\operatorname{sgn}\left(\sum_{i\in[n]}\ell_{i}^{\prime(t)}y_{i}j\left[(s_{i,11}^{(t)}+s_{i,21}^{(t)})y_{i}\boldsymbol{\mathbf{\mu}}+(s_{i,12}^{(t)}+s_{i,22}^{(t)})\boldsymbol{\mathbf{\xi}}_{i}\right]\right).$ $$

###### Proof.

This is by Lemma [D.2](#A4.Thmtheorem2).
∎

###### Lemma D.7 .

Consider the transformer model defined in Eq. ([1](#S2.E1)).
Let $\boldsymbol{\mathbf{w}}^{(t)}_{Q,s}$ and $\boldsymbol{\mathbf{w}}^{(t)}_{K,s}$ for $s\in[m_{k}]$ be the query and key parameters of the TF at the $t$-th iteration.
If we use sign gradient descent, the update formula of $\boldsymbol{\mathbf{w}}^{(t)}_{Q,s}$ and $\boldsymbol{\mathbf{w}}^{(t)}_{K,s}$ is

$$ $\displaystyle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s}=\boldsymbol{\mathbf{w}}^{(t)}_{Q,s}-\eta\operatorname{sgn}\left[\sum_{i\in[n]}{\ell}^{\prime(t)}_{i}y_{i}\langle\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,1}-\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,-1},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle(s_{11}^{(t)}s_{12}^{(t)}y_{i}\boldsymbol{\mathbf{\mu}}+s_{21}^{(t)}s_{22}^{(t)}\boldsymbol{\mathbf{\xi}}_{i})\right].$ $$

$$ $\displaystyle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s}=\boldsymbol{\mathbf{w}}^{(t)}_{K,s}-\eta\operatorname{sgn}\left[\sum_{i\in[n]}{\ell}^{\prime(t)}_{i}y_{i}\langle\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,1}-\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,-1},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},s_{11}^{(t)}s_{12}^{(t)}y_{i}\boldsymbol{\mathbf{\mu}}+s_{21}^{(t)}s_{22}^{(t)}\boldsymbol{\mathbf{\xi}}_{i}\rangle(y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i})\right].$ $$

###### Proof.

This is by Lemma [D.4](#A4.Thmtheorem4), Lemma [D.5](#A4.Thmtheorem5).
∎

###### Lemma D.1 (softmax derivative) .

###### Lemma D.2 (Gradient w.r.t. value parameters) .

###### Proof.

###### Lemma D.3 (Gradient w.r.t. pre-softmax quantity) .

###### Proof.

###### Lemma D.4 (Gradient w.r.t. query parameters) .

###### Proof.

###### Lemma D.5 (Gradient w.r.t. key parameters) .

###### Proof.

###### Lemma D.6 .

###### Proof.

###### Lemma D.7 .

###### Proof.

## Appendix E Proofs

In this section, we give a detailed analysis about the dynamics of sign gradient descent on our transformer models Eq. ([1](#S2.E1)) with our data model in Def. [2.1](#S2.Thmtheorem1).
These results are based on Condition [3.1](#S3.Thmtheorem1) and the conclusions in Appendix [C](#A3), which hold with high probability.
Denote by $\mathcal{E}_{\text{prelim}}$ the event that all the results in Appendix [C](#A3) hold (for a given $\delta$, we see $\mathbb{P}(\mathcal{E}_{\text{prelim}})\geq 1-10\delta-n^{-2}-nd^{-1/5}$ by a union bound).
For simplicity and clarity, we state all the results in this and the following sections conditional on $\mathcal{E}_{\text{prelim}}$.

### E.1 Technique Overview

For value, the dynamics is linear under sign gradient descent, which means the increment is exactly constant across different iterations and different neurons, and up to a constant multiplier across different samples.

For query and key, the dynamics is complicated. The key point is that we utilize the different increasing speed of different core quantities to divide the whole timeline into many stages where we can care about only one of or some of quantities in each stage while the dynamics of other quantities can be ignored or is very simple.
The main reason why the speed is different is attributed to the scale of hyperparameters including $\sigma_{p}s,\sigma_{p}\sqrt{s},\sigma_{0}$.

We will define some key timesteps in following analysis, including

$$ $\displaystyle T_{1}:=\tilde{\Theta}(\sigma_{0}\eta^{-1}s^{-1/2}m_{v}^{-1/2}),$ $\displaystyle T_{2}:=50\sqrt{2}n\beta_{\boldsymbol{\mathbf{\xi}}}\eta^{-1}\sigma_{p}^{-1}s^{-1}=\tilde{O}(\sigma_{0}s^{-1/2}\eta^{-1}),$ $\displaystyle T_{2}^{+}:=C_{1}\sqrt{\log(12m_{k}n/\delta)}\sigma_{0}ns^{-1/2}\eta^{-1}=\tilde{O}(\sigma_{0}s^{-1/2}\eta^{-1}),$ $\displaystyle T_{3}:=3\beta_{\boldsymbol{\mathbf{\mu}}}\eta^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{-1}=\tilde{O}(\sigma_{0}\eta^{-1}),$ $\displaystyle T_{3}^{+}:=C_{2}\sqrt{\log(12m_{k}/\delta)}\sigma_{0}\eta^{-1}=\tilde{O}(\sigma_{0}\eta^{-1}),$ $\displaystyle T_{4}^{-}\geq\tilde{T}_{4}^{-}:=\sqrt{\frac{0.99\pi}{2}}\sqrt{\log\left(\frac{\sigma_{p}s}{3\sqrt{2}n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\right)}\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1}=\tilde{O}(m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1}\eta^{-1}),$ $\displaystyle T_{4}^{-}\leq\hat{T}_{4}^{-}:=\sqrt{\frac{1.01\pi}{2}}\sqrt{\log\left(\frac{\sigma_{p}s}{\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\right)}\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1}=\tilde{O}(m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1}\eta^{-1}),$ $\displaystyle T_{4}^{+}\leq(1+\theta_{c})T_{2}^{-},$ $\displaystyle T_{4}:=C_{3}\log\left(\frac{C_{3}\sigma_{p}s}{\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\right)\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1}=\tilde{O}(m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1}\eta^{-1}),$ $$

where $C_{1},C_{2},C_{3}=\Theta(1)$ are large constants. The order among these timesteps is

$$ $\displaystyle T_{1}\ll T_{2}<T_{2}^{+}\ll T_{3}<T_{3}^{+}\ll T_{4}^{-}<T_{4}^{+}<T_{4}.$ $$

Our analysis basically consists of following parts:

- •
Mean value noise. Let
$\displaystyle T_{1}:=\tilde{\Theta}(\sigma_{0}\eta^{-1}s^{-1/2}m_{v}^{-1/2}).$
For $t\in[0,T_{1}]$, the mean value across all samples align quickly while other inner product barely move.
For $t\geq T_{1}$, we have the mean value noise is linear with time
$\displaystyle\sqrt{2}t\eta\sigma_{p}s\leq\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\leq 2t\eta\sigma_{p}s,$
for all $i\in[n]$, and mean value signal $\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ can be ignored compared with mean value noise $\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$.
Moreover, for $t\geq T_{2}$, we show the mean value noise across different samples are almost the same
$\displaystyle\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\sqrt{\frac{2}{\pi}}t\eta\sigma_{p}s(1+\tilde{O}(s^{-1/2})),$
for all $i\in[n]$.
- •
Useful bounds.
For $t\in[0,T_{3}^{+}]$, noise-signal softmax outputs are close to initialization, i.e., for all $i\in[n]$
$\displaystyle s_{i,21}^{(t)}=1/2+o(1).$
For $t\in[0,T_{4}]$, signal-signal softmax outputs are close to initialization, i.e., for all $i\in[n]$
$\displaystyle s_{i,11}^{(t)}=1/2+o(1).$
For $t\in[0,T_{2}^{+}]$, query/key signal are close to initialization, i.e., for all $s\in[m_{k}]$
$\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|=\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|(1+o(1)),$
$\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|=\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|(1+o(1)).$
For $t\in[0,T_{4}]$, the loss derivative are close to initialization, i.e., for all $i\in[n]$
$\displaystyle\ell_{i}^{\prime(t)}=1/2+o(1).$
For $t\in[0,T_{4}^{-}]$, the ratio of noise-signal softmax outputs are close to 1, i.e., for all $i,k\in[n]$
$\displaystyle\frac{s_{i,21}^{(t)}}{s_{k,21}^{(t)}}=1+o(1).$
- •
Query/Key noise dynamics part i. Let
$\displaystyle T_{2}^{+}:=C_{1}\sqrt{\log(12m_{k}n/\delta)}\sigma_{0}ns^{-1/2}\eta^{-1},$
$\displaystyle T_{2}:=50\sqrt{2}n\beta_{\boldsymbol{\mathbf{\xi}}}\eta^{-1}\sigma_{p}^{-1}s^{-1}=\tilde{O}(\sigma_{0}s^{-1/2}\eta^{-1}).$
where $C_{1}=\Theta(1)$ is a large constant.
For $t\in[0,T_{2}]$, we focus on the dynamics of query/key noise. As mentioned above, softmax outputs and query/key signal are stuck at initialization,
so we can ignore the variation of query/key signal and softmax outputs before $T_{2}$.
All query/key noise can be divided into two groups. Given $s\in[m_{k}]$ and $i\in[n]$,
1.
if $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, then they encourage each other to increase (in magnitude).
2.
if $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, then they first encourage each other to decrease in magnitude, but rebound quickly and reduce to the first case.
For $t\in[T_{2},T_{2}^{+}]$, we have some nice properties about query/key noise
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)}y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$
$\displaystyle t\eta\sigma_{p}s/2\leq\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq 2t\eta\sigma_{p}s,$
for all $s\in[m_{k}]$ and $i\in[n]$.
Also we have some nice properties about the sum of query/key noise $\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, $\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$
$\displaystyle\left|\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq 2n\beta_{\boldsymbol{\mathbf{\mu}}},$
$\displaystyle\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\frac{1}{\sqrt{2}}\eta\sigma_{p}s,$
$\displaystyle\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\frac{1}{\sqrt{2}}\eta\sigma_{p}s,$
for all $s\in[m_{k}]$.
- •
Query/Key signal dynamics part i.
Let
$\displaystyle T_{3}:=3\beta_{\boldsymbol{\mathbf{\mu}}}\eta^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{-1}=\tilde{O}(\sigma_{0}\eta^{-1}),$
$\displaystyle T_{3}^{+}:=C_{2}\sqrt{\log(12m_{k}/\delta)}\sigma_{0}\eta^{-1},$
where $C_{2}=\Theta(1)$ is a large constant.
For $t\in[T_{2},T_{3}]$, we focus on the dynamics of query/key signal. In this period, softmax outputs are stuck at initialization, and
the dynamics of query/key noise keep unchanged
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)}y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$
$\displaystyle\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\frac{1}{\sqrt{2}}\eta\sigma_{p}s,$
$\displaystyle\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\frac{1}{\sqrt{2}}\eta\sigma_{p}s,$
for all $s\in[m_{k}]$ and $i\in[n]$.
We show that the update direction of query/key noise is determined by the summation of query/key noise.
For $t\in[T_{2},T_{3}]$, for all $s\in[m_{k}]$ we have
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle-\operatorname{sgn}(\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert.$
For $t\in[T_{3},T_{3}^{+}]$, we have some nice properties about query/key signal
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle=_{\operatorname{sgn}}\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$
$\displaystyle t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/2\leq\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq 2t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$
for all $s\in[m_{k}]$.
- •
Query/Key signal and noise dynamics part ii.
Let $T_{4}^{-}\geq T_{3}$ be the first time the following condition does not hold, for all $s\in[m_{k}]$
$\displaystyle\left|\sum_{i=1}^{n}s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq\frac{1}{2}n\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|,$
$\displaystyle s_{i,21}^{(t)}s_{i,22}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq 2s_{i,11}^{(t)}s_{i,12}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|.$
Let $T_{4}^{+}\geq T_{4}^{-}$ be the first time the following condition holds, for all $s\in[m_{k}]$ and $i\in[n]$
$\displaystyle s_{i,21}^{(t)}s_{i,22}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq\frac{1}{2}s_{i,11}^{(t)}s_{i,12}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|.$
Let
$\displaystyle T_{4}=C_{3}\log\left(\frac{C_{3}\sigma_{p}s}{\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\right)\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1},$
where $C_{3}=\Theta(1)$ is a large constant
We discuss the dynamics by three parts.
For $t\in[T_{3},T_{4}^{-}]$, see stage IV.a in Appendix [E.7.1](#A5.SS7.SSS1).
For $t\in[T_{4}^{-},T_{4}^{+}]$, see stage IV.b in Appendix [E.7.2](#A5.SS7.SSS2).
For $t\in[T_{4}^{+},T_{4}]$, see stage IV.c in Appendix [E.7.1](#A5.SS7.SSS1).
For $t\geq T_{4}$, we have for all $s\in[m_{k}]$ and $i\in[n]$
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle-\operatorname{sgn}(\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)}y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)}\boldsymbol{\mathbf{\mu}}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)}\boldsymbol{\mathbf{\mu}}\rangle.$
Note that we have following bounds for $T_{4}^{-}$ and $T_{4}^{+}$
$\displaystyle T_{4}^{-}\geq\tilde{T}_{4}^{-}:=\sqrt{\frac{0.99\pi}{2}}\sqrt{\log\left(\frac{\sigma_{p}s}{3\sqrt{2}n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\right)}\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1},$
$\displaystyle T_{4}^{-}\leq\hat{T}_{4}^{-}:=\sqrt{\frac{1.01\pi}{2}}\sqrt{\log\left(\frac{\sigma_{p}s}{\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\right)}\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1},$
$\displaystyle T_{4}^{+}\leq(1+\theta_{c})T_{4}^{-}.$

### E.2 The dynamics of value

Write down the inner product between gradient of value parameters and signal & noise, we have

$$ $\displaystyle\langle\operatorname{sgn}\left(\nabla_{\boldsymbol{\mathbf{w}}_{V,j,r}}L\right),\boldsymbol{\mathbf{\mu}}\rangle=\operatorname{sgn}\left(\sum_{i\in[n]}l_{i}^{\prime(t)}j(s_{i,11}^{(t)}+s_{i,21}^{(t)})\boldsymbol{\mathbf{\mu}}[1]\right)\boldsymbol{\mathbf{\mu}}[1]=-\operatorname{sgn}(j)\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ $$

for all $j\in\left\{\pm 1\right\},r\in[m_{v}]$ and

$$ $\displaystyle\langle\operatorname{sgn}\left(\nabla_{\boldsymbol{\mathbf{w}}_{V,j,r}}L\right),\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=\sum_{k\in[d]}\operatorname{sgn}\left(\sum_{i^{\prime}\in[n]}l_{i^{\prime}}^{\prime(t)}y_{i^{\prime}}j(s_{i^{\prime},12}^{(t)}+s_{i^{\prime},22}^{(t)})\boldsymbol{\mathbf{\xi}}_{i^{\prime}}[k]\right)\boldsymbol{\mathbf{\xi}}_{i}[k],$ $$

for all $i\in[n],j\in\left\{\pm 1\right\},r\in[m_{v}]$, and simplified to sparse setting

$$ $\displaystyle\langle\operatorname{sgn}\left(\nabla_{\boldsymbol{\mathbf{w}}_{V,j,r}}L\right),\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=\sum_{k\in\mathcal{B}_{i}}\operatorname{sgn}\left(l_{i}^{\prime(t)}y_{i}j(s_{i,12}^{(t)}+s_{i,22}^{(t)})\boldsymbol{\mathbf{\xi}}_{i}[k]\right)\boldsymbol{\mathbf{\xi}}_{i}[k],$ $$

for all $i\in[n],j\in\left\{\pm 1\right\},r\in[m_{v}]$.

Consider the sparse property of the data, we have

$$ $\displaystyle\langle\operatorname{sgn}\left(\nabla_{\boldsymbol{\mathbf{w}}_{V,j,r}}L\right),\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=\sum_{k\in\mathcal{B}_{i}}\operatorname{sgn}\left(l_{i}^{\prime(t)}y_{i}j(s_{i,12}^{(t)}+s_{i,22}^{(t)})\boldsymbol{\mathbf{\xi}}_{i}[k]\right)\boldsymbol{\mathbf{\xi}}_{i}[k]$ $\displaystyle=-\operatorname{sgn}(y_{i}j)\sum_{k\in\mathcal{B}_{i}}\left|\boldsymbol{\mathbf{\xi}}_{i}[k]\right|$ $\displaystyle=-\operatorname{sgn}(y_{i}j)\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$ $$

where the last equality is due to $l^{\prime}<0$ and $s_{i,ab}>0$ for all $i\in[n],a,b\in[2]$.
Note the update to value parameters of is irrelevant to the magnitude of query, key, and attention weights. Then, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{V,y_{i},r}^{(t+1)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=\langle\boldsymbol{\mathbf{w}}_{V,y_{i},r}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle+\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{V,-y_{i},r}^{(t+1)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=\langle\boldsymbol{\mathbf{w}}_{V,-y_{i},r}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle-\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{V,y_{i},r}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=\langle\boldsymbol{\mathbf{w}}_{V,y_{i},r}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{V,-y_{i},r}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=\langle\boldsymbol{\mathbf{w}}_{V,-y_{i},r}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle-\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$ $$

As $t$ become large, both $\langle\boldsymbol{\mathbf{w}}_{V,y_{i},r}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{V,y_{i},r}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ increase.
Also, $\langle\boldsymbol{\mathbf{w}}_{V,y_{i},r}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ increases faster since we have $s\sigma_{p}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{-1}=\tilde{\Omega}(1)$.

### E.3 The dynamics of query and key: preparations

##### Additional Notations.

Let $\boldsymbol{\mathbf{v}}^{(t)}:=\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,1}-\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,-1}$.
When we talk about query and key parameters, we refer to $\boldsymbol{\mathbf{w}}_{Q,s},\boldsymbol{\mathbf{w}}_{K,s}$ for all $s\in[m_{k}]$.
When we talk about query/key signal (inner products), we refer to
$\langle\boldsymbol{\mathbf{w}}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle$, $\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle$ for all $s\in[m_{k}]$.
When we talk about query/key noise (inner products), we refer to
$\langle\boldsymbol{\mathbf{w}}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, $\langle\boldsymbol{\mathbf{w}}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ for all $s\in[m_{k}]$ and $i\in[n]$.
When we talk about (mean) value (noise), we refer to $\langle\boldsymbol{\mathbf{v}},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ for all $i\in[n]$.
We use $a=_{\operatorname{sgn}}b$ to denote $\operatorname{sgn}(a)=\operatorname{sgn}(b)$.
Let $\mathcal{B}_{i}$ be the support of sample $\boldsymbol{\mathbf{\xi}}_{i}$ in the training dataset.

##### Some Facts.

Based on gradient of query and key parameters and the sparse property of data, we firstly write down the inner product between gradient and signal & noise

$$ $\displaystyle\langle\operatorname{sgn}\left(\nabla_{\boldsymbol{\mathbf{w}}_{Q,s}}L\right),\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=\operatorname{sgn}\left(\sum_{i\in[n]}{l}^{\prime(t)}_{i}y_{i}\langle\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,1}-\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,-1},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle s_{i,11}^{(t)}s_{i,12}^{(t)}y_{i}\right),$ $\displaystyle\langle\operatorname{sgn}\left(\nabla_{\boldsymbol{\mathbf{w}}_{Q,s}}L\right),\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=\sum_{k\in\mathcal{B}_{i}}\operatorname{sgn}\left({l}^{\prime(t)}_{i}y_{i}\langle\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,1}-\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,-1},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle s_{i,21}^{(t)}s_{i,22}^{(t)}\boldsymbol{\mathbf{\xi}}_{i}[k]\right)\boldsymbol{\mathbf{\xi}}_{i}[k],$ $\displaystyle\langle\operatorname{sgn}\left(\nabla_{\boldsymbol{\mathbf{w}}_{K,s}}L\right),\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=\operatorname{sgn}\left(\sum_{i\in[n]}{l}^{\prime(t)}_{i}y_{i}\langle\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,1}-\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,-1},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},s_{i,11}^{(t)}s_{i,12}^{(t)}y_{i}\boldsymbol{\mathbf{\mu}}+s_{i,21}^{(t)}s_{i,22}^{(t)}\boldsymbol{\mathbf{\xi}}_{i}\rangle y_{i}\right),$ $\displaystyle\langle\operatorname{sgn}\left(\nabla_{\boldsymbol{\mathbf{w}}_{K,s}}L\right),\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=\sum_{k\in\mathcal{B}_{i}}\operatorname{sgn}\left({l}^{\prime(t)}_{i}y_{i}\langle\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,1}-\bar{\boldsymbol{\mathbf{w}}}^{(t)}_{V,-1},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},s_{i,11}^{(t)}s_{i,12}^{(t)}y_{i}\boldsymbol{\mathbf{\mu}}+s_{i,21}^{(t)}s_{i,22}^{(t)}\boldsymbol{\mathbf{\xi}}_{i}\rangle(-\boldsymbol{\mathbf{\xi}}_{i}[k])\right)\boldsymbol{\mathbf{\xi}}_{i}[k].$ $$

Before the formal analysis for the query/key dynamics, we have some observations to simplify the gradient update formula.
Firstly, we note that the update magnitude of query/key inner product is constant at all iterations. Formally, for all $t\geq 0$ we have

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|=\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ (8) $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|=\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$ (9) $$

which means that we only need to analyze the sign of update direction for all quantities of interest.

Secondly, we take a look at the key part in gradient formula for different quantities:

$$ $\displaystyle\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle-y_{i}\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\xi}}_{i}\rangle-y_{i}\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\xi}}_{i}\rangle,$ $\displaystyle\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}-\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\mu}}+\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+y_{i}\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\xi}}_{i}\rangle-y_{i}\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle-\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\xi}}_{i}\rangle.$ $$

By the subsequent analysis in the stage I, we have that
the $\left|\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|=o(1)\cdot\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle$ such that we can approximately neglect the effect of $\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$.
Also, we have $\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}\rangle\geq 0$.
Then, combined with the gradient formula, we have

- •
For $\langle\boldsymbol{\mathbf{w}}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, the sign of increment is aligned with
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$
$\displaystyle=_{\operatorname{sgn}}$
$\displaystyle-\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$
$\displaystyle=_{\operatorname{sgn}}$
$\displaystyle-\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle.$
(10)
- •
For $\langle\boldsymbol{\mathbf{w}}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle$, the sign of increment is aligned with
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$
$\displaystyle=_{\operatorname{sgn}}$
$\displaystyle\sum_{i\in[n]}(-l_{i}^{\prime(t)}s_{i,11}^{(t)}s_{i,12}^{(t)})\cdot\left(-\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right).$
(11)
- •
For $\langle\boldsymbol{\mathbf{w}}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, the sign of increment is aligned with
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$
$\displaystyle=_{\operatorname{sgn}}$
$\displaystyle-\left(-s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle-s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)$
$\displaystyle=_{\operatorname{sgn}}$
$\displaystyle s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle.$
(12)
- •
For $\langle\boldsymbol{\mathbf{w}}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle$, the sign of increment is aligned with
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$
$\displaystyle=_{\operatorname{sgn}}$
$\displaystyle\sum_{i\in[n]}(-l_{i}^{\prime(t)})\cdot\left(-s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle-s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right).$
(13)

### E.4 Stage I

In this part, we consider the dynamics of mean value signal and noise, i.e. $\langle\boldsymbol{\mathbf{v}}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\langle\boldsymbol{\mathbf{v}}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle$.
Roughly, at the end of this stage we want to show that mean value noise evolve fast and can be viewed as a linear variable with time, mean value signal is small and negligible afterwards, while query/key noise are very close to initialization.

From the analysis for value above, we directly have

$$ $\displaystyle\langle\boldsymbol{\mathbf{v}}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle+2\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ (14) $\displaystyle\langle\boldsymbol{\mathbf{v}}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+2\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$ (15) $$

Let

$$ $\displaystyle T_{1}:=4\beta_{\boldsymbol{\mathbf{\xi}}}m_{v}^{-1/2}\eta^{-1}\sigma_{p}^{-1}s^{-1}.$ $$

Note that we have $\left|\langle\boldsymbol{\mathbf{v}}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq\sqrt{2}\beta_{\boldsymbol{\mathbf{\xi}}}m_{v}^{-1/2}$, then we have

$$ $\displaystyle t\eta\sigma_{p}s\leq\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\leq 2t\eta\sigma_{p}s,$ $$

for all $t\geq T_{1}$ and $i\in[n]$.
However, the query/key noise do not deviate too much from the initialization at $T_{1}$ since the deviation is at most

$$ $\displaystyle\frac{\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{1})},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|}{\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|}\leq 1+\frac{T_{1}\eta\left\lVert\boldsymbol{\mathbf{\xi}}\right\rVert_{1}}{\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|}\leq 1+\frac{O(\beta_{\boldsymbol{\mathbf{\xi}}}m_{v}^{-1/2})}{\Omega(\beta_{\boldsymbol{\mathbf{\xi}}}m_{v}^{-1/4})}=1+o(1),$ $$

where the second step is by Lemma [C.11](#A3.Thmtheorem11) and magnitude of $T_{1}$.
Besides, the deviation of $\langle\boldsymbol{\mathbf{v}}^{(T_{1})},\boldsymbol{\mathbf{\mu}}\rangle$ from initialization is also small. Formally, we have

$$ $\displaystyle\frac{\left|\langle\boldsymbol{\mathbf{v}}^{(T_{1})},\boldsymbol{\mathbf{\mu}}\rangle\right|}{\left|\langle\boldsymbol{\mathbf{v}}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|}\leq 1+\frac{T_{1}\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\left|\langle\boldsymbol{\mathbf{v}}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|}\leq 1+\frac{\tilde{O}(\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert{s}^{-1/2}{m_{v}}^{-1/2})}{\Omega(\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert ns^{-1/3}m_{v}^{-1/2})}=1+o(1),$ $$

where the second step is by Lemma [C.12](#A3.Thmtheorem12)
Therefore, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle=o(\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle),$ $$

for all $i\in[n]$ and $t\geq T_{1}$, where this holds at $t=T_{1}$ due to Lemma [C.8](#A3.Thmtheorem8) and the definition of $T_{1}$.
Combined with the form of gradient of query/key parameters, this makes us able to ignore the effect of $\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ afterwards.

### E.5 Stage II

In this part, we consider the dynamics of all query/key noise. Roughly, at the end of this stage we want to show for all neurons $s\in[m_{k}]$ and samples $i\in[n]$, query/key noise have the same sign and linear with $t$.

Let

$$ $\displaystyle T_{2}^{+}:=C_{1}\sqrt{\log(12m_{k}n/\delta)}\sigma_{0}ns^{-1/2}\eta^{-1},$ $$

where $C_{1}=\Theta(1)$ is a large constant. In this section, we will analyze the dynamics of $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ for all $s\in[m_{k}]$, $i\in[n]$ and $t\leq T_{2}^{+}$.

We first consider the samples $i\in S_{s,K+,Q-}^{(0)}\cup S_{s,K-,Q+}^{(0)}$ at initialization.
Note that without loss of generality, we can assume that $i\in S_{s,K-,Q+}^{(0)}$.
Then, by Lemma [C.14](#A3.Thmtheorem14), we only need to consider two cases where the one is when we have

$$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\leq\frac{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\cdot(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle),$ $$

and the other one is when we have

$$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\geq\frac{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\cdot(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle).$ $$

Finally, we give an upper bound for $T_{2,s,i}^{3}$ (defined later) for all $s\in[m_{k}]$ and $i\in S_{s,K+,Q-}^{(0)}\cup S_{s,K-,Q+}^{(0)}$, and bound the magnitude of $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ before $T_{2}^{+}$.

We define following useful timesteps.
Let $T_{2,s,i}^{1}$ be the first time satisfying

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle.$ (16) $$

Let $T_{2,s,i}^{2}$ be the first time satisfying

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)}y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle.$ (17) $$

Let $T_{2,s,i}^{3}$ be the first time satisfying Eq. ([16](#A5.E16)),  ([17](#A5.E17)) and

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)}y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle.$ (18) $$

Let $T_{2}^{\prime}:=\sqrt{2}\beta_{\boldsymbol{\mathbf{\xi}}}\eta^{-1}\sigma_{p}^{-1}s^{-1}$.

The following lemma studies the order between $T_{2,s,i}^{1}$, $T_{2,s,i}^{2}$, $T_{2,s,i}^{3}$.

###### Lemma E.1 .

For all $s\in[m_{k}]$ and $i\in S_{s,K+,Q-}^{(0)}\cup S_{s,K-,Q+}^{(0)}$, we have

$$ $\displaystyle 0<T_{2,s,i}^{1}\leq T_{2,s,i}^{3},$ $\displaystyle 0<T_{2,s,i}^{2}\leq T_{2,s,i}^{3}.$ $$

###### Proof of Lemma E.1 .

WLOG, we assume that $i\in S_{s,K-,Q+}^{(0)}$. By the definition of $S_{s,K-,Q+}^{(0)}$, we have

$$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}1.$ $$

Then, by the definition of $T_{2,s,i}^{1}$, we have $0<T_{2,s,i}^{1}$.

By the definition of $T_{2,s,i}^{3}$, we naturally have $T_{2,s,i}^{1}<T_{2,s,i}^{3}$ and $T_{2,s,i}^{2}<T_{2,s,i}^{3}$.

At $t=0$, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex196))
and the second step is by Lemma [C.10](#A3.Thmtheorem10),
and

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex199))
and the second step is by Lemma [C.10](#A3.Thmtheorem10) and [C.13](#A3.Thmtheorem13),
which implies $0<T_{2,s,i}^{2}$.

∎

Note that by Lemma [E.14](#A5.Thmtheorem14), we have that for all $i\in[n]$ and $t\leq T_{2}^{+}$,

$$ $\displaystyle s_{i,11}^{(t)},s_{i,21}^{(t)}=s_{i,11}^{(0)}(1+o(1)),s_{i,21}^{(0)}(1+o(1))=1/2+o(1).$ $$

In the following analysis, we will use this fact many times.

###### Lemma E.1 .

###### Proof of Lemma E.1 .

###### When − ⟨ 𝐰 K , s ( 0 ) , y i ​ 𝝃 i ⟩ + ⟨ 𝐰 K , s ( 0 ) , 𝝁 ⟩ ≤ σ p ​ s − ∥ 𝝁 ∥ σ p ​ s + ∥ 𝝁 ∥ ⋅ ( ⟨ 𝐰 Q , s ( 0 ) , y i ​ 𝝃 i ⟩ + ⟨ 𝐰 Q , s ( 0 ) , 𝝁 ⟩ ) superscript subscript 𝐰 𝐾 𝑠 0 subscript 𝑦 𝑖 subscript 𝝃 𝑖 superscript subscript 𝐰 𝐾 𝑠 0 𝝁 ⋅ subscript 𝜎 𝑝 𝑠 delimited-∥∥ 𝝁 subscript 𝜎 𝑝 𝑠 delimited-∥∥ 𝝁 superscript subscript 𝐰 𝑄 𝑠 0 subscript 𝑦 𝑖 subscript 𝝃 𝑖 superscript subscript 𝐰 𝑄 𝑠 0 𝝁 -\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\leq\frac{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\cdot(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle) , the update direction are determined by query noise.

By the definition of $T_{2,s,i}^{2}$, we must have that

$$ $\displaystyle(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2,s,i}^{2}+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2,s,i}^{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\cdot(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(T_{2,s,i}^{2}+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(T_{2,s,i}^{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)>0.$ $$

Note that we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2,s,i}^{2}+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2,s,i}^{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}$ $\displaystyle\leavevmode\nobreak\ -\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{2})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{2})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ -\langle\boldsymbol{\mathbf{w}}^{(0)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(0)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+T_{2,s,i}^{2}\eta(\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\geq 0,$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(T_{2,s,i}^{2}+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(T_{2,s,i}^{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}$ $\displaystyle\leavevmode\nobreak\ s_{i,11}^{(T_{2,s,i}^{2})}s_{i,12}^{(T_{2,s,i}^{2})}\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{2})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(T_{2,s,i}^{2})}s_{i,22}^{(T_{2,s,i}^{2})}\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ \frac{1}{4}(1\pm o(1))(\langle\boldsymbol{\mathbf{w}}^{(0)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(0)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-T_{2,s,i}^{2}\eta(\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert))\geq 0,$ (19) $$

where the last steps in both lines are by the condition

$$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\leq\frac{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\cdot(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle).$ $$

We claim that the sign of update would not change for $t\in[T_{2,s,i}^{2},T_{2}^{+}]$. Suppose at $t\leq\tilde{t}\in[T_{2,s,i}^{2},T_{2}^{+}]$, the induction hypothesis holds, then we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$ $$

which imply

$$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq-\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq 0,$ $\displaystyle s_{i,11}^{(t^{\prime}+1)}s_{i,12}^{(t^{\prime}+1)}\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(t^{\prime}+1)}s_{i,22}^{(t^{\prime}+1)}\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq\frac{1}{4}(1\pm o(1))(\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\geq 0,$ $$

then the conclusion holds and thus $\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{3})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{3})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$.

Moreover, we have an upper bound for $T_{2,s,i}^{2}$ that

$$ $\displaystyle T_{2,s,i}^{2}\leq\frac{\langle\boldsymbol{\mathbf{w}}^{(0)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(0)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle}{\eta(\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)}\leq\frac{1.5\langle\boldsymbol{\mathbf{w}}^{(0)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle}{\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}}\leq 1.5T_{2}^{\prime},$ $$

where the first step is by Eq. ([19](#A5.E19)),
the second step is by Lemma [C.10](#A3.Thmtheorem10),
the last step is by the definition of $T_{2}^{\prime}$ and Lemma [C.4](#A3.Thmtheorem4).

###### When − ⟨ 𝐰 K , s ( 0 ) , y i ​ 𝝃 i ⟩ + ⟨ 𝐰 K , s ( 0 ) , 𝝁 ⟩ ≥ σ p ​ s + ∥ 𝝁 ∥ σ p ​ s − ∥ 𝝁 ∥ ⋅ ( ⟨ 𝐰 Q , s ( 0 ) , y i ​ 𝝃 i ⟩ + ⟨ 𝐰 Q , s ( 0 ) , 𝝁 ⟩ ) superscript subscript 𝐰 𝐾 𝑠 0 subscript 𝑦 𝑖 subscript 𝝃 𝑖 superscript subscript 𝐰 𝐾 𝑠 0 𝝁 ⋅ subscript 𝜎 𝑝 𝑠 delimited-∥∥ 𝝁 subscript 𝜎 𝑝 𝑠 delimited-∥∥ 𝝁 superscript subscript 𝐰 𝑄 𝑠 0 subscript 𝑦 𝑖 subscript 𝝃 𝑖 superscript subscript 𝐰 𝑄 𝑠 0 𝝁 -\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\geq\frac{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\cdot(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle) , the update direction are determined by key noise.

At $t=T_{2,s,i}^{2}$ we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2,s,i}^{2}+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2,s,i}^{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}$ $\displaystyle\leavevmode\nobreak\ -\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{2})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{2})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\leq$ $\displaystyle\leavevmode\nobreak\ -\langle\boldsymbol{\mathbf{w}}^{(0)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(0)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+T_{2,s,i}^{2}\eta(\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\leq 0,$ (20) $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(T_{2,s,i}^{2}+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(T_{2,s,i}^{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}$ $\displaystyle\leavevmode\nobreak\ s_{i,11}^{(T_{2,s,i}^{2})}s_{i,12}^{(T_{2,s,i}^{2})}\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{2})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(T_{2,s,i}^{2})}s_{i,22}^{(T_{2,s,i}^{2})}\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\leq$ $\displaystyle\leavevmode\nobreak\ \frac{1}{4}(1\pm o(1))(\langle\boldsymbol{\mathbf{w}}^{(0)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(0)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-T_{2,s,i}^{2}\eta(\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert))\leq 0,$ $$

where the last steps in both lines are by the condition

$$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\geq\frac{\sigma_{p}s+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\cdot(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle).$ $$

Then, similar to previous analysis, we have for $t\in[T_{2,s,i}^{2},T_{2}^{+}]$, we have

$$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle<0,$ $\displaystyle s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle<0,$ $$

which implies $\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{3})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{3})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle<0$,
and

$$ $\displaystyle T_{2,s,i}^{2}\leq\frac{\langle\boldsymbol{\mathbf{w}}^{(0)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle-\langle\boldsymbol{\mathbf{w}}^{(0)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle}{\eta(\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}+\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)}\leq\frac{1.5\langle\boldsymbol{\mathbf{w}}^{(0)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle}{\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}}\leq 1.5T_{2}^{\prime},$ $$

where the first step is by Eq. ([20](#A5.E20)),
the second step is by Lemma [C.10](#A3.Thmtheorem10),
the last step is by the definition of $T_{2}^{\prime}$ and Lemma [C.4](#A3.Thmtheorem4).

Next, we are ready to bound $T_{2,s,i}^{3}$.
Note that during $[0,T_{2,s,i}^{3}]$, only one of query and key noise change its update direction.
WLOG, assume query noise changes the update direction while key noise doesn’t.
For the key noise, we have with at most $T_{2}^{\prime}$ steps, its sign can align with the sign of its update.
For the query noise, we possibly have that the sign at $T_{2,s,i}^{2}$ can be contrary with the sign at initialization, but
with at most another $T_{2,s,i}^{2}$ steps, its sign can align with the sign of its update (and at initialization).
By these analyses, we have

$$ $\displaystyle T_{2,s,i}^{3}\leq\max\left\{T_{2}^{\prime},2T_{2,s,i}^{2}\right\}\leq 3T_{2}^{\prime},$ $$

for all $s\in[m_{k}]$ and $i\in S_{s,K+,Q-}^{(0)}\cup S_{s,K-,Q+}^{(0)}$.

In the main text, we define $T_{2}^{\text{SGN}}:=3T_{2}^{\prime}$.

The following lemma gives bounds for the magnitude of query/key noise before $T_{2}^{+}$.

###### Lemma E.2 .

For all $s\in[m_{k}]$, $i\in S_{s,K+,Q-}^{(0)}\cup S_{s,K-,Q+}^{(0)}$ and $t\in[3T_{2}^{\prime},T_{2}^{+}]$, we have

$$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$ $$

and for all $t\in[12T_{2}^{\prime},T_{2}^{+}]$

$$ $\displaystyle\leavevmode\nobreak\ t\eta\sigma_{p}s/2\leq\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq 2t\eta\sigma_{p}s.$ $$

###### Proof.

The first two lines are shown in the above discussion. For the last line,
for the lower bound, suppose $\langle\boldsymbol{\mathbf{w}}^{(T_{2,s,i}^{3})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$,
we have

$$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ (t-3T_{2}^{\prime})\eta\sigma_{p}s/\sqrt{2}+\langle\boldsymbol{\mathbf{w}}^{(3T_{2}^{\prime})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ (t-3T_{2}^{\prime})\eta\sigma_{p}s/\sqrt{2}$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ t\eta\sigma_{p}s/2,$ $$

where the first step is by Lemma [C.4](#A3.Thmtheorem4),
the second step is by $3T_{2}^{\prime}\geq T_{2,s,i}^{3}$,
the third step is by $t\geq 12T_{2}^{\prime}$.
For the upper bound, we have

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq\left|\langle\boldsymbol{\mathbf{w}}^{(0)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+t\eta\sigma_{p}s\leq 2t\eta\sigma_{p}s,$ $$

where the second step is by $t\geq T_{2}^{\prime}$.
∎

We next consider the samples $i\in S_{s,K+,Q+}^{(0)}\cup S_{s,K-,Q-}^{(0)}$ at initialization.
The following lemma studies the update direction and magnitude of these neurons.

###### Lemma E.3 .

For all $s\in[m_{k}]$, $i\in S_{s,K+,Q+}^{(0)}\cup S_{s,K-,Q-}^{(0)}$ and $t\in[0,T_{2}^{+}]$, we have

$$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$ $$

and for all $t\in[T_{2}^{\prime},T_{2}^{+}]$

$$ $\displaystyle\leavevmode\nobreak\ t\eta\sigma_{p}s/2\leq\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq 2t\eta\sigma_{p}s.$ $$

###### Proof.

The first two lines are shown in the above discussion.
For the last line, note that without loss of generality, we can assume that $i\in S_{s,K+,Q+}^{(0)}$.
We can show that the sign of update of $\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ are the same as their sign at initialization and would not change for $t\in[0,T_{2}^{+}]$. Formally

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}^{(0)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}^{(0)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$ $$

for all $0\leq t\leq T_{2}^{+}$.
Similar to the proof in Lemma [E.1](#A5.Thmtheorem1), the claim holds at $t=0$. Suppose at $t\leq t^{\prime}\in[0,T_{2}^{+}]$, the induction hypothesis holds, then we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\geq\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\geq\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$ $$

which imply

$$ $\displaystyle\leavevmode\nobreak\ -\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ -\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\eta(\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\geq 0,$ $$

and

$$ $\displaystyle\leavevmode\nobreak\ s_{i,11}^{(t^{\prime}+1)}s_{i,12}^{(t^{\prime}+1)}\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(t^{\prime}+1)}s_{i,22}^{(t^{\prime}+1)}\langle\boldsymbol{\mathbf{w}}^{(t^{\prime}+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ \frac{1}{4}(1\pm o(1))\cdot(\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle(1+o(1))+\langle\boldsymbol{\mathbf{w}}^{(t^{\prime})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\eta(\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert))\geq 0,$ $$

where the first step is by Lemma [E.14](#A5.Thmtheorem14).
Then we show the result for $t=t^{\prime}+1$ and prove the claim.
Therefore, for the lower bound, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(0)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+t\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\geq t\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\geq t\eta\sigma_{p}s/2,$ $$

where the first step and second step follow from the claim above,
the last step is by Lemma [C.4](#A3.Thmtheorem4),
and for the upper bound, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(0)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+t\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\leq T_{2}^{\prime}\eta\sigma_{p}s+t\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\leq 2t\eta\sigma_{p}s,$ $$

where the first step is by the claim above,
the second step is by the definition of $T_{2}^{\prime}$,
the last step is by $t\geq T_{2}^{\prime}$ and Lemma [C.4](#A3.Thmtheorem4).
∎

###### Remark E.4 .

Combining Lemma [E.2](#A5.Thmtheorem2) and [E.3](#A5.Thmtheorem3), we have for all $t\in[12T_{2}^{\prime},T_{2}^{+}]$, $s\in[m_{k}]$ and $i\in[n]$

$$ $\displaystyle t\eta\sigma_{p}s/2\leq\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq 2t\eta\sigma_{p}s.$ (21) $$

Note that the upper bound also hold for all $t\geq T_{2}^{+}$ since the magnitude of update of query/key noise is always $\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}$, i.e., Eq. ([9](#A5.E9)) always holds.

Next, for given neuron $s$, we study the magnitude of the summation of query/key noise over all samples. The following lemma studies the dynamics and magnitude of query/key noise summation.

###### Lemma E.5 .

Let $T_{2}:=50nT_{2}^{\prime}$.
We have

- 1.
(update direction). Let $\epsilon_{s,i}:=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(3T_{2}^{\prime})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)$.
We have the sign of the update of the summation of query/key noise would keep unchanged for $[3T_{2}^{\prime},T_{2}^{+}]$.
Formally, for all $s\in[m_{k}]$ and $t\in[3T_{3}^{\prime},T_{3}^{+}]$
$\displaystyle\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\eta\sum_{i=1}^{n}\epsilon_{s,i}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$
- 2.
(magnitude of update). With probability at least $1-O(m_{k}/\sqrt{n})$, for all $s\in[m_{k}]$
$\displaystyle\left|\eta\sum_{i=1}^{n}\epsilon_{s,i}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\right|\geq\eta\sigma_{p}s/\sqrt{2}.$
- 3.
(magnitude of query/key noise summation). For all $s\in[m_{k}]$ and $t\in[T_{2},T_{2}^{+}]$
$\displaystyle\operatorname{sgn}(\eta\sum_{i=1}^{n}\epsilon_{s,i}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1})\cdot\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,\operatorname{sgn}(\eta\sum_{i=1}^{n}\epsilon_{s,i}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1})\cdot\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq 2n\beta_{\boldsymbol{\mathbf{\mu}}}.$

To prove this lemma, we need the following lemma which studies the property of $\epsilon_{s,i}$.

###### Lemma E.6 .

Let

$$ $\displaystyle E_{s,+}:=\left\{i\in[n]:\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0\right\},$ $\displaystyle E_{s,-}:=\left\{i\in[n]:\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle<0\right\}.$ $$

Let

$$ $\displaystyle A_{s,i}:=\left\{\omega\in\Omega:-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle<\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right\}.$ $$

Let

$$ $\displaystyle X_{s,i}:=$ $\displaystyle\mathds{1}[i\in S_{s,K+,Q+}]+\mathds{1}[i\in S_{s,K-,Q+},A_{s,i}]+\mathds{1}[i\in S_{s,K+,Q-},A_{s,i}]$ $\displaystyle-$ $\displaystyle\mathds{1}[i\in S_{s,K-,Q-}]-\mathds{1}[i\in S_{s,K-,Q+},A^{c}_{s,i}]-\mathds{1}[i\in S_{s,K+,Q-},A^{c}_{s,i}].$ $$

Then we have

- 1.
(The distribution of $X_{s,i}$). For all $s\in[m_{k}]$ and $i\in[n]$,
$\displaystyle X_{s,i}=\begin{cases}1,\text{w.p. 1/2}\\
-1,\text{w.p. 1/2}\end{cases}.$
- 2.
(The probability of equal set size).
$\displaystyle\mathbb{P}[\left|E_{s,+}\right|-\left|E_{s,-}\right|=0]\leq O(n^{-1/2}).$
- 3.
(The probability of lucky neurons).
$\displaystyle\mathbb{P}[X_{s,i}>0,\sum_{i=1}^{n}X_{s,i}>0]\geq\frac{1}{4}-O(n^{-1/2}).$

###### Proof.

By the definition of $X_{s,i}$ and $\epsilon_{s,i}$, we have following equivalent formulations

$$ $\displaystyle\left|E_{s,+}\right|=\left|\left\{i\in[n]:\epsilon_{s,i}>0\right\}\right|=\sum_{i=1}^{n}\mathds{1}[i\in S_{s,K+,Q+}]+\mathds{1}[i\in S_{s,K-,Q+},A_{s,i}]+\mathds{1}[i\in S_{s,K+,Q-},A_{s,i}],$ $\displaystyle\left|E_{s,-}\right|=\left|\left\{i\in[n]:\epsilon_{s,i}<0\right\}\right|=\sum_{i=1}^{n}\mathds{1}[i\in S_{s,K-,Q-}]+\mathds{1}[i\in S_{s,K-,Q+},A^{c}_{s,i}]+\mathds{1}[i\in S_{s,K+,Q-},A^{c}_{s,i}].$ $$

Let
$Q_{s,i}:=\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$,
$K_{s,i}:=\langle\boldsymbol{\mathbf{w}}_{k,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$,
$\mu_{s}:=\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle$.
Then we have $\mu_{s}\sim N(0,\sigma_{0}^{2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2})$, and $Q_{s,i},K_{s,i}|y_{i}\boldsymbol{\mathbf{\xi}}_{i}\sim N(0,\sigma_{0}^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2})$.
Denote the distribution of $\boldsymbol{\mathbf{\xi}}$ and $\mu$ by $P_{\boldsymbol{\mathbf{\xi}}}$ and $P_{\mu}$, respectively.

Let $p(\mu,\boldsymbol{\mathbf{\xi}}_{i}):=\mathbb{P}[X_{s,i}=1|\mu_{s}=\mu,\boldsymbol{\mathbf{\xi}}_{i}]$,
$p_{\mu}:=\mathbb{P}[X_{s,i}=1|\mu_{s}=\mu]$.

Let $\phi(\mu,\boldsymbol{\mathbf{\xi}}_{i}):=\mathbb{P}[i\in S_{s,K-,Q+},A_{s,i}|\mu_{s}=\mu,\boldsymbol{\mathbf{\xi}}_{i}]$,
$\phi_{\mu}:=\mathbb{P}[i\in S_{s,K-,Q+},A_{s,i}|\mu_{s}=\mu]$.
By the symmetry of $S_{s,K-,Q+}$ and $S_{s,K+,Q-}$, we have $\mathbb{P}[i\in S_{s,K-,Q+},A_{s,i}|\mu,\boldsymbol{\mathbf{\xi}}_{i}]=\mathbb{P}[i\in S_{s,K+,Q-},A_{s,i}|\mu,\boldsymbol{\mathbf{\xi}}_{i}]$, which gives

$$ $\displaystyle p(\mu,\boldsymbol{\mathbf{\xi}}_{i})=\frac{1}{4}+2\phi(\mu,\boldsymbol{\mathbf{\xi}}_{i}),$ $\displaystyle p_{\mu}=\int p(\mu,\boldsymbol{\mathbf{\xi}}_{i})dP_{\boldsymbol{\mathbf{\xi}}_{i}}=\int\frac{1}{4}+2\phi(\mu,\boldsymbol{\mathbf{\xi}}_{i})dP_{\boldsymbol{\mathbf{\xi}}_{i}}=\frac{1}{4}+2\phi_{\mu}.$ $$

By the symmetry of the gaussian density, for any $\boldsymbol{\mathbf{\xi}}_{i}$, we have

$$ $\displaystyle\phi(\mu,\boldsymbol{\mathbf{\xi}}_{i})$ $\displaystyle=\int_{q>0,k<0,q+k>\mu}\exp(-\frac{q^{2}+k^{2}}{2\sigma_{0}^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2}})dqdk$ $\displaystyle=\int_{q>0,k<0,q+k<-\mu}\exp(-\frac{q^{2}+k^{2}}{2\sigma_{0}^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2}})dqdk$ $\displaystyle=\frac{1}{4}-\int_{q>0,k<0,q+k>-\mu}\exp(-\frac{q^{2}+k^{2}}{2\sigma_{0}^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert^{2}})dqdk$ $\displaystyle=\frac{1}{4}-\phi(-\mu,\boldsymbol{\mathbf{\xi}}_{i}),$ $$

which implies $\phi(\mu,\boldsymbol{\mathbf{\xi}}_{i})-\frac{1}{8}$ is an odd function, and thus

$$ $\displaystyle\mathbb{P}[i\in S_{s,K-,Q+},A_{s,i}|\boldsymbol{\mathbf{\xi}}_{i}]=\int\phi(\mu,\boldsymbol{\mathbf{\xi}}_{i})dP_{\mu}=\frac{1}{8},$ $$

thus $\mathbb{P}[i\in S_{s,K-,Q+},A_{s,i}]=\frac{1}{8}$ and $\mathbb{P}[X_{s,i}=1]=\frac{1}{2}$, which proves the first result.

Note that given $\mu$, $X_{s,i}$ are independent for $i\in[n]$, then

$$ $\displaystyle\mathbb{P}[\left|E_{s,+}\right|-\left|E_{s,-}\right|=0]$ $\displaystyle=\mathbb{P}[\sum_{i=1}^{n}X_{s,i}=0]$ $\displaystyle=\int\binom{n}{n/2}p_{\mu}^{n/2}(1-p_{\mu})^{n/2}dP_{\mu}$ $\displaystyle\leq\int\binom{n}{n/2}\frac{1}{2^{n}}dP_{\mu}$ $\displaystyle\leq\binom{n}{n/2}\frac{1}{2^{n}}\leq\frac{10}{\sqrt{n}},$ $$

which proves the second result.

Recall that

$$ $\displaystyle\leavevmode\nobreak\ \phi(\mu,\boldsymbol{\mathbf{\xi}}_{i})=1/4-\phi(-\mu,\boldsymbol{\mathbf{\xi}}_{i})\leavevmode\nobreak\ \forall\boldsymbol{\mathbf{\xi}}_{i}$ $\displaystyle\Longrightarrow$ $\displaystyle\leavevmode\nobreak\ \phi_{\mu}=1/4-\phi_{-\mu}$ $\displaystyle\Longrightarrow$ $\displaystyle\leavevmode\nobreak\ p_{\mu}=1-p_{-\mu}.$ $$

The quantity of interest can be written as

$$ $\displaystyle\leavevmode\nobreak\ \underbrace{\mathbb{P}[X_{s,i}>0,\sum_{j\neq i}X_{s,i}>0]}_{I_{1}}=\int\underbrace{\sum_{k=\lfloor(n-1)/2\rfloor+1}^{n-1}\binom{n}{k}p_{\mu}^{k}(1-p_{\mu})^{n-k}}_{p_{n-1,\mu}}\cdot p_{\mu}dP_{\mu},$ $\displaystyle\leavevmode\nobreak\ \underbrace{\mathbb{P}[X_{s,i}>0,\sum_{j\neq i}X_{s,i}<0]}_{I_{2}}=\int\underbrace{\sum_{k=\lfloor(n-1)/2\rfloor+1}^{n-1}\binom{n}{k}p_{\mu}^{n-k}(1-p_{\mu})^{k}}_{q_{n-1,\mu}}\cdot p_{\mu}dP_{\mu},$ $\displaystyle\leavevmode\nobreak\ \underbrace{\mathbb{P}[X_{s,i}<0,\sum_{j\neq i}X_{s,i}>0]}_{I_{3}}=\int\sum_{k=\lfloor(n-1)/2\rfloor+1}^{n-1}\binom{n}{k}p_{\mu}^{k}(1-p_{\mu})^{n-k}\cdot(1-p_{\mu})dP_{\mu},$ $\displaystyle\leavevmode\nobreak\ \underbrace{\mathbb{P}[X_{s,i}<0,\sum_{j\neq i}X_{s,i}<0]}_{I_{4}}=\int\sum_{k=\lfloor(n-1)/2\rfloor+1}^{n-1}\binom{n}{k}p_{\mu}^{n-k}(1-p_{\mu})^{k}\cdot(1-p_{\mu})dP_{\mu}.$ $$

By the symmetry of $P_{\mu}$ and $p_{\mu}$, we have

$$ $\displaystyle\leavevmode\nobreak\ \mathbb{P}[X_{s,i}>0,\sum_{j\neq i}X_{s,i}>0]=\mathbb{P}[X_{s,i}<0,\sum_{j\neq i}X_{s,i}<0],$ $\displaystyle\leavevmode\nobreak\ \mathbb{P}[X_{s,i}>0,\sum_{j\neq i}X_{s,i}<0]=\mathbb{P}[X_{s,i}<0,\sum_{j\neq i}X_{s,i}>0].$ $$

Since

$$ $\displaystyle I_{1}+I_{4}-I_{2}-I_{3}=\int(p_{\mu}-(1-p_{\mu}))(p_{n-1,\mu}-q_{n-1,\mu})dP_{\mu}\geq 0,$ $$

where the last step is by $p_{n-1,\mu}>q_{n-1,\mu}$ when $p_{\mu}\geq\frac{1}{2}$, then we have

$$ $\displaystyle\mathbb{P}[X_{s,i}>0,\sum_{i=1}^{n}X_{s,i}>0]\geq I_{1}\geq\frac{1}{4}(1-\mathbb{P}[\sum_{j\neq i}X_{s,i}=0])\geq\frac{1}{4}-O(n^{-1/2}),$ $$

where the second step is by $I_{1}=I_{4}\geq I_{2}=I_{3}$, the third step is by the second result above, which concludes the proof.
∎

###### Proof of Lemma E.5 .

By Lemma [E.2](#A5.Thmtheorem2) and [E.3](#A5.Thmtheorem3),
we have the update direction of each element $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ is unchanged for $t\in[3T_{2}^{\prime},T_{2}^{+}]$, which implies the first result.

By Lemma [E.6](#A5.Thmtheorem6), with probability $1-O(m_{k}/\sqrt{n})$, we have for all $s\in[m_{k}]$,
$\left|E_{s,+}\right|-\left|E_{s,-}\right|\neq 0$.
Below, suppose this fact holds, and all the results are based on this fact which hold with high probability.

WLOG, for fixed neuron $s$, suppose $\left|E_{s,+}\right|-\left|E_{s,-}\right|\geq 1$, then by Lemma [C.5](#A3.Thmtheorem5), we have

$$ $\displaystyle\eta\sum_{i=1}^{n}\epsilon_{s,i}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\geq\eta\sigma_{p}s/\sqrt{2}.$ (22) $$

Finally, we note that at $t=12T_{2}^{\prime}$, we have

$$ $\displaystyle\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(12T_{3}^{\prime})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(12T_{3}^{\prime})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq-24T_{3}^{\prime}n\eta\sigma_{p}s.$ $$

By the definition of $T_{2}$, we have $T_{2}\geq(12+24\sqrt{2}n)T_{2}^{\prime}+2n\beta_{\boldsymbol{\mathbf{\mu}}}\eta^{-1}\sigma_{p}^{-1}s^{-1}$.
Then, for $t\in[T_{2},T_{2}^{+}]$, we have

$$ $\displaystyle\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq 2n\beta_{\boldsymbol{\mathbf{\mu}}}.$ (23) $$

Similarly, if $\left|E_{s,+}\right|-\left|E_{s,-}\right|\leq-1$ for some other neuron $s$, then for $t\in[T_{2},T_{2}^{+}]$, we have

$$ $\displaystyle\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\leq-2n\beta_{\boldsymbol{\mathbf{\mu}}}.$ $$

∎

###### Lemma E.2 .

###### Proof.

###### Lemma E.3 .

###### Proof.

###### Remark E.4 .

###### Lemma E.5 .

###### Lemma E.6 .

###### Proof.

###### Proof of Lemma E.5 .

### E.6 Stage III

Let

$$ $\displaystyle T_{3}^{+}:=C_{2}\sqrt{\log(12m_{k}/\delta)}\sigma_{0}\eta^{-1},$ $$

where $C_{2}=\Theta(1)$ is a large constant. In this section, we will analyze the dynamics of $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ for all $s\in[m_{k}]$, $i\in[n]$ and $T_{2}\leq t\leq T_{3}^{+}$.
Generally, we already have controlled the direction and magnitude of the quantity $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ at the timestep $T_{2}$, we can show the update direction of query/key noise is determined by these quantities.

For $t\leq T_{2}$, we want to show the change of $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ is not much.
Actually, we can show that for $t\leq T_{2}^{+}$, $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ are close to initialization. Formally, we have the following lemma.

###### Lemma E.7 .

For all $s\in[m_{k}]$ and $t\leq T_{2}^{+}$, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\cdot(1\pm o(1)),$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\cdot(1\pm o(1)).$ $$

###### Proof.

For key signal, we have

$$ $\displaystyle\frac{\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|}{\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|}\leq\frac{\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|+t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|}\leq 1+\frac{t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\Omega(\sigma_{0}ns^{-1/3}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)}=1+o(1),$ $$

where the first step is by Eq. ([8](#A5.E8)),
the second step is by Lemma [C.12](#A3.Thmtheorem12),
and the last step is by the definition of $T_{2}^{+}$.
This also imply the sign of $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ is the same as initialization for all $t\leq T_{2}^{+}$. The proof for query signal is similar.
∎

For $t\geq T_{2}$, we want to show the update direction of $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ or $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$) would be determined by the sign of $\sum_{i}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ or $\sum_{i}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, respectively.
To this end, we need to control the weighting term to make the weighted summation of query/key noise similar to the ordinary summation, which enables us to apply the previous analysis.
These weighting terms include: (1) The softmax outputs; (2) The loss derivative; (3) The magnitude of mean value noise $\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$.
We want to show that these quantities are almost same across all samples, or even keep constant before $T_{4}^{+}$.

The following lemma studies the magnitude of $\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ after $T_{2}$.
This lemma also shows $\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ are close between different samples.

###### Lemma E.8 .

For all $i\in[n]$ and $t\geq T_{2}$,
we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=2\sqrt{\frac{2}{\pi}}t\eta\sigma_{p}s(1+\tilde{O}(s^{-1/2})).$ $$

###### Proof.

We can show

$$ $\displaystyle\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=\langle\boldsymbol{\mathbf{v}}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+2t\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}$ $\displaystyle=2t\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(1+o(1))$ $\displaystyle=2\sqrt{\frac{2}{\pi}}t\eta\sigma_{p}s(1+\tilde{O}(s^{-1/2}))(1+o(1)).$ $$

where the second step is by $\langle\boldsymbol{\mathbf{v}}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=O(\beta_{\boldsymbol{\mathbf{\xi}}}m_{v}^{-1/2})=o(T_{2}\eta\sigma_{p}s)$,
the third step is by Lemma [C.4](#A3.Thmtheorem4).
∎

By Lemma [E.14](#A5.Thmtheorem14), we have that softmax outputs are concentrated at $1/2$ for all $i\in[n]$ and $t\leq T_{3}^{+}$.
By Lemma [E.15](#A5.Thmtheorem15), we have that the derivative of loss of all samples are concentrated at $1/2$ for all $i\in[n]$ and $t\leq T_{3}^{+}$.

With these properties, we are ready to show that the query/key noise dynamics keep unchanged for $[T_{2},T_{3}^{+}]$,
and the query/key signal dynamics are dominated by the sign of the summation of query/key noise, i.e., $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$.
The next two lemmas characterize the dynamics of query/key noise and query/key signal, respectively.

###### Lemma E.9 .

For all $s\in[m_{k}]$, $i\in[n]$ and $T_{2}\leq t\leq T_{3}^{+}$, we have

$$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$ $$

and for all $t\in[T_{2},T_{3}^{+}]$

$$ $\displaystyle\leavevmode\nobreak\ t\eta\sigma_{p}s/2\leq\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq 2t\eta\sigma_{p}s.$ $$

###### Proof.

WLOG, given neuron $s\in[m_{k}]$ and sample $i\in[n]$, suppose $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$, $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$, we claim that $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ are increasing for $T_{2}\leq t\leq T_{3}^{+}$. At $t=T_{2}$, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2}+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq\tilde{\Omega}(\sigma_{0}\sigma_{p}\sqrt{s}n)-\tilde{O}(\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)>0,$ $$

and

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(T_{2}+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}$ $\displaystyle s_{i,11}^{(T_{2})}s_{i,12}^{(T_{2})}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(T_{2})}s_{i,22}^{(T_{2})}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=$ $\displaystyle(\frac{1}{4}+o(1))(\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)$ $\displaystyle=$ $\displaystyle(\frac{1}{4}+o(1))(\tilde{\Omega}(\sigma_{0}\sigma_{p}\sqrt{s}n)-\tilde{O}(\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert))>0.$ $$

Suppose the induction holds for $t\leq\tilde{t}-1$, then at $t=\tilde{t}$, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(\tilde{t}+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(\tilde{t})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}^{(\tilde{t})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(\tilde{t})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}^{(T_{3})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(T_{3})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+(\tilde{t}-T_{3})\eta(\left\lVert\boldsymbol{\mathbf{\xi}}\right\rVert_{1}-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)>0,$ $$

and

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(\tilde{t}+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(\tilde{t})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}$ $\displaystyle s_{i,11}^{(\tilde{t})}s_{i,12}^{(\tilde{t})}\langle\boldsymbol{\mathbf{w}}^{(\tilde{t})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(\tilde{t})}s_{i,22}^{(\tilde{t})}\langle\boldsymbol{\mathbf{w}}^{(\tilde{t})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=$ $\displaystyle(\frac{1}{4}+o(1))(\langle\boldsymbol{\mathbf{w}}^{(\tilde{t})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(\tilde{t})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)$ $\displaystyle\geq$ $\displaystyle(\frac{1}{4}+o(1))\left(\langle\boldsymbol{\mathbf{w}}^{(T_{3})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(T_{3})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+(\tilde{t}-T_{3})\eta(\left\lVert\boldsymbol{\mathbf{\xi}}\right\rVert_{1}-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\right)>0.$ $$

Finally, the bounds about magnitude follow from Lemma [E.2](#A5.Thmtheorem2), [E.3](#A5.Thmtheorem3) and [C.4](#A3.Thmtheorem4).
∎

###### Remark E.10 .

Lemma [E.9](#A5.Thmtheorem9) implies that the dynamics of $\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$,
$\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ keep unchanged for $T_{2}\leq t\leq T_{3}^{+}$ and
for all $s\in[m_{k}]$.

###### Lemma E.11 .

For all $s\in[m_{k}]$, $i\in[n]$ and $T_{2}\leq t\leq T_{3}^{+}$, we have

$$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle-\operatorname{sgn}(\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert.$ $$

###### Proof.

WLOG, given neuron $s\in[m_{k}]$, suppose we have $\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$,
thus $\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$.
Therefore, for $T_{2}\leq t\leq T_{3}^{+}$, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}(-l_{i}^{\prime(t)}s_{i,11}^{(t)}s_{i,12}^{(t)})\cdot\left(-y_{i}\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}\frac{\sqrt{2}}{8\sqrt{\pi}}t\eta\sigma_{p}s(1+o(1))\cdot(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-n\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle\geq$ $\displaystyle\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-n\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+(t-T_{2})(\eta\sigma_{p}s/\sqrt{2}-\eta n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)>0,$ $$

where the first inequality is due to the lower bound Eq. ([22](#A5.E22)),
and the second inequality is due to the lower bound Eq. ([23](#A5.E23)) and $n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert=o(\sigma_{p}s)$. Similarly, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}(-l_{i}^{\prime(t)})\cdot\left(-y_{i}s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle-s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}\frac{\sqrt{2}}{8\sqrt{\pi}}t\eta\sigma_{p}s(1+o(1))\cdot(-\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle-\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-n\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle\leq$ $\displaystyle-\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-n\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle-(t-T_{2})(\eta\sigma_{p}s/\sqrt{2}-\eta n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)<0.$ $$

∎

###### Lemma E.12 .

Let $T_{3}:=3\beta_{\boldsymbol{\mathbf{\mu}}}\eta^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{-1}=\tilde{O}(\sigma_{0}\eta^{-1})$, then for $T_{3}\leq t\leq T_{3}^{+}$, we have

$$ $\displaystyle t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/2\leq\operatorname{sgn}(\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\cdot\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\leq 4t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/3,$ (24) $\displaystyle-4t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/3\leq\operatorname{sgn}(\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\cdot\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\leq-t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/2.$ (25) $$

###### Proof.

WLOG, given neuron $s\in[m_{k}]$, suppose we have $\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$.
We only prove Eq. ([24](#A5.E24)). The proof of Eq. ([25](#A5.E25)) is identical.
For the upper bound, for $t\in[T_{3},T_{3}^{+}]$ we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\leq\beta_{\boldsymbol{\mathbf{\mu}}}+t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\leq(t+T_{3}/3)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\leq 4t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/3,$ $$

where the first step is by Eq. ([8](#A5.E8)),
the second step is by the definition of $T_{3}$.
For the lower bound, for $t\in[T_{3},T_{3}^{+}]$ we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\geq-\beta_{\boldsymbol{\mathbf{\mu}}}-T_{2}\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert+(t-T_{2})\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\geq(t-T_{3}/3-2T_{2})\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\geq t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/2,$ $$

where the first step is by Eq. ([8](#A5.E8)),
the second step is by definition of $T_{3}$,
the third step is by $T_{2}T_{3}^{-1}=o(1)$.
∎

###### Remark E.13 .

Note that the upper bound in Eq. ([24](#A5.E24))
and lower bound in Eq. ([25](#A5.E25))
hold for all $t\geq T_{3}$ since the magnitude of update of query/key signal is always $\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert$, i.e., Eq. ([8](#A5.E8)) always holds.

###### Lemma E.7 .

###### Proof.

###### Lemma E.8 .

###### Proof.

###### Lemma E.9 .

###### Proof.

###### Remark E.10 .

###### Lemma E.11 .

###### Proof.

###### Lemma E.12 .

###### Proof.

###### Remark E.13 .

### E.7 Stage IV

This section study the concentration behavior of softmax outputs and loss derivative.
We first show the magnitude of softmax outputs at initialization $s_{i,11}^{(t)}$ and $s_{i,21}^{(t)}$ are concentrated at 1/2 for all $i\in[n]$ before $T_{4}$ and $T_{3}^{+}$, respectively.

###### Lemma E.14 .

Let

$$ $\displaystyle T_{4}:=C_{3}\log\left(\frac{C_{3}\sigma_{p}s}{\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\right)\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1}$ (26) $$

where $C_{3}=\Theta(1)$ is a large constant. Then, for all $i\in[n]$, we have

$$ $\displaystyle s_{i,11}^{(t)}=1/2\pm o(1),\leavevmode\nobreak\ \forall t\leq T_{4},$ $\displaystyle s_{i,21}^{(t)}=1/2\pm o(1),\leavevmode\nobreak\ \forall t\leq T_{3}^{+}.$ $$

###### Proof of Lemma E.14 .

Note that for $t\geq 0$,

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|$ $\displaystyle\leq\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+2t\eta\sigma_{p}s$ $\displaystyle\leq 4\max\left\{\beta_{\boldsymbol{\mathbf{\xi}}},t\eta\sigma_{p}s\right\},$ $$

for all $i\in[n]$ and $s\in[m_{k}]$ and

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|$ $\displaystyle\leq\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle\right|+t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert$ $\displaystyle\leq 4\max\left\{\beta_{\boldsymbol{\mathbf{\mu}}},t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\right\},$ $$

for all $s\in[m_{k}]$. Similarly, we also have

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|$ $\displaystyle\leq 4\max\left\{\beta_{\boldsymbol{\mathbf{\mu}}},t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\right\},$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|$ $\displaystyle\leq 4\max\left\{\beta_{\boldsymbol{\mathbf{\xi}}},t\eta\sigma_{p}s\right\}.$ $$

When $t\leq T_{3}^{+}$, we have for all $i\in[n]$

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq O(\sqrt{\log(12m_{k}/\delta)}\sigma_{0}\sigma_{p}s),$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq O(\sqrt{\log(12m_{k}/\delta)}\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert).$ $$

Then we have for all $i\in[n]$,

$$ $\displaystyle s_{i,21}^{(t)}$ $\displaystyle=\frac{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right)}{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right)+\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}$ $\displaystyle\leq\frac{\exp\left(O(m_{k}\log(12m_{k}/\delta)\sigma_{0}^{2}\sigma_{p}s\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\right)}{\exp\left(O(m_{k}\log(12m_{k}/\delta)\sigma_{0}^{2}\sigma_{p}s\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\right)+\exp\left(-O(m_{k}\log(12m_{k}/\delta)\sigma_{0}^{2}\sigma_{p}^{2}s^{2})\right)}$ $\displaystyle=\frac{1}{1+\exp\left(-O(m_{k}\log(12m_{k}/\delta)\sigma_{0}^{2}(\sigma_{p}^{2}s^{2}+\sigma_{p}s\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert))\right)}$ $\displaystyle\leq\frac{1}{2}+O(m_{k}\log(12m_{k}/\delta)\sigma_{0}^{2}(\sigma_{p}^{2}s^{2}+\sigma_{p}s\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)),$ $$

and

$$ $\displaystyle s_{i,21}^{(t)}$ $\displaystyle=\frac{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right)}{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right)+\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}$ $\displaystyle\geq\frac{\exp\left(-O(m_{k}\log(12m_{k}/\delta)\sigma_{0}^{2}\sigma_{p}s\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\right)}{\exp\left(-O(m_{k}\log(12m_{k}/\delta)\sigma_{0}^{2}\sigma_{p}s\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\right)+\exp\left(O(m_{k}\log(12m_{k}/\delta)\sigma_{0}^{2}\sigma_{p}^{2}s^{2})\right)}$ $\displaystyle=\frac{1}{1+\exp\left(O(m_{k}\log(12m_{k}/\delta)\sigma_{0}^{2}(\sigma_{p}^{2}s^{2}+\sigma_{p}s\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert))\right)}$ $\displaystyle\geq\frac{1}{2}-O(m_{k}\log(12m_{k}/\delta)\sigma_{0}^{2}(\sigma_{p}^{2}s^{2}+\sigma_{p}s\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)).$ $$

Similarly, when $t\leq T_{4}$, we have for all $i\in[n]$

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq O(\log(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)m_{k}^{-1/2}),$ (27) $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|,\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq O(\log(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert).$ (28) $$

Then we have for all $i\in[n]$,

$$ $\displaystyle s_{i,11}^{(t)}$ $\displaystyle=\frac{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right)}{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right)+\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}$ $\displaystyle\leq\frac{\exp\left(O(\log^{2}(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\sigma_{p}^{-2}s^{-2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2})\right)}{\exp\left(O(\log^{2}(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\sigma_{p}^{-2}s^{-2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2})\right)+\exp\left(-O(\log^{2}(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\sigma_{p}^{-1}s^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\right)}$ $\displaystyle=\frac{1}{1+\exp\left(-O(\log^{2}(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\sigma_{p}^{-1}s^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1+\sigma_{p}^{-1}s^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert))\right)}$ $\displaystyle\leq\frac{1}{2}+O(\log^{2}(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\sigma_{p}^{-1}s^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1+\sigma_{p}^{-1}s^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)),$ $$

and

$$ $\displaystyle s_{i,11}^{(t)}$ $\displaystyle=\frac{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right)}{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right)+\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}$ $\displaystyle\geq\frac{\exp\left(-O(\log^{2}(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\sigma_{p}^{-2}s^{-2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2})\right)}{\exp\left(-O(\log^{2}(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\sigma_{p}^{-2}s^{-2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2})\right)+\exp\left(O(\log^{2}(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\sigma_{p}^{-1}s^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\right)}$ $\displaystyle=\frac{1}{1+\exp\left(O(\log^{2}(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\sigma_{p}^{-1}s^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1+\sigma_{p}^{-1}s^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert))\right)}$ $\displaystyle\geq\frac{1}{2}-O(\log^{2}(C_{3}\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\sigma_{p}^{-1}s^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1+\sigma_{p}^{-1}s^{-1}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)).$ $$

∎

The next lemma shows that loss derivative $\ell_{i}^{\prime(t)}$ are concentrated at 1/2 for all $i\in[n]$ before $T_{4}$.

###### Lemma E.15 .

Recall the $T_{4}$ defined in Eq. ([26](#A5.E26)).
For all $i\in[n]$, we have

$$ $\displaystyle\ell_{i}^{\prime(t)}=1/2\pm o(1),\leavevmode\nobreak\ \forall t\leq T_{4}.$ $$

###### Proof of Lemma E.15 .

Note that $\ell(z)=\log(1+\exp(-z))$ and $-\ell^{\prime}(z)=1/(1+\exp(z))$. For $t\leq T_{3}^{+}$ we have

$$ $\displaystyle\ell_{i}^{\prime(t)}=$ $\displaystyle\frac{1}{1+\exp\left((s_{i,11}^{(t)}+s_{i,21}^{(t)})\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle+(s_{i,12}^{(t)}+s_{i,22}^{(t)})\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}$ $\displaystyle=$ $\displaystyle\frac{1}{1+\exp\left(y_{i}(1\pm o(1))(\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\right)}$ $\displaystyle=$ $\displaystyle\frac{1}{1+\exp\left(y_{i}(1\pm o(1))\tilde{O}(\sigma_{0}\sigma_{p}s)\right)}$ $\displaystyle=$ $\displaystyle\frac{1}{2}\pm o(1),$ $$

where the first step is by definition,
the second step is by Lemma [E.14](#A5.Thmtheorem14),
the third step is by Lemma [E.8](#A5.Thmtheorem8),
the last step is due to $\sigma_{0}\sigma_{p}s=o(1)$.
For $t\leq T_{4}$, suppose $y_{i}=1$, we have $\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle,\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$, then

$$ $\displaystyle\ell_{i}^{\prime(t)}$ $\displaystyle=\frac{1}{1+\exp\left((s_{i,11}^{(t)}+s_{i,21}^{(t)})\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle+(s_{i,12}^{(t)}+s_{i,22}^{(t)})\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}$ $\displaystyle\leq\frac{1}{2}+o(1),$ $$

where the first step is by definition,
the second step is by $\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle,\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$, and

$$ $\displaystyle\ell_{i}^{\prime(t)}=$ $\displaystyle\frac{1}{1+\exp\left((s_{i,11}^{(t)}+s_{i,21}^{(t)})\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle+(s_{i,12}^{(t)}+s_{i,22}^{(t)})\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}$ $\displaystyle\geq$ $\displaystyle\frac{1}{1+\exp\left(y_{i}(3/2\pm o(1))(\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\right)}$ $\displaystyle=$ $\displaystyle\frac{1}{1+\exp\left(y_{i}(3/2\pm o(1))\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}$ $\displaystyle=$ $\displaystyle\frac{1}{1+\exp\left(y_{i}(3/2\pm o(1))\tilde{O}(m_{k}^{-1/2})\right)}$ $\displaystyle\geq$ $\displaystyle\frac{1}{2}-o(1),$ $$

where the first step is by definition,
the second step is by Lemma [E.14](#A5.Thmtheorem14),
the third step is by Lemma [E.8](#A5.Thmtheorem8),
the fourth and last step is by $\langle\boldsymbol{\mathbf{v}}^{(T_{4})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\tilde{O}(m_{k}^{-1/2})=o(1)$.
∎

###### Lemma E.14 .

###### Proof of Lemma E.14 .

###### Lemma E.15 .

###### Proof of Lemma E.15 .

#### E.7.1 Stage IV.a

This stage studies the dynamics in the interval $[T_{3},T_{4}^{-}]$, show the fast decay of post-softmax weights and their impact.
Roughly, everything keeps unchanged in this stage.
The idea is that everything keeps going on as usual before $T_{4}^{-}$, which is defined later.
And we don’t control the behaviour between the interval $[T_{4}^{-},T_{4}^{+}]$, since $T_{4}^{-}$ is not too far from $T_{4}^{+}$, we still can keep the good property at $T_{4}^{+}$ even in the worst case.
Finally, after $T_{4}^{+}$, the good behaviour continues.
(Also at $T_{4}^{+}$, we have the softmax ratio is basically concentrated at 1.)

Let $T_{4}^{-}\geq T_{3}$ be the last time the such that for all $t\in[T_{3},T_{4}^{-}]$, $s\in[m_{k}]$ and $i\in[n]$, following conditions hold,

$$ $\displaystyle\left|\sum_{i=1}^{n}s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq\frac{1}{2}n\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|,$ (29) $\displaystyle s_{i,21}^{(t)}s_{i,22}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq 2s_{i,11}^{(t)}s_{i,12}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|.$ (30) $$

Note that at $T_{3}$ we have

$$ $\displaystyle\sum_{i\in[n]}s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\frac{1}{2}n\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=$ $\displaystyle\frac{1}{4}\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\frac{1}{2}n\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle\geq$ $\displaystyle\frac{1}{4}\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\frac{1}{2}n\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle-(t-T_{2})\eta(\frac{\sigma_{p}s}{4\sqrt{2}}-\frac{n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{2})>0,$ $$

and

$$ $\displaystyle s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-2s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=$ $\displaystyle\frac{1}{4}(1+o(1))\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\frac{1}{2}(1+o(1))\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle\geq$ $\displaystyle\frac{1}{4}(1+o(1))\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\frac{1}{2}(1+o(1))\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle-(t-T_{2})\eta(\frac{\sigma_{p}s}{4\sqrt{2}}-\frac{\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{2})>0,$ $$

which implies $T_{4}^{-}$ is well-defined.

We use induction on $t$ to simultaneously prove the following properties for all $t=T_{3},\dots,T_{4}^{-}$:

- •
$\mathcal{A}(t)$, the linear-with-$t$ estimation: for all $i\in[n]$ and $s\in[m_{k}]$
$\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|=t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),$
$\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|=t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})).$
- •
$\mathcal{D}(t)$, query noise is increasing at $t$: for all $i\in[n]$ and $s\in[m_{k}]$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$
- •
$\mathcal{E}(t)$, key noise is increasing at $t$: for all $i\in[n]$ and $s\in[m_{k}]$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$
- •
$\mathcal{F}(t)$, query/key noise have the same sign at $t$: for all $i\in[n]$ and $s\in[m_{k}]$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle.$

###### Claim E.16 .

$\mathcal{F}(t)\Longrightarrow\mathcal{E}(t)$.

###### Claim E.17 .

$\mathcal{D}(T_{3}),\mathcal{E}(T_{3}),\dots,\mathcal{E}(t-1),\mathcal{F}(T_{3}),\dots,\mathcal{F}(t-1)\Longrightarrow\mathcal{D}(t)$.

###### Claim E.18 .

$\mathcal{A}(t),\mathcal{D}(t),\mathcal{E}(t)\Longrightarrow\mathcal{A}(t+1)$.

###### Claim E.19 .

$\mathcal{F}(t),\mathcal{D}(t),\mathcal{E}(t)\Longrightarrow\mathcal{F}(t+1)$.

Note that we have $\mathcal{A}(T_{3})$, $\mathcal{D}(T_{3})$, $\mathcal{F}(T_{3})$ hold.
$\mathcal{D}(T_{3})$ and $\mathcal{F}(T_{3})$ are already shown in the stage III.

###### Proof of $\mathcal{A}(T_{3})$ .

At $t=T_{3}$, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=(t-T_{2})\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}+\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=t\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(1\pm O(T_{2}t^{-1}))$ $\displaystyle=t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2}))(1\pm O(T_{2}t^{-1}))=t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),$ $$

where the first step follows from the monotonicity of $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ for $[T_{2},T_{3}]$,
the second step follows from Eq. ([21](#A5.E21)), the estimate of the magnitude of $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$,
the third step follows from Lemma [C.4](#A3.Thmtheorem4), the concentration of $\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}$,
the last step follows from $T_{2}/T_{3}=\tilde{O}(s^{-1/2})$. The analysis is similar for $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$.
∎

###### Proof of Claim E.16 .

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex199)), the second step is by the definition of $T_{4}^{-}$, i.e., Eq. ([30](#A5.E30)).
∎

###### Proof of Claim E.17 .

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle>$ $\displaystyle-\langle\boldsymbol{\mathbf{w}}^{(T_{3})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(T_{3})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+(t-T_{3})\eta(\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}-\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)>0,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex196)), the second step is due to
$\mathcal{E}(T_{3}),\dots,\mathcal{E}(t-1),\mathcal{F}(T_{3}),\dots,\mathcal{F}(t-1)$, the third step is by $\mathcal{D}(T_{3})$.
∎

###### Proof of Claim E.18 .

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|$ $\displaystyle=\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}$ $\displaystyle=t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2}))+\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2}))$ $\displaystyle=(t+1)\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),$ $$

where the first step is by $\mathcal{D}(t)$, the second step is by $\mathcal{A}(t)$ and Lemma [C.4](#A3.Thmtheorem4). The proof is same for $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$.
∎

###### Proof of Claim E.19 .

Suppose $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$, then by $\mathcal{D}(t)$ we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert>0.$ $$

By $\mathcal{F}(t)$, we have $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$, then by $\mathcal{E}(t)$ we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert>0,$ $$

which gives $\mathcal{F}(t+1)$.
∎

Now we study the ratio of softmax outputs across different samples. The following lemma is crucial to prove the lower bound of $T_{4}^{-}$.

###### Lemma E.20 .

Let

$$ $\displaystyle t^{*}:=\min\left\{T_{4}^{-},T_{4}\right\}.$ (31) $$

Then, for all $T_{3}\leq t\leq t^{*}$ and all $i,k\in[n]$, we have

$$ $\displaystyle\frac{s_{i,21}^{(t)}}{s_{k,21}^{(t)}}=1\pm o(1),\leavevmode\nobreak\ \forall i,k\in[n].$ $$

###### Proof of Lemma E.20 .

By $\mathcal{A}(t)$ and $\mathcal{F}(t)$ we have for $t\in[T_{3},T_{4}^{-}]$ and all $i\in[n]$

$$ $\displaystyle\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle=\frac{2}{\pi}m_{k}t^{2}\eta^{2}\sigma_{p}^{2}s^{2}(1\pm\tilde{O}(s^{-1/2})).$ (32) $$

Then, for all $T_{3}\leq t\leq t^{*}$ and $i\in[n]$

$$ $\displaystyle s_{i,21}^{(t)}$ $\displaystyle=\frac{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle\right)}{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle\right)+\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}$ $\displaystyle=\frac{1\pm o(1)}{1\pm o(1)+\exp\left(\frac{2}{\pi}m_{k}t^{2}\eta^{2}\sigma_{p}^{2}s^{2}(1\pm\tilde{O}(s^{-1/2}))\right)}$ $\displaystyle=\frac{1\pm o(1)}{1\pm o(1)+\exp(\frac{2}{\pi}m_{k}t^{2}\eta^{2}\sigma_{p}^{2}s^{2})\exp(\pm\tilde{O}(\frac{2}{\pi}m_{k}t^{2}\eta^{2}\sigma_{p}^{2}s^{3/2}))}$ $\displaystyle=\frac{1\pm o(1)}{1\pm o(1)+\exp(\frac{2}{\pi}m_{k}t^{2}\eta^{2}\sigma_{p}^{2}s^{2})(1\pm o(1))}$ $\displaystyle=\frac{1\pm o(1)}{1+\exp(\frac{2}{\pi}m_{k}t^{2}\eta^{2}\sigma_{p}^{2}s^{2})},$ (33) $$

where the first step is by definition, the second step is by Eq. ([28](#A5.E28)), ([27](#A5.E27)), and Eq. ([32](#A5.E32)), the fourth step is by $m_{k}T_{4}^{2}\eta^{2}\sigma_{p}^{2}s^{2}=o(1)$. Then we have for all $T_{3}\leq t\leq t^{*}$ and all $i,k\in[n]$

$$ $\displaystyle\frac{s_{i,21}^{(t)}}{s_{k,21}^{(t)}}$ $\displaystyle=(1\pm o(1))\frac{(1\pm o(1))(1+\exp(\frac{2}{\pi}m_{k}t^{2}\eta^{2}\sigma_{p}^{2}s^{2}))}{(1\pm o(1))(1+\exp(\frac{2}{\pi}m_{k}t^{2}\eta^{2}\sigma_{p}^{2}s^{2}))}=1\pm o(1),$ $$

where the first step is by Eq. ([E.7.1](#A5.Ex389)).
∎

The following lemma studies the lower bound for $T_{4}^{-}$.

###### Lemma E.21 .

We have

$$ $\displaystyle T_{4}^{-}\geq\tilde{T}_{4}^{-}:=\sqrt{\frac{0.99\pi}{2}}\sqrt{\log\left(\frac{\sigma_{p}s}{3\sqrt{2}n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\right)}\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1}.$ (34) $$

###### Proof.

Suppose $T_{4}^{-}<\tilde{T}_{4}^{-}$, thus we have $T_{4}^{-}+1<\tilde{T}_{4}^{-}+1<T_{4}$.
By Claim [E.18](#A5.Thmtheorem18) and [E.19](#A5.Thmtheorem19), we have $\mathcal{A}(T_{4}^{-}+1)$ and $\mathcal{F}(T_{4}^{-}+1)$ hold, which implies Eq. ([32](#A5.E32)) holds at $t=T_{4}^{-}+1$. Then, for all $i\in[n]$, we have

$$ $\displaystyle s_{i,21}^{(T_{4}^{-}+1)}$ $\displaystyle=\frac{1\pm o(1)}{1\pm o(1)+\exp\left(\frac{2}{\pi}m_{k}(T_{4}^{-}+1)^{2}\eta^{2}\sigma_{p}^{2}s^{2}(1\pm\tilde{O}(s^{-1/2}))\right)}$ $\displaystyle=\frac{1\pm o(1)}{1+\exp(\frac{2}{\pi}m_{k}(T_{4}^{-}+1)^{2}\eta^{2}\sigma_{p}^{2}s^{2})}$ (35) $\displaystyle\geq\frac{1\pm o(1)}{1+\exp(\frac{2}{\pi}m_{k}(\tilde{T}_{4}^{-}+1)^{2}\eta^{2}\sigma_{p}^{2}s^{2})}$ $\displaystyle\geq\frac{1\pm o(1)}{1.01\exp(\frac{2}{\pi}m_{k}(\tilde{T}_{4}^{-}+1)^{2}\eta^{2}\sigma_{p}^{2}s^{2})}$ $\displaystyle\geq\frac{1\pm o(1)}{1.01}\left(\frac{3\sqrt{2}n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s}\right)^{0.99+o(1)}$ $\displaystyle\geq\frac{(1\pm o(1))3\sqrt{2}n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{1.01\sigma_{p}s},$ (36) $$

where the first step is by Eq. ([28](#A5.E28)), ([27](#A5.E27)), and Eq. ([32](#A5.E32)),
the second step is by $T_{4}^{-}+1<T_{4}$,
the fifth step is by $\eta=o(m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1})$,
the last step is by $n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\sigma_{p}^{-1}s^{-1}=o(1)$,
Then we estimate the magnitude of $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$.
WLOG, given neuron $s$, suppose $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$.
For $T_{3}<t\leq T_{4}^{-}+1$

$$ $\displaystyle\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq\frac{1}{\sqrt{2}}(t-T_{2})\eta\sigma_{p}s+\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq\frac{1}{\sqrt{2}}(t-T_{2})\eta\sigma_{p}s+2n\beta_{\boldsymbol{\mathbf{\mu}}}$ $\displaystyle\geq\frac{1}{\sqrt{2}}t\eta\sigma_{p}s(1\pm O(T_{2}t^{-1}))=\frac{1}{\sqrt{2}}t\eta\sigma_{p}s(1\pm o(1)),$ (37) $$

where the first step follows from $\mathcal{D}(t)$ and Eq. ([22](#A5.E22)),
the second step follows from the lower bound of $\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ Eq. ([23](#A5.E23)),
the last step follows from $T_{2}T_{3}^{-1}=o(1)$.
The analysis is similar for $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$.
Now we can show that

$$ $\displaystyle\sum_{i=1}^{n}s_{i,21}^{(T_{4}^{-}+1)}s_{i,22}^{(T_{4}^{-}+1)}\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-}+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=\leavevmode\nobreak\$ $\displaystyle\frac{1\pm o(1)}{1+\exp(\frac{2}{\pi}m_{k}(T_{4}^{-}+1)^{2}\eta^{2}\sigma_{p}^{2}s^{2})}\cdot(1-\frac{1\pm o(1)}{1+\exp(\frac{2}{\pi}m_{k}(T_{4}^{-}+1)^{2}\eta^{2}\sigma_{p}^{2}s^{2})})\cdot\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-}+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq\leavevmode\nobreak\$ $\displaystyle\frac{(1\pm o(1))3\sqrt{2}n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{1.01\sigma_{p}s}\cdot 0.5\cdot\frac{1}{\sqrt{2}}(T_{4}^{-}+1)\eta\sigma_{p}s(1\pm o(1))$ $\displaystyle\geq\leavevmode\nobreak\$ $\displaystyle 1.4n(T_{4}^{-}+1)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert$ $\displaystyle\geq\leavevmode\nobreak\$ $\displaystyle\frac{1}{2}n\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{4}^{-}+1)},\boldsymbol{\mathbf{\mu}}\rangle,$ $$

where the first step is by Eq. ([35](#A5.E35)),
the second step is by Eq. ([36](#A5.E36)), ([37](#A5.E37)) and $T_{4}^{-}\geq T_{3}$,
the last step is by the upper bound in Eq. ([24](#A5.E24)). Similarly, we can show that for all $s\in[m_{k}]$ and $i\in[n]$

$$ $\displaystyle s_{i,21}^{(T_{4}^{-}+1)}s_{i,22}^{(T_{4}^{-}+1)}\left|\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-}+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\geq 2s_{i,11}^{(T_{4}^{-}+1)}s_{i,12}^{(T_{4}^{-}+1)}\left|\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-}+1)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|,$ $$

which contradicts with the definition of $T_{4}^{-}$. Therefore, we have $T_{4}^{-}\geq\tilde{T}_{4}^{-}$.
∎

###### Remark E.22 .

Note the constant $0.99$ in the definition of $\tilde{T}_{4}^{-}$ is crucial and its impact is on the exponent.
When this constant is greater than 1, then the bound does not hold. Furthermore, we have $s_{i,21}^{(t)}\leq 0.1$ for $t\geq T_{4}^{-}$.

Next, the following lemmas study the behaviour of query/key signal.

###### Lemma E.23 .

The query signal are monotonic for $t\in[T_{3},t^{*}]$. Formally, for all $t\in[T_{3},t^{*}]$ and $s\in[m_{k}]$,

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ $$

where $t^{*}$ is defined in Eq. ([31](#A5.E31)).

###### Lemma E.24 .

The key signal are monotonic for $t\in[T_{3},t^{*}]$. Formally, for all $t\in[T_{3},t^{*}]$ and $s\in[m_{k}]$,

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ $$

where $t^{*}$ is defined in Eq. ([31](#A5.E31)).

###### Proof of Lemma E.23 .

WLOG, given neuron $s$, suppose $\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle>0$ and $\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle<0$. For query signal, similar to previous analysis, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}(-\ell_{i}^{\prime(t)}s_{i,11}^{(t)}s_{i,12}^{(t)})\cdot\left(-y_{i}\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}\frac{\sqrt{2}}{8\sqrt{\pi}}t\eta\sigma_{p}s(1\pm o(1))\cdot(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-n\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle\geq$ $\displaystyle\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-n\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+(t-T_{2})(\eta\sigma_{p}s/\sqrt{2}-\eta n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)>0,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex198)),
the second step is by the concentration of $\ell_{i}^{\prime}$ and $s_{i,11}^{(t)}$ for $t\leq T_{4}$, i.e., Lemma [E.14](#A5.Thmtheorem14) and [E.15](#A5.Thmtheorem15),
the fourth step is by $\mathcal{E}(t)$ and the lower bound Eq. ([22](#A5.E22)),
and the last step is by lower bound Eq. ([23](#A5.E23)).
∎

###### Proof of Lemma E.24 .

For key signal, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}(-l_{i}^{\prime(t)})\cdot\left(-y_{i}s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle-s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}\frac{\sqrt{2}}{2\sqrt{\pi}}t\eta\sigma_{p}s(1\pm o(1))\cdot(-s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\frac{1}{4}(1\pm o(1))\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle-\sum_{i\in[n]}s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\frac{1}{4}(1\pm o(1))n\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle-\sum_{i\in[n]}s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle-s_{21}^{(t)}s_{22}^{(t)}(1\pm o(1))\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle<0,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex198)),
the second step is by the concentration of $\ell_{i}^{\prime}$ and $s_{i,11}^{(t)}$ for $t\leq T_{4}$, i.e., Lemma [E.14](#A5.Thmtheorem14) and [E.15](#A5.Thmtheorem15),
the fourth step is by the definition of $T_{4}^{-}$, i.e., Eq. ([29](#A5.E29)),
the fifth step is by the $\mathcal{B}(t)$,
the last step is due to the upper bound Eq. ([23](#A5.E23)) and $\mathcal{D}(T_{3}),\dots\mathcal{D}(t-1)$.
∎

At the end of this section, we give an upper bound for $T_{4}^{-}$.

###### Lemma E.25 .

We have

$$ $\displaystyle T_{4}^{-}\leq\hat{T}_{4}^{-}:=\sqrt{\frac{1.01\pi}{2}}\sqrt{\log\left(\frac{\sigma_{p}s}{\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}\right)}\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1}.$ (38) $$

###### Proof.

Suppose $T_{4}^{-}>\hat{T}_{4}^{-}$.
First, we estimate the magnitude of $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ for $T_{3}\leq t\leq T_{4}^{-}$.
WLOG, given neuron $s$, suppose $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$.
For all $s\in[m_{k}]$ we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\geq\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},\boldsymbol{\mathbf{\mu}}\rangle+(t-T_{3})\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\geq t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/2,$ (39) $$

where the first step is by Lemma [E.23](#A5.Thmtheorem23) and $\hat{T}_{4}^{-}<t^{*}$,
the second step is by the lower bound in Eq. ([24](#A5.E24)) at $t=T_{3}$.
For all $s\in[m_{k}]$ we also have

$$ $\displaystyle\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\leq\sum_{i=1}^{n}\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq tn\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),$ $$

where the second step is by $\mathcal{A}(t)$.
By the definition of $T_{4}^{-}$, at $t=\hat{T}_{4}^{-}$ we have

$$ $\displaystyle tn\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/4$ $\displaystyle\leq\frac{1}{2}n\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle\leq\left|\sum_{i=1}^{n}s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|$ $\displaystyle\leq\frac{1\pm o(1)}{1+\exp(\frac{2}{\pi}m_{k}t^{2}\eta^{2}\sigma_{p}^{2}s^{2})}\cdot tn\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2}))$ $\displaystyle\leq\frac{1\pm o(1)}{\exp(\frac{2}{\pi}m_{k}t^{2}\eta^{2}\sigma_{p}^{2}s^{2})}\cdot tn\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),$ $$

which implies that

$$ $\displaystyle\hat{T}_{4}^{-}\leq\sqrt{\frac{\pi}{2}}\sqrt{\log\left(\frac{4\sqrt{2}\sigma_{p}s}{\sqrt{\pi}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}(1+o(1))\right)}\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1}s^{-1},$ $$

which gives a contradiction.
∎

###### Remark E.26 .

The upper bound for $T_{4}^{-}$ also proves that $t^{*}=T_{4}^{-}$, which gives that for all $t\in[T_{3},T_{4}^{-}]$, $s\in[m_{k}]$, and $i,k\in[n]$

$$ $\displaystyle\frac{s_{i,21}^{(t)}}{s_{k,21}^{(t)}}=1\pm o(1),$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert.$ $$

Moreover, for neuron $s$, suppose $\langle\boldsymbol{\mathbf{w}}^{(T_{3})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle>0$, then at $t=T_{4}^{-}$, for all $s\in[m_{k}]$, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=(t-T_{3})\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert+\langle\boldsymbol{\mathbf{w}}^{(T_{3})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1\pm o(1)),$ (40) $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle=-(t-T_{3})\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert+\langle\boldsymbol{\mathbf{w}}^{(T_{3})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle=-t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1\pm o(1)),$ (41) $$

where for query signal the first step is by Lemma [E.23](#A5.Thmtheorem23)
and the second step is by the bound Eq. [24](#A5.E24).
The proof of statement for key signal is similar.

###### Claim E.16 .

###### Claim E.17 .

###### Claim E.18 .

###### Claim E.19 .

###### Proof of 𝒜 ​ ( T 3 ) 𝒜 subscript 𝑇 3 \mathcal{A}(T_{3}) .

###### Proof of Claim E.16 .

###### Proof of Claim E.17 .

###### Proof of Claim E.18 .

###### Proof of Claim E.19 .

###### Lemma E.20 .

###### Proof of Lemma E.20 .

###### Lemma E.21 .

###### Proof.

###### Remark E.22 .

###### Lemma E.23 .

###### Lemma E.24 .

###### Proof of Lemma E.23 .

###### Proof of Lemma E.24 .

###### Lemma E.25 .

###### Proof.

###### Remark E.26 .

#### E.7.2 Stage IV.b

This stage studies the dynamics in the transition interval $[T_{4}^{-},T_{4}^{+}]$.
For $T_{4}^{-}\leq t\leq T_{4}^{+}$,
given neuron $s$, for those samples where query/key noise align with the query/key signal, query/key noise dynamics keep unchanged, which will be defined later.
On the other hands, for those samples where query/key noise do not align with query/key signal, query noise dynamics keep unchanged.
However, key noise start to decrease but are still greater than zero at $T_{4}^{+}$.

Let $T_{4}^{+}$ be the first time the following condition holds

$$ $\displaystyle s_{i,21}^{(t)}s_{i,22}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq\frac{1}{2}s_{i,11}^{(t)}s_{i,12}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|,$ (42) $$

for all $s\in[m_{k}]$ and $i\in[n]$. Note that at $T_{4}^{+}$, we have

$$ $\displaystyle\sum_{i=1}^{n}s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\leq\sum_{i=1}^{n}s_{i,21}^{(t)}s_{i,22}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq\frac{1}{2}\sum_{i=1}^{n}s_{i,11}^{(t)}s_{i,12}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|$ $\displaystyle=\frac{1}{8}(1+o(1))ns_{i,11}^{(t)}s_{i,12}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|<\frac{1}{2}ns_{i,11}^{(t)}s_{i,12}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|,$ $$

which implies $T_{4}^{+}\geq T_{4}^{-}$, $T_{4}^{+}$ is well-defined.

Let

$$ $\displaystyle E_{i}^{(T_{3})}:=\left\{s\in[m_{k}]:\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},\boldsymbol{\mathbf{\mu}}\rangle\right\},$ $\displaystyle E_{s}^{(T_{3})}:=\left\{i\in[n]:\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},\boldsymbol{\mathbf{\mu}}\rangle\right\}.$ $$

all neurons with $s\in[m_{k}]$ and $i\in E_{s}^{(T_{3})}$ are called aligned neurons.
The next lemma studies the concentration behaviour about $\left|E_{i}^{(T_{3})}\right|$. This lemma is crucial to upper bound $s_{i,21}^{(t)}$ for $T_{4}^{-}\leq t\leq T_{4}^{+}$, and further upper bound $T_{4}^{+}$.

###### Lemma E.27 .

Suppose that $\delta>0$
Then with probability at least $1-\delta$,

$$ $\displaystyle\left|E_{i}^{(T_{3})}\right|\geq m_{k}/2-\tilde{O}(m_{k}n^{-1/2}+m_{k}^{1/2})\geq m_{k}/2-\tilde{O}(\sqrt{m_{k}}),$ $$

for all $i\in[n]$.

###### Proof.

We can write $E_{i}^{(T_{3})}$ as

$$ $\displaystyle E_{i}^{(T_{3})}=\left\{s\in[m_{k}]:\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right\},$ $$

where this is by the sign of $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ and $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ does not change for $[T_{2},T_{3}]$ and the sign of $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},\boldsymbol{\mathbf{\mu}}\rangle$ is the same as that of $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$, i.e., Lemma [E.9](#A5.Thmtheorem9) and [E.11](#A5.Thmtheorem11).
Then we have

$$ $\displaystyle p$ $\displaystyle:=\frac{1}{m_{k}}\mathbb{E}[\left|E_{i}^{(T_{3})}\right|]$ $\displaystyle=\mathbb{P}[\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0,\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0]+\mathbb{P}[\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle<0,\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle<0]$ $\displaystyle=\mathbb{P}[X_{s,i}>0,\sum_{i=1}^{n}X_{s,i}>0]+\mathbb{P}[X_{s,i}<0,\sum_{i=1}^{n}X_{s,i}<0]$ $\displaystyle\geq\frac{1}{2}-O(n^{-1/2}),$ $$

where the third step by that the sign of $X_{s,i}$ determines the sign of $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{2})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$,
the last step is by Lemma [E.6](#A5.Thmtheorem6).
Then, by Hoeffding’s inequality, with probability at least $1-\delta/n$

$$ $\displaystyle\left|\left|E_{i}^{(T_{3})}\right|-pm_{k}\right|\leq\sqrt{\frac{m_{k}}{2}\log(2n/\delta)}.$ $$

Finally, by the condition $n=\Omega(m_{k}^{2})$, we get the conclusion.
∎

We do not directly characterize the dynamics for $[T_{4}^{-},T_{4}^{+}]$.
Instead, we characterize for $[T_{4}^{-}],(1+\theta)T_{4}^{-}$, where $\theta$ is a large constant in $(0,1)$, and show that the definition of $T_{4}^{+}$ are satisfied, which characterize he dynamics before $T_{4}^{+}$ and give an upper bound for $T_{4}^{+}$ simultaneously.
Specifically, we use induction on $t$ to simultaneously prove the following properties for all $t=T_{4}^{-},\dots,(1+\theta)T_{4}^{-}$:

- •
$\mathcal{A}(t)$, the monotonicity of query noise aligned with signal: for any $s\in[m_{k}]$ and $i\in E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$
- •
$\mathcal{B}(t)$, the monotonicity of key noise aligned with signal: for any $s\in[m_{k}]$ and $i\in E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$
- •
$\mathcal{C}(t)$, the monotonicity of query noise unaligned with signal: for any $s\in[m_{k}]$ and $i\in[n]/E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$
- •
$\mathcal{D}(t)$, the sign does not change for unaligned key noise: for any $s\in[m_{k}]$ and $i\in[n]/E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle.$
- •
$\mathcal{E}(t)$, the monotonicity of query signal: for any $s\in[m_{k}]$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert.$
- •
$\mathcal{F}(t)$, the sign does not change for key signal: for any $s\in[m_{k}]$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle.$
- •
$\mathcal{G}(t)$, the linear-with-$t$ estimation: for all $s\in[m_{k}]$ and $i\in[n]$
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),\forall i\in E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),\forall i\in E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=-\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),\forall i\in[n]/E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1\pm o(1)).$

###### Remark E.28 .

Before stating the results in this section, we first recall and summarize the results until $T_{4}^{-}$.
WLOG, given neuron $s$, we suppose that $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$ in this section below.
Then we have $\langle\boldsymbol{\mathbf{w}}^{(T_{3})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle>0$ and $\langle\boldsymbol{\mathbf{w}}^{(T_{3})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle<0$.
Then at $t=T_{4}^{-}$ we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),\forall i\in E_{s}^{(T_{3})}$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),\forall i\in E_{s}^{(T_{3})}$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=-t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),\forall i\in[n]/E_{s}^{(T_{3})}$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=-t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),\forall i\in[n]/E_{s}^{(T_{3})}$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1\pm o(1)),$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle=-t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1\pm o(1)).$ $$

This gives that $\mathcal{A}(T_{4}^{-})$,
$\mathcal{B}(T_{4}^{-})$, $\mathcal{C}(T_{4}^{-})$, $\mathcal{D}(T_{4}^{-})$, $\mathcal{E}(T_{4}^{-})$, $\mathcal{F}(T_{4}^{-})$ hold.

To prove these properties, we need some further properties after $T_{4}^{-}$.

###### Lemma E.29 (Properties after $T_{4}^{-}$ ) .

Let $\theta$ be a large constant in $(0,1)$.
For all $t\geq T_{4}^{-}$, we have

$$ $\displaystyle(2T_{4}^{-}-t)\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1-\tilde{O}(s^{-1/2}))\leq\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\cdot\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\leq t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1+\tilde{O}(s^{-1/2})),$ $\displaystyle(2T_{4}^{-}-t)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1-o(1))\leq\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle)\cdot\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle\leq t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1+o(1)),$ $$

and for all $T_{4}^{-}\leq t\leq(1+\theta)T_{4}^{-}$

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle,$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle.$ $$

Particularly, For all $0\leq\theta^{\prime}\leq\theta$, $s\in[m_{k}]$ and $i\in[n]$,
setting $t=(1+\theta^{\prime})T_{4}^{-}$, we have

$$ $\displaystyle(1-\theta^{\prime})T_{4}^{-}\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1-\tilde{O}(s^{-1/2}))\leq\left|\langle\boldsymbol{\mathbf{w}}^{((1+\theta^{\prime})T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq(1+\theta^{\prime})T_{4}^{-}\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1+\tilde{O}(s^{-1/2})),$ $\displaystyle(1-\theta^{\prime})T_{4}^{-}\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1-o(1))\leq\left|\langle\boldsymbol{\mathbf{w}}^{((1+\theta^{\prime})T_{4}^{-})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle\right|\leq(1+\theta^{\prime})T_{4}^{-}\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1+o(1)).$ $$

###### Proof of Lemma E.29 .

Suppose $\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}\rangle>0$, then we have for all $t\geq T_{4}^{-}$

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\leq\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+(t-T_{4}^{-})\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\leq t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),$ $$

where the first step is by that the magnitude of single update of query/key noise is always $\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}$,
the second step is by summary above and Lemma [C.4](#A3.Thmtheorem4). Similarly, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\geq\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-(t-T_{4}^{-})\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}\geq(2T_{4}^{-}-t)\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})).$ $$

Set $t=(1+\theta^{\prime})T_{4}^{-}$ gives the conclusion.
Also, the lower bound gives

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq(2T_{4}^{-}-t)\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2}))$ $\displaystyle\geq(1-\theta)T_{4}^{-}\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2}))>0,$ $$

where the last step is by $\theta$ is a constant.
Similarly, we can prove the statement for $\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle$.
∎

###### Remark E.30 .

Note that Lemma [E.29](#A5.Thmtheorem29) proves the $\mathcal{D}(t)$ and $\mathcal{F}(t)$ for all $T_{2}^{-}\leq t\leq(1+\theta)T_{4}^{-}$.

###### Claim E.31 .

$\mathcal{G}(t),\mathcal{F}(t)\Longrightarrow\mathcal{A}(t)$.

###### Claim E.32 .

$\mathcal{G}(t)\Longrightarrow\mathcal{B}(t)$.

###### Claim E.33 .

$\mathcal{D}(t),\mathcal{F}(t)\Longrightarrow\mathcal{C}(t)$.

###### Claim E.34 .

$\mathcal{B}(t)\Longrightarrow\mathcal{E}(t)$.

###### Claim E.35 .

$\mathcal{G}(t),\mathcal{A}(t),\mathcal{B}(t),\mathcal{C}(t),\mathcal{E}(t)\Longrightarrow\mathcal{G}(t+1)$.

###### Proof of Claim E.31 .

We have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex196)),
the second step is by $\mathcal{F}(t)$,
the last step is by $\mathcal{G}(t)$.
∎

###### Proof of Claim E.32 .

We have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex199)),
the second step is by $\mathcal{G}(t)$.
∎

###### Proof of Claim E.33 .

We have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex196)),
the second step is by $\mathcal{F}(t)$,
the last step is by $\mathcal{D}(t)$.
∎

###### Proof of Claim E.34 .

We have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}(-\ell_{i}^{\prime(t)}s_{i,11}^{(t)}s_{i,12}^{(t)})\cdot\left(-y_{i}\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{v}}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}\frac{\sqrt{2}}{8\sqrt{\pi}}t\eta\sigma_{p}s(1\pm o(1))\cdot(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-n\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=$ $\displaystyle-n\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\sum_{i\in E_{s}^{(T_{3})}}\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\sum_{i\in[n]/E_{s}^{(T_{3})}}\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle-n\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\sum_{i\in E_{s}^{(T_{3})}}\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+(t-T_{4}^{-})\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}+\sum_{i\in[n]/E_{s}^{(T_{3})}}\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-(t-T_{4}^{-})\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}$ $\displaystyle\geq$ $\displaystyle-n\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+(t-T_{4}^{-})\eta\sigma_{p}s/\sqrt{2}$ $\displaystyle\geq$ $\displaystyle-n\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\sum_{i\in[n]}\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+(t-T_{4}^{-})\eta(\sigma_{p}s/\sqrt{2}-n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)>0,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex198)),
the second step is by the concentration of $\ell_{i}^{\prime}$ and $s_{i,11}^{(t)}$ for $t\leq T_{4}$, i.e., Lemma [E.14](#A5.Thmtheorem14) and [E.15](#A5.Thmtheorem15),
the fifth step is by $\mathcal{B}(t)$,
the sixth step is by the lower bound Eq. ([22](#A5.E22)),
the last step is guaranteed by the dynamics in $[T_{3},T_{4}^{-}]$.
∎

###### Proof of Claim E.35 .

This claim is straightforward.
∎

The following lemma gives an upper bound for $T_{4}^{+}$.

###### Lemma E.36 .

Let $\theta_{c}$ be a small constant satisfying $0.99(1+\theta_{c})-O(\theta_{c}/\sqrt{m_{k}})>1$,
then we have

$$ $\displaystyle T_{4}^{+}\leq(1+\theta_{c})T_{4}^{-},$ (43) $$

and for all $t\in[(1+\theta_{c})T_{4}^{-},(2-\theta_{c})T_{4}^{-}]$ we have

$$ $\displaystyle s_{i,21}^{(t)}s_{i,22}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq\frac{1}{2}s_{i,11}^{(t)}s_{i,12}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|.$ $$

###### Proof.

For fixed $\theta\in[\theta_{c},1-\theta_{c}]$, at $t=(1+\theta)T_{4}^{-}$, for all $i\in[n]$, we have

$$ $\displaystyle\leavevmode\nobreak\ \sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \sum_{s\in E_{i}^{(T_{3})}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle+\sum_{s\in[m_{k}]/E_{i}^{(T_{3})}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ \left|E_{i}^{(T_{3})}\right|\frac{2}{\pi}t^{2}\eta^{2}\sigma_{p}^{2}s^{2}(1\pm\tilde{O}(s^{-1/2}))+(m_{k}-\left|E_{i}^{(T_{3})}\right|)\frac{2}{\pi}t(2T_{4}^{-}-t)\eta^{2}\sigma_{p}^{2}s^{2}(1\pm\tilde{O}(s^{-1/2}))$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ (\left|E_{i}^{(T_{3})}\right|t^{2}+(m_{k}-\left|E_{i}^{(T_{3})}\right|)t(2T_{4}^{-}-t))\cdot\frac{2}{\pi}\eta^{2}\sigma_{p}^{2}s^{2}(1\pm\tilde{O}(s^{-1/2}))$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ (\left|E_{i}^{(T_{3})}\right|(1+\theta)^{2}T_{4}^{-2}+(m_{k}-\left|E_{i}^{(T_{3})}\right|)(1+\theta)(1-\theta)T_{4}^{-2}\cdot\frac{2}{\pi}\eta^{2}\sigma_{p}^{2}s^{2}(1\pm\tilde{O}(s^{-1/2}))$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ ((\frac{m_{k}}{2}-o({m_{k}}))(1+\theta)^{2}+(\frac{m_{k}}{2}+o({m_{k}}))(1-\theta^{2})\cdot T_{4}^{-2}\frac{2}{\pi}\eta^{2}\sigma_{p}^{2}s^{2}(1\pm\tilde{O}(s^{-1/2}))$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ ((1+\theta)m_{k}-(\theta^{2}+\theta)o(m_{k}))\cdot T_{4}^{-2}\frac{2}{\pi}\eta^{2}\sigma_{p}^{2}s^{2}(1\pm\tilde{O}(s^{-1/2})),$ (44) $$

where the second step is by $\mathcal{G}(t)$ and Lemma [E.29](#A5.Thmtheorem29),
the fifth step is by Lemma [E.27](#A5.Thmtheorem27).

Then, we have

$$ $\displaystyle s_{i,21}^{(t)}$ $\displaystyle=\frac{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle\right)}{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\mu}}\rangle\right)+\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}$ $\displaystyle\leq\frac{1\pm o(1)}{1\pm o(1)+\exp\left(((1+\theta)m_{k}-(\theta^{2}+\theta)O(\sqrt{m_{k}}))\cdot T_{4}^{-2}\frac{2}{\pi}\eta^{2}\sigma_{p}^{2}s^{2}(1\pm\tilde{O}(s^{-1/2}))\right)}$ $\displaystyle=\frac{1\pm o(1)}{1\pm o(1)+\exp(((1+\theta)m_{k}-(\theta^{2}+\theta)O(\sqrt{m_{k}}))\cdot T_{4}^{-2}\frac{2}{\pi}\eta^{2}\sigma_{p}^{2}s^{2})\exp(\pm\tilde{O}(\frac{2}{\pi}m_{k}T_{4}^{-2}\eta^{2}\sigma_{p}^{2}s^{3/2}))}$ $\displaystyle=\frac{1\pm o(1)}{1\pm o(1)+\exp(((1+\theta)m_{k}-(\theta^{2}+\theta)O(\sqrt{m_{k}}))\cdot T_{4}^{-2}\frac{2}{\pi}\eta^{2}\sigma_{p}^{2}s^{2})(1\pm o(1))}$ $\displaystyle=\frac{1\pm o(1)}{1+\exp(((1+\theta)m_{k}-(\theta^{2}+\theta)O(\sqrt{m_{k}}))\cdot T_{4}^{-2}\frac{2}{\pi}\eta^{2}\sigma_{p}^{2}s^{2})}$ $\displaystyle\leq\frac{1\pm o(1)}{\exp(((1+\theta)m_{k}-(\theta^{2}+\theta)O(\sqrt{m_{k}}))\cdot T_{4}^{-2}\frac{2}{\pi}\eta^{2}\sigma_{p}^{2}s^{2})},$ (45) $$

where the first step is by definition, the second step is by Eq. ([28](#A5.E28)), ([27](#A5.E27)), and ([E.7.2](#A5.Ex476)), the fourth step is by $T_{4}^{-}\lesssim T_{4}$ and $m_{k}T_{4}^{2}\eta^{2}\sigma_{p}^{2}s^{2}=o(1)$.

Therefore, we have

$$ $\displaystyle\leavevmode\nobreak\ s_{i,21}^{((1+\theta)T_{4}^{-})}s_{i,22}^{((1+\theta)T_{4}^{-})}\left|\langle\boldsymbol{\mathbf{w}}^{((1+\theta)T_{4}^{-})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|$ $\displaystyle\leq$ $\displaystyle\leavevmode\nobreak\ \frac{1+o(1)}{\exp(((1+\theta)m_{k}-(\theta^{2}+\theta)O(\sqrt{m_{k}}))\cdot T_{4}^{-2}\frac{2}{\pi}\eta^{2}\sigma_{p}^{2}s^{2})}\cdot 1\cdot(1+\theta)T_{4}^{-}\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1+\tilde{O}(s^{-1/2}))$ $\displaystyle\leq$ $\displaystyle\leavevmode\nobreak\ (1+\theta)T_{4}^{-}\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1+\tilde{O}(s^{-1/2}))\cdot\left(\frac{3\sqrt{2}n\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert}{\sigma_{p}s}\right)^{0.99(1+\theta)-O(\theta/\sqrt{m_{k}})}$ $\displaystyle\leq$ $\displaystyle\leavevmode\nobreak\ \frac{1}{9}(1+\theta)T_{4}^{-}\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert$ $\displaystyle\leq$ $\displaystyle\leavevmode\nobreak\ \frac{1}{2}s_{i,11}^{((1+\theta)T_{4}^{-})}s_{i,12}^{((1+\theta)T_{4}^{-})}\left|\langle\boldsymbol{\mathbf{w}}^{((1+\theta)T_{4}^{-})}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|,$ $$

where the first step is by Eq. ([E.7.2](#A5.Ex482)) and $\mathcal{G}(t)$,
the second step is by Lemma [E.21](#A5.Thmtheorem21),
the third step is by $0.99(1+\theta)-O(\theta/\sqrt{m_{k}})>1$,
the last step is by $\mathcal{G}(t)$ and
the concentration results in Lemma [E.14](#A5.Thmtheorem14),
thus the conclusion holds.
∎

###### Lemma E.27 .

###### Proof.

###### Remark E.28 .

###### Lemma E.29 (Properties after T 4 − superscript subscript 𝑇 4 T_{4}^{-} ) .

###### Proof of Lemma E.29 .

###### Remark E.30 .

###### Claim E.31 .

###### Claim E.32 .

###### Claim E.33 .

###### Claim E.34 .

###### Claim E.35 .

###### Proof of Claim E.31 .

###### Proof of Claim E.32 .

###### Proof of Claim E.33 .

###### Proof of Claim E.34 .

###### Proof of Claim E.35 .

###### Lemma E.36 .

###### Proof.

#### E.7.3 Stage IV.c

$T_{4}$ is a timestep where $s_{i,11}^{(t)}=1/2+o(1)$,
we characterize the dynamics in this interval.

WLOG, given neuron $s$, we suppose that $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$ in this section below.

Specifically, we use induction on $t$ to simultaneously prove the following properties for all $t=(1+\theta_{c})T_{4}^{-},\dots,T_{4}$:

- •
$\mathcal{A}(t)$, the monotonicity of query noise aligned with signal: for any $s\in[m_{k}]$ and $i\in E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$
- •
$\mathcal{B}(t)$, the monotonicity of key noise aligned with signal: for any $s\in[m_{k}]$ and $i\in E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$
- •
$\mathcal{C}(t)$, the fate of query noise unaligned with signal: for any $s\in[m_{k}]$ and $i\in[n]/E_{s}^{(T_{3})}$,
for $t\leq(2-\theta_{c})T_{4}^{-}$ we have
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{+})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$
and for $t\geq(2+3\theta_{c})T_{4}^{-}$ we have
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle-\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{+})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$
- •
$\mathcal{D}(t)$, the monotonicity of key noise unaligned with signal: for any $s\in[m_{k}]$ and $i\in[n]/E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}.$
- •
$\mathcal{E}(t)$, the monotonicity of query signal: for any $s\in[m_{k}]$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert.$
- •
$\mathcal{F}(t)$: the monotonicity of key signal: for any $s\in[m_{k}]$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert.$
- •
$\mathcal{G}(t)$, the linear-with-$t$ estimation: for all $s\in[m_{k}]$ and $i\in[n]$
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),\forall i\in E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)t\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2})),\forall i\in E_{s}^{(T_{3})}$
$\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)t\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert(1\pm o(1)).$
- •
$\mathcal{H}(t)$: for any $s\in[m_{k}]$ and $i\in[n]$
$\displaystyle s_{i,21}^{(t)}s_{i,22}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq\frac{1}{2}s_{i,11}^{(t)}s_{i,12}^{(t)}\left|\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|.$

###### Remark E.37 .

By Lemma [E.36](#A5.Thmtheorem36), we have $\mathcal{H}(t)$ holds for $t\in[(1+\theta_{c})T_{4}^{-},(2-\theta_{c})T_{4}^{-}]$.
The $\mathcal{A}(t)$, $\mathcal{B}(t)$, $\mathcal{E}(t)$ and $\mathcal{G}(t)$ for $t\in[(1+\theta_{c})T_{4}^{-},(2-\theta_{c})T_{4}^{-}]$ are already shown in the induction statement in the last section. Actually, the proof of $\mathcal{A}(t)$, $\mathcal{B}(t)$, $\mathcal{E}(t)$ and $\mathcal{G}(t)$ for all subsequent $t$ are same as those in section [E.7.2](#A5.SS7.SSS2), thus we omit them in this section.

###### Claim E.38 .

$\mathcal{H}(t)\Longrightarrow\mathcal{D}(t),\mathcal{F}(t)$.

###### Proof of Claim E.38 .

We have

$$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\leavevmode\nobreak\ s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle+s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex199)),
the second step is by $\mathcal{H}(t)$ and

$$ $\displaystyle\leavevmode\nobreak\ \langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\leavevmode\nobreak\ \sum_{i\in[n]}(-l_{i}^{\prime(t)})\cdot\left(-s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle-s_{i,21}^{(t)}s_{i,22}^{(t)}\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\leavevmode\nobreak\ \sum_{i\in[n]}(-l_{i}^{\prime(t)})\cdot(-s_{i,11}^{(t)}s_{i,12}^{(t)}\langle\boldsymbol{\mathbf{v}}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle)$ $\displaystyle=_{\operatorname{sgn}}$ $\displaystyle\leavevmode\nobreak\ -\langle\boldsymbol{\mathbf{w}}^{(t)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex201)),
the second step is by $\mathcal{H}(t)$.
∎

The next lemma gives an one side bound for unaligned key noise for $t\geq(1+\theta_{c})T_{4}^{-}$.

###### Lemma E.39 .

Let $t\geq(1+\theta_{c})T_{4}^{-}$, suppose $\mathcal{D}(t)$ holds for $t=(1+\theta_{c})T_{4}^{-},\dots,\tilde{t}$,
then for all $s\in[m_{k}]$, $i\in[n]/E_{s}^{(T_{3})}$ we have

$$ $\displaystyle\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\cdot\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}\rangle\leq(2(1+\theta_{c})T_{4}^{-}-t)\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1+\tilde{O}(s^{-1/2})).$ $$

###### Proof.

Since $\sum_{i=1}^{n}\langle\boldsymbol{\mathbf{w}}^{(T_{2})}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle>0$,
we have $\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle<0$ for $s\in[m_{k}]$, $i\in[n]/E_{s}^{(T_{3})}$.
We can show

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}\rangle$ $\displaystyle=\langle\boldsymbol{\mathbf{w}}^{(1+\theta_{c})T_{4}^{-}}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}\rangle+(t-(1+\theta_{c})T_{4}^{-})\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}$ $\displaystyle\geq\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}\rangle+(t-(1+\theta_{c})T_{4}^{-}-\theta_{c}T_{4}^{-})\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}$ $\displaystyle=-T_{4}^{-}\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1\pm\tilde{O}(s^{-1/2}))+(t-(1+\theta_{c})T_{4}^{-}-\theta_{c}T_{4}^{-})\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}$ $\displaystyle\geq-(2(1+\theta_{c})T_{4}^{-}-t)\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1+\tilde{O}(s^{-1/2})),$ $$

where the first step is by $\mathcal{D}(t)$ for $t\in[(1+\theta_{c})T_{4}^{-},\tilde{t}]$,
the second step is by Lemma [E.29](#A5.Thmtheorem29),
the third step is by the argument in Remark [E.28](#A5.Thmtheorem28),
the last step is by Lemma [C.4](#A3.Thmtheorem4).
∎

Now we only need to and are ready to show $\mathcal{H}(t)$ holds for $t=(2-\theta_{c})T_{4}^{-}+1,\dots,T_{4}$.
Suppose $\mathcal{A}(t^{\prime}),\dots,\mathcal{H}(t^{\prime})$ hold for $t^{\prime}\leq t$, we want to show $\mathcal{H}(t+1)$ holds.
We split to two stages to prove where the first one is $[(2-\theta_{c})T_{4}^{-},(2+3\theta_{c})T_{4}^{-}]$. In this stage, we can bound the magnitude of unaligned key noise $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ by Lemma [E.29](#A5.Thmtheorem29) and [E.39](#A5.Thmtheorem39).
Note that

$$ $\displaystyle\leavevmode\nobreak\ \sum_{s\in E_{i}^{(T_{3})}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \sum_{s\in E_{i}^{(T_{3})}}(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1})(\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1})-\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \sum_{s\in E_{i}^{(T_{3})}}\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|)$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \left|E_{i}^{(T_{3})}\right|(\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+2t\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}\eta^{2}(1\pm o(1))),$ (46) $$

where the last step is by $\mathcal{G}(t)$,
and

$$ $\displaystyle\leavevmode\nobreak\ \sum_{s\in[n]/E_{i}^{(T_{3})}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \sum_{t\geq T_{4,s,i}^{Q,\text{flip}}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle+\sum_{t<T_{4,s,i}^{Q,\text{flip}}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \sum_{t<T_{4,s,i}^{Q,\text{flip}},\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle>0}-\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}-\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle)$ $\displaystyle\leavevmode\nobreak\ +\sum_{t<T_{4,s,i}^{Q,\text{flip}},\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle<0}-\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}-\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle)$ $\displaystyle\leavevmode\nobreak\ +\sum_{t\geq T_{4,s,i}^{Q,\text{flip}},\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle>0}\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle)$ $\displaystyle\leavevmode\nobreak\ +\sum_{t\geq T_{4,s,i}^{Q,\text{flip}},\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle<0}\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}-\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle+\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle)$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ \sum_{s\in[n]/E_{i}^{(T_{3})}}-\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}-\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|).$ $$

Furthermore, when $(2-\theta_{c})T_{4}^{-}\leq t\leq(2+3\theta_{c})T_{4}^{-}$, for $s\in[m_{k}]$ and $i\in[n]/E_{s}^{(T_{3})}$, we have

$$ $\displaystyle-3\theta_{c}T_{4}^{-}\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1+\tilde{O}(s^{-1/2}))\leq\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}^{(T_{4}^{-})}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}\rangle)\cdot\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}\rangle\leq 3\theta_{c}T_{4}^{-}\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1+\tilde{O}(s^{-1/2})),$ (47) $$

where the upper bound is by Lemma [E.39](#A5.Thmtheorem39), and the lower bound is by Lemma [E.29](#A5.Thmtheorem29).
Therefore, for $(2-\theta_{c})T_{4}^{-}t\leq(2+3\theta_{c})T_{4}^{-}$ we have

$$ $\displaystyle\leavevmode\nobreak\ \sum_{s\in[n]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ \sum_{s\in E_{i}^{(T_{3})}}\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|)$ $\displaystyle\leavevmode\nobreak\ -\sum_{s\in[n]/E_{i}^{(T_{3})}}\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|)$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \left|E_{i}^{(T_{3})}\right|(\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+2t\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}\eta^{2}(1\pm o(1)))-(m_{k}-\left|E_{i}^{(T_{3})}\right|)(\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+(t+3\theta_{c}T_{4}^{-})\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}\eta^{2}(1\pm o(1)))$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ (\frac{m_{k}}{2}-O(\sqrt{m_{k}}))(\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+2t\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}\eta^{2}(1\pm o(1)))-(\frac{m_{k}}{2}+O(\sqrt{m_{k}}))(\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+(t+3\theta_{c}T_{4}^{-})\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}\eta^{2}(1\pm o(1)))$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ -O(\sqrt{m_{k}})\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+(\frac{m_{k}}{2}(t-3\theta_{c}T_{4}^{-})-O(\sqrt{m_{k}})(t+\theta_{c}T_{4}^{-}))\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}(1\pm o(1)))>0,$ $$

where the second step is by Eq. ([E.7.3](#A5.Ex515)) for aligned part,
Lemma [E.29](#A5.Thmtheorem29) for upper bound of unaligned query noise and
Eq. ([47](#A5.E47)) for upper bound of unaligned key noise,
the third step is by concentration for $\left|E_{i}^{(T_{3})}\right|$ in Lemma [E.27](#A5.Thmtheorem27).
Combined with $\mathcal{H}(t^{\prime})$ for $(1+\theta_{c})T_{4}^{-}\leq t^{\prime}\leq t$, we have

$$ $\displaystyle s_{i,21}^{(t+1)}s_{i,22}^{(t+1)}\left|\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|$ $\displaystyle\leq s_{i,21}^{((1+\theta)T_{4}^{-})}\left|\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|$ $\displaystyle\leq\frac{1}{9}(1+t)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\leq\frac{1}{2}s_{i,11}^{(t+1)}s_{i,12}^{(t+1)}\left|\langle\boldsymbol{\mathbf{w}}^{(t+1)}_{Q,s},\boldsymbol{\mathbf{\mu}}\rangle\right|,$ $$

where the first step is by induction, the second step is similar to the proof in Lemma [E.36](#A5.Thmtheorem36),
which implies $\mathcal{H}(t+1)$ holds.

Note that now we have $\mathcal{D}(t^{\prime})$, $\mathcal{F}(t^{\prime})$ hold for $t=(2-\theta_{c})T_{4}^{-},\dots,(2+3\theta_{c})T_{4}^{-}$. Then
for $t\geq(2+3\theta_{c})T_{4}^{-}$,
suppose $\mathcal{F}(t)$ and $\mathcal{D}(t)$ hold for $t^{\prime}\leq t$,
for all $s\in[m_{k}]$ and $i\in[n]/E_{s}^{(T_{3})}$, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)}-\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},\boldsymbol{\mathbf{\mu}}\rangle+\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq\langle\boldsymbol{\mathbf{w}}^{(t)}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq(t-2(1+\theta_{c})T_{4}^{-})\eta\sqrt{\frac{2}{\pi}}\sigma_{p}s(1+\tilde{O}(s^{-1/2}))>0,$ $$

where the first step is by Eq. ([• ‣ E.3](#A5.Ex196)),
the second step is by $\mathcal{F}(t)$,
the last step is by Lemma [E.39](#A5.Thmtheorem39).
Combined with $\mathcal{D}(t)$, when $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle<0$ (note that we suppose $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T_{3})},\boldsymbol{\mathbf{\mu}}\rangle>0$, in more general case this is $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle$), this gives

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|=\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|.$ (48) $$

Note that at $t=(2+3\theta_{c})T_{4}^{-}$, we have

$$ $\displaystyle\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{((2+3\theta_{c})T_{4}^{-})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{((2+3\theta_{c})T_{4}^{-})},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|\leq(2+6\theta_{c})T_{4}^{-}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}\eta^{2}(1\pm o(1)),$ (49) $$

where the upper bound for query noise is by Lemma [E.29](#A5.Thmtheorem29) and the upper bound for key noise is by Lemma [E.39](#A5.Thmtheorem39).
Then we have

$$ $\displaystyle\leavevmode\nobreak\ \sum_{s\in[n]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},\boldsymbol{\mathbf{\xi}}_{i}\rangle-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ \sum_{s\in E_{i}^{(T_{3})}}\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|)$ $\displaystyle\leavevmode\nobreak\ -\sum_{s\in[n]/E_{i}^{(T_{3})},\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle}\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}(\left|\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|+\left|\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\xi}}_{i}\rangle\right|)$ $\displaystyle=$ $\displaystyle\leavevmode\nobreak\ \left|E_{i}^{(T_{3})}\right|(\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+2t\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}\eta^{2}(1\pm o(1)))-(m_{k}-\left|E_{i}^{(T_{3})}\right|)(\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+(2+6\theta_{c})T_{4}^{-}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}\eta^{2}(1\pm o(1)))$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ (\frac{m_{k}}{2}-O(\sqrt{m_{k}}))(\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+2t\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}\eta^{2}(1\pm o(1)))-(\frac{m_{k}}{2}+O(\sqrt{m_{k}}))(\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+(2+6\theta_{c})T_{4}^{-}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}\eta^{2}(1\pm o(1)))$ $\displaystyle\geq$ $\displaystyle\leavevmode\nobreak\ -O(\sqrt{m_{k}})\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}+(m_{k}(t-(1+3\theta_{c})T_{4}^{-})-O(\sqrt{m_{k}})(t+(1+3\theta_{c})T_{4}^{-}))\eta^{2}\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1}^{2}(1\pm o(1)))>0,$ $$

where the second step is by Eq. ([E.7.3](#A5.Ex515)) for aligned part,
Eq. ([48](#A5.E48)) and ([49](#A5.E49)) for unaligned part,
the third step is by concentration for $\left|E_{i}^{(T_{3})}\right|$ in Lemma [E.27](#A5.Thmtheorem27).
Combined with $\mathcal{H}(t)$, this implies $\mathcal{H}(t+1)$ holds.
Then we have $\mathcal{H}(t)$ holds for all $(1+\theta_{c})T_{4}^{-}\leq t\leq T_{4}$.

###### Remark E.40 .

When $t=(4+7\theta_{c})T_{4}^{-}$, we have that for all $s\in[m_{k}]$ and $i\in[n]$

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle.$ $$

Finally, for all $s\in[m_{k}]$, $i\in[n]$ and $t\geq T_{4}$, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\xi}}_{i}\right\rVert_{1},$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t+1)},\boldsymbol{\mathbf{\mu}}\rangle=\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle+\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle)\eta\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert,$ $\displaystyle\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle=_{\operatorname{sgn}}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle=_{\operatorname{sgn}}-\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle.$ $$

###### Remark E.37 .

###### Claim E.38 .

###### Proof of Claim E.38 .

###### Lemma E.39 .

###### Proof.

###### Remark E.40 .

### E.8 Proof of Theorem 3.2

In this section, we aim to prove the Theorem [3.2](#S3.Thmtheorem2).

###### Proof of Theorem 3.2 .

By Lemma [E.41](#A5.Thmtheorem41) and [E.42](#A5.Thmtheorem42), we complete the proof.
∎

###### Lemma E.41 (Training and Logistic Test Loss) .

For any fixed $\epsilon>0$,
let $T=2\log(\epsilon^{-1})\eta^{-1}\sigma_{p}^{-1}s^{-1}$, we have $L_{S}(\boldsymbol{\mathbf{W}}^{(T)})\leq\epsilon$ and $L_{\mathcal{D}}(\boldsymbol{\mathbf{W}}^{(T)})\geq 0.1$.

###### Proof.

Note that

$$ $\displaystyle\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle$ $\displaystyle=\langle\boldsymbol{\mathbf{v}}^{(0)},\boldsymbol{\mathbf{\mu}}\rangle+\eta T\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert$ $\displaystyle=\pm\tilde{O}(\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert m_{v}^{-1/2})+2\log(\epsilon^{-1})\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\sigma_{p}^{-1}s^{-1}$ $\displaystyle=\Theta(\log(\epsilon^{-1})\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\sigma_{p}^{-1}s^{-1}),$ $$

where the last step is by $\tilde{O}(\sigma_{0}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert m_{v}^{-1/2})=o(\eta T\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)$.
Then, for all $i\in[n]$, we have

$$ $\displaystyle y_{i}f(\boldsymbol{\mathbf{W}}^{(T)},\boldsymbol{\mathbf{X}}_{i})$ $\displaystyle=(s_{i,11}^{(T)}+s_{i,21}^{(T)})\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle+(s_{i,12}^{(T)}+s_{i,22}^{(T)})\langle\boldsymbol{\mathbf{v}}^{(T)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq(s_{i,12}^{(T)}+s_{i,22}^{(T)})\langle\boldsymbol{\mathbf{v}}^{(T)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq 0.9\langle\boldsymbol{\mathbf{v}}^{(T)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ $\displaystyle\geq\log(\epsilon^{-1}),$ $$

where the first step is by the definition of $f$,
the second step is by $s_{i,11}^{(T)}\geq 0$, $s_{i,21}^{(T)}\geq 0$ and $\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle\geq 0$,
the third step is by $s_{i,12}^{(T)}\geq 0$ $s_{i,22}^{(T)}\geq 0.9$,
the last step is by Lemma [E.8](#A5.Thmtheorem8) and the definition of $T$.
This implies $\ell_{i}^{(T)}=\log(1+\exp(-y_{i}f_{i}(\boldsymbol{\mathbf{W}}^{(T)},\boldsymbol{\mathbf{x}}_{i})))\leq\epsilon$
and thus $L_{S}(\boldsymbol{\mathbf{W}}^{(T)})\leq\epsilon$.

On the other hand, for a new data point $\boldsymbol{\mathbf{x}}=(y\boldsymbol{\mathbf{\mu}},\boldsymbol{\mathbf{\xi}})^{\top}$, let event $\mathcal{E}$ be the event that $\boldsymbol{\mathbf{\xi}}$ has disjoint support with $\boldsymbol{\mathbf{\xi}}_{1},\dots,\boldsymbol{\mathbf{\xi}}_{n}$. Then when $\mathcal{E}$ holds, we have

$$ $\displaystyle yf(\boldsymbol{\mathbf{W}}^{(T)},\boldsymbol{\mathbf{X}})$ $\displaystyle=(s_{11}^{(T)}+s_{21}^{(T)})\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle+(s_{12}^{(T)}+s_{22}^{(T)})\langle\boldsymbol{\mathbf{v}}^{(T)},y\boldsymbol{\mathbf{\xi}}\rangle$ $\displaystyle\leq 2\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle+2\langle\boldsymbol{\mathbf{v}}^{(T)},y\boldsymbol{\mathbf{\xi}}\rangle$ $\displaystyle\leq\tilde{O}(\log(\epsilon^{-1})\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\sigma_{p}^{-1}s^{-1}+\sigma_{0}\sigma_{p}s^{1/2}m_{v}^{-1/2})$ $\displaystyle\leq 1,$ $$

where the second step is by $s_{11}^{(T)},s_{12}^{(T)},s_{21}^{(T)},s_{22}^{(T)}\leq 1$,
the last step is by $\sigma_{0}\sigma_{p}s^{1/2}=o(1)$.
And when $\mathcal{E}$ does not hold, we have

$$ $\displaystyle yf(\boldsymbol{\mathbf{W}}^{(T)},\boldsymbol{\mathbf{X}})$ $\displaystyle=(s_{11}^{(T)}+s_{21}^{(T)})\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle+(s_{12}^{(T)}+s_{22}^{(T)})\langle\boldsymbol{\mathbf{v}}^{(T)},y\boldsymbol{\mathbf{\xi}}\rangle$ $\displaystyle\leq 2\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle+2\langle\boldsymbol{\mathbf{v}}^{(T)},y\boldsymbol{\mathbf{\xi}}\rangle$ $\displaystyle\leq O(\log(\epsilon^{-1})\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\sigma_{p}^{-1}s^{-1}+\log(\epsilon^{-1}))$ $\displaystyle\leq O(\log(\epsilon^{-1})).$ $$

Note that $\mathbb{P}[\mathcal{E}]\geq 1-n^{-2}$, which gives

$$ $\displaystyle\mathbb{E}[\ell(yf(\boldsymbol{\mathbf{W}}^{(T)},\boldsymbol{\mathbf{X}}))]$ $\displaystyle=\mathbb{E}[\mathds{1}(\mathcal{E})\ell(yf(\boldsymbol{\mathbf{W}}^{(T)},\boldsymbol{\mathbf{x}}))]+\mathbb{E}[\mathds{1}(\mathcal{E}^{c})\ell(yf(\boldsymbol{\mathbf{W}}^{(T)},\boldsymbol{\mathbf{x}}))]$ $\displaystyle\geq(1-n^{-2})\log(1+e^{-1})\geq 0.1,$ $$

which completes the proof.
∎

###### Lemma E.42 (Attention layer all attends to noise patch) .

Let
$T_{\mathrm{attn}}=\Theta(\log(\sigma_{p}s/\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\eta^{-1}m_{k}^{-1/2}\sigma_{p}^{-1/2}s^{-1/2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{-1/2})$, we have
$s_{i,11}^{(T_{\mathrm{attn}})}=o(1)$, $s_{i,21}^{(T_{\mathrm{attn}})}=o(1)$, for all $i\in[n]$.

###### Proof.

By the previous analysis, we already have $s_{i,21}^{(t)}=o(1)$ for $t\geq T_{4}^{-}$.
For $t\geq T_{\mathrm{attn}}$

$$ $\displaystyle s_{i,11}^{(t)}$ $\displaystyle=\frac{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right)}{\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\right)+\exp\left(\sum_{s\in[m_{k}]}\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(t)},\boldsymbol{\mathbf{\mu}}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s}^{(t)},y_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle\right)}$ $\displaystyle\leq\frac{\exp\left(2m_{k}t^{2}\eta^{2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2}\right)}{\exp\left(2m_{k}t^{2}\eta^{2}\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert^{2}\right)+\exp\left(m_{k}t^{2}\eta^{2}\sigma_{p}s\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert/4\right)}$ $\displaystyle\leq\frac{1}{\exp\left(\Omega(m_{k}t^{2}\eta^{2}\sigma_{p}s\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert)\right)}$ $\displaystyle\leq o(1),$ $$

where the first step is by definition,
the second step is by $\mathcal{H}(t)$ in Stage IV.c, Remark. [E.40](#A5.Thmtheorem40), and the concentration in Lemma [E.27](#A5.Thmtheorem27),
the last step is by $t\geq T_{\mathrm{attn}}$.
∎

###### Proof of Theorem 3.2 .

###### Lemma E.41 (Training and Logistic Test Loss) .

###### Proof.

###### Lemma E.42 (Attention layer all attends to noise patch) .

###### Proof.

## Appendix F Discussion on Extensions

### F.1 0-1 Test Loss

In this section, we talk about the magntiude of 0-1 test loss.
In binary classification,
constant logistic test loss doesn’t necessarily mean constant 0-1 test loss.
And actually, we find it depends on the network initialization parameter $\sigma_{0}$.

The 0-1 Test Loss is defined as

$$ $\displaystyle L_{\mathcal{D}}^{0-1}(\boldsymbol{\mathbf{W}}):=\mathbb{E}[\mathds{1}(yf(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})>0)]=\mathbb{P}[yf(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})>0].$ $$

The value of $yf(\boldsymbol{\mathbf{W}},\boldsymbol{\mathbf{X}})$ depends on two components, the signal part $\langle\boldsymbol{\mathbf{v}},\boldsymbol{\mathbf{\mu}}\rangle$, which continuously increases during training, and the noise part $\langle\boldsymbol{\mathbf{v}},\boldsymbol{\mathbf{\xi}}\rangle$, which is random for the unseen noise $\boldsymbol{\mathbf{\xi}}_{i}$. With probability at least $1-n^{-2}$, the unseen $\boldsymbol{\mathbf{\xi}}$ is disjoint with all training data $\boldsymbol{\mathbf{\xi}}_{i}$.

If $\sigma_{0}$ is small enough, i.e.,
$\sigma_{0}=o(\sigma_{p}^{-2}s^{-3/2}m_{v}^{1/2})$, then the noise part can be ignored:
By the disjoint property, we can use the estimate at initialization for unseen noise.
By Lemma [C.6](#A3.Thmtheorem6), with probability at least $1-\delta$, we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle=\Theta(\log(\epsilon^{-1})\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\sigma_{p}^{-1}s^{-1})\gg O(m_{v}^{-1/2}\sigma_{0}\sigma_{p}s^{1/2})=|\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\xi}}\rangle|.$ $$

This is similar for
$\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle$ v.s. $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T)},\boldsymbol{\mathbf{\xi}}\rangle$, and
$\langle\boldsymbol{\mathbf{w}}_{K,s}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle$ v.s. $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(T)},\boldsymbol{\mathbf{\xi}}\rangle$,
which gives $s_{11}^{(T)},s_{21}^{(T)}=\Omega(1)$. Then, we have

$$ $\displaystyle yf(\boldsymbol{\mathbf{W}}^{(T)},\boldsymbol{\mathbf{X}})$ $\displaystyle=(s_{11}^{(T)}+s_{21}^{(T)})\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle+(s_{12}^{(T)}+s_{22}^{(T)})\langle\boldsymbol{\mathbf{v}}^{(T)},y\boldsymbol{\mathbf{\xi}}\rangle$ $\displaystyle\geq\Theta(\log(\epsilon^{-1})\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\sigma_{p}^{-1}s^{-1})-\tilde{O}(m_{v}^{-1/2}\sigma_{0}\sigma_{p}s^{1/2})$ $\displaystyle>0,$ $$

which implies a small 0-1 loss
$\mathbb{P}[yf(\boldsymbol{\mathbf{W}}^{(T)},\boldsymbol{\mathbf{X}})>0]\geq 1-n^{-O(1)}$.

On the other hand, if $\sigma_{0}$ is large, i.e., $\sigma_{0}=\Omega(\sigma_{p}^{-2}s^{-3/2}m_{v}^{1/2})$ but still satisfies Condition 4.1, then
we have

$$ $\displaystyle\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle=\Theta(\log(\epsilon^{-1})\left\lVert\boldsymbol{\mathbf{\mu}}\right\rVert\sigma_{p}^{-1}s^{-1})\ll\Omega(m_{v}^{-1/2}\sigma_{0}\sigma_{p}s^{1/2})=|\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\xi}}\rangle|.$ $$

This is similar for
$\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle$ v.s. $\langle\boldsymbol{\mathbf{w}}_{Q,s}^{(T)},\boldsymbol{\mathbf{\xi}}\rangle$, and
$\langle\boldsymbol{\mathbf{w}}_{K,s}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle$ v.s. $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(T)},\boldsymbol{\mathbf{\xi}}\rangle$,
which gives $s_{11}^{(T)},s_{21}^{(T)}=1/2\pm o(1)$ since the symmetry of unseen noise. This means that the softmax outputs are random. Then,

$$ $\displaystyle yf(\boldsymbol{\mathbf{W}}^{(T)},\boldsymbol{\mathbf{X}})$ $\displaystyle=(s_{11}^{(T)}+s_{21}^{(T)})\langle\boldsymbol{\mathbf{v}}^{(T)},\boldsymbol{\mathbf{\mu}}\rangle+(s_{12}^{(T)}+s_{22}^{(T)})\langle\boldsymbol{\mathbf{v}}^{(T)},y\boldsymbol{\mathbf{\xi}}\rangle$ $\displaystyle=\Omega(1)\langle\boldsymbol{\mathbf{v}}^{(T)},y\boldsymbol{\mathbf{\xi}}\rangle,$ $$

which implies $|\mathbb{P}[yf(\boldsymbol{\mathbf{W}}^{(T)},\boldsymbol{\mathbf{X}})>0]-1/2|\leq 1-n^{-O(1)}$.
Note that the condition regarding $\sigma_{0}$ here is not included in our Condition 4.1.

### F.2 Extension to Longer Contexts

In this section, we talk about how to extend our current theory to longer context lengths with additional assumptions.

We consider two case of extension:

Noise patches are the same within one sample.
Note that many works use similar settings, e.g., Tarzanagh et al. (); Vasudeva et al. (); Sheen et al. () assume all non-optimal tokens are the same in some way.
If the data $\boldsymbol{\mathbf{X}}$ has $L$ patches $\boldsymbol{\mathbf{X}}=[\boldsymbol{\mathbf{x}}^{(1)},\dots,\boldsymbol{\mathbf{x}}^{(L)}]$ where the first $L/2$ patches are signal vectors $\boldsymbol{\mathbf{\mu}}$ and the latter $L/2$ patches are noise vectors $\boldsymbol{\mathbf{\xi}}^{(l)}$, and all noise patches are the same, i.e., $\boldsymbol{\mathbf{\xi}}^{(l)}=\boldsymbol{\mathbf{\xi}}^{(l^{\prime})}$ for $L/2<l,l^{\prime}\leq L$, then the training dynamics are the same as in the $L=2$ setting in the paper. If we change the ratio of number of signal patches to noise patches, e.g., 1 signal patch and $L-1$ noise patches, it remains essentially the same except that the signal/noise-signal softmax is around $1/L$ while (the sum of) signal/noise-noise softmax is around $(L-1)/L$ at initialization. Note that in this case the softmax of all noise patches have the same value (no single patch becomes dominant).

Noise patches are different within one sample.
We remark that WLOG we can consider $1$ $\boldsymbol{\mathbf{\mu}}$ and $(L-1)$ $\boldsymbol{\mathbf{\xi}}$ vectors. With high probability, the $n(L-1)$ noise patches are disjoint. The main difficulty here is the competition between noise patches, which can be seen in the gradient and increase speed. To address or avoid this issue, one solution is to assume some form of sparsity in the dimension of context length. For example, Tarzanagh et al. (); Vasudeva et al. (); Sheen et al. () assume there is only one optimal token in one sample, which has a large gap from the non-optimal ones.
Jiang et al. () assumes the second noise vector $\boldsymbol{\mathbf{\xi}}_{2}$ is greater than remain noise $\boldsymbol{\mathbf{\xi}}_{3},\dots,\boldsymbol{\mathbf{\xi}}_{n}$ in one data point.
In our case, we need to assume for each sample there is a noise patch $L$ (WLOG, we assume it is the last patch) such that $\sum_{l^{\prime}=2}^{L-1}\|\boldsymbol{\mathbf{\xi}}^{(l^{\prime})}\|_{1}=o(\|\boldsymbol{\mathbf{\xi}}^{(l^{\prime})}\|_{L})$. We can achieve this by applying large sparsity $s$ and/or large std $\sigma_{p}$ on this patch. Note that this large patch across different samples can be in different positions. We provide a detailed example here: First, we still have stage I since the gradient of value does not change. The gradient of query is

$$ $\displaystyle\nabla_{\boldsymbol{\mathbf{w}}_{Q,s}}L_{S}({\bf W})=\frac{1}{n}\sum_{i=1}^{n}\ell_{i}^{\prime}y_{i}\sum_{l=1}^{L}{\bf x}_{i}^{(l)}\sum_{a<b}s_{i,la}s_{i,lb}\langle\bar{\boldsymbol{\mathbf{w}}}_{V,j},{\bf x}_{i}^{(a)}-{\bf x}_{i}^{(b)}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s},{\bf x}_{i}^{(a)}-{\bf x}_{i}^{(b)}\rangle,$ $$

and the inner product is

$$ $\displaystyle\langle\operatorname{sgn}(\nabla_{\boldsymbol{\mathbf{w}}_{Q,s}}L_{S}),y_{i}\boldsymbol{\mathbf{\xi}}_{i}^{(l)}\rangle$ $\displaystyle=-\|\boldsymbol{\mathbf{\xi}}_{i}^{(l)}\|_{1}\operatorname{sgn}(\sum_{a<b}s_{i,la}s_{i,lb}\langle\boldsymbol{\mathbf{v}},{\bf x}_{i}^{(a)}-{\bf x}_{i}^{(b)}\rangle\langle\boldsymbol{\mathbf{w}}_{K,s},{\bf x}_{i}^{(a)}-{\bf x}_{i}^{(b)}\rangle)$ $\displaystyle\approx-\|\boldsymbol{\mathbf{\xi}}_{i}^{(l)}\|_{1}\operatorname{sgn}(\langle\boldsymbol{\mathbf{w}}_{K,s},y_{i}\boldsymbol{\mathbf{\xi}}_{i}^{(L)}-\sum_{a<L}{\bf x}_{i}^{(a)}\rangle).$ $$

We approximate the inner sum with only one term by 1) concentration of softmax to $1/L$, 2) the assumption $\sum_{l^{\prime}=2}^{L-1}\|\boldsymbol{\mathbf{\xi}}^{(l^{\prime})}\|_{1}=o(\|\boldsymbol{\mathbf{\xi}}^{(l^{\prime})}\|_{L})$, which implies the dominance of $\langle\boldsymbol{\mathbf{v}},\boldsymbol{\mathbf{\xi}}_{i}^{(L)}\rangle$ after stage I, and the dominance of $\langle\boldsymbol{\mathbf{w}}_{K,s}^{(0)},\boldsymbol{\mathbf{\xi}}_{i}^{(L)}\rangle$. The approximation for the update of key noise is essentially the same, thus ensuring the alignment between the largest query and the largest key, thus extending our result for $L=2$. Subsequently, due to the gap in the $L_{1}$ norm, the softmax at the position of the largest noise patch will converge to 1, while the others will go to zero.

### F.3 Extension to Joint Training of Linear Layer

In this section, we talk about how to extend our main theoretical results, which fix the second linear layer, to the joint training of softmax attention layer and linear layer.

Firstly, we would like to remark that the value matrix in our model setting already acts as a linear layer while many works Tarzanagh et al. (; ) on analyzing training dynamics of transformers fix the value matrix.
In our specific case, we can extend the current theory to the joint training of the softmax attention layer and linear layer.

We can write the gradient with respect to the parameters in the second layer as follows:

$$ $\displaystyle\nabla_{\theta_{j,r}}L_{S}(W)=\frac{1}{n}\sum_{i=1}^{n}y_{i}\ell_{i}^{\prime}j[(s_{i,11}+s_{i,21})\langle\boldsymbol{\mathbf{w}}_{V,j,r},y_{i}\boldsymbol{\mathbf{\mu}}\rangle+(s_{i,12}+s_{i,22})\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{\xi}}_{i}\rangle],$ $$

where $\theta_{j,r}$ denotes the parameter for the second layer $F_{j}$ at the $r$-th entry, initialized with $j/m_{v}$, for $1\leq r\leq m_{v}$. During stage I, it moves at most $T_{1}\eta$ which is negligible compared to the initialization $1/m_{v}$, and will not affect our analysis for value in stage I. After stage I, the terms $\langle\boldsymbol{\mathbf{w}}_{V,j,r},jy_{i}\boldsymbol{\mathbf{\xi}}_{i}\rangle$ become positive for all $1\leq i\leq n$ and larger than the value signal term $\langle\boldsymbol{\mathbf{w}}_{V,j,r},\boldsymbol{\mathbf{\mu}}\rangle$, which leads to that
$\operatorname{sgn}(\nabla_{\theta_{j,r}}L_{S}(\boldsymbol{\mathbf{W}}))=-j$. Consequently, all $\theta_{1,r}$ keep increasing while all $\theta_{-1,r}$ keep decreasing.
Until then, all $\theta_{j,r}$’s sign are the same as initialization and will not change during the rest of training.
Therefore, it will not affect the analysis for other parameters.

In summary,
the key to the joint training extension lies in the different magnitudes of initialization between the second layer weights and the value matrix, which is $1/m_{v}$ and $\sigma_{0}$, respectively. When value converges to pattern specified in stage I, the second layer weights have almost no change.
But after stage I, the sign of gradient of second layer become fixed and aligned with its initialized sign, so that it will have no effect on other parameters.

## Peer Reviews

confidence: 2
contribution: 3
initial_rating: 8
presentation: 2
questions: 1. In Figure 2(d), the test loss for the less noisy case (SNR=1.0) is higher than for the noisier case (SNR=0.5) when using SignGD, which contradicts theoretical expectations about the benefits of higher-quality data. Could you provide an analysis of the noise characteristics or data distribution across different SNR levels that might explain this behavior? Additional experiments or analyses could also be valuable in investigating this phenomenon further. For instance, it’s possible that entirely clean data may reduce generalization performance compared to slightly noisy data, though too much noise also significantly hinders generalization.

2. Clarity and Detail in Training Dynamics Analysis (Section 3.1): Could you offer a clearer mapping between training steps and each stage of the training dynamics? More detailed explanations of the transitions between each of the four stages would be helpful. Additionally, including data plots that capture the state of relevant variables before and after critical transition points would improve clarity (rather than focusing solely on within-stage details, like in Figure 1(a)).

3. Real-world applications of Transformers, such as language modeling, demonstrate robustness to noise. How do the noise levels assumed in your theoretical model compare to those typically encountered in practical Transformer applications? Could you discuss any factors or experiments that might explain the gap between your theoretical findings and the observed robustness of Transformers to noise in real-world scenarios?
soundness: 3
strengths: 1. The paper offers a comprehensive four-stage analysis of transformer training dynamics, which is a promising contribution to understanding the behavior of transformers. 

2. The demonstration of fast convergence but poor generalization with SignGD on noisy data is important. These findings are beneficial for practical applications as they guide the selection of optimizers and the preparation of datasets in real-world tasks.

3. The theoretical approach is novel.
summary: This paper studies the learning dynamics of two-layer transformers optimized with SignGD. It introduces a theoretical framework that categorizes the learning process into four distinct stages, providing an analytical tool to study these dynamics. The study demonstrates that while SignGD achieves fast convergence of the transformer model, it leads to poor generalization on noisy datasets.
weaknesses: 1) The paper claims that both SignGD and Adam require high-quality data to perform well. However, an apparent contradiction arises in Figure 2(d), where the model with less noisy data (SNR=1.0) performs worse than that with noisier data (SNR=0.5). This discrepancy challenges the theoretical claims and needs to be addressed or explained to validate the model’s consistency across different noise levels.

2)  The description of the training dynamics in Section 3.1 lacks clarity, making it difficult to follow through the stages. It would be beneficial to explicitly map the training steps of the analysis to each stage in the diagrams (e.g., steps 1-2 corresponding to Stage 1). Additionally, providing data plots before and after switching steps of different stages, such as providing the mean noise values before and after the transition steps for 1st stage and 2nd stage in Fig 1(a), would greatly enhance understanding.

3) In many real-world applications of Adam in transformer, such as those used in vision and language tasks, often exhibit robustness to noise in practical applications, which is contrary to the findings presented. Understanding whether this discrepancy is due to lower real-world noise levels or other factors that contradict the assumptions made in the study is essential for aligning theoretical insights with empirical observations.

---

confidence: 3
contribution: 3
initial_rating: 5
presentation: 3
questions: Can the framework be extended to deeper transformers or multi-head attention mechanisms? I ask that because no method is currently using two layers transformer, so its extension to SOTA networks could provide broader applicability. 

From the current work what strategies could the authors envision can be used to improve generalization of SignGD and consequently Adam in data noisy settings? Identifying such strategies would be valuable for increasing robustness in real-world applications, where noise is often unavoidable.
soundness: 3
strengths: The paper provides a multi-stage framework to analyze the transformer training using SignGD, making this complex behaviour into small interpretable stages.

By establishing that SignGD is a proxy for Adam, the paper is capable of offering new perspectives on the reasons why Adam present some generalization problems.

The combination of theoretical proofs and experimental validation strengthens the proposed analysis and its overall message.
summary: The paper investigates the optimization and generalization properties of Sign Gradient Descent (SignGD) for transformers, focusing on binary classification tasks involving noisy, linearly separable data. The authors claim that SignGD is an effective surrogate for the Adam optimizer. They identify four stages in the training dynamics of SignGD and observe that it achieves fast convergence in training loss but generalizes poorly, especially on noisy data.
weaknesses: The main weakness of the method is its limited applicability to real-world scenarios. The reliance on assumptions such as linearly separable data and the use of only two-layer transformers restricts its effectiveness when dealing with more complex datasets and modern, state-of-the-art transformer architectures.

---

confidence: 4
contribution: 3
initial_rating: 6
presentation: 4
questions: Please see the weaknesses.
soundness: 3
strengths: To the best of my knowledge, this manuscript is the first to theoretically analyze the detailed training dynamics of the transformer and softmax attention layer with trainable query-key parameterization using signSGD. It breaks down the process into four stages, capturing the complex and subtle behaviors, which deepens our theoretical understanding of the optimization process for Transformers.
summary: This manuscript provides a deep theoretical analysis of the training dynamics of a simplified two-layer transformer using signSGD on a synthetic dataset. It explicitly identifies four complex but distinct stages in the training dynamics. Furthermore, it empirically uncovers that Adam exhibits similar behavior under the same conditions as signSGD. The analysis demonstrates that while signSGD typically converges quickly, it generalizes poorly and requires high-quality data compared to SGD.
weaknesses: - Some assumptions in the analysis are overly strong and unrealistic. For example, the manuscript assumes that the data dimension satisfies  $ d = \Omega(\text{poly}(n)) $, and that the variance of noise satisfies $\sigma_p = \Omega(d^{\frac{1}{4}} n^3)$, where $ n $ is the number of data samples, which is typically much larger in practice.

- Although the manuscript attempts to analyze the complex softmax attention with trainable query-key parameterization, it simplifies the attention block to a single-head variant rather than using the original multi-head version. Additionally, the formulation of $ v $ is somewhat unusual, defined as $v = \bar{w} _{v,1} - \bar{w} _{v,-1}$ which undermines the theoretical insights in relation to real-world scenarios. 

- While the empirical results validate the theoretical analyses for training a two-layer transformer using signSGD, it would be valuable to investigate whether the empirical training dynamics of a multi-layer transformer, or even the original transformer, align with the theoretical findings from the simplified two-layer model.  I would like to see whether the results  can extend to deeper transformers, or implement more experiments to test if the key behaviors you identified persist in more complex models.

**Minor Issues**  
- At Line 168, the variable $ i $ in the equation is not explicitly defined, though it can be inferred that $ i$ represents the index of data samples.

- In Figure 2, the manuscript suggests distinct differences between the training dynamics of signSGD and Adam, yet does not provide sufficient explanations. For instance, in (a) and (b), it is noted that the negative term $\langle w_{Q,s}, y_i \epsilon_{i} \rangle$ for signSGD ultimately becomes positive, while for Adam, it remains consistently negative. In (c), signSGD converges at a linear rate, whereas Adam converges at a sublinear rate. I think it needs to provide more detailed explanations for these differences, and to discuss their implications for understanding Adam's behavior in practice.

Meta-review: This work presents a theoretical analysis of the training dynamics of a simplified two-layer transformer using sign gradient descent on a synthetic dataset. They identify four stages in the training dynamics observe that it achieves fast convergence in training loss but generalizes poorly. The paper received a mix of reviews, but some reviewers raised their score after rebuttal. I recommended an acceptance.

Decision: Accept (Spotlight)