On the Optimization and Generalization of Two-layer Transformers with Sign Gradient Descent
Authors: Bingrui Li
Wei Huang
Andi Han
Zhanpeng Zhou
Taiji Suzuki
Jun Zhu
Jianfei Chen
Venue: ICLR 2025 Spotlight (2025)
Keywords: Sign Gradient Descent; Transformer; Training Dynamics; Theory
Primary area: learning theory
TL;DR: We study the optimization dynamics and generalization properties of a two-layer transformer trained on a signal-noise dataset.

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.

Published as a conference paper at ICLR 2025
ON THE OPTIMIZATION AND GENERALIZATION OF TWO-LAYER TRANSFORMERS WITH SIGN GRADIENT DESCENT
Bingrui Li', Wei Huang", Andi Han", Zhanpeng Zhou', Taiji Suzuki Jun Zhu', Jianfei Chen 'Dept. of Comp. Sci. and Tech., Institute for AI, BNRist Center, THBI Lab, Tsinghua-Boss Joint MI. Center, Tsinghua University
"RIKEN AIP University of Tokyo 'Shanghai Jiao Tong University
1br22@mai 1s. ts sainghua.... edu. en; jianfeic@tsinghual edu. on
ABSTRACT
The Adam optimizer is widely used for transformer optimization in practice, which makes understanding the underlying optimization mechanisms important prob- lem. However, due to the Adam's complexity, theoretical analysis of how it optimizes transformers remains 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 softmax attention layer with trainable query-key parameterization followed by linear layer on 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. Addition- ally, we find that the poor generalization 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.
SignGD is not solely due to data noise,
INTRODUCTION
The transformer architecture (Vaswani et al., 2017) has become ubiquitous across various domains, achieving state-of-the-ard results in areas such as language modeling (Devlin et al., 2019; Brown et al., 2020), computer vision (Dosovitskiy et al., 2021; Peebles Xie, 2023), and reinforcement learn- ing (Chen et al., 2021). Regardless of the specific task data modality, the Adam optimizer (Kingma & Ba, 2015) is typically employed to train large transformer models, making it the de facto choice in practice. This widespread highlights that understanding the inner mechanism optimizes transformers an important problem. However, the complexity of Adam's 's formulation presents significant challenges for rigorous analysis. Many of the underlying mechanisms of how Adam optimizes transformers are still poorly understood.
how Adam
Recent theoretical works (Jelassi al., 2022; Tarzanagh al., 2023a; Tian et al., 2023) study the training dynamics of transformers across various datasets and objectives. Training dynamics analysis allows us trace the evolution of model parameters throughout the training process. In doing so, enables 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, 2023; Xu & Du, 2023) and CNNs (Cao et al., 2022; Kou et al., 2023), 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., 2024b) unrealistic initialization (Li et al., 2023c). 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. (2023)). While this assumption simplifies the analysis, it remains

---

Published as a conference paper at ICLR 2025
unrealistic in practice. Moreover, existing analyses focus primarily 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 effective surrogate for understanding Adam. SignGD is 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, 2018; Balles et al., 2020; Bernstein et al., 2018; 2019), and has inspired the development of optimizers like Adam (Kingma & Ba, 2015) and Lion (Chen al., 2023). More importantly, SignGD shares many similarities with Adam, making it an effective proxy for gaining insights into Adam's 's optimization behavior (Balles Hennig, 2018; Bernstein al., 2018; Kunstner et al., 2023; 2024; Wu et al., 2020; Zou et al., 2023). For example, Kunstner et al. (2023) has shown that while 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 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. 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 linear convergence rate in training loss but maintaining high constant test loss, leading to 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
Stage II
The query & key noise align their sign to each other. Majority voting determines the sign of query & key signals.
Lemma 4.2, 4.3
Stage III
Lemma 4.4
The noise-signal softmax outputs decay fast exponentially, then the query & key noise align their sign to signals.
Lemma 4.5
Stage IV
Lemma 4.6, 4.7
Technical novelties. We use the feature learning framework (Allen-Zhu & Li, 2023; Cao et al., 2022; Zou et al., 2023; Huang et al., 2023a) 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 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, 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 SignGID Our findings offer valuable insights into the inner workings of both SignGD and Adam, advancing our theoretical understanding of transformers and their optimization.
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 v = [e]...., va] we denote the le and e, norm by ||v|| and ||v|]]- respectively. For two fixed non-negative sequences {x.)
2

---

Published as a conference paper at ICLR 2025
and {yn }. we denote nn, = (y)) if there exist some absolute constant C' > and N > 0 such that [x,| - C' yal for all n 2 N. We say 2n = (2yn.) ) if yn = ( and say In = ( In = ((n)) and En = (((n). We use O()), (2)) and (-) to hide logarithmic factors in these notations, respectively. Moreover, we denote In poly if In O ) for some constant D > 0, and Zn = polylog(vn) if In = poly(log(иa (log()n)) )) We use [d] to denote the set (1.2...., d}. We use sg)(z)) xr | when æ o and sgn(() O. We denote a n-dim all-ones vector by I,,
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.
Definition 2.1. Ler μ E Rd be fixed vector representing the signal contained in each data point. For each data point (X. y), the predictor X = x E RRdx2
We assume l = [1.0......0)
consists of two patches (or tokens, vectors), where x(1), x(2) TRe, and the label y is binary, i.e., y E {+1]. The data is generated from distribution 'D, which we specify as follows:
1. The label y generated as Rademacher random variable.
2. Randomly select s coordinates [d]; {1} uniformly, denoted as a vector s E 10.. Generate each coordinate in E, from distribution N N(0, of)) and then mask off the first coordinate and other
d - s - 1 coordinates, i.e., =
3. One of x(1), (2) is randomly selected and then assigned as yıs, representing the signal, while the other is designated as &. representing noise.
The design of the signal patch yu and the noise patch can be viewed as 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, 1/1 represents useful, label-correlated features (referred to as signal), whereas & represents non-informative fcatures or irrelevant noise (referred to noise).
Remarks on data assumptions. We make several assumptions regarding sparsity. orthogonality, and context length. Specifically, we assume fu is 1-sparse (i.e., it has only one non-zero entry), is s-sparse, and that 1a and & are orthogonal. The sparsity assumption is essential for analysing optimizers that not invariant under orthogonal transformations, such as Adam and SignGI. Our results can be easily extended to any C-sparse signal vector μu, 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 for details). We also assume context length of L 2 for technical simplification. With additional appropriate assumptions, our analysis can be extended to data with longer contexts (see discussion in Appendix F.2). We empirically validate our theoretical results for non-sparse, non-orthogonal, and multi-patch data in Appendix B. Data models comprised of signal and noise patches with similar assumptions have also been studied in recent works (Allen-Zht & Li, 2023; Cao et al., 2022; Jelassi et al., 2022; Zou et al., 2023; Huang et al., 2023b; Han et al., 2024; Huang et al., 2024a).
Two-layer transformers. Motivated by vision transformers (Dosovitskiy et al., 2021), we consider 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 W := (Wq. q, W K, Wv,), where W Q. W E R xd and W Wv.j R x for j E (=1}. The parameters of the second layer are fixed as 1/m and -1/mg respectively. We also talk about learnable linear head in Appendix F3. Then, the network can be written as f(W,X) == FiW,,X FW where (W,, X) and F-1(W,
(W.X) X) are defined as:
(W X): = my Σ
T my
Wv, X.oftmxx (xT (TWWWxx"))
(,,).)) E Rd WV W Wii E Rd be the s-th or r-th row of the parameter Wq. WK. W..j. respectively. Let Wy. Ere(m.) W....// /mm be the mean value in Fz. Let Wv.1 Wy,-1 1 be the mean value. We can write the model in simpler form:
Let WQ,s W (.).s)) E Rd. WK.a = W Wk....)
(W, X)
Σ my
[sn1 + s21) (wV.mm((")) + (s12 s22) (WV.).x(1)) (1)
rE(ma)
3

---

Published as a conference paper at ICLR 2025
where Sta == softmax (211, .............. 2) and 21a Exe(me) (WQ.sxx (.,,,x(()) Unless other- wise 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 nignel softmax outputs. Formally, for data point (X,, Mk), it is defined as
exp
E Eee(ma)
(w Q.e Ex) (t) K.WKK, K. Vidl)
(0) 1,21
exp [Ee(ma) (wl))
(t) Qia E.) (w Kis vi(1)) + exp (Eeemma) (
E.V K,a E,)
Similarly, we refer to 5,,11 $,,12 $122 as Signal-signal, signal-noise and noise-noise softmax outputs, respectively. When == 2, a key fact about the softmax function is that <( s 1, for I E [2]- The subscript is used only to distinguish between signal and noise, without imposing any restrictions on the permutation of patches. Although use the symbol for sparsity, softmax outputs, and the indices of query and key neurons simultaneously, the context clearly indicates which one is being referred to.
(t) (6) (1)
Training algorithm. We train our transformer model by minimizing the empirical cross-entropy loss function I.g (W: == 1 Ei=r Eyy J(W, X.)). where {(x) == log (1 + exp(-x)) is the logistic loss function, and 5 = {(x,,) (..36)) ())) ;==1 the training dataset. We further define the test loss
LD(W) == E<x,y>~D((( /(W,)))
We study SignGD starting from Gaussian initialization, where each entry of WQ W K. Wy. for j E {+1} is sampled from a Gaussian distribution 5(0,o8). The SignGD update for the parameters can be written as
Widrr w (t)
(t+1)
Vjr.r 17 sgn V Twv L. (s(((())),
(++1) (1)
w Q. sgnn V w... Ls( [s()())), w (++1) (0)
WK,a Wk. sgn V Ww,, Ls( (W))),
Q.
for all [ma)..j E {+1] and E [m.]. The expanded gradient formulas and update rules can be seen in Appendix D.
MAIN RESULTS
In this section, present our main results. Firstly, we provide 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 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 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
[Data dimension and Sparsity] Data dimension sufficiently large with ((poly(n)).
Sparsity satisfies: s (2
2. [Noise strength] The standard variance of noise op satisfies: o, 12(d-1/44pp)
3. [Network width and initialization] Network width of value m. and of query & key satisfy:
my; m. ((pllylog(d)) Network initialization on satisfies: do o 11- 1mm )
4. [Training dataset size] The training sample size satisfies: n = ((m
5. [Learning rate] The learning rate 17 satisfies: 17 = O(poly(d-1)) is sufficiently small.
Remarks on Condition 3.1. Our Condition 3.1 is frequently used the literature and realistic in practice (Cao al., 2022; Chatterji & Long, 2021; Frei et al., 2022). The conditions on d and make sure the different noise patches have disjoint support with high probability. The condition on op implies a v = Q(n [|д])), which indicates the noise in the dataset is strong. The conditions on m.s m&, n, n technical and mild. my and m, affect the convergence of mean value noise and softmax outputs, respectively. The size of n affects the concentration of ||E]] Finally, the conditions on do ensures the network weights are small enough at initialization, which makes the learning process fall into the feature learning regime. Note we set my = (polyyogg(d) (d), then our condition become oo = ((2d) which is realistic in practice.
4

---

Published as a conference paper at ICLR 2025
Stage L, value
Stage key noise
Stage III: query key
Stage query and key
Stage IV: softmax
0.00
0.6
0.02
1.
0.03
0.01
0.02
0.00
0.
0.01
-0.01
0.-
-0.02
0.0
-0.00
100
Stegs
Staps
Staps
Staps (d)
Steps
(8)
(19))))4))
(
(
(1
(
positive (v
positive ( (126))55))
50
(119.p)
negative (2..)))))
negative (18))))))
10
The training dynamics of two-layer transformers with SignGD. (a) Dynamics of mean value noise and mean value signals in Stage I, and II. (b) Dynamics of key noise in Stage I, and
Figure 1:
II. We mark different key noise in different colors. D: w(O Kia E stu)
x+.Q+ == {wh. K.s K.s U.E,) > (w K,et U(E,)
(w(o)
O (w Vi,)) > 0}. 2): w Kia E S SK-Q- = [w (t) K2* w Kia E (SK+ Q+ USS K- )". (c) Dynamics of query noise, key noise, query signals, key signals in Stage II and III. The dotted lines represent positive (query and key) noise at f = 40, and the solid lines represent negative noise at the same point. (d) Dynamics of query noise, key noise, query signals, key signals in Stage III and IV. The dotted lines and solid lines have the same mcanings in (c). (e) Dynamics of softmax outputs in four stages. The dynamics over the whole time horizon is provided in Fig. 18. An illustration explaining the behaviors of all quantities in all stages is provided in Fig. 19.
10)
(0)
0, (w (.) < 0}. 3:
3.1 TRAINING DYNAMICS ANALYSIS: FOUR-STAGE ANALYSIS
Based on Condition 3.1, we aim 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 summarized in Tab. 1. 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.
To better understand the four-stage training dynamics, the dynamics of key quantities key timesteps or during the entire dynamics are illustrated in the Fig. and Tab. 2.
From / 0 to f is the Stage I of the dynamics, which shows the early shift and stabilization of mean value noise, i.e.. (v) y.E.). The mean value noise increases monotonically from the random initialization and becomes positive, then stabilizes into a linear relationship with by the stage's 's end (Fig. (a)). This period is so rapid that other quantities, including value signals and query & key noise all remain close to their initialization (Fig. (b) illustrates key noise as an example). This linear behavior of mean value noise make it ordered in the gradients of query and key noise
From t 2 to = 10 is the Stage II of the dynamics, which illustrates the sign alignment between query & key noise, i.e., ( and (w K. y.E.). 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 continuous to grow subsequently. Additionally, the number of positive and negative neurons is nearly equal. Fig. I (b) shows how key noise aligns their signs. Tab. 2 provides statistics on the signs of query and key noise at initialization and at the end of Stage II (t 10).
From f = 10 to t = 40 is the Stage III of the dynamics, which shows how the sign of query & key signals, i.c. (w)) pe) and (wk)... μ1), are determined by majority voting. Before Stage III, query and key signals remain close to the initialization, and their gradients are disordered. However, at the start of Stage III, the sum of key noise En (w K.s V.C.) 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 this mechanism effectively act majority voting. The sign of the key signals determined symmetrically by the sum of the query noise, with an opposite sign to the query signals. The behaviors in Stage III are shown in Fig. (c).
5

---

Published as a conference paper at ICLR 2025
s(t) Table 2: Sign alignment between query and key noise in Stage II. S++ defined as +K+QQ+ ((s, [mx] [n] : viS..) > 0, (w &. VvE,) > 0}, represents the number of neurons and samples having positive query noise and positive key noise. The definitions for S q $ Sk- Q+
s(t)
s Q- 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, | ms x+ The signs of query and key noise are independent at initialization but aligned at t 10, which can be seen as an estimate of T2
486.
init(t = 0) 1 10
|S
k+++2++
IS
SK.Q--l
((())
K--Q-
Q--1
Row sum
*++Q+
486
1
0
25
512
[9k K*.Q-
244
4
9
250
507
| 2+
223
10
4
221
458
37
2
3
481
523
Column sum
990
17
16
977
2000
From & = 40 to f 2000 is the Stage IV of the dynamies. During this stage, noise-signal softmax outputs, i.e. s1,21 decay exponentially fast in Stage IV. The duration from Stage to Stage IIII can be relatively short, with all softmax outputs (t) + concentrated at 1/2. However, in Stage IV, $1.21 decreases exponentially to zero, while = remains stuck at 1/2. The concentration and fast decay are shown Fig. (e). The difference between ( and 1,21 due to the varying rates of increase between query signals and query noise.
$1,11.
Furthermore, "negative"" query and key noise align their sign to signals in Stage IV. We focus on single neuron with a positive query signal, shown in Fig. (d), though this applies to all neurons. Before the noise-signal softmax outputs decay to zero, the dynamics all signals and noise remain unchanged. At a critical point (1 = 150 in Fig. I (d)). 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 (around t = 300 in Fig.1 (d)). 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. The sign alignment of negative key and query noise is shown in Fig. (d).
Overall, the dynamics exhibit complex and intriguing sign alignment behaviors within the attention layer, and they ultimately reach stabilization.
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 >> under Condition 3.1, with probability at least n -1/3, there exists 1/2 -1/2 s -1/2 |||||| -1/2) such that "m-1/"
= Olog(c " 1)) and Tatta = Ö(1 'm, Op
1. [Traiiiin loss] The training loss converges
Ls (T() S..
[Test loss] The trained transformer has constant order test loss: ( 0(1).
3. [Noise memorization of value] The value matrix in attention layer memorizes noises in the training data: For all E [n], (T)(T)) E.). Q(1), ((v(T), 11] 0(0='s-1)
4. [Noise memorization of query keyl The softmax outputs of attention layer attends all the (Toaa) (Tauu) = o(1), 8,21 = o(1).
weights to the noise patch in the training data: For all E [n]] 8
Theorem 3.2 outlines the training and test loss 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 also presents new results on noise
"To be more accurate, the "negative" key noise means all the key noise with the sign opposite of query signals. When focus neuron with positive query signal, the "negative" key noise is exactly negative key noise.

---

Published as a conference paper at ICLR 2025
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., 2023; 2024; Li et al., 2023c). In contrast to prior works (Cao et al., 2022; Zou et al., 2023) 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 is in Appendix E.8.
Remark. Theorem 3.2 focus on the logistic test loss. We further give a result about final 0-1 test loss in Appendix F.1 since bad logistic loss doesn't 't necessarily imply bad 0-1 loss in binary classification task. Interestingly, the size of 0-1 test loss depends on the network initialization On-
Fast convergence but poor generalization contradeits the 'train faster, generalize better' ar- gument. Algorithmic stability (Bousquet & Elisseeff, 2002; Hardt et al., 2016) is a widely used technique in generalization analysis. One typical argument of algorithmic stability is "train faster, generalize better" (Hardt et al., 2016). Our results provide a counterexample to this argument by showing SignGD trains fast but generalizes poorly. Notably, Teng et al. (2023) introduced 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 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 (b). Further similarities in other stages, and the results across different 31 values are given in Appendix B.3. 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 when 31 and (32 in Adam are close to 1, which is commonly used in practice, SignGD does not always behave like Adam. In our experiments, 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 in Adam.
proxy for understanding Adam. However,
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., 2022; Kou et al., 2023) or transformers for different tasks (Nichani et al., 2024; Huang et al., 2024b). Empirically, our experiments demonstrate that SignGD trains faster than GD (Fig. 2 (c)), but GD generalizes better consistently across different levels of data noise (Fig. 2 (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 SignGID, 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.
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 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.
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, 2023; Cao et al., 2022), 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 (WQ... p4). (WQ..r J.C.), (WK.., p)) (WK.,, E.).). (Wv.j.r. p4). (wv. E,E,,, and 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 (v. ((v, ((,)). Conditional these quantities, the output f(W, X,) is independent of parameters W.
7

---

Published as a conference paper at ICLR 2025
Training dynamics w/ SignGD
Training dynamics wi Adam
Training loss on synthetic data
Test loss on noisy MNIST
10
2.00
1.23
1.00
10
0.75
1.25
0.50
0.75
0.0
0.50
-0.25
0:25
500
1000
1500
2000
1900
1500
2000
300
1500
2000
0.6
1.0
Steps
Steps
Steps
Scaled SNR
(b)
(c)
(d)
positive (16)) positive (.g.)118)
(M) H)
SignGD
Adami -0)
negative ((i)
negutive (16.))15))
(Mox)))
Adamif; -0.5)
Adamxp -0.9)
Figure 2:
Comparison of SignGD with Adam and GD on synthetic and real-world datasets. (a)
Dynamics of query noise, key noise, query signals, and key signals with SignGD on the synthetic dataset. (b) Dynamics of the same quantities with Adam((; = 0.9). (c) Training loss curve (log scale) on the synthetic data for different optimizers. The training loss with SignGI decays exponentially. Note that the training losses for Adam(3 = 0.9), Adam(3; = 0.5), Adam(3 = 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.
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 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.
Secondly, we perform a multi-stage analysis for transformers by breaking down the complex dy- namies 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 Tp. To. 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 j, r, and i, with high probability, we have:
(w (++1) (e)
W.s. u0)
(w (e+1) (e) K.s Kist 10)/ //w (((tt)) (t) w 12)
17 |Aall
((+1) (w ViJ
W.S. Mr
(1+1)
K. vrE.) (w (w@
(0+1) w. vrE.) = |1E.l1
K.s
This property implies that each iteration's 's update is determined solely by its sign, allowing us to simplify the analysis by focusing only on the update direction's '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. (v and (v p). Let ((
(0) 3e = max a. E.Y. (w K.. E.)l, ((
(0) E.)l and T3 := 43em. 1/2 " oy : V.j.r:
s-1
102mm 1/2), we call E [0, T2] Stage I. The following lemma is the main result in Stage I.
Lemma 4.1 (Stage D) We have (I) (Magnitude). For all i [n] and t 2 Tı, (v(t), ,1E,) @(tnops).
(2) (Negligibility of mean value signals). For all i E [n] and 2 Tı, (v(), (1) o((v(", (.Ed))
(3) (Query & Key noise barely move). For all E [ma], i E [m] and 1 S Ti, w. E.) (0)
(w]0) Q.s: Ex) (1 do o(1)), (w19, K...
(...::E) E.) +1+ o()).
Lemma 4.1 (1) shows the mean value noise quickly stabilize into linear relationship with t.
Lemma 4.1 (2) allows to disregard the influence of the mean value signals on the query and key gradients, as is negligible compared to the impact of the mean value noise. Lemma 4.1 (3) states that the query & key noise remain close to their initialization before Ti, allowing us to approximate them as still being at their initial values at Ti. Lemma 4.1 (1)() jointly show the early shift and stabilization of mean value noise. The proof of this section is in Appendix E.4.
8

---

Published as a conference paper at ICLR 2025
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., (w Q.a: ME,)) and (w (i.. BiE.)- Let T3 = 50/2n opp' s-' = ( 1/112,,, we call t E [Ti, T2] Stage II.
Lemma 4.2 (Sign Alignment Between Query and Key Noise). Let 12Th == oo s
Then, with probubility 1 - 6, sgn((w
(.1,,)) = sgn(
(T K,s
v..)).). Specifically: (1)
If sgu(( (0)
Q.a M=E))) = sgu((wX) k. M.E))),, then sgu((w (T -.EE)) sgu((w ((. v.Ei)) (2)
(0)
If sgn((w a.a g.E.))
-sgn( K.s v(S.)), then the conditional probability of the event Q.a (.,, = j} are at least 1/2 O for {+1).
10)
{sgn( (w
Lemma 4.2 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 3,, == max. (1 (w (0)
Q.a: 1.) K. 1)). KW V.J.r: MD] We have: (I)
(0)
(Magnitude) At for all s E [ms]] and i E [n]. ( Q.a M M. (
(t)
Re s(E,)1 (w
(t)
= (tnnops). (2) (Con-
centration of softmax) (t)
$2,11 = 1/2 (1). 1221 = 1/2 + (11) for all i E [n]] and t S Tz. (3) (Sum of noise) With probability at least 1 for all s E [mm] / q.a U. 1E (w (Ta)
(72)
K,s MeEi)) 2
2m (4) For all s E [ma]. and t S T2, (w a. = (w Q p) (1 = o(1)), (w K.s H4) =
(w K. p1) -1++
(0)
Lemma 4.3 (1) shows that query & key noise become linear with t at T3. However, according
to Lemma 4.3 (2), the softmax outputs of attention layer remain concentrated around 1/2 at TT. Lemma 4.3 (3) estimates the magnitude of the sum of the query & key noise, i.e., Σm= (w M+E,) and ( (w... k. U:E.) at T2. which plays a role in Stage III. Lastly, Lemma 4.3 (4) states that the query & key signals remain near their initialization before T2. The proof of this section is in Appendix E.5.
4.3 STAGE III. MAJORITY VOTING DETERMINES THE SIGN OF QUERY AND KEY SIGNALS.
In Stage III. we focus on query & key signals, i... (w uu) and (w K..:
14). Recall By = Vamn W))). Let Ta 382n |||||| @(oun-"), "), we call
(0) max..jr (((w a. 1)),,(WW, kw K.s 14) (( E [T2,7, Stage III.
(0)
Lemma 4.4 (Stage III). We have: (1) (Noise determines the sign of signals by majority vot- ing) For all s E [ma] and t E [T2, Tz), sgn((w!(t)) (t+1)
(1) (wQ,)))) = sgu(E (W UiE)), K.. H4) - (w.., = sga( ( (w Q.at v.E)). (2) (Magnitude) At Ts. for all
(t)
and sgu((wk. (e+1)
E [mx]. (w(Th) pa) = sgn(E, (w(Th)
(Ta) Q. (ad.))-((TT))) [llll), and (w - =sgn(- --, Q. -VEE))
(Ta)
(1) ((Ts)) [|a|])- (3) (Concentration of softmax) $ 1/2 + o1), $1,21 1/2 + o(1) for all i [n]]
(t)
and E [T2, Tz].
Lemma 4.4 (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 E [n]. the sign dictation can be viewed as a majority voting mechanism. Lemma 4.4 (2) shows that query & key signals become linear with / at T3. Although much smaller than query & key noise, they will become dominant in Stage IV when noise-signal softmax outputs approach zero. Lemma 4.4 (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.
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 sofimax outputs (1) and all query & key signals and noise that could be affected by 8,211 Let T4 == C, log(ppa ||||| "m=m"""mm"" "
0(m²) -1/2,, o, s1)) where C3- = +(1) is a large constant, we call & [Ts, Ti] Stage IV.
9

---

Published as a conference paper at ICLR 2025
Lemma 4.5 (Exponentially Fast Decay of Noise-Signal Softmax Outputs), Let TT Ty be the
( last time such that for all t E [T,, b. s E [mm]] and i E [n], Eis 1221 1222 ( vc.,) 2
(w (0
e.a: u)) and (1) (0)
$1,21 $,22 (w a. ME,) 2 2 2 (w (1) Then, we have T4
(1 m
pp ) and for all E [m]: (1) 8,21 = (mpp-0mme12))) 22))) for & E [Ts, T
and
8 1
exp(OOlogg(
|| /ops))) = o(). (2) 91,11 1/2 + ffr. [T.,
(0)
Lemma 4.5 states that the noise-signal softmax outputs decay exponentially and approach zero during
[T.T; while other softmax outputs stay around 1/2. All signals and noise continue to grow as in Stage III until just before T4- Shortly after TT. the final sign alignment of query and key noise begins.
Lemma 4.6 (Sign Alignment of Key Noise). There exists a small constant 0g E (0, 1) such that for (t+1)
all E [ma]] and i E [n]. we have: (1) sgn( (en((w/) K,s W:E) (ww (t)
K.s: WeE) = sgn( (w (t)
K, y:E)), (.),.ff for
t E [s. T. (2) sgn((g (( (t+1)
K.s + .3:E) (w (t)
K. UrE) = sgn( (Ta)
Q. 1)) for t 2 (1 + 0.)TT.
Negative key noise alignment with signals. Consider a single neuron with a positive query signal. Lemma 4.6 states that after (1 +1+68 0c)TT all negative key noise flip direction and begin aligning with query signal. Before T the dominant gradient terms for key noise are related to both query noise and noise-signal softmax outputs. However, after (1 + 0,)TT, due to the exponential decay of (t) the query signal term becomes dominant.
1,21
Lemma 4.7 (Delayed Sign Alignment of Query Noise). With the same O2 in Lemma 4.6, for (t+1)
all E [mx] and é E [n], we have: (I) sgn (w((++))
Q. (EE)) (w Q.o y=E)) = sgn( (onw(w)) (w((t Q, y:E)), for (t+1) Q.s .M - (w Q. ViE) sgn((w Q.s p))), 2 (2+30.)TT-
t E [T3, (2-0)) (2) sgu((w
(Ts)
Different alignment times for negative query and key noise. Lemma 4.7 indicates a time gap between the alignment of negative key noise and negative query noise. Negative query noise continues to grow until (2- 0,)T, since 2/ 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-0.)Tī 0.)TT is the point where key noise still has significant magnitude, while (2 + 30.)TT indicates the completion of key noise alignment.
Lemma 4.8 (End of Stage IV). (I) For all 2 TA, 8 E [ma]. and i E [n] sgn( (w (0) Q. MEE)) =
sgn((v (w.,s, MnE,)) = sgn((w (w. p)). (2) Loss does not converge. (s( O(1)
Kist
Lemma 4.8 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.
CONCLUSION AND LIMITATIONS
In conclusion, we present theoretical analysis of the training dynamics of two-layer transformers using SignGD for binary classification task involving dataset with both signal and noise. We iden- tify 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 deeper theoretical understanding and aids in the design of more efficient optimization methods for transformers.
L 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 in our main theory, creating 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 analyzed includes only a angle-head self-attention layer, whereas deeper transformers and multi-head attention are more commonly used in practice. We discuss some potential extensions in Appendix F, and discuss differences from real-world setup in Appendix H.
10

---

Published as a conference paper at ICLR 2025
ACKNOWLEDGMENT
This work was supported by the NSFC Project (No. 62376131), Tsinghua Institute for Guo Qiang, and the High Performance Computing Center, Tsinghua University. J.Z is also supported by the XPlorer Prize.
REFERENCES
Kwangjun Ahn, Xiang Cheng, Hadi Daneshmand, and Suvrit Sra. Transformers learn implement preconditioned gradient descent for in-context learning, Advances in Neural Information Processing Systems, 36, 2023.
Kwangjun Ahn, Xiang Cheng, Minhak Song, Chulhee Yun, Ali Jadbabaie, and Suvrit Sra. Linear attention is (maybe) all you need (to understand transformer optimization). International Conference on Learning Representations, 2024. URL https: http::: /tpp////ppee/eeee/ee/eee net wwt.norttttnnaaeeeee /formmm 1d=0uI5415ry1.
Zeyuan Allen-Zhi and Yuanzhi Li. Towards Understanding Ensemble, Knowledge Distillation and DittDissillation in Deep Learning. In International Conference on Learning Representations, 2023.
Lukas Balles and Philipp Hennig. Dissecting adam: The sign, magnitude and variance of stochastic gradients. In International Conference on Machine Learning, pp. 404-413. PMLR, 2018.
Lukas Balles, Fabian Pedregosa, and Nicolas Le Roux. The geometry of sign gradient descent. arXiv preprint arXiv:2002.080556 2020.
Jeremy Bernstein, Yu-Xiang Wang. Kamyar Arizzadenesheli, and Animashree Anandkumar. signsgd: Compressed optimisation for non-convex problems. In International Conference on Machine Learning, pp. 360-569. PMLR, 2018.
Jeremy Bernstein, Jiawei Zhao, Kamyar Azizzadenesheli, and Anima Anandkumar. signSGD with majority vote is communication efficient and fault tolerant. In International Conference on Learning Representations, 2019, URL https: h/p///peeeveeeeeeee/eeeeeeeeeeee www.ereeeiieeeeeeeeeeeeee.ee.eee wtw.oooouuiddmiie jACY7.
Olivier Bousquet and André Elisseeff. Stability and generalization. The Journal of Machine Learning Research, 2:499-526, 2002.
Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 13:1877-1901, 2020.
Yuan Cao, Zixiang Chen, Misha Belkin, and Quanquan Gu. Benign overfitting in two-layer convolu- tional neural networks. Advances in Neural Information Processing Systems, 35:25237-25250, 2022.
Niladri Chatterji and Philip M. Long. Finite-sample analysis of interpolating linear classifiers in the overparameterized regime. Journal of Machine Learning Research, 22 22(129) 2021. URL /itt:://jmis http://mlwooorr../eeeeee wapgg//apprrs. /v22/222222222/// 974 html.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in Neural Information Processing Systems, 11115588449999 2021.
Siyu Chen, Heejune Sheen, Tianhao Wang, and Zhuoran Yang. Training dynamics of multi-head softmax attention for in-context learning: Emergence, convergence, and optimality. arXiv preprint arXiv:2402. 2024.
Xiangning Chen, Chen Liang, Da Huang, Esteban Real, Kaiyuan Wang, Hieu Pham, Xuanyi Dong, Thang Luong, Cho-Jui Hsich, Yifeng Lu, et al. Symbolic discovery of optimization algorithms. Advances in Neural Information Processing Systems, 36, 2023.
11

---

Published as a conference paper at ICLR 2025
Aakanksha Chowdhery, Sharan Narang, Jacob Devlin, Maarten Bosma, Gaurav Mishra, Adam Roberts, Paul Barham, Hyung Won Chung. Charles Sutton, Sebastian Gehrmann, et al. Palm: Scaling language modeling with pathways. Journal of Machine Learning Research, 24(240)1--13,, 2023.
Michael Crawshaw, Mingrui Liu, Francesco Orabona, Wei Zhang, and Zhenxun Zhuang. Robustness to unbounded smoothness of generalized signsgd. Advances in neural information processing systems, 35:9955-9968. 2022.
Tri Dao, Dan Fu, Stefano Ermon, Atri Rudra, and Christopher Ré. Flashattention: Fast and memory- efficient exact attention with io-awareness. Advances Neural Information Processing Systems, 35:16344-16359 2022.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. In Proceedings of NAACL-HLT, pp. 4171- 4186, 2019.
Yihe Dong, Bean-Baptiste Cordonnier, and Andreas Loukas. Attention not all you need: Pure attention loses rank doubly exponentially with depth. In International Conference on Machine Learning, pp. 2793-2803. PMLR, 2021.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recogmition at scale. In International Conference on Learning Representations, 2021. URL https: /pwnnevieeeeeeeeeeeeeeeee net/ het/forunt Eorum? i id=Yi cbFdNTTy.
Spencer Frei, Niladri S Chatterji, and Peter Bartlett. Benign overfitting without linearity: Neural network classifiers trained by gradient descent for noisy linear data. Conference on Learning Theory, pp. 2668-2703. PMLR, 2022.
Andi Han, Wei Huang, Yuan Cao, and Difan Zou. On the feature learning in diffusion models. arXiv preprint arXiv:24/2.01021, 2024.
Moritz Hardt, Ben Recht, and Yoram Singer. Train faster, generalize better: Stability of stochastic gradient descent. In Maria Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of The 33rd International Conference on Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 1225-1234, New York, New York, USA, 20-22 Jun 2016. PMLR. URL https: http::s/pprcceeeieg
mir mwr..pree press / v48 htw.rttrlet 6 html.
Wei Huang, Yuan Cao, Haonan Wang, Xin Cao, and Taiji Suzuki. Graph neural networks provably ben- efit from structural information: A feature learning perspective. arXiv preprint arXi::2306.13926 2023a.
Wei Huang, Ye Shi, Zhongyi Cai, and Taiji Suzuki. Understanding convergence and generalization in federated learning through feature learning theory. In The Twelfth International Conference Learning Representations, 2023b.
Wei Huang, Andi Han, Yongqiang Chen, Yuan Cao, Zhiqiang Xu, and Taiji Suzuki. On the comparison between multi-modal and single-modal contrastive learning. arXiv preprint 2rXiv:2411.02881 2024a.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers. arXiv preprint 1Xiv:2310.00249, 2023c.
Yu Huang, Zixin Wen, Yuejie Chi, and Yingbin Liang. Transformers provably learn feature-position correlations in masked image modeling. arXiv preprint arXiv:2403.02233, 2024b.
Samy Jelassi, Michael Sander, and Yuanzhi Li. Vision transformers provably learn spatial structure. Advances in Neural Information Processing Systems, 35:37822-37883 2022.
Jiarui Jiang, Wei Huang, Miao Zhang, Taiji Suzuki, and Liqiang Nie. Unveil benign overfitting for transformer in vision: Training dynamics, convergence, and generalization. arXiv preprint arXiv:2409. 19345, 2024.
12

---

Published as a conference paper at ICLR 2025
Kaiqi Jiang, Dhruv Malik, and Yuanzhi Li. How does adaptive optimization impact local neural
network geometry? Advances in Neural Information Processing Systems, 36, 2023.
Juno Kim and Taiji Suzuki. Transformers learn nonlinear features in context: Nonconvex mean-field dynamics on the attention landscape. arXiv preprint arXiv:2402.01258, 2024.
Diederik Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In International Conference Learning Representations, 2015.
Yiwen Kou, Zixiang Chen, Yuanzhou Chen, and Quanquan Gu. Benign overfitting in two-layer ReLU convolutional neural networks. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 17615-17659, PMLR, 23-29 Jul 2023. URL https: http::::/ppppceeeeeeeeee .mw.mmlr....... pres.press// v202/ kou23a htmll
Frederik Kunstner, Jacques Chen, Jonathan Wilder Lavington, and Mark Schmidt. Noise is not the main factor behind the gap between sgd and adam on transformers, but sign descent might be. In International Conference on Learning Representations, 2023. URL https: /ppnrevvee. net/fome forum? id a
Frederik Kunstner, Robin Yadav, Alan Milligan, Mark Schmidt, and Alberto Bietti. Heavy-tailed class imbalance and why adam outperforms gradient descent on language models. arXiv preprint arXiv:2402.19444, 2024.
Bingrui Li, Jianfei Chen, and Jun Zhu. Memory efficient optimizers with 4-bit states. Advances in Neural Information Processing Systems, 36: 16136-111,,,, 2023a.
Haochuan Li, Alexander Rakhlin, and Ali Jadbabaie. Convergence of adam under relaxed assumptions.
Advances in Neural Information Processing Systems, 36, 2023b.
Hongkang Li, Meng Wang, Sijia Liu, and Pin-Yu Chen. A theoretical understanding of shallow
vision transformers: Learning, generalization, and sample complexity. In International Confer- ence on Learning Representations, 2023c. URL https: /tttt:::///ppeevieeww net htt/moormmmm/ jCiGv30jhb.
Hongkang Li, Meng Wang, Songtao Lu, Xiaodong Cui, and Pin- Yu Chen. Training nonlinear transformers for efficient in-context learning: A theoretical learning and generalization analysis. arXiv preprint arXiv:2402. 2024.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding. In International Conference on Machine Learning, pp. 19689-19729. PMLR, 2023d.
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. In Proceedings of the 2020 Conference on Empirical Methods Natural Language Processing (EMNLP), pp. 5747-5763, 2020.
Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is provably the optimal in-context learner with one layer of linear self-attention. In International Conference on Learning Representations, 2024. URL https: http////peereee..eee/e net forum?
1d=8p3fu5611s 6cc..
Eshaan Nichani, Alex Damian, and Jason D Lee. How transformers learn causal structure with gradient descent. arXiv preprint 140.35 2024.
Lorenzo Noci, Sotiris Anagnostidis, Luca Biggio, Antonio Orvieto, Sidak Pal Singh, and Aurelien Lucchi. Signal propagation transformers: Theoretical perspectives and the role of rank collapse. Advances in Neural Information Processing Systems, 15:27118-27211 2022.
Samet Oymak, Ankit Singh Rawat, Mahdi Soltanolkotabi, and Christos Thrampoulidis. On the role of attention in prompt-tuning. In International Conference on Machine Learning, pp. 26724-26766 PMLR, 2023.
13

---

Published as a conference paper at ICLR 2025
Yan Pan and Yuanzhi Li. Toward understanding why adam converges faster than sgd for transformers. arXiv preprint arXiv:2306.0022 2023.
William Peebles and Saining Xie. Scalable diffusion models with transformers. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 4195 4205, 2023.
Sashank J. Reddi, Satyen Kale, and Sanjiv Kumar. On the convergence of adam and beyond. In International Conference Learning Representations, 2018. URL /uRLLttps://ppeneev.ee net met/forumm forum? id= redd-rryuu--BBe
Clayton Sanford, Daniel Hsu, and Matus Telgarsky. Representational strengths and limitations of transformers. Advances in Neural Information Processing Systems, 36, 2023.
Heejune Sheen, Siyu Chen, Tianhao Wang, and Harrison H Zhou. Implicit regularization of gradient flow on one-layer softmax attention. arXiv preprint MxXV:2403308699 2024.
Davoud Atace Tarzanagh, Yingcong Li, Christos Thrampoulidis, and Samet Oymak. Transformers as support vector machines. arXiv preprint arXiv:2308. 2023a.
Davoud Atace Tarzanagh, Yingcong Li, Xuechen Zhang, and Samet Oymak. Max-margin token selection in attention mechanism. Advances in Neural Information Processing Systems, 36:48314- 48362, 2023b.
Jiaye Teng, Bohang Zhang, Ruichen Li, Haowei He, Yequan Wang, Yan Tian, and Yang Yuan. Finding generalization measures by contrasting signal and noise. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Hnngelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 33983-34010. PMLR, 23-29 Jul 2023, URL https: httpttss://ppooeeeeeee /ptppeeeioeg..... mlr.press v21//222222///22222 htt////enn22.. html.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon S Du. Scan and snap: Understanding training dynamies and token composition in I-layer transformer. Advances in Neural Information Processing Systems, 36:71911-7199 2023.
Yuandong Tian, Yiping Wang. Zhenyu Zhang, Beidi Chen, and Simon Shaolei Du. JoMA: De- mystifying multilayer transformers via joint dynamics of MLP and attention. In International Conference Learning Representations, 2024. URL https: /ttp:////peeeeeeeeee/m net meee/fae www.foruumttttm LdGDbUGRG
Bhavya Vasudeva, Puneesh Deora, and Christos Thrampoulidis, Implicit bias and fast convergence rates for self-attention. arXiv preprint wrXiv:2402.05738, 2024.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017
Bohan Wang, Jingwen Fu, Huishuai Zhang, Nanning Zheng, and Wei Chen. Closing the gap between the upper bound and lower bound of adam's iteration complexity. Advances Neural Information Processing Systems, 36, 2023.
Mingze Wang and Chao Ma. Understanding multi-phase optimization dynamics and rich nonlinear behaviors of relu networks. Advances in Neural Information Processing Systems, 36, 2023.
Zixuan Wang, Stanley Wei, Daniel Hsu, and Jason D Lee. Transformers provably learn sparse token selection while fully-connected nets cannot. In Forty-first International Conference on Machine Learning, 2024.
Yikai Wu, Xingyu Zhu, Chenwei Wu, Annie Wang, and Rong Ge. Dissecting hessian: Understanding common structure of hessian in neural networks. arXiv preprint 00X1::2010.04261, 2020.
Weihang Xu and Simon Du. Over- paramttrrizzaioonn exponentially slows down gradient descent for learning a single neuron. In The Thirty Sixth Annual Conference on Learning Theory, pp. 1155-1198. PMLR, 2023.
14

---

Published as a conference paper at ICLR 2025
Jingzhao Zhang, Tianxing He, Suvrit Sra, and Ali Jadbabaie. Why gradient clipping accelerates training: A theoretical justification for adaptivity. International Conference on Learning Representations, 2020a. URL
Jingzhao Zhang, Sai Praneeth Karimireddy, Andreas Veit, Seungyeon Kim, Sashank Reddi, Sanjiv Kumar, and Suvrit Sra. Why are adaptive methods good for attention models? Advances in Neural Information Processing Systems, 33:15383- 15393, 2020b.
Jintao Zhang, Haofeng Huang, Pengle Zhang. Jia Wei, Jun Zhu, and Jianfei Chen. Sageattention2: Efficient attention with thorough outlier smoothing and per-thread int4 quantization. arXiv preprint arXiv:2411.10958, 2024a.
Jintao Zhang, Jia Wei, Pengle Zhang, Jun Zhu, and Jianfei Chen. Sageattention: Accurate 8-bit attention for plug-and-pin inference acceleration. In International Conference Learning Representations, 2025a.
Jintao Zhang, Chendong Xiang, Haofeng Huang, Jia Wei, Haocheng Xi, Jun Zhu, and Jianfei Chen. Spargeatm: Accurate sparse attention accelerating any model inference. arXiv preprint 1rXii:2502.18137, 2025b.
Ruiqi Zhang, Spencer Frei, and Peter L. Bartlett. Trained transformers learn linear models in- context. Journal of Machine Learning Research, 25(49):1-55, 2024b. URL http:// /tmtr///mmr.
http/papers/v2555 v2t//225//2//22 22 html.
Yushun Zhang, Congliang Chen, Tian Ding, Ziniu Li, Ruoyu Sun, and Zhi-Quan Luo. Why trans- formers need adam: A hessian perspective. arXiv preprint 1688::2800.,,788, 2024c.
Yushun Zhang, Congliang Chen, Ziniu Li, Tian Ding, Chenwei Wu, Yinyu Ye, Zhi- QuiiQuan Luo, and Ruoyu Sun. Adam-mini: Use fewer learning rates to gain more. International Conference on Learning Representations, 2025c.
Difan Zou, Yuan Cao, Yuanzhi Li, and Quanquan Gu. Understanding the generalization of adam in learning neural networks with proper regularization. In International Conference on Learning Representations, 2023. URL https: / //tp///vvveeeeeeeeeeeeeeeeeeeeee www.eeeiieeeee..... www.eeeoorummmiiiiiio 4qjTF.
15

---

Published as a conference paper at ICLR 2025
Appendix
CONTENTS
Detailed Related Work
17
B Experimental Details and More Experiments
19
B.1
Experimental Settings
19
B.2
Comparison with Non-sparse and/or Non- Non-orthogonal Data
19
Comparison with Adam and Gradient Descent
20
B.4
Comparison with Greater Context Length
21
B.5
Comparison with Multi-head Attention
23
B.6
Comparison with More Complex Attention Models
24
B.7
Explanations for Differences between SignGD and Adam
26
C Preliminary Lemmas
29
Gradient
37
Proofs
40
E.1
Technique Overview
40
E.2
The dynamics of value
43
E.3
The dynamics of query and key: preparations
44
E,4
Stage
45
E.5 Stage II
45
E,6
Stage III
53
E.7
Stage IV
56
E.7.1
Stage IV..
58
E.7.2
Stage IV.b
64
E.7.3
Stage IV.c
69
E.8 Proof of Theorem 3.2
72
F
Discussion Extensions
75
F1
()-1 Test Loss
75
F2
Extension Longer Contexts
75
F3 Extension to Joint Training of Linear Layer
76
F.4 Extension to Multi-head Attention
77
G More Detailed Explanations for the Four-Stage Dynamics
78
Discussion on Differences from Real- World Setup
79
16

---

Published as a conference paper at ICLR 2025
A DETAILED RELATED WORK
Training dynamics of transformers. Currently, the theoretical analysis of transformers training dynamies 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. (2022) 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. (2023) 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 al. (2024) analyzed the joint dynamies of an attention and MLP layer. Li et al. (2023d) studied the training dynamics of a single-head attention layer on the &g loss with the data modeled by topic modeling. Li et al. (2023c) studied a three-layer vision transformer with query and key parameterization trained by SGD on the hinge loss. Oymak et al. (2023) 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. (2023a;b) studied the dynamics of a single-head attention layer and tunable token, and connects the training dynamics to a certain SVM problem. While Tarzanagh et al. (2023a;b) only presented an asymptotic convergence result, Vasudeva et al. (2024) showed the global convergence, provided a convergence rate / 1-3/4 for GD, as well as removing the restriction on a fixed linear head. Also extending Tarzanagh et al. (2023a;b), Sheen et al. (2024) studied this problem with a query-key parameterized transformer, which gives different implicit regularization compared with the single attention matrix parameterization. Huang et al. (2024b) studied the GD dynamics of single-head attention layer with a self-supervised learning objective. Nichani et al. (2024) studied the GD dynamics 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. (2024) studied the data model introduced by Sanford et al. (2023), showing that one-layer transformer can efficiently learn this task via GD, while fully-connected networks cannot express the task. Jiang et al. (2024) studied the GD dynamics of a two-layer transformer with single-head self-attention layer and fixed linear head dataset with signal and noise, which has the similar data and model setting to our work. However, Jiang et al. (2024) 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. (2023); Mahankali et al. (2024) 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. (2024b) studied the Gradient Flow (GF) of single-head linear attention layer on in-context linear regression task. Huang et al. (2023c) studied how single-head attention layer trained by GD solves an in-context linear regression task where the the input data are orthogonal. Chen et al. (2024) 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 taaki-taskk structure. Kim & Suzuki (2024) studied the mean-field dynamics of a transformer with one linear attention layer and one MLP layer
in-context regression task. Li et al. (2024) studied the SGD dynamics of one-layer transformer with a single-head self-attention layer and a two-layer MI.
in-context classification task.
Comparison with other works using query-key parameterization. Practically, the modern trans- former architecture uses query-key parameterization in the attention module (Vaswani al., 2017; Dao et al., 2022: Zhang et al., 2024a; 2025a;b). However, most of works mentioned above take simplifications that replaces the query and key matrices with a single attention matrix. Among those works, only Li et al. (2023c; 2024); Sheen et al. (2024); Jiang et al. (2024) studied the softmax attention with trainable query-key parameterization, but there are also some limitations. Li al. (2023e; 2024) introduced relaxed assumptions about initialization, which are too stringent and make softmax outputs not concentrated at 1/2 at initialization anymore. Sheen et al. (2024) 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. (2024) 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 focus on SignGD.
17

---

Published as a conference paper at ICLR 2025
Understanding of Adam on transformers. While Adam may fail converge in convex objec- tive (Reddi et al., 2018), it performs so well on transformers and better than SGD, which means Adam converges faster and achieves lower training loss (Ahn et al., 2024; Jiang al., 2023; Kunstner et al., 2023; 2024; Pan & Li, 2023; Zhang et al., 2024c; 2020b).
Previous works tried to give explanation about this fact from different perspectives. Liu al. (2020) observed unbalanced gradients in transformers and Adam can give uniform parameter update. Zhang al. (2020b) 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 (2023) 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 still better than GD. Jiang et al. (2023) showed that Adam could bias the trajectories towards regions where Hessian has relatively more uniform diagonals while SGD cannot. Zhang al. (2024c) 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. Zhang et al. (2025c) and Li et al. (2023a) found that second moment of the Adam optimizer is not sensitive which indicates the advantage over SGD is robust. Pan & Li (2023) showed that Adam can lead to smaller directional smoothness values which may imply better optimization. Kunstner et al. (2024) 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 al. (2020a) and Crawshaw (2022) sought alternative relaxed assumptions for Adam. Li et al. (2023b) and Wang et al. (2023) improved the analysis of
Adam under those assumptions.
18

---

Published as a conference paper at ICLR 2025
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. For data hyerparameters, they can be uniquely determined by one row in Tab. 3 and the value of d. In the Fig. and Tab. 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 1) le-4. For Adam, we use the learning rate 1) le-4, B1 0.9, (3 0.999, and c= 1e-15. For gradient descent, we use the learning rate le-1.
Also, we use neuron 0 in Fig. 1 by default and in following experiments. We sign gradient descent with the learning rate y = le-7 in 2000 iterations to simulate sign gradient descent with the learning rate 7) le-4 in iterations in Fig. (b). and in following experiments. We use learning rate of lc-4 for all optimizers in Fig. 2 (c) for fair comparison.
Table 3: Experimental settings of data model. the training sample size and there are always equal samples in both classes. s is the noise sparsity level. op is the standard deviation of noise. do the network initialization standard deviation. m, is the value dimension. my is the query and key dimension. 'orthogonal" means whether the noise patch and signal patch are orthogonal. If they are not orthogonal, the coordinates are selcted form [d] instead of [d]' {1}. 'iters' is the total iteration/epoch number in one run. Signal patch μ is always [1, 0, ...0...000 0]
Pp
orthogonal
00
my
m
iters
(a)
0.01d
0.04d
2.0/ V
True
0.1/Vd
0.01d
0.05d
2000
(h)
0.01d
0.04d
2.0/ Vs
False
0.1/Va
0.01d
0.05d
2000
(c)
0.01d
0.4d
2.0/ VB
True
0.1/v7
0.01d
0.05d
2000
(d) 0.01d
0.4d
250/V5
False
0.1/va
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 X, which we call "scaled SNR", and then add gaussian random noises with standard deviation 1 to the outer regions with a width of 7. Additionally, we only the class 3 and 7 for classification, to make binary classification task which is consistent with theoretical settings. To input the data into the transformers, we patchify the data with a size of 7 x 7.
We train two-layer transformer model consistent with our theoretical setting. We set d 49, m. my 10, ou 0.1/Vd. We only use 2000 training data points and use deterministic optimizers to train the network. For SignGD and Adam across different values of 3,, we use learning rate 7 le 2, and train the models for 200 epochs. For GD, we use different learning rates for different SNR since one learning rate of GD cannot adapt to data with noise in different magnitude. Specifically, we use 17- 1e 1 for SNR in [0.9. 1.0]. 17 3e for SNR [0.7, 0.8]. 17 = 6e for SNR in [0.5, 0.6].
Le0 for SNR in [0.2. 0.3, 0.4]. We train the model for 500 epochs for GD. In all these settings, our training setup can guarantee 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, we run the experiments with the same setting as Fig. 1, but use a different legend. In the experiments with non-orthogonal sparse data, orthogonal non-sparse data, and
Fig. 4, 5, 6, non-orthogonal non-sparse data, respectively. The legends follow Fig. 3. The figures show that our theoretical results hold empirically in those data settings.
19

---

(1) t = 150, coinciding with $1.1 approaching zero, while delayed sign alignment of query noise begins Published as a conference paper at ICLR 2025
V13.) Stage
(v,11()) in Stage!
1,11
and
",21
(Was. (1 and (wasiP)
0.03
0.8
0.02
0.006
0.4
0.00
0.004
0.5
0.002
-0.1
0.000
0.0
-000
0.0
-0.3
150300
Steps
Steps
Steps
Stops
(4)
(b)
(d)
(16.))
positive (WKx..
positive ( M(k)
9121
(Wes.))
negative (11):)16))
negative (Mo...):6))
Figure 3: Data setting (a) in Tab. with d = 2000. (a) Key noise dynamics over = 0 to = 2. f While mean value noise stabilizes into linear relationship with early, key noise remains close to initialization. (c) Softmax output dynamics over (1)
(b) Mean value noise dynamics over
t= 0 to 900. The softmax outputs decay exponentially. At t 150, $1,21 approaches zero, while
i1 remains close 1/2. (d) Dynamics of query noise, key noise, and query signals over ( to = 900: The dotted lines represent positive query and key noise at = 100, and the solid lines represent negative noise 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
around 300, about twice late as the key noise.
0.6
<v.66>
0.02
0.5
0.6
0.4
0.004
0.3
0.4:
train
0.01
0.002
0.2
0.0
0. 0.000
0.1
-02-
-0.002
0.0
-0.4
Figure 4: Data setting (b) in Tab. with d = 2000.
B.3 COMPARISON WITH ADAM AND GRADIENT DESCENT
In Fig. 8, 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, is actually very small in magnitude, which implics the small gradients be casily dominated by the momentum Adam. But otherwise, symbolic gradient descent is almost similar 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. is also noted that gradient descent can leads small generalization gap in this setting.
20

---

Published as a conference paper at ICLR 2025
0.03
0.6
< >
A/
0.5
0.005
10-
0.001
train
0.00
0.007
0.3
test
0.000
0.2
0.0
-0.02
0.002
0.03
230
800
400
400
0.0
1000
2000
Figure 5: Data setting (c) in Tab. with d 2000.
0.03
0.4
<1
2
<M0 -M6
0.02
f=0
0.5-
10-
0.6
0.005
0.00
0.004
0.5
Rest
frain
-0.1
0.0
0.2
-003
-0.007
200
400
600
-05
200
400
600
0.0
1000
2000
Figure 6: Data setting (d) in Tab. 3 with d 2000.
B.4 COMPARISON WITH GREATER CONTEXT LENGTH
In this section, we investigate when will happen when context length is greater than 2, i.c.. L > 2. At this time, the model is defined as
FJ (W,X) =
T W Wv,jXsoftne Xsoftmax
x W (Wxx
1=1
For each data point (X, y). predictor X x((2) x( Radx have L patches (or tokens),
where x(1) x (2) x())) E Rd, and label is binary, i... y E {+1}. The data generation is similar to the L 2 case, except that randomly select L/2 patches and assign them by yμ signal patches, while the remaining I.2 patches are noise patches. The noise patches in one data sample mutually independent.
In our experiments, we use L= 10, and Fig. 10, 11, 12 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 E [r)]. 0g is defined as
arg max (T) 7,1
We empirically observe that for all data point, only one element in each line of L x L post-softmax attention matrix is activated, while other elements in the line are almost zero, which means that the attention attends to only patch. This patch for each line is also uniform in different lines and therefore uniquely corresponds to one data point, which is exactly a;- In the = case, we have a, 2 for all i [n]. But when L. and there many noise patches in data sample, 0; can be varied across samples but ( must be noise patch.
SignGD
Adam(8;=0)
Adem(6,=0.5)
Adem(8;=0.9)
0.8
0.75
0.2
ME Hi
1.0
B.5
0.4
0.1
0.00
-0.5
-0.25
-0.25
-03
230
750
502
500
230
500
3000
2000
Figure 7: Sign alignment of signal to noise by majority voting in Stage III and sign alignment of negative noise to query signal by decay of niise-signal softmax outputs in Stage IV, We use data setting (d) in Tab. 3 with d = 2000. We always use 3r 0.99 and le-15 in Adam.
21

---

Published as a conference paper at ICLR 2025
SignGD
Adamla:-0)
AdamiB;-0.81
Adam(8)=009)
0.3
0.3
0.3
0.3
0.8
0.1
0.1
0.1
0.3
0.0
0.0
0.0
0.5
250
750
750
500
750
500
750
5.0
1000
2000
Figure 8: The dynamics of softmax outputs. We use data setting (d) in Tab. 3 with 2000. We
always use 32 = 0.99 and = le-15 15 in Adam.
SignGD
Adem(B;=0)
AdamiB:=0.51
Adam(8;=0.9)
GD
train
0.4-
0.0
0.0
0.0
0.0
1000
2000
1000
2000
1000
2000
1000
2000
1000
2009
Figure 9: Training and test loss. We use data setting (d) in Tab. 3 with d = 2000. We always use
$2 = 0.99 and le-15 in Adam.
The greater context length is consistent with the L 2 case, and thus consistent with theoretical results the sense that (I) For E [n]. the dominated query noise 9: and
key noise (WK.s X/.)( are the same in sign with the query signal (Wq.s 14). (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 basically not
learned, even with learning rate 1.
SignGD
AdamlA:=0)
Adam(B:=0.5)
Adam(8;=0.9)
GD
1.25
0.6
0.6
1.00
<W <W
0.4
0.4
0.04
0.2
0.2
0.00
300
500
520
300
73U
1000
2000
Figure 10: The dynamics of query noise, key noise and query signal. We use data setting (d) in Tab. with d 2000. We use B₂ 0.99 and le-15 in Adam and 7 1.0 for gradient descent.
22

---

Published as a conference paper at ICLR 2025
SignGD
AdamiB:=0)
Adami8;=0.3)
Adamn(e
0.100
0.8-
0.8
0.8-
02-
0.098
0.6-
0.6
0.6
0.098
0.088
0.2
02
0.2
406
500
500
1000
1500
500
1000
1500
500
1000
19:0
1000
2000
Figure 11: The dynamics of softmax outputs. We data setting (d) in Tab. with d 2000. We
use 5 = 0.99 and c le-15 in Adam and we use = 1.0 for gradient descent.
SignGD
Adem(B;=0)
AdamiB:=0.31
AdamiB;=0.99
GD
0.8
brain
0.7
0.2
0.2
0.2-
0.2
0.0
0.0
0.0
0.0-
0.0
1000
2000
0
1000
2000
1000
2000
1000
2000
1000
2009
Figure 12: Training and test loss. We use data setting (d) in Tab. 3 with d = 2000. We use 82 0.99 and c = le-15 in Adam and we use 1) 1.0 for gradient descent. In this case, gradient descent converges much faster than sign gradient descent and Adam since we a large learning rate = 1.0, which is times the learning rate used sign gradient descent and Adam.
B.5 COMPARISON WITH MULTI-HEAD ATTENTION
See Fig. 13 for the full dynamics of simplified transformer model with 4 attention heads using SignGD. We observe that the softmax outputs of each head exhibit the same behaviors, and all dynamics are consistent with the single-head case.
Figure 13: Equivalent of Fig. 18 but using 4 head in the simplified transformer model. The subfigure 3-6 shows the dynamics softmax outputs in each attention head.
23

---

Published as a conference paper at ICLR 2025
B.6 COMPARISON WITH MORE COMPLEX ATTENTION MODELS
We have conducted additional experiments on deeper transformers using our synthetic dataset with SignGD, exploring various settings. Specifically, we extend our analysis to models with additional attention layers. MLP layers, and residual connections, which are essential components of modern transformer architectures. Since our theory primarily predicts the behavior of parameerreeter inner products, for transformers with multiple attention layers, we focus on the dynamics of the first layer.
To examine how well the key behaviors identified by our theory persist in more complex models, we performed an ablation study. We provide the full dynamics of all relevant quantities in Fig. 14 and Fig. 15 and augment these results with Tables 4-11, which illustrate the sign alignment behavior during Stage II.
Transformers with Residual Connections. Firstly, on transformers with residual connections, across all model configurations we tested including 2-layer transformers without MLPs, 3-layer trans- formers without MLPs, 2-layer transformers with MLPs, and 3-layer transformers with MLPs observe the following behaviors, consistent with our theoretical predictions:
Stage I: Value noise increases faster than query and key noise, and the value signal remains small relative to the value noise.
Stage II: Query and key noise exhibit sign alignment behavior early in training.
Stage III: The query and key signals have opposite signs, determined by query (and key) noise via a majority-voting mechanism.
Stage IV:
Koise-feature softmax outputs firstly decay and decay exponentially, and both negative query and key noise align with the query signal.
However, we remark that in more complex models, the final alignment observed in Stage IV-iee.. the flip of negative query and key noise-often halts midway. This phenomenon becomes more pronounced with the addition of MLP layers, where the final alignment stops earlier. We attribute this behavior to the rapid shrinking of query and key gradients. This is partly driven by the decay of softmax outputs (as shown in Lemma D.7). Furthermore, as the number layers increases and/or MLP layers introduced, additional layers significantly contribute to this gradient shrinkage, as illustrated in the last column of Figure 14. It is worth noting that this gradient shrinking is a numerical precision issue unrelated to our theory. In theory, the sign operation maps gradients to #1 regardless of their magnitude. However, in practice, extremely small gradients are rounded to zero, disrupting the alignment process. Despite this, conclude that the key behaviors predicted by our theory persist in deeper transformers with residual connections.
Transformers without Residual Connections. On the other hand, in deeper transformers lacking residual connections, the dynamics become erratic. While some short-term behaviors (e.g., sign alignment between query and key noise in Stage II, and the opposing signs between query and key signal) are preserved (see Tables 8-11, and Figure 15), long-term term behaviors deviate significantly from theoretical predictions. For instance:
Feature-feature softmax outputs start to increase instead of decreasing.
The dynamics of positive key noise become non-monotonic.
Value noise exhibits irregular patterns rather than increasing consistently.
Additionally, we remark that the training dynamics of transformers without residual connections are less stable and more irregular compared to those with residual connections. This instability may be linked to the phenomenon of rank collapse in transformers, as discussed in prior works (Dong et al., 2021; Noci et al., 2022).
Based on these findings, we conclude that the key behaviors predicted by our theory persist in deeper transformers with residual connections. Without the residual connections, the key behaviors outlined in our theory are only partially preserved. Understanding the behaviors in subsequent layers and why deeper models make the gradient of first layer shrink faster could be future direction.
24

---

Published as a conference paper at ICLR 2025
softmax
gradient
1.02
0.75
g000
1.
04-
0.25
0.6
00.25
500
1000
1000
2000
500
1000
1900
7000
520
1000
1500
2000
505
1000
1500
2000
500
3000
3500
2000
1.5
2m--8
1.0
1.50
0.3
0.50
21
0.0
8.00
500
1000
1300
2000
500
1000
1500
2002
320
1000
1500
2000
502
1020
1520
2000
6
350
1000
2000
2000
0.3
0.2
0.00
0.
-0.23
0.8
5.0
3.5
0.0-
0.2
0.0
1000
2000
1000
320
2000
502
2000
300
0.5
0.0
0.20
0.3
0.35
-05
0.30
0.2
0.-
-10
500
1000
1500
2000
500
1000
1500
2002
520
1000
3500
2000
505
1000
1500
2000
500
3000
3500
2000
Figure 14: Equivalent of Fig. 18 but with four additional configurations for transformer models:
1) 2 layer attention-only transformer blocks with residual connections; 2) layer attention-onf, transformer blocks with residual connections; 3) layer attention+MLP transformer blocks with residual connections; 4) 3 layer attention+ transformer blocks with residual connections;
25

---

Published as a conference paper at ICLR 2025
soltmax outbuts
gradient
1,5
1.20
grod
1.0
0.1
065
0.4
0.50
6.5
0.2
0.0
6:3
0.43
2300
2000
1000
1500
500
2500
2000
1502
300
1000
2000
2.0
1.25
1.00
0.0
3.5
0.9
1.0
0.55
0.2
-0.
2.0
6.6
2.0
1.5-
2.6
1.5
3.0-
0.0
0.5-
-55
2000
500
1550
0.15
0.30
6.2
0.6-
0.00
0.00
-02
0.2
533
1020
1500
2000
1050
1906
3000
020
1050
1908
3000
1.000
1905
500
1000
1930
2000
16-2
0.1
1.0
0.2
0.5
015
0.000
0.005
0.0
0.2
0.4
0.000
-0.2
0.3
0.0
500
1000
1.00
2000
500
1050
3500
3000
520
1000
2502
2200
10000
1002
500 1000
1500
2000
Figure 15: Equivalent of Fig. 18 but with four additional configurations for transformer models:
2 layer attention-only transformer blocks without residual connections; 2) 3 layer attention-ong transformer blocks without residual connections; 2 layer attention+MLP transformer blocks without residual connections; 4) 3 layer attention+ transformer blocks without residual connections; The last line represents the mean gradient Wa and W x parameters in each iteration.
Table 4:
Sign alignment between query and key noise in Stage II for model: 2L, w/o MLP, w/ residual. The notation in this table and following tables the same as Tab.
init(1 0)\ = 10
((t)) K+.0+
((t)
K+.Q
s(t) K
Q+1
[S(t]
K--QQ-
Row sum
ss())
S*+0+
-+
483
2
26
512
|S(0)
|K+-Q-
242
3
9
253
507
| -+
224
10
3
221
458
---QQ-
37
3
3
480
523
Column sum
986
18
16
980
2000
B.7 EXPLANATIONS FOR DIFFERENCES BETWEEN SIGNGD AND ADAM
Although SignGD can serve as a proxy for understanding Adam, our experiments reveal notable differences between the two. In Figure 2(a) and (b), SignGD causes the negative query eventually
26

---

Published as a conference paper at ICLR 2025
Table 5: residual.
Sign alignment between query and key noise in Stage II for model: 3L, w/o MLP, w/
init(t 0)
10
k++0++
[
K+.Q-
-x+
|
K-.QQ
Row sum
482
2
0
28
512
243
4
7
253
507
[S Q+
222
12
3
221
458
[S
39
3
480
523
Column sum
986
19
13
982
2000
Table 6: Sign alignment between query and key noise in Stage for model: 2L, w/ MLP, w/ residual.
init(I = 0)\ = 10
k+++Q+]
(s())
s(0)
SK-Q++
x
Row sum
|s
479
0
29
512
Q+
-K+-Q-
241
14
6
246
507
| +++
217
9
11
221
458
| -5--Q-
39
3
2
479
523
Column sum
976
30
19
975
2000
Table 7: Sign alignment between query and key noise in Stage for model: 31, w/ MLP, w/ residual.
init(t 0)\t 10
[s() K+.0+
(s
|s x---0-
Row sum
(
K++0+
481
3
1
27
512
|s K++Q0-
232
25
16
234
507
213
9
20
216
458
K-
36
3
1
483
523
Column sum
962
40
38
960
2000
Table 8: Sign alignment between query and key noise in Stage II for model: 2L, w/o MLP, w/o residual.
init(t 0 = 10
(s(t) 1
(s(t)
K+.Q-
s(t)
S
(s(t)
Row sum
486
2
0
24
512
249
4
4
250
507
|S --
224
4
223
458
[
33
2
487
523
Column sum
992
15
9
984
2000
Table 9; Sign alignment between query and key noise in Stage II for model: 3L, w/o MLP, w/o residual.
init(t = 0)\t = 10
[s(t) K+.00
sit)
K4.9-1
((t)
K --.Q+
.+11
|s()
K-.0Q-
Row sum
(
483
2
26
512
++++++
|s()) PktQ-1
238
9
13
247
507
222
8
9
219
458
K Q+
1S--Q-
34
4
5
480
523
Column sum
977
22
29
972
2000
27

---

Published as a conference paper at ICLR 2025
Table 10: residual.
Sign alignment between query and key noise in Stage for model: 2L, w/ MLP, w/o
init(t 0)
10
k++00+
|
((t))
PK+Q-1
K
Xx.Q+I
(s(t)
K-.Q--
Row sum
15K+.4
484
1
0
27
512
[S -
245
5
12
245
507
[S
[K--04
Q+
213
12
5
228
458
41
2
480
523
Column sum
983
18
19
980
2000
Table 11: residual.
Sign alignment between query and key noise in Stage II for model: 3L, w/ MLP, w/o
init(t = 0) = 10
(s(t)
K++4++
(s(t)
KK.Q-1
|
Row sum
| Q+
473
3
32
512
|S -
233
12
14
248
507
|Sk- Q+
211
14
9
224
458
| PK--Q-
36
2
6
479
523
Column sum
953
32
32
983
2000
become positive, whereas remains negative with Adam. Additionally, in Figure 2(c), the training loss of SignGD converges linearly, while Adam exhibits sublinear convergence. While previously suggested that these differences might arise from Adam's momentum term, we did not provide detailed evidence. Here, we try explain these differences in terms of training dynamies and convergence rates.
To investigate factors influencing Adam's behavior, vary its f parameters and conduct experiments under the same model and dataset as in Figure 2. In Figure 2, we observe that 31 values ranging from 0 (no first moment) to 0.9 (commonly used in practice) do not significantly impact training speed. Similarly, in Figure 7, changes in 31 have little effect on training dynamics. Thus, our focus shifts to the role of $2.
Convergence rate. In Figure 16,we we observe that when 32 > 0.9, the training loss exhibits sublinear convergence rate. We remark that when the 32 0.9, the loss curve closely resembles that of SignGD, thus we use a range of [0.9, 0.999] for 8g. Since the training loss convergence is primarily driven by the growth of mean value noise, we believe this behavior can be approximated through the analysis of linear model fitting the noise.
Training Dynamics. Figure 17 (first row) shows that only small values of 3a prevent the negative query noise from turning positive. As 32 increases, the dynamics become smoother, and the evolution of query noise halts earlier.
To understand this, examine the mean gradient and update magnitude in the second and last rows of Figure 17. Unlike deeper transformers, the query and key gradicnts do not shrink faster. Instead, Adam's s update magnitude for query parameters decays zero before the gradients approach zero. This early decay of the update magnitude (or effective step size) can be attributed to 83. As 82 increases, the update magnitude decreases earlier, while the gradient shrinkage occurs at the same point.
These observations suggest that 32 plays crucial role in both the convergence rate and training dynamics of Adam, highlighting key differences from SignGD.
28

---

Published as a conference paper at ICLR 2025
Adamip.1
Adamx9t=0.55
Adamm9t-
30"
0.92
31
0.90
310
0.90
10
10
10
Training loss curve (log scale) on the synthetic dataset for different optimizers. In each plot, $1 is fixed to 0, 0.5, 0.9, respectively, while 82 varies from 0.9 to 0.999. The colorbar next to each plot represents the value of 32. This range for 8g is chosen because we observe that the training loss of Adam with 32 values below 0.9 very similar to that of SignGD.
Figure 16:
8,-0.3)
1.2
1001 DAB 1865
1583
1921
2500
500
1003
3900
12:00
2.00
2.75
Kirdd
10-
075-
0.70
1.00
8.03
500
3000
2008
300
1000
1503
1550 3000
500
1800
330
2003
200
3300
346
2004
1203
3200
Figure 17: Equivalent of Fig. 18 but using Adam with 31 0 and varying 33. The last represents the mean update magnitude, i.e., l norm between next the current iterates, of parameters Wo. WK all iterations.
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 /2d)/2n -2). Let {(X,, ((X,,)))== ==........ be the training dataset sampled according Definition 2.1. Moreover, let B, = supp(E.) be the support of E., Then with probability at least -2 B,nB, for all i E [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 (22 d = (n)y(n). Suppose that the training datasets are generated forratte.foolwww following Definition. 2.1 except that the non-zero coordinates of noise vectors are uniformly selected from [d] instead of [d] \ {1}. Then, with probability at least
1 (11 (nv/))) have μ is orthogonal to E, all [n].
29

---

Published as a conference paper at ICLR 2025
Proof. We have
P(3i E [n]] (E:)0 + 0] =
(1- 8 2)
I exp(
2ns/d
(L)(nvV)).
Note that this lemma can be extended to the signal vectors with constant non-zero entries without modifying the proof idea.
Let $1 fily, 1} and S-1 {ilm 1}. We have the following lemmas characterizing their sizes.
Lemma C.3. Suppose that 8 and 8log(4/6). Then with probability least 1 ö, |S-1, E [n/4,30 //3/].
Proof. Since [11 EiEM) (ys 1). |S-1| Eienn] [ -1), we have E||S[]]
B[|S-1]] = n/2. By Hoeffding's inequality, for arbitrary t >> the following holds:
2t
2t
P||S|| E|S 2 2 exp
P E-S ||| 2 1 2xxp
n
Setting
(nn/2) and taking a union bound, follows that with probability at least 1
we have
n
n
log ) ||S-1| <
4
|81|
2
V 2
2
log
Therefore, as long as 2 8 logg448)), we have v log 4/8)/2 < n/4 and hence n/4 [511,18-11 < 3n/4.
The following lemma estimates the norms of the noise vectors &, for all i [n]. Lemma C.4. Suppose > 0 and s = (4obgg(4)))...)) Then with probability at least
// < ||5:12 - //
ops/v2 || ops,
1/2), |3:11 = v o.s + 0( ((Vlog(4n/5)s 11/2),
for all i E [n].
Proof. By Bernstein's inequality, with probability at least 6/(2n), we have
| 07 O log(4n/))),
Therefore, if we set appropriately s = 22(llg((1))))) (4m/5)), we get
0/3/2 | < 3/
Let l E B, Since [k) is Gaussian, we have [A]] is sub-gaussian satisfying
[A]] Elle. [+]] 2 ||| [+]]l
- CCop
By sub-gaussian tail bounds, with probability at least
- 8/2n, we have
2
|,
V ops
=0 4n (n))))
30

---

Published as a conference paper at ICLR 2025
Therefore, if we set appropriately = ((log(4n/0)), get
aps/v2 S || S 0p8.
By union bound, we complete the proof.
Lemma C.S. Suppose ò > 0 and (2 (og(4n/0)) Let S1, S2 C [n]] satisfying S1 U S2 [m]] S and S2 are disjoint, and |S1] S2 = c for some constant c, then
ca,8/ /2
Σ
1s:hh
ElEl caps.
iES:
Proof. By sub-gaussian tail bounds, with probability at least 1 - 6/2n, we have
2 V ops
|5:11
/(opvslog(41))
π
Then have
Σ
||E 2 18il (V= 0ps 000o00((0gggl/////
iES,
||& S | (V2->>
Σ RE
-ps
Jps s + O(oov slog(44))) slog(4n/0))
which imply
2
Σ || Σ 115.112
JpS ops-O(nin 00nop v 2 V x
(og(4n/))).
ES
iES.
If we set appropriately
Qnn²² log( (An/6)), we have
Σ
||
||& 2 c0/8//
1/20//2. /22
RES,
VESu
Similarly, have
||
Σ VE
1kE:11 S cops.
RES
Now turning to network initialization, the following lemmas study the inner product between a randomly initialized value parameter vector Wy.ir for j (+1} and [m,,] or query/key parameter vector wQ.,/wx, K.s for 8 E [mL] in attention layer, and the signal/noise vectors in the training data. The calculations characterize how the neural network initialization randomly captures signal and noise information.
Lemma C.6. Suppose that ((log(m.m/o)) m. 20llg(11000000000 18)). Then with probability least
1 8,
d/ -
Vid.r d0/
(0)
(w (0)
V.inm 14)
5 V2log(12m,/05 to ||
(w
Vijar E2)
< 22m1og((2m,nn/5 agopy/s.
for all r E [m.), j E {+1] and E [n]. Lemma C.7. Suppose that Q(log(m.n/// m. = 2(logg(/66). Then with probability at least 1-8, 6,
(v(0) m))
$2 11mm./0) /0) m,.1/2 ||
(v(0)) ..:) 22 2og2log (1mm, m 000 0000V5 /s,
for all i E [n).
31

---

Published as a conference paper at ICLR 2025
Lemma C.8. Suppose that 22(oo(mmn/60)) my ((log((/6)). Then with probability least
1 6,
ogd/2 <
[
30/d/2,
max 0o 12
(w (0)
Q::: mb)
5 v 2200 oo 1|
* =E(mp]
max 000PV8/4
(w Q Ei)
(0)
$ (12mgg122 oudpVs.
mmim
odd/2 -
(0)
d//
K.s
00
||||| /2
max eE(mp]
W Kist mb)
(0)
- (2log(12ma/0) oo |112
000p <
max
(w (W Ex)
(0)
S 22/1og (00mm1m))/0) ooo,Vo.
mm
for all s E [m and E [n].
Proof.
The proof of Lemma C.6, C.7 and C.8 are the same as the Lemma B.3 in Cao et al. (2022).
Define
BE = max
{ (,
(w Q.::
Ex)
(w
(0)
k Ei)
(w:0)
V.j.r Ex)
(2)
Bp max a3"
Qu
(w k.a
(0) w1)
1w 14)
(0)
The following lemmas study the anti-comcenttatioo behaviour in randomly initialized neural network. Lemma C.9. Let N(0, o²). Then
2ca
[|4 S c) S erf
< 1
expep
/20
TO
Proof. The probability density function for x given by
1
x
f(x)
exp(
V2To
20
Then we know that
x² exp(
dd.
V2TO
20
By the definition of erf function
2
/ expl-x² 2d)de,
er((c)
and variable substitution yields
/ exp( x²
erf( c
202 dr.
v /2 V2To
Therefore, we first conclude P([rr| () 2erl( c=).
/20
Next, by the inequality crf(x) < /1 - exp( -xx /T) we finally obtain
2c expl
P[|2| S c) < 22
-2
32

---

Published as a conference paper at ICLR 2025
Lemma C.10. Suppose the results in Lemma C.8 and Lemma C.4 hold. Suppose that ogs 2 32 3|l| log(12ma/0r Then with probability at least 1 8, have
lo 54/4mgn")
(w s.) 2 28p; (w K.s E.) 2 28
for all s [mm]] and i [n].
Proof. Given En (w (w Ex) ~ N0, [1E,11). By Lemma C.9, we have
(0) (w
Q.st E2) S sc)
exp(
2ca тO ||
Let c 2. By Lemma C.8, 3r can be upper bounded by /2 (12gm11) oo soll By Lemma C.4,
/ < | S //
for all E [n]. Therefore, if we set appropriately
32 || log( 1 1
o
π log(1 - 82/ /mm4 n2)
then we have
< 5
P (w Q.a E.)
S 238
2m
The analysis for (w K.s: E.) exactly the same. By union bound, we get the final result.
Lemma C.11. Suppose the results in Lemma C.8 and Lemma C.4 hold. Suppose that m. 2
n² log (1 (1-82 (1-82 34 Then with probability at least 6, we have
256 log" (12mgn/8))))
(wo (0) Q, E.) 2 m, 1/4 (w K.a E.)) 2 Bem 1/4
for all s [ma]] and i [n].
Proof. Given En (w 2 Ex) ~ N0, 11),13). By Lemma C.9, we have
( 2. Ex) S sc) S exp( To ||
(0)
2ca
Let c= = 8em Lemma C.4,
1/4 By Lemma C.8, Bg can be upper bounded by 2 wlgg/o/mm//////// 0oop/s. By
// 1| - 30
for all i E [n]. Therefore, if we set appropriately
256 log" 2 (12mmn/6)
my 2 x² logg (1 82 / n)
then we have
P ((
(1) Ex) $ Bem 1/4 <
2m
(0) The analysis for (w K.s: E,) is exactly the same. By union bound, we get the final result.
33

---

Published as a conference paper at ICLR 2025
Lemma C.12. Suppose that n ( log(1 2
(=
3/2
Then with probability least
have
(w Q. 10) //
(w K.a (m))
(0)
oons 1//3-113 ||
for all [mz)]
Proof. Note that (w Q.:: 14) ~ N(0, og [1111 By Lemma C.9, we have
2c
P ( qiat m4) sc) -
exp(
TO |44
Let = oyns 1/3 ||ll If we set appropriately
3/2
8 2 n²
-T log(1 82 02/4mm))
then we have
IP ( (Ww
(0)
S oonss 1/33 || S
2m
(U) The analysis for (w K.a p)) is exactly the same. By union bound, get the final result.
The following lemma studies the magnitude of output of softmax operation. Lemma C.13. Suppose that (0lo(mmnn00)), least 8,
((log((/8)). Then with probability at
81.1
(0)
1/2
16mz log (00ggm1mmmmm)) oe max {ozs, ||
$121
(0)
1/2
16m, log 12m of max {ojs, ||
for all i [n]. When my log(1mmen/5) into
of max {ozs | o(1), we can simplify the results
111 121 1/2 + o(1.
for all i E [n].
Proof. We directly have
me
(w a (1)
(0)
(v(w
(0)
K.,: 11)
s mx
mi
w (0)
qia: Mi
(w (0)
K. y:E.)
S mp:3e
me
mk
10)
11)
(w
K. viEe)
Σ
(w (0)
Q.+ MrEe)
(w (0)
K. H1)
S 3p
a
for all E [n]. Then we have
exp (m838
(0)
$
exp (mx 3ppp)) + exp (-m+.33
1 exp (-mr3e(Be + 3,,)
1/22 3 Bp))
34

---

Published as a conference paper at ICLR 2025
and
expl (-m
(0)
$1,21
exp ( m. Be Sp) + exp mm:88))
+ exp (m Be(BE + Bn))
1/2 2m 1e Bm),
for all i E [n]. Similarly, we have
1/2 - 2m + 3p) - $1,11 S 11/2+ 2m + 1pm
(0)
for all i E [n]. By Lemma C.8, with probability 1 - 6, 3g can be bounded by 2 /og (2mmmn/6) oo and Br can be bounded by (2log(12mg//5) du || and
max (. }g} - 2/ log (12mmnm oo max {opv5, 1|
which completes the proof.
Lemma C.14. Suppose that /,8/ | 12(d1/6). Let (0)
X; (w o y:E,) + (w (0)
Ht)
Y; = - (w) k. yiEi) + (w (0)
"W H).
k.
Then with probability at least 1 nd- / for all [n], we have
Tp.- ||,a||
((w (0)
Q. veE
(WQ,.:
Q. 11) (w
(0)
(w (0)
K. (iE
(W..s 14),
(0)
o2S + ||1|||
or
0,5 + ||||||
(ww v1Et)) + (w (0)
K,st pb)
((w (. V.E. V:E.)
(wol (0)
12)
S,S |||||
Proof.
We only need show that for all
E [n] the probability of following event can be controlled.
0 ||
X; Y2 TpS + ||
ops - || X,.
o,S + ||
Conditioning on yıE, we have X; d Yi ~ N(0, of (114) + 11), this is because (E., 11) = 0.
Denote
2² = 2 + ||E,
D (0,)) IR. Tps |e
Jp + |
< y <
oys + ||
0nS ||
Therefore, we have
Op8 |eal|
8,8 + ||,a||
X, < Y, <
X:l1iE.
0p8 + |al||
8p8
||
JJ 1 ddPxx.Pr,
JI, 2TO2 exp
x² + y²
dady
2a²
22
2TO2 P exp
dr
20
π
Tp5
||
2 arctan
2
S +
|A
2 |||a|||
3 3 + ||
1/6).
35

---

Published as a conference paper at ICLR 2025
where the fifth step is by arctan(1 + 2) I + - 1) when a 1 - o(1), and the last step is by the condition sps/ | ()d Then, by integrating over yıE; and union bound, we have
P 3 E [n], 0ps - ||0l||
Jpst |||| X - O(nd 1/5).
X1 < Yi
oys+ + ||
0pS
|0||
Next, let
((0)
EE [n]
(w K.s E.) 0, (w
Q. E.) < 0, - 3
(4)
*,1,K+QQ-
+.
s.0)
8 E [maa] K. +.)
4.) < 0)
(5)
{ [mx] (W K. k) 0, k)) < <0)
(6)
uk++Q-
((0)
i E [n] Kw K.st M.E,) > 0, (w a yrE,)) < 0,
(0)
(0)
=,,,K+,00 U.:0)
(())
S,K+QQ
-
36

---

Published as a conference paper at ICLR 2025
D GRADIENT
In this section, present the gradients of the loss Ls with respect to Wyj.r, WQ,, and WK,, for s E [ma], j E {=1} and E [my]- Our analysis starts from more general form:
1
I (WW,, v.y. X.oftmax (x W W ))
(W X)
and finally focus L = 2 and o(x) case, which is exactly the model defined in Eq. (1). Lemma D.1 (softmax derivative). Let ($1..... . Sn)
softmax (21...., ,n), then we have
0s: 9; (1 (i = j} -5)).
0zj
Lemma D.2 (Gradient w.r.t. value parameters). Consider the transformer model defined in Eq. (1).
Let (X.y) y) be a data point. The derivative of (W, X) with respect to Wviir
1 mh
V VWV.. P,W,X)
[$L,1 + 82,1) x( + (81,2 + 82 (2)
=
Proof. We have
$t,0o ((wrsmx((e))) x(aa)
V Vwy, (,, X
my
[E(L]
a=1
(81,1 + 82,1) of ((Wvsn)x()) (1) (81,2 + +222) ((wrsso)
x(x)) x(2)
my
(01,1 + 82,1) x (1) + (81,2 + 82,2) x(2)]
where the second step is by L = 2, the third step is by (xx) =
Lemma D.3 (Gradient w.r.t pre-sofimax quantity). Consider the transformer model defined in Eq. (1), Let (X, y) be a data point. The derivative of F,(W, X) with respect to 21, is:
OF, (F,(W,X)
Σ my rE(ma)
81,1812 [( xa)) ((WVijrr (.VVj,r:x--a)))) Xz-a))]
Proof. We have
OF,(W,
Osi.b
( (Wvssnx((1)
my
dzi,a
rE m..]
$1,0 ( ((WKsmx(()) - ((mrsrr.(0))
my
rE(mp] bitt
31.a Σ 81b (o((()smm)())) -o ((((rsmm
my
bE(L]
Σ
$1,1 la ((WV,,r) - (V-W)i)-----a))
my
rE(ma]
$118t,2 ([ - ((WV
my
rE[ma]
where the penultimate step is by L = 2, the final step is by oxx)
37

---

Published as a conference paper at ICLR 2025
Lemma D.4 (Gradient w.r.t query parameters), Consider the transformer model defined in Eq. (1). Let (X,9) v) be a data point. The derivative of FW with respect to WQ, is:
V wa.. (W,
1
Σ
(WVS.r:()) -x(e)) (WK,; -x(1) (81181,2(() $2,1 $2,182,2300 $2,2x(
me re(mu]
=(WV,),x( (1) -x(() (wKioxx() (31191225)) + (22119222(1)
Proof. By Lemma D.3, we have
OF, (W, X
V Vw... ( = Σ
Ww.. 21
La
$1.0 Σ $L.6 [o ((Mvsnx(oo) -o -o((msss>x()))
})]]]
( (Wo,ar x(1)
myy
La
me|de]
DEILI
51,1 [o (( ((mvsm))to)) -
(WKKon x(0)
my
La
rE(me]
1 my
81,181,2 17 (WWr).x(") - ((((
(W,,ox((1)) x(1)
rE(ma]
WV Vir WV. x(2)
( (wV..x(")))
Wk.,xx(2) x(1)
+ $1,1
+ 82,182,
WV...rs'
(( (W....r x2)) (WK,,, x(2) x(2)
(WV... -0
((Wrg.
(WK,,, x()
(2)
x(2)
+ 92,1
1
0 (wrx.r,(()) -0 -0((>>) (WWV,3r,,
me rE(mml
W x -x(3)) (51,111,2() + 82,1 (2.12222((2))
1
(WVs)ri(()) - -x()) (WK,.. x(1)
xxx) -x))xx(xxx)x))xx)))))))))))))- (.2r91t55)) + $2,1 22.1)))
my rE(mu]
=WWVj)(())
- -xx)))
(WK. x(1) -x(1) (IL191eIs)) + (* +2,1)2222())
where the third step is by L L= 2, the last step is by a(x)= = x.
Lemma D.5 (Gradient w.r.t. key parameters), Consider the transformer model defined in Eq. (1).
Let (X,y) be a data point. Similarly, The derivative of F, X) with respect to WK,, is:
Vwg. FW, (Wr,,x"- (wQ,s) $2,1 s2.192,2x()) (x() -x(2)
Proof. By Lemma D.3, we have
Σ 51.6 (o (( ((mvaa))(/o))) -o ((mrsss (wec) x(o)
V Ww., FW,,
$L
my
1,a
rE(me]
BE[L.
1
my
rElmal
(W..: + $2,1 x()) (x(1) x
=(WV,)) (WVJsx(") xxxxxxxxxx -x(2) (wQ..r +111,,x1 + (211)22(118
(x()-xx x(2).
where the second step is by I. -= 2, the third step is by (x:)
38

---

Published as a conference paper at ICLR 2025
Lemma D.6. Consider the transformer model defined Eq. (1). Let W... Jor j E {+13, r [m.] be the value parameters of the TF at the t-th iteration. If we use sign gradient descent, the update formula of W (0)
(t)
Vaja is
W (1+1) Vijer
W (1)
Σ iE/n]
4u yi
[0 + 8,21 Da
10)
(1) 26.22)
V.j.r 1g8gh
Proof. This by Lemma D.2.
(1) Lemma D.7. Consider the transformer model defined in Eq. (1). Let wQ and W.s be the query and key parameters of the TF the t-th iteration. If we use sign gradient descent, the (t)
for for [ma]
(0)
update formula of Q,e
and w K.a
(t)
WQ,a
(++1)
W..
(1)
778g1
Σ (()) (o
w 1. E,) (w)
K. Biu E,)(5 (9) 91 v. + (9
22)
iEiml
W K.8 Kis
(1+1)
(t) (2)
1) sgn
w(O)
vi E.)(w (0)
0. 811 12 9e + 521 22 E,) Ex)
(0)
(E(M)
Proof. This is by Lemma D.4, Lemma D)5.
39

---

Published as a conference paper at ICLR 2025
E PROOFS
In this section, we give a detailed analysis about the dynamics of sign gradient descent on our transformer models Eq. (I) with our data model in Def. 2.1. These results are based on Condition 3.1 and the conclusions in Appendix C, which hold with high probability. Denote by Epretim the event that all the results in Appendix C hold (for given 8, we see P(Eprelim) 2 1 106 n 2 nd 115 by union bound). For simplicity and clarity, we state all the results in this and the following sections conditional on prelim.
E.1 TECHNIQUE OVERVIEW
For value, the dynamics is linear under sign gradient descent, which means the increment exactly constant across different iterations and different neurons, and up to constant multiplier across different samples.
For query and key, the dynamics is complicated. The key point is that utilize the different increasing speed of different core quantities to divide the whole timeline into many stages where we care about only one of some of quantities in each stage while the dynamies 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 pp8, op fyy/s,o oo-
We will define some key timesteps in following analysis, including
T Ô(onn 15-1/2m,) 1/2),
Ta == 50V/n 00p 5 = 0(008-1121111)
T2 Ci vlog(1 (10m1mmm///m //)ooons 1/2 0(00s -1/3,111111
Ts = 3, -1 = Ō(oon
T5 C2 C2vloge 10gg(1m4//5org
Ō(aun-1)
T 2 f =
1022 0.99T
log
0,8 2v ||
17 1 - 1/2 1 1-=00mm 1/2 op -- 1-
m
0
2
TI $ T = V
1.01
2
log
||
Tp8
1m
17
m
112 -1 Ö(mm
0p
1/2 1-1-- 11 11
0p
T S (1 + 0.)T2,
C3ops
T4 = C. log
"1 m
0, -1/20-11--1
1/2
Ō(m
-1/2 1 11
||,a|||
where C1, C2, C3 = (1) are large constants. The order among these timesteps is
T < To T < To < & TT <T < T4.
Our analysis basically consists of following parts:
Mean value noise. Let
Ti : Olonn 1= 1/2)..
For t [0, TiJ. the mean value across all samples align quickly while other inner product barely move. For t 2 Ti, we have the mean value noise is linear with time
V2lma (v(0) y.E..) < 211
for all E [n]], and mean value signal (v(t), (u) can be ignored compared with mean value noise (v((), V.).). Moreover, for t > Tz, we show the mean value noise across different samples are almost the same
- /2 1np + (0s--(11))
(v(1) ..E)))
for all E [n).
40

---

Published as a conference paper at ICLR 2025
Useful bounds. For E [0, 13 noise- softmax outputs close to initialization, i.e., for all E [n]
(e) 1/2 + o(1).
Fort f E Ta), signal-signal softmax outputs are close to initialization, i.e., for all i E [v]
( 1,11 = 1/2 + o(1).
For / E [0, query/key signal are close to initialization, i.e., for & [ms]]
(w (e) K,sz pa)
(w (0) K. 14)
(1 (1+)
Q..t w))
(0) (w
.. u1)
(1 (11)
For /. E Ta). the loss derivative are close to initialization, i.e., for all i E [n)
eto = 1/2+
For t ]. the ratio of niise-signal softmax outputs are close 1, i.e. for all i, [n]
21
o(1).
k
Query/Key noise dynamics part i. Let
T == C1 (ov/lg(122mpp)) (0)ogm/m0)/ons1111 11: T2 50 o
s- a
where C₁ = (1) is large constant. For L E T2). 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 TT. All query/key noise can be divided into two groups. Given E [ma] and E [r)].
1. if Q. B:E,)) =ggu (w(O) tude).
(0)
(0) K.st E:E,), then they encourage each other to increase (in magni-
2. if (w! (0)
Q.a M.E Fegn -ww K.a M.E.) then they first encourage each other to decrease in
(0)
magnitude, but rebound quickly and reduce to the first case.
For 1 E [T2, J, we have some nice properties about query/key noise
(wil) (. VeEa) =gg (w K...
(t)
(w (w. .VEE.) Q.:
(++1)
Q M1E + sgn((w()) (t)
(e)
Q.a: vrE (
(w K.s .yEE,) (w k. UrE,) + sgn((w (t)
(t+1)
K.s y,E.)) |l
m S (w ( yxE,) (w K. )
2b1o,s
for all [ma] and i [n]] Also have some nice properties about the of query/key noise Eie(n (0)
Q.a (iE.) V.E.) Eie(n) (w k. ViE,)
(1) a yrE.)
Σ
w/ K.s B1E.)
2 22
t=1
w
(1+1) Q,a ,yiE,) >
(w (c)
2.a viEt)
+ sgn
(t)
Q.ay B:E,))
V2" no
7
(w Kis ,yiE)) 2
(++1)
(() (t)
Kist mnEE)) + sgu Σ
w K. 31
v2" 11
for all [ma].
41

---

for all s [mm]]. Published as a conference paper at ICLR 2025
Query/Key signal dynamics part i. Let
T3 33,07
| Ö(oon )
T3 = (22/lool120000 (12mx/05)o0
where C2 = (1) a large constant. For & (72, Tz], 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
( UoE.) agn (w (1)
K...Bi.)),
(wilt)) -VGE.) (ww)) Q M.S,) - sgn((w (1)
Q Me ) 11
(w+1))
K, y:E,) -(w Ki v:E.) + sgu( (w K,st y& ||E:]1
(t)
(t+1) (e, .y(E,)
. N
(w()) (t) Q.a MrE)) + sgn( Σ
2. B1E,)) 110p8,
(t)
V2"
Σ
(wk (++1)
K,s yeE,)) 2
(wk) (4) K.s M:E,) + sgu
(w/)
(0) K 3:E,))
V2 110
for all E [ma]] and E [n]. We show that the update direction of query/key noise is determined by the summation of query/key noise. Fort E [T2, Tz], for all s E [ma]] we have
(1+1) w q. [1] = (w Q.a: 14) + sgn(
(Ta) Q. V:S.))) ||
(w
(w K.s 12) = (w (w..st (1) sgn
(t+1)
(w
Q.a
(Ta)
BEE)) [|||
For t E (Ts, 1. we have some nice properties about query/key signal
(w Q.st (1) =agn
(ww!!)
11)
K,st M.E.).
K.
=gg
iE[n]
bny |||||||| /2 S (w(()
Qio: 12))
//(w// K. 12)
S 26y ||||
(w (++1) 11) (w(() (1)
2.00 11) sgn(( (w/ (1)
Q.a 14) ||
(w K. (1) (w (1)
(1+1)
K. sgn K. p))m ||]l]
(1)
Query/Key signal and noise dynamics part ii Let T7 2 T3 be the first time the following condition does not hold, for all
[me)
(w (2)
(w (t)
(. mi)
2,21 2.22
2 2
(w v:E,))
> 23
(w e. -1)-
2.1 1,22
Let T 2 TT be the first time the following condition holds, for all s E [mx] and i [m]
8,21 9,22 (w a viEa)
(2) (t)
1 (t) (6)
(w (w (1)
2 1,
Let
Caans ||pa||
T4 = C3 log
1 m
1/2
0p
where C3- O(1) a large constant We discuss the dynamics by three parts. For E [T,, stage IV.a in Appendix E.7.1. For i E (Tt) stage IV.b in Appendix E.7.2. For
42

---

Published as a conference paper at ICLR 2025
t [TT,, Ta], see stage IV.c in Appendix E.7.1. For t 2 Ta, we have for all [mg] and i E [n]]
(1+1) (w 0, (e) = -(we. 14) + sgn
w (7s)
Q 315,)) |eell
(W,,s K,s
((+1)
.10 (w (t)
K,s 11) sgn
Q.
(Ta) 3:E.)) |l
(wQ.s ()t,,) (w (1)
(f+1)
Q.st VIE,)) sgu (w (1)
Q y )) ||Ell
K,s K.s VvE,) (w (t)
(++1)
K. UvE,) sgn( (w K.st v. 1
(t)
(w Q. y:E,) =ugn (w v(C, agn (w (Q.lM) =gn -w H)
Note that we have following bounds for T4 and TT7
T 2 1
0.99
log
3//n 0 10
11 mmm
mk
-1/2
1/2 1 18_
0
2
Ti < t =
1.01
log
JPS
m
m
1/2 -1 Op
2
|
Tf ~ (1 + 0.)T
E.2 THE DYNAMICS OF VALUE
Write down the inner product between gradient of value parameters and signal & noise, have
(sgn (V Vww.s L) (4) -ggn
Σ iE|ml
+ + (6) ((a))
[[ll] -sgn(j) |ell
for all j e {#1}, [mm] and
(sgn (V (Vw... L) ,E))
Σ sgn
Σ 1 we
10) E2) E,(,)),
kE(d]
MEln
for all i [n].j [=1} [mo]. and simplified to sparse setting
(sgn [V L) E.) = Σ sgn
(1.0 u: + (2) (..2)) E.
REB:
for all E im].j (=1}, r [m,]-
Consider the sparse property of the data, we have
(sgn V VVV... L) EE:)
Σ sgn keB,
(1 un + 2222) E.1A]
sgn()
Σ [5,
kEB.
- sgg(().)
|11
where the last equality is due to l' and st,ab 0 for all i [n].a, a, 6 [2]. Note the update to
value parameters of is irrelevant to the magnitude of query, key, and attention weights. Then, we have (t)
(w Va Vat1) (w Van.r Map1) n ||
(I+1)
W.- Vipt) (w (W r Bipl) n ||l4|]
(1+1)
(0)
(w ((+1)
(Wi.r {,) (Wrarrr
(0)
{.) 17 ||Eilli
(w (wV.-u..r- E,) (w (WV..vrr E,) n 1l
(t-1)
(1) V. Vifte) and (w( V.g..r E,) increase. Also, (w Vige.r E,) increases faster since have sap |||=|| 2(1)
As become large, both (w (t)
43

---

Published as a conference paper at ICLR 2025
E.3 THE DYNAMICS OF QUERY AND KEY: PREPARATIONS
Additional Notations. Let v(t == WV (2)
-1- When we talk about query and key parameters. we refer to WQ.s; WK,s for all s [m,]. When we talk about query/key signal (inner products), we refer to (Wo,s) (1), (WK,s; (1) for all E [ma]. When we talk about query/key noise (inner products), we refer to (Wo.ss S:C.), (WK.s; W&E;)) for all 8 E [ma] and i E [n]] When we talk about (mean) value (noise), we refer to y.C) for all E [n]. We use aga b to denote sgn(a) sgn(b). Let B, be the support of sample E; in the training dataset.
Some Facts, Based on gradient of query and key parameters and the sparse property of data, firstly write down the inner product between gradient and signal & noise
(sgn (V "wo. T.) :14) = sgm
Σ iE|n
Yi
V11
W (1)
V. E.) (w K Beft 6.)8,,11
(sgn (Vwg.oT)) [.)) = Σ
sgn
(()))
Vi
V,1
(1)
-w
-. Vift - -Ei)
K, Vitk -E) 21.25.111) E,(A),
k B.
(sgn (Vwx.. (o) sgn
Σ
((t)
Ye No(W V.1
(1)
w/ 10
1. 4:14 E.)
Q,at 81 12 + (6)
(t) (4)
iE|ni
(sgn (Vwx, ,.E,)
Σ B.
y/ (w (1)
(0 W.. Vil (.)
10 0 2
$ 221 (22))) (-E,1K)))) E.
sgn
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 2 O we have
[w K.s
(e+1)
w
(c)
K.st 11)
(w (1+1) w (1) Q,s: 11) 17 ||
Q
(8)
(w K.s w (1)
(e+1)
K E
(
(e+1)
w 2 yrE,) = 1 ||Eall.
(2)
Q.
which means that we only need to analyze the sign of update direction for all quantities of interest.
Secondly, we take look at the key part in gradient formula for different quantities:
(v(2, Vilt (w (1)
K. Will E.)
(1) = =(v(0, (H) K. 1o) vo(v!9, pa) (w K.s Ei) 10 (w) K. ()) + (w)
(1)
(..... E.),
(v(1),
yipe 1it (,) (1)
Q. Vol4 + Ei)
= -(v(0) (w) o p) + ye (v(1)) 11) (w 2. E.) (v((() E.) (w (t)
(t)
a. 14) - (g) Q. E.)-
By the subsequent analysis in the stage I, we have that the (v(0) wo) o(1) (v 3 such that we can approximately neglect the effect of (v(t), p1). Also, we have (v(") y,E) 2 0. Then, combined with the gradient formula, we have
For (WQ.s: yeEa), the sign of increment is aligned with
(w (t+1)
(t)
Q.at Q. UiE.)
(2)
Q,e
-agn (v())
ViE,) (w K. 1) + (v.
(c)
(w U.E,)
sgn K,s 14)
(w12 K. weEl).
(1)
(10)
For (WQ. ui), the sign of increment is aligned with
(w (w(+t))
(++1)
w :
(t)
(1)) (t)
1,11 (112)
W. 1)) + (v (v... (R.. (11)
sgn
( m2]
44

---

Published as a conference paper at ICLR 2025
For (wK.ss yrE.), the sign of increment is aligned with
(w
(E+1)
(e) w K V:S.))
K.8
sgn
(1) (1) (v(0) y. (w (1)
Q. pa)
1.21
(c) (2)
9,22 (v"))
(.EE)) UiE.))
(1)
14) + 1,21 1.22 (w Q, V:C.).
(2) (t)
(12)
1111 1,12 Q.
For (WK,s, p1), the sign of increment is aligned with
(++1) (w K2s
w (f) K, 11)
(1)) (- (0) 10)
8,,1 1,12 (v ,:E,)) Q... 14) - - (2) Q.
6,22 (v
M. (w
Q.a viE,))
sgn
vE[n]
(13)
E.4 STAGE I
In this part, we consider the dynamics of mean value signal and noise, i.e. (v(") :&E,)) and (v(0), 14). Roughly, at the end of this stage we want to show that mean value noise evolve fast and can be viewed as 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
(y((+)), (²) = (v("), 11) + 2n || (v((+1), = ((('), + 2n |lE:
(14)
(15)
Let
T 4.8mm, 1/2 p s-1
Note that we have (v y:E.) S 223mm. 1/2 then we have
t S (v(") .v:E.) 2tmops,
all Z Tj and i [n]. However, the query/key noise do not deviate too much from initialization at Ti since the deviation is at most
(w (11) Q. Ee) (0)
Tin 11
51+
< 1
+ o(1)
2 ((Bmm. 1/4
Q.st {,)
{,)
where the second step is by Lemma C.11 and magnitude of Ti. Besides, the deviation of (v(Ti), .1) from initialization also small. Formally, we have
(T), p4) (v(0), p4)
Tin ||]|
< Ö( ||||| 1/2 m/2m" -1/2)
51+
oo(),,
[v(0)) p1)
2(ao |||||| ns ns-nnnn/n 1m3 1/2
where the second step is by Lemma C.12 Therefore, have
(v(), 11) ((,))
for all i E [n]] and > Ti, where this holds f = Ti due Lemma C.8 and the definition of Ti- Combined with the form of gradient of query/key parameters, this makes us able ignore the effect of (v(t), p1) afterwards.
E.5 STAGE
In this part, consider the dynamics of all query/key noise. Roughly, the end of this stage want to show for all neurons E [mg] and samples E [n]. query/key noise have the same sign and linear with t.
45

---

Published as a conference paper at ICLR 2025
Let
T=: :- O1 -- Iog (1 2mma/0)oonns-/777
(0) where C = (1) is large constant. In this section, will analyze the dynamics of (w K, V:E,) and (w (t) V. Y:E,) for all E [ma]] E [n] and T
We first consider the samples E S +. U S
sK at initialization. Note that without loss of generality, we can assume that E s, .Q+ Then, by Lemma C.14, we only need to consider two cases where the one is when we have
((0)
0p8 ||
(0) K Me + (w (0)
(0) (w.. V:E.) (WQ.st Q 14),
(0)
K,s 11)
0p8 + |
and the other one when we have
(w o viE,) + (w K, 11) 0p8 + ||
((w ((wQ.s
Q, y&E.) + (w (0) q. 1i)).
(0)
0pS - ||||||
Finally, we give an upper bound for T2,. (defined later) for all E [ma] and ( ((U) q. and bound the magnitude of (w!!) (!)
+.+
K. M&E,) and (w v&E.) before T
We define following useful timesteps. Let T2, be the first time satisfying
(w q. ViE.) =aga (w W. UiE.).
(16)
Let T2.,, be the first time satisfying
(w 1++1) w a U.E.)) Fagn (wkkte - w K...
(17)
Q
Let T3 be the first time satisfying Eq. (16), (17) and
(w ((+1) (1)
Q.+ Q. v(E agg (
(18)
Let T2 = 28g1 0p
The following lemma studies the order between T2, T2, 2,,,; T2 Lemma E.1. For all 8 E [ma] and i ((0)
- U ((( - we have
<< 2, S T
0 < T S TT,
Proof of Lemma E. WLOG, assume that i - By the definition of .Q+ have
(0)
(w K. V.E.) =ggu (w (0) a. 31E,) =sgn 1.
Then, by the definition of T2,,. we have 0 < T2,,.
By the definition of T 2, we naturally have T2 < T3 2,,, and T2,0, < T
At / = 0, we have
(wQ.e (t+1)
Ve....15 v.E,) =gn -w n)) + (w K.s U1E,) =gn (w Kiet MiE.),
w Q.
where the first step by Eq. (10) and the second step is by Lemma C0 and (1+1) (w K.s w (0)
Q. u) + s $ (w Q. v&E.)) =ago (w.) U.E.), where the first step is by Eq. (12) and the second step is by Lemma C. and C.13, which implics 0 T2 T2
K.s. B.E.) =agn $1,11
(1) (0) (w
Note that by Lemma 14, have that for all E [n] and < T
211 8 = o (1 + o(1), 2 (1 + o(1) = 1/2 + o(1).
In the following analysis, will use this fact many times.
46

---

Published as a conference paper at ICLR 2025
WHEN - (WW2s (0)
K.st MeEd) + (w (w
K k
(0) -ل-_]|1
||ااا|||
(( Q vvE.) + (0) (Q.s THE UPDATE DIREC-
(0)
TION ARE DETERMINED BY QUERY NOISE.
By the definition of T we must have that (T)..))
(1) ((W (WQQ.* Q.
(T
(1)
w Q ,VEE))
((w)
K.8
w K
v.+))) O.
Note that we have
(w (T2 +1)
(2)...)
(2)..))
.p) + (w (T2.as)
W Q.
U:E.) =gn
(WK,s K,
K. 9:S,))
2
(w (0)
K 1i) + (w K,st ViE.)) + T2 7.. - | 2 O,
(0)
(w (2 +1)
(1 (25,..)
(1
11) + .. (2g,..)
(22,a.)
K28
w K.
Sun 8 $1,11
S $1,12
(W Q
S1,21 $,,2
(w
.yE.,) Q.+
4 (1= 2. 11) + (wo, y1E,) T2 M + || 2 0,
10)
(19)
where the last steps in both lines are by the condition
K. pa) S TpS ||al|| (0)
K. MEE + (w (0) K.
(U)
TpS + || 1, U.E.) + ( (0)
11).
We claim that the sign of update would not change for t E [T Suppose at t S 8 [T2 T2 the induction hypothesis holds, then we have
(t'+1) (w K. ..E.))
(w K.a V(E,),
(We. Q,a .2ih,) > (w Q.st viLi),
(t'+1)
which imply
(1+1)) .D) + (Wiss , .vrE)) 2 -(w (t')
(1++1)
K, 1e) + (w (4')
-(wk,
K, yeE,) 2 0,
2111 1,.2 (WQ. (r) + 8,21 8,22
(1+1)) (141) (1'+1)
(1+11) (1-1)
(w Q, -.EE)) 2 1 (1 + o(1)) (t')
(1'+1)
(4) (w))) + (ww (0)
w vrE.)) 0,
(2) V:E.), (w WQ.s -1E.)) >0
(I)...)
then the conclusion holds and thus (w
K,s
Moreover, we have an upper bound for T3 2,,,, that
(w (0) Q. 1)) + (w q. B.E.)
(0)
1.5(w
Q v:E.) S 1.5T2,
T2.
S
n( + [11l)
"7 &
where the first step is by Eq. (19), the second step is by Lemma 1.10, the last step is by the definition of TV and Lemma C.4.
WHEN- - (w K. M.E + (w (0) K. p1)
3,5-114
((w
Q. M=E,) + (. THE UPDATE DIREC-
(0)
(0)
TION ARE DETERMINED BY KEY NOISE. At ttt= -t? 2 we have
(TH (w
(T)..))
.vEE)) =gg (w Kis .p) (w (T
Q.9 -Q..
K.S v.E,)
S K 12) + (w Ki yeEt) + T2
(0)
(0)
((1lall + +||11) S 0,
(20)
(w (T2...+)) +1)
(T2...) w:S)) =ggn
(T2 (T2...) (T
(T2 (TZ...)
(T (T) (T2...)
K.S
W K.s
8,11 $,12
(Wo. 1/) + 51,21
81,22
(Wo. .3iE.)
(1 (1== O. 14) + (w Q,at ( v.k.) - T T2,, m.llsl (|ml)) S 0.
where the last steps in both lines are by the condition
K. m) > TpS + ||ا|||
-wwihee K.e MeEa) + (W/) (0)
(0)
(0) ((w 2, v.E.) + (w:0) (0)
H)).
Tn8 ||
47

---

Published as a conference paper at ICLR 2025
Then, similar to previous analysis, we have for t [T2,,,: 1, we have
(wkk (t)
Ki 14) + (w
Kist viSi) < 0, (t)
1,11 (w 2 14) + 1,21 1,22
(t) Q,st MiS.) < 0,
(T aT
(T (T2...)
which implies K,s
y.E.). ( yBE,)) 0, and
(0)
K.st (1) (w K ))E,)
(0)
1.5( n 16:
K. viE.) S 1.572
T 22,,5
5111 + |
where the first step is by Eq. (20), the second step is by Lemma C. the last step is by the definition of T and Lemma C.4.
Next, we are ready to bound T change its update direction. WLOG, assume query noise changes the update direction while key noise doesn't. t. For the key noise, we have with at most Tř steps, its sign can align with the sign of its update. For the query noise, we possibly have that the sign at T2,s, can be contrary with the sign at initialization, but with at most another T2,,, steps, its sign can align with the sign of its update (and at initialization). By these analyses, we have
22,,5" Note that during [0, T2. only one of query and key noise
T 2. S max (T2,
2T2 2.,,) S 337,2
for all 8 E [ms]] and E + U (.(0)
((0)
-
In the main text, we define T,SON 3T2-
The following lemma gives bounds for the magnitude of query/key noise before T Lemma E.2. For all s E [ma], i 5(0)
S.K+ U US(())
s and E [3T2, T.l. we have
(w Q.s yiE.) = (w Q U:E.) + sgn( (g Q.at S(E:)) l|E:la
(t+1)
(++1) (w K,s
y.E,) = (w k.s U&Ei) + sgn (w K. v.E,)) )n |llll
for all E [12Tg, 7
tnoss/2 2
10 ( MEE
(1) (W o MEE S 2tna,s.
Proof. The first two lines are shown in the above discussion. For the last line, for the lower bound, (T...) ..EE)) > we have
suppose (Wo.a
(w (1) q. ):E.)
mo /2
(3TH) (w .vE.))
(t mop
8/v2 //2
11o,8/22
where the first step is by Lemma C.4, the second step is by 3T3 2 T 2 the third step is by 12T2. For the upper bound, we have
(w W. viEa
(w (w o. mEEi) + tmops S 2tnoys,
where the second step by / 2 Tgg
We next consider the samples E S00 ++ Q+ U Q- initialization. The following lemma
studies the update direction and magnitude of these neurons.
48

---

Published as a conference paper at ICLR 2025
Lemma E.3. For all 8 E [me]] i E (0)
+.++ U s(0)
S and E [0, we have
(WQ.*
(w (t+1)
Q.+ :.iE)) (w (t)
Q,a V:E.) + sgn(( Q..: ve 1
(0)
(w (t+1)
Kis MeEi) (w Kis veLl) + sgn(( (( k. v:E,)) ( |
for all [T2,
(w (0)
/ ( viEl)
S 21naas.
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 E S **K++ We can show that the sign of update of (w (. yrE,) and (w K.st ViE,) are the same as their sign at initialization and would not change for / E [0, Tzt Formally
(e)
(++1) (w K.s -.E.)) (( (8)
K,s M.E.)) Fagu (w (0)
K.s U:E.), (0)
(1+1) (w Q.s M.G.)) (( Q.5 M.E.) Fagn q,st V:C.))
for all S t S Tzt. Similar to the proof Lemma E.I, the claim holds at t = 0. Suppose at t I t' E [0, 125 the induction hypothesis holds, then we have
(f'+1) (w Kio M:E.) = (w Kist M.E + ||E 2 (w (e')
Kist yoE.),
(w Q.+ yiE,) = (w Q.s M.E + || 2 (w (t)) Q.a v:E.),
(1'+1)
which imply
(1+4) .1) (W K.s MEE
(f+11)
2 - (w K.st (4) + (w K. (.E.) + nllEilll -1|el)) 2 0,
and
(1++1) (1--1) (1++)((1(11)
1.11 1.12
(w
0,s
(t'+1)
14) + (t'+1) (t'+1) (1'11)
8,21 $,22 (.. ..E.))
2 7 (1 + o()) ((w
(t)
Q (1 + (1)) + (w
innnlllllll (|xl1) 2 0, a
U.
where the first step is by Lemma F. 14. Then we show the result for L = l' + 1 and prove the claim. Therefore, for the lower bound, we have
(wf)) K.s V.E.) (w K,s U.E,) + tn [| 2 tn |1 2 tnopss/2
where the first step and second step follow from the claim above, the last step is by Lemma C.4, and for the upper bound, we have
(1) (w K. U:E,) (.... K. MrE,) + tn ||&/ - Tano,s + tn |1 S 2tnops,
(0)
where the first step is by the claim above, the second step is by the definition of Ty the last step by t 2 T3 and Lemma C.4.
Remark E.4. Combining Lemma E.2 and E.3, we have for all t [121H, TT E [mx] and i
[n] (21)
(w (t) tmops/2 (w Q. U1E,)
(t) (wW.. K. v:E.)) /
2tnops.
Note that the upper bound also hold for all 2 Tst since the magnitude of update of query/ky noise is always 77 ||E: 11, i.e., Eq. (9) 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 F.S. Let T3 50nTz. We have
49

---

Published as a conference paper at ICLR 2025
(update direction). Lel 6x,i = sgn((w Q, -.E,)) We have the sign of the update of the summation of query/ky noise would keep unchanged for [3T2 T.l. Formally, for all s [me] and t E [3Ty, 1
(3T2)
Σ w (1+1) R.a
U:E.))
2(w K.a V.S.)
( (1+1) U.S.))
Σ
(w/ O.or MiS.) = Σ cou |16:11
Q,
2. (magnitude of update). With probability at least 1 O vn), for all E [mm]]
S.s ||eill 2 nops/V2.
3. (magnitude of query/ke noise summation). For all E [mx]] and E [T2, Tzt
T1
sgn 7
Σ Ca.
( U.E.).
sgn(n)
Ca (IE E V.E.) 2 2n
To prove this lemma, we need the following lemma which studies the property of cs,y. Lemma E.6. Let
Ea,+ == i E [n] (w (T))
MEE .1E.) > 0)
Qa
Ea,- = i E [m] (w
Q. v(E,)) < =0.
(T2)
Let
Asi = (w EΩ:: :::::=(w:> (0)
-(w2 MrE,) + (w (0)
K.s: 14) < (w Q. ViEi) + (w/o) w. 1.)
Let
1/i I S. + 1/1 S
-1 S. 11i S. Q+ 1 - 1 E S A
Then we have
1. (The distribution of X,,)). For all E [mx] and E [n]]
X 1, w.p. 1/2
1, w.p. 1/2
2. (The probability of equal set size).
P P([Es,t] 1B,-1= 0] 5 (1221/2).
3. (The probability of lucky neurons).
n
P[X., 0,
Σ X,,
2 4
11/2))))
Proof. By the definition of X and cs,i, we have following equivalent formulations
n
|E,+I (i E [m] : C.,; > 031 Il S.k+.0+] + 1/i S..e -.. 1li S,xt,ou;
1,,-1= ((i E [m] G,. < 03|
Σ ==1
I S.,k - l S.K - -Q+: AEd + l1 E S.,k+.Q, 4%,]
50

---

Published as a conference paper at ICLR 2025
Let Qa,i Qw:. MrS.), Ka (Wa, viSi), Ms (w K,,, K,s wQ. 14). Then we have jaa
(0)
(0)
(0)
(0)
(00, ob [|)l))), and Qs,i, Ka,lBbiE (0, og [|E,1(). Denote the distribution of & and p by Pe and Pa. respectively.
Let p(p, Ex) == P(X.,t 1/ = U. Pp P(Xs,t 1/4= = jl.
Let (1, == Pii E 5..- = u, E.E. @n = Pi E Sak-Q+ Agilus = ud. By the symmetry of -+ and Sa,K+ we have Pli e SaK--Q+1 Aoilu, Eil Pi E
Ss,K+ As, E.l. which gives
p(u, 6.) +2)(,,E)).
Ps. = / plu, S:)dPe - 1 - - 20(,,, E.) 4
1 + 2
By the symmetry of the gaussian density, for any E., have
q² + h22
exp(
dqdk
14>0,k<0.0tk50
20 | 1
9² + k2
exp(
ddqkk
2 E,
g² + k2
exp(
20 | dddkk
4 -(-H.E)),
which implies 6(4, {, - $ is an odd function, and thus
Pli E S..- A..lel = f( E.)dP, S(dP. g'
thus Pii Ss,K-QQ, A.,) and P(X,, = 1] 1. which proves the first result.
Note that given μr, Xs,, are independent for E [n]] then
Ea - |E,_- 0] = X = 0]
n1
p
Pa
/dPPp
1//2/
71
1
dPp
T//2
n
n1
10
<
n/2
2
which proves the second result.
Recall that
o(,, E,) = 1/4 - 6(-pppp u. E.) VE,
= 0₂ = 1/4 8-p
Pp = p-
51

---

Published as a conference paper at ICLR 2025
The quantity of interest can be written
P[X,,; > 0. >
Σ
X,,,
0]= /
E
(w) pu)" Pu
jti
11/2/+1
11
Pr L.2
P[Xa,i > 0, Σ
Xa,i <0] /
Σ
(x) Pp Pp-k((1 Duu) PrdP
jz1
k= 11/2/+1
12
@n-1.p
P(X <0, Σ
Na,1 > 0] =
Σ
(m)
(k)
(0) - Pa) pm)m- (1-PP) Pp)dP
1#1
k= 11/2]+1
1g
P(X., a, < 0,
Σ
Na, <0]
Σ
pp - -Pa)* (1 -P
(n
By the symmetry of Pa. and Pa. we have
p| > 0, Σ
X.. > = <0, <0.Σ Ex.,
X..i < 0),
iti
jtti
pf > 0, Σ
N..i < <0] = PP(3,, < 0,
x
Σ
EX,, > 0].
X
Since
11 + 1414-12 13 /
(p)))(n-11)) Qn- -n- )dP, 2 0,
where the last step is by Pn-1,, > 9n- 1, when p/, 2 2. then we have
F|X.. O, Σ Xa,> 0) > 11 - Xx, 0)) On 1/2
where the second step is by 1 1 2 1n 1g. the third step is by the second result above, which concludes the proof.
Proof of Lemma E.S. By Lemma E.2 and E.3, we have the update direction of each element (w (t) K.s yiE,) and (w B1E,) is unchanged for t E [3T; T5"), which implies the first result.
By Lemma E.6, with probability 1 - Om2// /v), we have for all [mg]. |E,+1 - |E,,-1 * 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 |E,,t - -IE,, 1 2 1, then by Lemma C.5, we have
G |1Eil1 2 nops/V2.
(22)
Finally, we note that at / = 12T2, we have
Σ w Q ..E.)),
(121g)
(12Tg) w K. ..E.)) 2 -24T
By the definition of T2, we have T3 2 (12 + 24 2n + 2n 1 Then, for t E Tt), we have
Q.e v:E.))
(1)
Σ
w
k. wcEi) > 2n
(
(23)
52

---

Published as a conference paper at ICLR 2025
Similarly, if Ea,+l - E -1 for some other neuron 8, then for t [T2, T2
we have
Σ (w Qi yiE:),
w K. M:Ei) S -
E.6 STAGE III
Let
T3 = C2 (C2Vool(22mg)) (09g(1mm//)0mm) 18)on'
where C2 = (1)) is a large constant. In this section, we will analyze the dynamics of (w (t) (w (1
Kis V.C.).
(t) q. V:E,), (w Q,st (4) and (wx.s K.s: 14) for all E [m=]. E [n]] and T2 S t S Tg- Generally, already have controlled the direction and magnitude of the quantity Σ (w Q. yrE,) and (t)
(i=(w k. K.ez UrE,)) at the timestep T2, can show the update direction of query/key noise is
determined by these quantities.
Fort < T2, we want show the change of (w (1) much. Actually, we can show that for (t) t S Tz (w q.s (e) and (w K.
(t)
(W,, (4) are close to initialization. Formally, we have the following lemma.
(t)
Lemma E.7. For all 8 E [mm]] and l S TTt, we have
(w (6) K,st 11) = (w (t)
Kist (1) (1 o(1)),
(w (4) = (w Q u)) (1 + o(1).
Proof. For key signal, we have
(t) (w K.s: 1e)
(w(o) (0)
k. pe)
bn |laa||
tn |||
< 1 +
1 + o()),
(w)) (0)
Ki pe)
(w
K.e m))
M(oons-1/3 pa]))
where the first step is by Eq. (8), the second step is by Lemma C.12, and the last step is by the (f) definition of Tz". This also imply the sign of (WK,st (e) the same as initialization for all T The proof for query signal is similar.
For 2 T2, we want to show the update direction (w Que 14) or (w K.s (4)) would be determined by the sign of >: (w K.s MIE.) or Ex(w ViEa), 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
(e)
apply the previous analysis. These weighting terms include: (1) The softmax outputs; (2) The loss derivative; (3) The magnitude of mean value noise (v((), y.E.). We want to show that these quantities are almost same across all samples, or even keep constant before TT.
The following lemma studies the magnitude of (v(9, .1E,) after T2. This lemma also shows (vl9), y:E.) are close between different samples.
Lemma E.S. For all [n] and 2 T2, we have
(v(1), v.E,)) 2/ 2/2 tnops(1 [(1 ((2-))))
Proof. We can show
(v()) .v..e) (y(0),EE + 2tn ||5 221 |1lll (1 + o(1))
=2/ bn + ((ss-1/2) o(1))
53

---

Published as a conference paper at ICLR 2025
where the second step is by ,y0,,,E)) OBBemi o(T2nj,,)), the third step is by Lemma C.4.
By Lemma 1.14, we have that softmax outputs are concentrated 1/2 for all E [m] and I Ts By Lemma E.15, we have that the derivative of loss of all samples are concentrated at 1/2 for all i E [n] and t t T3.
With these properties, we are ready to show that the query/key noise dynamics keep unchanged for (T2. ]. and the query/key signal dynamics are dominated by the sign of the summation of query/key noise, i.e., Q. (t) and (w
K.s: 14). The next two lemmas characterize the dynamics of query/key noise and query/key signal, respectively.
Lemma E.9. For all s [mmx] i [n] and T2 Sts T3. we have
(w (t+1)
Q.a -.E,.) (w Qia U.E.)) + sgn( (0)
Q se ) 1 .WE.)) = (w k veE + sgn(( Kess v:E,)) n |
(w (t+1)
Kis
for all t T1)
110ps/2 < (w!2 Q.st viEd ( viEi) S 26naps.
Proof. WLOG, given neuron s E [ma] and sample i E [n]. suppose (w V:E,) > 0,
(w (1)
K. M:E,) > O, we claim that (w Q waE, and (w K,s w&E are increasing for T2 < t < T At
t t= T2. we have
(w Q,a
(12+1)
(T2) w Q,a v:Ei)) =sgn - -(w (T2)
K,s 11) + (w (T2)
K.s (.E,) 2 (n) - ||ll > 0,
and
(w K. K,s y:E,)
(T++1)
(T2)
agn 1,11
(72) (T:)
1,12 Q.0 1) +
(T2)
s1,2 $1,2 1. 3.E,))
(Tb) (Tb) (T2)
o(1))(
(Tb)
mb) + (w (T2)
Q,5 MiEEi)
4
0,
4
(1)) Von) - ( [| > 0.
Suppose the induction holds for t < s---1 1, then at t = i, we have
(w Q.
(i+1)
w (1)
Qui UcE =agn - -(w Kist (u) + (w K.s M.Ei)
(
2 (w 2- (Th)
Kie p1) + (w Kis ..E + (i Ta)n(1EE1, -1lll) > O,
(T5)
and
(f+2) (w K.s
w/O 1
WK V1E,)
1 !1) =ggn a
2,11 - 1,12 Q. 1b) + 8 8122 (w (:E.)
2 (
+
11)
+ (w (iE.))
2( + (w
Q, p1) + (w (Ts)
(Tb)
Q UrEi) + (f Ta)n(1el) - -llell) > 0.
Finally, the bounds about magnitude follow from Lemma E.2, E.3 and C.4.
Remark E.10. Lemma E.9 implies that the dynamics of Eietn (w a v.E.). Eieiny (whe U.E,) keep unchanged T2 <t tt S T and all s E [ma].
Lemma E.11. For all E [maa] i E [n]] and T2 St Tx., we have
(wQ.) .14) =(w (w (1) sgn(
(t+1)
(t)
(T)) Q.s vn
Q.5t
ie(n]
(WK,s K.s .14) =(w (t)
(t+1)
K 14) sgn(
(w (T2)
Q,a y. )) ||]l
iE|n]
54

---

Published as a conference paper at ICLR 2025
Proof. WLOG, given neuron E [ms]]. suppose have Σie(n) (w (Tp)
Kis w:E.) 0, thus
Eie[n] (W (T2)
V.S.)) > O. Therefore, for T2 S i T3. we have
(1+1)
(t)
(()) 1,11 (.12) (-vi(v(o), (- K. 14) + (.) E.. E.)
agn
icIn
v2
8V tmy ps + +o(1) (WW2.. UEEi) - (w/2.
gg
in|n
(w Kist 1/i nn(w K.s
(t)
iE(n]
(Ta) (w (WKss K.s W:Ei) n( (Ta) K. ui) + (t T2)(no,8/22 17 || > 0,
sein]
where the first inequality is due to the lower bound Eq. (22), and the second inequality is due to the lower bound Eq. (23) and n |e| = o(pps). Similarly, we have
1++1) (w
K. 14)
K,s (-1
(1) (1) (-V: (1 Ex) ( (1) (1) (Q 8 2,22
(.) (.. E.)
kE(m]
8 1110 1 + o(1)) (- (-(w @ V...) (w (w))
5gm
XE [n]
8gm
(wilt) (t)
Q.st V:E,) -n(w (t)
-n. 14)
Q.a
/E|n]
(w (T2)
Q,a M.E,)) n(w
(T2) (1) - - (T2)) - mm (mll < O.
kE(n)
Lemma E.12. Let T3 == =33
|||0|1==0000000 =Ō(oon then for T3 S S Tz. we have
in |||l||| /2 < sgn(
(w (Ta)
Q.a M.Ei)) (w!1) (1)
Q. 1²) < 41m ||||||
(24)
iEIn]
- -tm |||||||| /3 < sgn
(w (Tp)
Q.a y:E,)) (w K.s 13) S -tn || /2.
(25)
sE/n|
Proof. WLOG, given neuron s E [ma]. suppose we have Σie(a (w (T2)
K,s yrE,) > 0. We only prove
Eq. (24). The proof of Eq. (25) is identical. For the upper bound, for t [73, T75] we have
(t) (w (w 14) S 3p + by |||||| s (I + Ta/3)n || S 40n ||]||| /3,
Q
where the first step is by Eq. (8), the second step is by the definition of Tz. For the lower bound, for t [Ts, TTy1] we have
(w 1) 2 -3p- - Ten ||||| + (t T2)n ||l|||| 2 (t - - || 2 bn || /2,
where the first step is by Eq. (8), the second step is by definition of Ta. the third step is by ToT o(1).
Remark E.13. Note that the upper bound in Eq. (24) and lower bound in Eq. (25) hold for all t 2 T3 since the magnitude of update of query/key signal is always ||laa], iee, Eg. (8) always holds.
55

---

Published as a conference paper at ICLR 2025
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 i,11 and 1,21 are concentrated at 1/2 for all i [n] before T4 and 7 respectively. Lemma E.14. Let
(t) (t)
Csaps |
1/2 " 8-1 where C3 = (11) a large constant. Then, for all E [n], we have
T4 = C2 log
11 m
(26)
$1.11 1/2 + o11, Vt < T4,
1,21 = 1/2 + o(1), 1I S T
Proof of Lemma 14 Note that for 2 0,
(w (t) Q.at miEa)
S (w ( MiS.) + 20110
S 4 max {se, tnopss,
for all E [n] and 8 [ms]] and
(w Q m4)
- (w (0) 12) tm ||
I 4 max (0pm tn ||
for all s [ma]] Similarly, we also have
(w (1)
K 12)
I 4 max (8ps by ||
(w (0)
Kist MCE,)
s 4 max {,8, tno
When t $ Tz", we have for all i [n]
(w Q.a: viEi)
(w (2)
k. viE,)
s O( D(vToo(12me///o00pp))
(wl2 Q.st 10)
(w R.
O((/log(12mn//)) ||
Then we have for all i E [n].
(1) exp (Enelme) ( (w()
Q.o: yiE,) (w(
Kist 11)
1,21
exp
Eme(mu)
Q. EiE (w (w K. 12) + exp
Exe(me) (w (1)
Q. U:E,) (w (t)
K.s (w,s viE,)
exp (0(m² log( 2 (8)0://0)0a0sss/ 1 ||11))))1)1))
exp (0(ma
(1 + exp (-0 (o22:o/0000/0000000//5
+ exp 10
+ aps [|x)1)))
< + O(m (ogmmo/000mm000000000 +ops ((]1)) (0n
2
and
exp (Ene (Exe(m)) (w (t)
Qu (w yiE,) (w (w (t) K.st (t)
$,21
exp
Exe(my) (w (1)
Q.st U.E,) (w K 12) + exp [xe(mm)) (w (1)
(t)
Q. U:Ei) (w (t)
K. y:E,))
exp (-0 log 12mm 18) (m))))
exp (-0(m log (1mma/0)00 4 ()) exp (0(m) log (12m² (1)mg/0)000333)))
+ exp (O(m, log (1mml (0 + ons ([))))))
1
>
O(mi log (0m) Tn | [()1))
2
56

---

Published as a conference paper at ICLR 2025
Similarly, when t I Ta, we have for all E [mx]
(w s m
(1)
(w (t)
Ki MiE,)
s SO(log(Cooppp Iull)m 1/33)
(27)
(w (1)
12)
(w k 12)
S0(log(C30pp// ll/ps//llllmmmmm -1 1/2 op 1 |1ll))
(28)
Then we have for all E [n],
exp E Exe(ma) (w Qio: (w) (2)
10
k.
exp
(Eee(ms
(w Qio H4) (wk.ee (t)
k. Hi) + exp [Ee(ma)) (w() Qiot He) (w(l)
kiet B1Es))
exp (((log((aapps)) (Ca0ps/ 1111)pp 2-2 11)13))
exp (O(log"(Cso))) (C3aps/ 1m)1) 2s -- [|[|)() exp (-O (Caaps/ [1 s-1 11
1 + exp (-O(log* (Caops// [1ll)ap s-1 || (1 + -
s-1 (w
1
+ O(log"(Gaapp hatps///llllo/ [ul) s-1 || + o ||
2
and
exp (Znelma) (w ( (e) (t)
(0)
$1,11
exp
(Exefm)) (w (w
Qio 14 (w (t)
k. (22.0 14) + exp (Zne(me) (w(()
Q.or H4) (w
(t)
K MEE))
exp (-O(log3 (C2ops/ herss/I||11/ [1 -2 1m113))
(-O(log² (ec00s//llul))g (u1))ppp -- 2 [|[21(3)) 22
+ exp (Ollog" (Caops/ [1al)) t --1 (|<al))
exp
1 1 exp (O(log" (C3aps/ (Csop/ H 1)0 s-1 |||||| (1 + or o s-1 [(]))))
1 (Csops/ |I 5- ||l||| (1+op's s- (|))))
2
2
The next lemma shows that loss derivative e(0) are concentrated at 1/2 for all i [n] before T4. Lemma E.15. Recall the Ta defined in Eq. (26). For all i [n], we have
e00 1/2 + o(1), 14 < T4.
Proof of Lemma 15. Note that ((z) = log(1 + exp(-z)) and -E(2) = 1/(1 + exp(z)). For t < T we have
(t) (s/1)) + (v(1), yepe) + (o! + (4) (v(() Ei)
1 1 exp expp
1 exxp ( m1) + (v(1), v.E)))
1 exp yi eexp
1 o
where the first step is by definition, the second step is by Lemma 141 the third step is by Lemma E.8, the last step is due 0 o(1). For & < T4, suppose yi 1, have (v(", ,yt1), (v(1) E.) >
57

---

Published as a conference paper at ICLR 2025
then
(00
1
+ exp ((s + (0) (2(((((()) yipe) + (0 + (v(1) E.)
1
S o(1):
where the first step is by definition, the second step is by (v(*), VaPl)) (v(t), E.) 0. and
1 + exp (0%)) (t) + (v(()) yepe) + (0 + ( (v(0) E.)
1
1 + cxp (1:33/2 ((1))((v()) 14) + (v(0,E))))
1 exep (3/3222 + o())( y1E,))
1 exp 1x13/2 + (m)) 1/3)
1
o(1),
2
where the first step is by definition, the second step is by Lemma E.14, the third step by Lemma E.8, the fourth and last step is by (v(fa), y:E.)) = Ö(mp
11/2) -(12212
o(1)
E.7.1 STAGE IV.A
This stage studies the dynamics in the interval [Tg, Tí 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 TT, which is defined later. And we don't control the behaviour between the interval [T [T) since T7 not too far from T,t, we still can keep the good property at Ty even in the worst case. Finally, after Ti², the good behaviour continues. (Also at Ti", we have the softmax ratio is basically concentrated at 1.)
Let T 2 T3 be the last time the such that for all t E [T3,TT] s E [mm]] and i E [n]]. following conditions hold,
Σ 2,21 9222
yrE,)
1
(w (t)
2: 10/
(29)
2,21 6,22
(
e viE.))
(1) (1)
(w (2)
Q. 113
(30)
Note that at T3 we have
2.2 9.22
Q.a:
vcEi)
2 nn(w!l. (t)
1
e. 14)
[E(n]
1
( Q.at viE
n((w())
Q.a p4)
2
EE/m]
(w (Ta)
Q:a v(E.))
1
n
s 14) (1 T2)n(( Tp8
(T))
a1|111
0,
5
4
2
:E|0
and
1,21 8 Q. v(E
(0) (1)
(0)
28
1.11 w...
(1) (1) (1)
(1 + o (4)
1
(1
((1)
(t)
2
2
1
(1+
0. v:E.))
(1))
2
1
+
o(1))
(
Q,
(Ta)
(a) (t TT2)m(
Tp8 4v2
|
0,
58

---

Published as a conference paper at ICLR 2025
which implies T4 is well-defined.
We use induction on / to simultaneously prove the following properties for all t Ts...., T7:
A(t), the linear-with-I estimation: for all E [n]] and E [mm]]
(w q.a vcEi) = Lm
0 pp + 0ss (112))
(w22aa K vcL.) = Mu
- o + ((s-1
2
Dt)), query noise is increasing at for all i [n] and E [m+]
(wQ,s) U.S,) = (w (f) Q, y&C,) + sgn ( (t)
Q.at 3e |E:ll
E(), key noise is increasing t: for all E [n] and [mm]]]
(1+1) (w (W',a U.E.)) (w K,st y sgn (w K,st M.E)) [&ll
(c)
(0)
F(t), query/key noise have the same sign at for all E [m] and s E [me]
(w Q.s: M1E,) Fggu (w 31E.)
Claim E.16. F(t) = E(t).
Claim E.17. D(7a), ETT)),.... E(1 ().). - D(1).
Claim E.18. A(t), D(t), 8(t) A(t + 1).
Claim E.19. F(t), D(t), = F(t+
Note that we have A(Ta), D(Ta), F(T3) hold. D(T3) and F(Tz) are already shown the stage II
Proof of A(Ts). Att Ts. we have
(w()) (1) 3. V:E = ( (2) || 1 + (w Q. M&E)) = = E (1 (1+ O
(Th)
=tmy 2 0 =1=
2
2 O(T)t tn op + ((s-1(()))) (1)
(T2) where the first step follows from the monotonicity of (wo, 3.E.) for [T2, Tz}. the second step
follows from Eq. (21), the estimate of the magnitude of (wo..) y:E.). the third step follows from
Lemma C.4, the concentration of ||&.l] the last step follows from T2/T3 ((8-1/2) The analysis
similar for (w K. V.E.).
Proof of Claim 1.
(w K.s w (1)
(((1)
K, yiE,)
sgn 1,11 111 1,12 ( Q.a1 1) + (0) (1)
(t)
9,21 (,2 ( MiE
Q.
=agn (w Q,a B:E.):
where the first step is by Eq. (12), the second step is by the definition of Tí. i.e., Eq. (30).
Proof of Claim E.17.
(w ((+1)
w Q.a7 U.E.)
0.
egn (w(()
K, 14) + (w K. 3:E (T3)
(w(2)) Kis 11) + (w (Ta)
K.s wrE.) + (t - 73)m(1Ell - ||/l]) > 0,
where the first step is by Eq. (10), the second step is due to ((T))...,,E( F(T,), F(t 1). the third step is by (Ts).
59

---

Published as a conference paper at ICLR 2025
Proof of Claim E.18.
(w vrEe)
= =( (1)
Q.o: MiEi)) +17
= typ - p o + ((5-1/1)) m/ 2 o - ((-11
2
= (1+ 1Mw 0 + ((1))
where the first step is by D(t), the second step is by A(t) and Lemma C.4. The proof is same for (wl) K.s M.E.)
Proof of Claim E.19. Suppose (w (wQ.s: Q.a: MeE,) > O, then by D(1) we have
(w (t+1)
Q μ.) = (we Q.a 14) + sgn( Sggi(wel Q.a 1)) || > 0.
By F(t), we have (w Kis ))E > 0, then by E(t) we have
(0)
((+1)
11) (w (1) K.s
|al| 0, K,st 14) agg((w (w K,s p2)
which gives F(L+1)
Now study the ratio of softmax outputs across different samples. The following lemma is crucial to prove the lower bound of TT.
Lemma E.20. Let
;= min ; Ta]
(31)
Then, for all T3 Stst to and all i, k E [n]], we have
(t)
$1.1
= 1 + o(1), Vi e [m].
(E)
k.21
Proof of Lemma E.20. By A(t) and F(t) we have for t [Ta, T and all E [n]
Σ
(w (wl Q.s: Ex) (w Kist E.)
2
11 (1-1/
(32)
SE [m,))
Then, for all T3 S S 1r and i E [n]
exp Exe(me)) (w/))
Q.o: (i)
Kis (ipte)
1121
exp
E [Ee(mm))
( Ex) (W. K. will))
(1)
(2)
exp EEe(me) (W(t) (t)
2.00 Ci) R. E.)
1
1 + exp 1m ((8-112))
1 exp + ( (r
o
1 +
+ o
o
(33)
1+ ++exppmma_pp0333
where the first step is by definition, the second step is by Eq. (28), (27), and Eq. (32), the fourth step is by m 2032 o(1). Then we have for all T3 - t ts and all i, [m]
(4)
(1 (()) (xm33
1,21
(1 o(1)
+ o(1)
(I (xp
+
where the first step is by Eq. (33).
60

---

Published as a conference paper at ICLR 2025
The following lemma studies the lower bound for T4- Lemma E.21. We have
T4 2 14
0.99
log
SpS
1/2 1 m
(34)
2
3//nn ||
Proof. Suppose T3= < TT. thus we have T7 +1 <
T7 +1 < T4. By Claim E.18 and E.19, we
have A(TT + 1) and F(T + 1) hold, which implies Eq. (32) holds at / = Tí + 1. Then, for all i [n], we have
(T +1)
1
1 o1) + exp (m + 1)307)303(1 1/2))
1 + o(1)
(35)
1 exp( 2 mm (T7 + 1)21120333)
1 (1)
1 + exp( f m4 (T. + 1) 10
1 + (1)
1.01expo 2 mmx(Ti (TT + 1)2120233)
0.99+a(1)
11)
2n |ll
1.01
Tps
(1 = o(1))3 2n
||
(36)
1.01ops
where the first step is by Eq. (28), (27), and Eq. (32), the second step by Tj + 1 < T4, the fifth
p s 1), the last step is by n || op 1 = o(1), Then we estimate the magnitude of Σi= (w
step is by = o(m= -1 11/2
2.5t V.E.). WLOG, given neuron suppose Σ; (w (Ts) Q, y:E.) > O. For
i=1
T3 <t< T4 +1
Σ
(t) Q.or viEa) 2 1/11 (1 To +
Σ
(w (T2) Q,s M.E.) 2 (1 - (1- T2) + 2n,8,
1
V2
/2
1=1
z V2 tno O = v2 tmo s( (1 o(1)).
(37)
where the first step follows from D(1) and Eq. (22), the second step follows from the lower bound of
Eie(n) (w (Ta) (w))
Q,s :iE,)) Eq. (23), the last step follows from T2T = o(1). The analysis is similar for
(w (1)
K. S...). Now we can show that
(TT +1)
122
$,22
(wo.a Q.a
1=
1 (1)
o
(T, +1) w
(1
,)E,,)
1 ++exp( mx
(T +)21100382)
1 1 exp((@mmeTT 1)2120352)
Q,a
1=1
(1 o(1))3/2n ||||||
1
0.5
(TI
11 +1 o())
1.01aps
V2
2 1T + +1)
1 (T++1) +1)
2 5 0.*
14)
where the first step is by Eq. (35), the second step by Eq. (36), (37) and 7J 2 Ts. the last step by the upper bound in Eq. (24). Similarly, we can show that for all E [mL] and E [n]
$.21
(T +1) (TT+1)
$,22
+1)
(w Q v.E.))
(TT+1)
2
(TT +1) (TT +1)
7,12
( 14)
(T)-1)
61

---

Published as a conference paper at ICLR 2025
which contradicts with the definition of T4. Therefore, have TT TT.
Remark E.22. Note the constant 0.99 in the definition of T 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 (0)
1121 0.1 for & 2 T
Next, the following lemmas study the behaviour of query/key signal. Lemma E.23. The query signal are monotonic for E [T's tn]. Formally, for all t E [Ta, t²] and E [mg].
(wQ Q.+ #2) = (wQ,. q. po) + sgn((w Q, u))) ||ll,
(t+1)
where L is defined in Eq. (31).
Lemma E.24. The key signal monotonic for t E [T,,t]] Formally, for all t E [Tg, (*] and
8 E [mg].
K.s .p) K.s (4) + sgu((w (9 K.s ))) ||llll
(t+1)
where 12 is defined in Eq. (31).
Proof of Lemma E.23. WLOG, given neuron s, suppose (w pe) 0 and W. 14) O. For query signal, similar to previous analysis, we have
(t+1) (w 0,8
w (t)
qq.:
(00
(. R. u) + (v) $1,119;112
(1)
(-y((v
(w) (... Ei)
agn
(E(n)
V2
l1 11 + d()) k. k. v1a,) (w (.)
8VT
(E(n)
(w K. ViEi) - nn(w K.
sgu
iE(n]
(w (7a)
K.s y:E,) n( (T2) K.s pe) (t- (2)22(ppps//22 my (1 > 0,
RE(M]
where the first step is by Eq. (11), the second step is by the concentration of eg and s for 1. S T4. i... Lemma E. and E.15, the fourth step is by (t) and the lower bound Eq. (22), and the last step by lower bound Eq. (23).
(2)
Proof of Lemma E.24. For key signal, we have (t)
(w
(t+1)
w K.st 14)
K.s
((t))
(-vis(t 2,11 81,12 (0 (.) (4) - (9) (1) (v
-,,21 (v (, (w (.. E.i
kgn
xE[m]
22
tno
p + o())
(-s (2) (t)
2,22 Q.at v:E.) 1 + o(1)) (w.) H))
2/T
XE [n]
(t)
Ww Q.at M.E
(1 (w)) H1)
8go
1,21
1,22
(E(n)
agm
$,21 22
(8)
e. v1Ei))
(E[n)
(1 + ((1)
Σ
(w Q. v.E < 0,
8gu
$21 822
xE(m]
where the first step is by Eq. (11), the second step is by the concentration of E and (1) i... Lemma 1. and 1., the fourth step is by the definition of TT i.e., Eq. (29), the fifth step is by the (t), the last step is due the upper bound Eq. (23) and D(T3),... D(t 1).
in for t S T4.
0
62

---

Published as a conference paper at ICLR 2025
At the end of this section, we give upper bound for T4
Lemma E.25. We have
T4 T
1.01 2
log
TpS
17 mk
1/2
0p 1-18-1
(38)
|||||
Proof. Suppose T4 > fr. First, we estimate the magnitude of Ei= (w (t)
Q. M&S,)) and (w
Q 14)
for T3 St TT. WLOG, given neuron s, suppose E (w (Ta) have
Q. 31E,) > 0. For all E [mg] we
ww Qu p2) 2
( (Ta) Q, pe) + (t Ta)n ||lll 2 tr) p /2,
(39)
where the first step is by Lemma E.23 and T/ t", the second step is by the lower bound in Eq. (24)
at t = T3. For all s E [mx] we also have
w Q. v6C
(we.. a. UnEi) S kmmy o =
1=
where the second step is by A(t). By the definition of Tī at t = Pn we have
tnn ||| /4 <
n2 W (e) 2
Q.
Σ
i121 (.2
Q. BrE
(1)
i=1
1 (11)
<
- Ops(( ( + ((s-11())
exp( m
o
2
tmm
o + (123)
22
2
which implies that
T
S
log
4v 20 ps
(1+ (1+)
-1/2 112 mm 00 1--1
V (
which gives contradiction.
Remark E.26. The upper bound for T7 also proves that t TT which gives that for all t e [T3 T.). [mu], and i, [n]
(0)
1,21
1 + o(1).
(witt)) (++1)
Q.s [1) = (w Q.5 14) + (gg((w[) (0)
Q.: 4)) Illll
(.s .H) = (w K.s: 1) + sgn( (( K.s )m ||ll
(++1)
Moreover, for neuron 8, suppose (w (T')
(w..2) 11) 0, then at = Tī for all s E [ma] we have
(w (1)
Q.s:t 14) = (t--Tss) (t-))==tt)p)) (tt-Ta)n + (w (T5)
1. p1) =t |||| (1 o(1)),
(40)
(w K.87 11) --(t-Ta(
(1)
(w(Ta) (Ta)
K.s x1) --n (1 o(1));
(41)
where for query signal the first step is by Lemma E.23 and the second step is by the bound Eq. 24. The proof of statement for key signal is similar.
63

---

Published as a conference paper at ICLR 2025
E.7.2 STAGE IV.B
This stage studies the dynamics in the transition interval [T] T. For I 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
<tSTI given neuron
Let T be the first time the following condition holds
(c) Q.a
v(E.) - 1,12
(w Q.o: 14)/
(42)
1,21
1,22
for all 8 E [me)] and E [m). Note that at TTt, we have
1 n
1,21 1,22 (w (0)
(8)
Q.at MiEi) S
(e) (t)
(w 3. B
(t)
$
Σ
s
(1) (t) (t)
1,1181,12
2
+ (1)) 11
(0)
(wol.
(w
:1 11)
(0) (1)
<w
2 12)
which implies T 2 To. T is well-defined.
Let
ET(Ta = 8 E [mx]: : (w. (Ta)
q,s U:E.) sgg (w (T5) Q. 14)
E = i [m]: (w (Ts)
Q. m.E,) =ago (w
(Ts) Q,s 11)
all neurons with E [ma] and i E (TTa) are called aligned neurons. The next lemma studies the
concentration behaviour about E [e(Ts) This lemma crucial to upper bound (() and further upper bound TT. Lemma E.27. Suppose that 8 > 0 Then with probability at least - 6,
s,21 for T7 S S Tzt
E (Ts)
2 mx/2 (mmgn-1/// 1/2 + m/ m/22 2 m²/2 Ö /mz);
for all i [n].
Proof. We can write E p((bb) as
- { s [mm] : (w(()) ,.Ei =ggn Σ
n
(
(Tb) Q,a U
w
where this is by the sign of (w viE,) and Σi= (w) (t) Q. U.E,) does change for [72, T3] and the sign of (wQ,s) (u) is the same as that of Σ Ei_( (.s M.E.)) i.e., Lemma E.9 and E.11. Then we have
(Ts)
(w (T))
1 mi
El (Tb)
=
n
(T2) =P((w Q yIEa) 0, Σ
w (T))
Q.e . > 0 + Fl(w (T))
Q v.E,) 0, Σ
(T2) Q, M.E,) 0]
w
1=1
= =P|X,j > 0, Σ
N > 0]+ P|X,, < 0, Σ X 0]
1=1
1=1
1
2, O(n 1/2),
where the third step by that the sign of X..; determines the sign of (w (Tz) y:E.), the last step is by Lemma E.6. Then, by Hoeffding's 's inequality, with probability least 1 5/n
[T(T
pmg S
71 2
log(2n/8))
Finally, by the condition R(mî), we get the conclusion.
64

---

Published as a conference paper at ICLR 2025
We do not directly characterize the dynamics for 0)T' where 0 is large constant in (0, 1), and show that the definition of II are satisfied, which characterize he dynamics before T7 and give upper bound for T4 simultaneously. Specifically, we use induction on / to simultaneously prove the following properties for all t= ........ ..... (1 0
Instead, characterize for [T4 (1
A(t), the monotonicity of query noise aligned with signal: for any [mx] and e E!
(w (+++1) Q. .vE,)) = (w (t) Q, MAE,) + sgn( (w (t) a ye Eill.
Bt), the monotonicity of key noise aligned with signal: for any [ms]] and e ETTs)
(1+1) WKs v (w K,st yiE.) #gg((w22. K,st V.E,)) 16.111-
C(t), the monotonicity of query noise unaligned with signal: for any E [ma] and e [n]/E]
(++1) (we,a MrE,) ( y:E.) sgu((w @ (E))))
(t)
Dt),, the sign does not change for unaligned key noise: for any E [ma] and E [m]/E!Th
(w K.s yaE,) Fagn (w (T4 K,s MeE)) =agn - (w (t)
Q. (w VeE.).
Et), the monotonicity of query signal: for any s E [mm] (++1)
(we,s ()) (w..s Q,st 14) + sgn( ((wl(( (t)
Q.a u) ||
F(),, the sign does not change for key signal: for any [m=]
K. =gg (TT)
Kis (4) =gn - (w (t)
(.. ViE.).
(t),, the estimation: for all s E [mg] and E [n]
(w(() a.a: v.E, = sgn((w Q. 14) 14)
2
00 + + (1-)) vi E E(Ta)
(w K. v1E,) = sgn((w (1)
Q, 11) 1))
2
pps(1 + 0(s (111) Vi E E(
(w qiat M&E,) sgn( (w
(1)
Q.at u)) )tn
2
Op (1 0(s 1/3)), ((- Vi E (m]/E(T))
(w Q. 11)= =
(1)
Q 14) sgn((ww(
(1))
Remark F.28. Before stating the results in this section, we first recall and summarize the results (Tt)
Q. a y:E,) > 0 in this section below. until TT. WLOG, given neuron s, we suppose that Σi= (
(Ts) Then we have (w Q. 11) > and (w Kis p4)
(Ts)
< O. Then at t TT we have
2
1. U.E,) =tn
s
1 Ö 1/2)), vi EXTs)
π
2
(w K.st M.E,) t/
op (1 0 1/2 Vie E!Ts)
-
π
(w Que y:E,) = /
0 s1 + Ō ((- Vi E (n)/E(T))
-
(w K,st viE.)) = MMn
- p + - vie (
(w Q. 14) bn || (1 + o()),
(w (w.... K.s 11) -t |||||| (1 o
This gives that A(TT), C(TT), D(TT), E(TT), F(T hold.
65

---

Published as a conference paper at ICLR 2025
To prove these properties, need some further properties after T4. Lemma E.29 (Properties after T ). Let 0 be large constant in (0, 1). For all 2 T we have
2
(2T - V Ops(1 - ((s ((-1)(1))) & sga( (w Kis (.E.))) (w (w.,.r 3.E _ tmy opo(1 1 + (()-11)))
(2T4 -t)n ( (1 o(1)) S sgn((w (TT) for all II St - (1+
Kia p.) (w (t)
K.s: 12) S bn || (1++
(w (t) K,8 MiEit) =ggu (w
(WK,s K.s
(2T
U.S.),,
(w k 14) =yn (w
(t)
(T. K.s 14)
Particularly, For all 0 S 0' v 0, 8 E [ma)] and E [n]. setting t = +1+ 0')TT, we have
2 (1 1'-0' VV o sp
1/2)) - (WK,s
(11-0'xTT] yyE.))
- +1+ "V ops(1 + ((5-1/2)),
2
π
(1 - 0'yTT 7 ||l|||||| (1 o(1)) s (wkis (1+10)77) mi) s (1+ (1+
(TT)
Proof of Lemma E.29. Suppose (WK,e) y:E) > 0, then we have for all 2 TT
(TT) (w K M&E S (w K,s M. + (1 - (T TT Illl S bn / - + 0(s (12-)1
where the first step is by that the magnitude of single update of query/key noise is always v || the second step is by summary above and Lemma C.4. Similarly, we have
2 (.E,) - (t (t- T4 1l 2 (2T4 - -tny ops(1 = 0(s--113) (w Kist U.E,) 2 (w K.a
(TT)
Set t (1 + T') gives the conclusion. Also, the lower bound gives
(w K.s y.E 2 (2T - -tyn pp =s--111))
2 2 (1 - O)- m op (1 + ((0ss113)) > 0,
where the last step is by 0 is a constant. Similarly, we can prove the statement for (w (t)
Kis 1!).
Remark E.30. Note that Lemma E.29 proves the D(t) and F(t) for all Tz <t (1 8)T.
Claim E.31. F(I) = A(1..
Claim E.32. G(() = B(t).
Claim E.33. D(t), = C(t)-
Claim E.34. B(t) = E(t).
C(t) Claim E.35. A(t),)(t), ((t),
= G(t + 1).
Proof of Claim E.31. We have
(w (t+1)
Q. Q.5 v:E,) agn -(w K.:: 14) + (w K. v.E.) (w(2. K.s ):
(t)
> 0,
where the first step is by Eq. (10), the second step is by F(t), the last step is by (()).
Proof of Claim E.32. We have
(w K.s w K.s V:E.) =agn $ 8,12 2. p1)
(t+1)
(1)
(t) (t)
(w (1)
8:21 $4.2 (w
(6) (2)
Q.a: v.E.)) 0,
where the first step is by Eq. (12), the second step is by (()).
66

---

Published as a conference paper at ICLR 2025
Proof of Claim E.33. We have
(w Q.
(t+1)
w Qua MiEi) =agn
-(w K.st pi) + (w K,a y1E,) (w K.s y:E,) > O,
where the first step is by Eq. (10), the second step is by F(t), the last step is by (tt).
Proof of Claim E.34. We have
(w
Q
(1+1)
Q.
(t)
13)
(e) (1)
(-uitoll)
(E) (e 1) + (.),), (0) E.) E.)
agn
2,11
xE(n]
V2
tno (1 + o()) ((w (W....V.. UiEEt - (wh.er 11)
8 8
Ki
(wll
K,st ViE.)) n( K.. 14)
agn
kE|n]
n( (1)
K. 14)
Σ
(w (1) K,st yiE,))
Σ
(w
K, ViEi)
(T)
vE(n)/E(7a)
--n(wl)
K pL)
(w (TT) K,s
(.EE)) + - Tn-TTTnll 11 +
Σ
(w (T)
(W,s ,.EE,) -(t -t- TTT )n |
TEF
n((w (1)
K.s 14)
+
K.s y.Et)) + (1 - --TT)n
sE|n]
(TT) n(
(TT) y.E.)) + MTTpp//p// ml/l)))))))))))))))))) > 0,
(1)
K.s
<E[m]
where the first step is by Eq. (11), the second step is by the concentration of & and (t) for t < TT. i.e. Lemma 1. and E.15, the fifth step is by (t)). the sixth step is by the lower bound Eq. (22), the last step is guaranteed by the dynamics in [T's ].
Proof of Claim E.35. This claim is straightforward.
The following lemma gives an upper bound for T Lemma E.36. Let 0g be a small constant satisfying 0.99(1 + 0.) - 000./ >
then we have
T s (1 0
(43)
and for all fllt E [[1+ 0.)TT. (2 0 we have
$,218,,2
(9) (t)
(w (t)
Q B1Et))
81,118,112
(1) (1)
(w
Q. Hi)
67

---

Published as a conference paper at ICLR 2025
Proof.
For fixed 0 E [0., 1 0z]. at (1 + B)Tr , for all i E [n). we have
Σ #E[ma]
a. E2) (w E.)
Σ
(W Q.a (.).w/..r E
Σ
(w (w 9,o E.) (.)(w (w Et)
SE FE
2 E
(T5)
2
1) p
2 112)) + (mx |E(Tbi|
2 t(T -t)y (1 0 0(s 1/3)
π
π
2 + (mx | )0(27) - -)) 11
2
0
π
D [3)
(1 + 0 2 + (m+ - (T) |1 + (1 2 21 (1 0(s 1/3)
m
o(m)) ))(1 + 0) + mk
+ + - 0²) 2 oo + (15
2 ((
2
2
((1 + 0) (02 0))(m))) T 2 2
o 0
(44)
π
where the second step is by (t)) and Lemma E.29, the fifth step is by Lemma E.27.
Then, we have
exp (Er ( Q.s: Ex) (e/(w)) K. RiPe
$1.11
exp
E Eoelme)
(w (1)
Q.o: (w) (t)
K. vi(e))
exp
EEE|muu 2.) E.)
(8)
(t)
K.o Ei)
1 o(1)
1 1 + exp ((1+0)m= - (02 + ()0(()(()))) mm))
TT=22 2 2 Ō
T4 22
1 o()
o + exp(((1 0)ma (02 (0*+
2 exp exp(+ (
1 1
o + exp(((1 + 0)m= (02 + (=00((mmm))) I7 - (1 o
1 (1))
1 + 0)m; - (02 + (0)(0((m))))
T 2 13332
o()
(45)
exp(((1 + 0)mz (02 + 0)0( T 22,20222
where the first step is by definition, the second step by Eq. (28), (27), and (44), the fourth step by TT & T4 and m aa3 =(())
Therefore, we have
((+0)TT) ((++)TT)) (w
(1+0)TT) q.s
4,21 $,22
o1)
< exp(((1 + 0) (02 + (0)((mm)))
.T- 2==2 22232 11 11(1 mV ops(1 + Ō(s-111))
0.99(1-0) (0(1--))0(0//mm
2 < (1+ (1++)TT "/ ops(1 + ((s-1/2))
(3/2gs
(3vEn ||||||
@ps ops
1 (1 + 0)T n ||al1
9
1 (1+0011) (1+0)TT)
(11+0)14) (w Q
2 1,11
$,12
where the first step is by Eq. (45) and (t), the second step is by Lemma E.21, the third step is by
0.99(1 + 8) - 000/ h00//mm, / > 1, the last step is by G(t) and the concentration results in Lemma 1.14.
thus the conclusion holds.
68

---

Published as a conference paper at ICLR 2025
E.7.3 STAGE IV.c
T4 timestep where 1,11 1/2 + o(1), we characterize the dynamics in this interval.
WLOG, given neuron s, we suppose that Σx=
(Ta) Q. y.C.) in this section below.
Specifically, use induction on simultaneously prove the following properties for all t (1 + 0e) ....,,
A().. the monotonicity of query noise aligned with signal: for any s E [ma]] and E (T3
(wQ.s 1++1) q. y.,,)) = (w Q, U=E,)) + sgn((w (t) Qu U. [E:ll
. Bt), the monotonicity of key noise aligned with signal: for any [ma] and e E.(T))
(w', (1+1)
K,s U.E.) (w K,st yvE,) + sgu((w (t) K,s M.E)) |6:11
Ctt), the fate of query noise unaligned with signal: for any s e [mg] and i [n]/ mn]/E!"a for
(2 0 have
(t-1) (w.. viEi) = (w
(2) Q.o vvE.) + sgn((w q. 1lE.l)
(T))
and fort (2+ TT we have
(Tt) (w m.EE)) = (w y1E,)) sgn( (w Q, y.E.) [|E:ll
D(t), the monotonicity of key noise unaligned with signal: for any e [mm]] and i [m]/Fg ())
(t+1) (w K.s y.E.) = (. k... UeE + sgu((w a. u))n Ilell
(1)
{l), the monotonicity of query signal: for any E [ma)]
(1+1) (w 0,s .1) (w Qua: 14) + sgn((w 10) Q.a 1))n ||
F(t): the monotonicity of key signal: for any s E [m²]
((+1) (w K.s pp) (w(t) K. (1) + sgn(( k. p1) ||lll
(()), the linear-with-t estimation: for all s [ma] and i E [n]
2 (w Q.a 31E.) = sgn( Q,a 12)
o + 1 Vi e E(Ts)
(w/)) K,s viE,) = sgn(( (gn((w/]) (t)
Q.a7 )
2
sp + 0 vi F!
(w!)) 2. 14) -sggg((w Q.at ())) |||||| (1 o
(1)
H((t:: for any s E [my]] and i E [n]
(0) 8 8.21
(w()) (t)
(w (t) Q. U.E.)) S 2 s Qu 11)
(1) (t)
S 1,22
Remark E.37. By Lemma E.36, have H(t) holds for
T.)TT ). The A(t),
B(t), E(t) and (tt) for [[1 + 0.)TT (2 TT are already shown in the induction statement in the last section. Actually, the proof of A(L), B()), E(I) and (()) for all subsequent are same as those in section 2.7.2, thus omit them in this section.
Claim E.38 H(1) = F(1..
69

---

Published as a conference paper at ICLR 2025
Proof of Claim E.38. We have
(w ((+1)
w()) (t)
Ki ViS.))
Kis
ag 1,11 (w piar 12)
(1) (1)
(2)
1,22 Q... v.S.)
mgm (w (t)
Q.st u!),
where the first step is by Eq. (12), the second step is by H(t and
(w (++1) w (e)
)())) (-8 (v(1), M(E (w Q. 14)
8 (,22 (v (w viE,)
eE|n]
((()))
1-, i (v()) w:E.) (w (w,ss
(t) (0)
sgn
iein]
(w (f)
Q.st (1) =ago (w Kist u))
sgm
where the first step is by Eq. (13), the second step is by H((t)
The next lemma gives an one side bound for unaligned key noise for 1. 2 (1 + 0.)TT
Lemma E.39. Let t 2 (1 + 0.)TT, suppose D(t) holds for t = (1 + o.)TT E, then for all
E [mx]. i E [n]/ (B/EETTa (Ta) we have
sgn((w (w K.s :.EE.))
(wR, K.s V&E) S (2(1 + +.)TT - -tym Jps(1 + p ((--11))).
Proof. Since Di.1 (w.2.
(Tp) Q,0 M:E,)) > O, we have (w (Ta K,s (1E,) < for s E [ma], E [m]/ (b). We
can show
(w K.st McE) = (WK,
(2)
(1+0.TT
vvE) + (1- (14 (1+ |Eal|
2 (W
(TT)
.1iE) + (t (t-(1- 11-(1+0))2 - -02T )n 111i11
= - pV: ops(1 (1 + 0(5-111) + (t - -11 + 0.)TT - O.TT )n ||&,l|
(t-(1-
2 - (2(1 + 0.)TT - 1mm - o (1 + 01--1)))
where the first step is by (tt) for E [[1 + 0o)TT 8 the second step is by Lemma E.29, the third step is by the argument in Remark E.28, the last step is by Lemma C.4.
Now we only need to and are ready to show H((t) holds for t = (2 - 8.)TT 1..... Tr. Suppose A(t')...., H(')) hold for t' < f, want to show H(t + 1) holds. We split to two stages to prove where the first one is [(2 0.)TT (2 + 30. In this stage, we can bound the magnitude of unaligned key noise (w K.s {,) by Lemma E.29 and E.39. Note that
Σ
(w (1+1)
Q.a (w K. ...)- - (w/2 Q. (.)
((+1)
(0) 2.. Ei)
(5)
Σ
(wwl)) E.) + + + ( (w 22.0 E.)
+ + > - 2. (w E/d
SEE!TS
Σ
n²
|| 2 + 11 ( (w 1 E +
(w (t)
k. (
TTE (3)
(n² || + 2t 11. n + o(1))),
(46)
70

---

Published as a conference paper at ICLR 2025
where the last step is by ((t), and
Σ sEIm)/E(7 TD
(1+1) (w
Q E.) (w K,s EE) - (w (1)
((+1)
(wQ..: (.).W..W.... 2 E.)
Σ
(111)
E.) K,s .E,) - (w Si> W K. Ei) + Σ (wo E.)( (tt0 ..:) -(w2
((+1)
Ex) (w
K.s. E.)
Q
12T
Σ
17 11 ? - 1115:11 (- (w (t)
Q.a E.) + (w (w..e: 51)
ETT (w (0)
E
-1² [E&lll - 7 |1& ((w
(1)
2. Ei) (w Ei)
tCT (w!!)
n² |L,, [ + n |11.111 ((w Q.a: E.) + (w (t)
K.s E.))
127 4,d Q. 14)
(ww
Σ
v ||,, i |1:111 (( Q.a E.) + (w K.a E.)
127 dip ((wl))
10000 <0
Σ
1p |1Eil1 17 || ( (w E.) + (w 8.)
Furthermore, when (2 0.) St (2 30.)TT for s [mmx] and i [n]/E!Ta), we have
2 -30,TI " π 0 < sgn(( ( (TT) K,s ..EE) (. V)E) S 30 "y opo(1 + ((--113),,
2
(47)
where the upper bound is by Lemma E.39, and the lower bound is by Lemma E.29. Therefore, for (2 0.) / S (2 + 30,)TT we have
Σ
(w
(++1)
Q. .., (w (1+1)
K,8 .E) - (w E,) (w K.at E,)
=E|m]
Σ
n | + 1 I1lll ( (( @/.. E2> + (w R. E.>
AEE! (T)
Σ
n² | + 7 11 ( ( {/ي + <w 2 E
EiMlEE/Ta T
E E(TA
(n7 11: + 2t 2t 11 77 + (())) (m* E | 1072 + ( (1+3 (1:001000) o()))
mh -O(mm)))( (1² 1| || (1 o(1))
-1/ mm +0 (1² ||E.1] (0.30.TT) 1 + o()))
O(Vma)n |s.11i² + ( (1 10-30.TT - (0.m))/mm)))))))))))) +.TT )m | (1 (1 o > 0,
where the second step is by Eq. (46) for aligned part, Lemma E.29 for upper bound of unaligned query noise and Eq. (47) for upper bound of unaligned key noise, the third step is by concentration
for E (5) in Lemma E.27. Combined with H((t)) for (1 + 0.) < t' < t, we have
(++1) (t+1)
1,21 $,22 (w (c+1)
(WQ. (.EE))
S 8 $,21
(11+0)1
(w
Q.+ .BAEA)
(t+1)
1 (1+ t)m ||||||
1 (+11) (e+1)
(w .1)
(t+1)
- 2 $1,2
where the first step is by induction, the second step is similar to the proof in Lemma E.36, which implies H(t + 1) holds.
71

---

Published as a conference paper at ICLR 2025
Note that have D(t"), F(t') hold for t (2 - (.)TTT ...... (2 + 30.)TT Then for t 2
(2 + T suppose F(t and (t) hold for I t, for all & E [mm]] and i E [n|/ ETTa), we have
(w((-)) w (0)
v:E.) Fagn -(W (1)
K,s: 14) (w!)) K...
(w K, U:E,)
(0)
2 (1 2(1 + O.)TT Iy y go +
0,
where the first step is by Eq. (10), the second step is by F(t), the last step is by Lemma E.39. Combined with D(t), when (WQ.s (f+1)
(.E.)) < 0 (note that we suppose (we, p)) 0, in more (++1)
general case this is (WQ.s (.E,) sgn -(wQ. pe)), this gives
(++1)
(w@. (++1) MEEt) + (wk, (e+1)
3nE = (w .. srEi) + K.o (
(48)
Note that at t = (2 + 38c) T
4. we have
(2443.))77) Q.
+ (w (2+336))T))
m. < (2 60.)7T 115, (1 + o(1)),
(49)
K.s
where the upper bound for query noise is by Lemma E.29 and the upper bound for key noise is by Lemma E.39. Then we have
(w (++1)
Q E.) R.s
(++1) Ei) - (w (w.. Ei) E.)( W.K. E.)
K.s.
=E|n]
2
Σ
r 11 + 7 11ll ( Q... S + ( K, 5.)1
.E (T)
Σ
v/ ||&.11 + +1 111/l11 @ Eil + R. {
(t)
sEin)/E (3) (w (++1)
2/ E,) sug
(. .14
(T)T (n² 115:117 | + 2t || 12 + o(1)) (m+ - E (T(TT) M07 1| + (2 60. Tz 11E.112 + o())))
mik
O( (( ||&.1 + 21 || || 7² + o(1))) - ( + + |1 + (2 60.)TT || (1 o()))
mi
2
2
O(/m))n² 15.11 + (m (1 30.)TT) O ((/mm) (1 + 38. 38.)TT || (1 + o())) > 0,
where the second step is by Eq. (46) for aligned part, Eq. (48) and (49) for unaligned part, the third
step is by concentration for (Ta) in Lemma E.27 Combined with H()), this implies H( + 1)
holds. Then we have H((t) holds for all (1 + 0.)TT Sts T4.
Remark E.40. When t = (4 + 70.)TT we have that for all s E [mx] and i e [m]]
(w y:E,) =sgn (w()) v.E.) =gn (w (. 14).
Finally, for all E [mx], E [n]] and 2 T4, we have
(w(tt)) W:G,) (w(1) Q, ) + (gg((w) (t)
Q. 3:E (ill.
(witt))) (++1) y:E,) (w 3.E,) + sgn( (w K yrE,)) )m ||Eall
(1)
(1)
(w(tttl) me) (w 14) +gu (w o. (1) ||]l)
(1)
(w K.s p1) (w()) (c)
(1+1)
K.s: 14) sgn((w() K.s 14) ||lall
(t)
(w()) ()) Q.at v:E.) =gn (w K. viEa) = (w!2 Q. 11) =sgn - -(wl R. 14).
E.8 PROOF OF THEOREM 3.2
In this section, we aim prove the Theorem 3.2.
72

---

Published as a conference paper at ICLR 2025
Proof of Theorem 3.2. By Lemma E.41 and E.42, complete the proof.
Lemma E.41 (Training and Logistic Test Loss), For any fixed > 0, ler I' = 2 2log(E ( s- we have (s(W(T)) S e and Lp( (W(T) 2 0.1.
Proof. Note that
(v(T) 11) (v("), (4) + T/T ||
= =0 || m. 1/2) + 2 2log(e
|| 1p
O(log(c || 07 ))
where the last step is by Ö(0 |||||| my 1/2) = o(n]T' (|l)) Then, for all E [n]]. we have
f 3is(W(2) (2) X;) = (8),11 + + n²) + (0 + (2) (v(0 v.E,))
(0 + (2)
+1,22) (v .1iE.)
0.9(v(T), (T),
,,9(v(T)y;EE)) y)E
log(e-) '),
where the first step is by the definition of f. the second step is by $1,1 2 0, $,21 2 and (v(T), 14) 2 O, the third step is by $1,12 2 0 $1,2 2 0.9, the last step is by Lemma E.8 and the definition of T' This implies (T) log( 1 + exp( - (W x.))) S e and thus Ls (s(((T)) < e.
On the other hand, for a new data point x = (y/u. T) let event & be the event that & has disjoint support with , Then when & holds, have
(3) (V yE Vf(W(T), x) = (5 + $21 (v(T), 14) + (812 + +
(T)
2(v(T), , 2(v(T), ,yE)
s Öllog( |||||| op -1 0oo 1/2 m 112
1/2)
$1
where the second step is by sn $12 821 522 E does not hold, we have
(T) (T) (T) (T)
< 1, the last step is by ooops 1/2 1/2 = o(1). And when
VT(
x) (i + 821 ( (((()(())) (T) 11) + (8 + 822 ( (( yE
< 2(v(T), ,H) + 2(v(T),yE)
< Ollog(E " |||1|| op 3 log(e-1)
Ollog(e -1))) 1))
Note that P[8] 2 1 -2, which gives
El((y(((()) X))] = (E(()E((((((T))))) 1|
> (1 - -2 log(1 + 0.1,
which completes the proof.
Lemma E,42
(Attention layer all attends to noise patch). Let Tattn (Taan)
O(log(aps/ 1|pl1)mmmmmm 1 1/2 for all E [n).
my
1/,,1//101000/0 -12 1//2-11/1100 11/2), have 8,11 o(1), 8
op
1/2
(Taun)
o(1),
73

---

Published as a conference paper at ICLR 2025
Proof. By the previous analysis, we already have 81,1
(0)
o(1) for L 2 TT. For 2 Tattn
(w/ (w Qia H4) (w K pe)
exp
(2)
8,11
exp [Exe(m)) (w
Kist ui) cxp 22 12 (1
(w (1)
Q,at u1) (w (6)
+ +exp (Zne(me) (w Qua u) (0)
Kis B:E,))
22222222 [|(l))) + exp (mat21fapps) | /4)
exp
exp (ppm (pll)) S o1),
where the first step is by definition, the second step is by H(t) in Stage IV.c, Remark. E.40, and the concentration in Lemma E.27, the last step is by t Tastn
74

---

Published as a conference paper at ICLR 2025
F DISCUSSION ON EXTENSIONS
F.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 I necessarily mean constant 0-1 test loss. And actually, we find it depends on the network initialization parameter du-
The 0-1 Test Loss is defined as
T (W) == F X) Ply/( 00.
The value of yf X) depends two components, the signal part (v, (u), which continuously increases during training, and the noise part (v, E), which random for the unseen noise &,- With probability least 1 2, the unseen & is disjoint with all training data E.-
oo is small enough, i.e., op = o 3/2 m. 3/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, with probability at least 1 - 6, we have
(v(T) (a) = 0(log(( ||||||| " 11 >> O(m 10001/0000p0//// = ((v(T),))
This is similar for (w (T) 2(1). Then, we have
Q.a: H4) v.s. (w. E). and (w K,st 1u) v.s. (w (w E), which gives (17) .22
(T)
VT(W(T), X) ($ + S21 ( (v( (T) m1) + (8
(2)
(1) + $22 (T (7) .EE
2 Olog(( ) || op - - 0m,, 1/2 ooo 212)
0,
which implies a small 0-1 loss Plys(W(T), .X) 0] 2 -(1))
On the other hand, if oo is large, i.e., ou = 2(0 2s 3/2 m/2 1/2 but still satisfies Condition 4.1, then have
(T) 11) = O(log(E 11) |||1||| op 1s 1 < 0(m 1/2 100000/)//3))00110 ((v(7), ,))))))))
This is similar for (w (T)
Q.st 14) v.s. (w E), and (w K.s U) v.s. (w K,s
(T)
W. E). which gives $11 $2
(T)
1/2 + o(1) since the symmetry of unseen noise. This means that the softmax outputs are random. Then.
(f(W(T), x) = (sff) + si (v()) 1)) + (8 + $22 (V .
=((1)(((T), v
which implies Ply/(w(aa) X) > - 1/2 S 1 - n (11), Note that the condition regarding do 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. (2023a); Vasudeva et al. (2024); Sheen al. (2024) assume all non-optimal tokens are the same in some way. If the data X has L signal vectors 4 and the latter L/2 patches are noise vectors g ( and all noise patches are the same, i.e., (!)) = g!d)) for L/2 < 1, < L. then the training dynamics
patches X (x) x2)] where the first L/2 patches
the same 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 noise patches, it remains essentially the same except that the signal/noise-signal softmax is around 1/L. while (the sum of)
75

---

Published as a conference paper at ICLR 2025
signal/noise-noke 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 μ and (L-1) & 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 al. (2023a); Vasudeva et al. (2024); Sheen et al. (2024) assume there is only one optimal token in one sample, which has a large gap from the non-optimal ones. Jiang et al. (2024) assumes the second noise vector &2 is greater than remain noise fz..... E,, in one data point. In our case, we need to assume for each sample there is a noise patch I. (WLOG, we assume it the last patch) such that Σt=2 [((h))1 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 since the gradient of value does not change. The gradient of query is
((11)(()1))) We can achieve this by applying large sparsity and/or large std op
(0) (b) xxx)) ) -x ,,
V Vwa., Ls( (W)
81,1a8,,t0
x;
72
1=1
a<b
and the inner product is (sgn(V Wo. Ls), (( = -1E
1a
sgn
Si,la Si,bb x(.
(WK.)xx ((0)
a<b
~ -
wi
Σ
a<L
We approximate the inner sum with only one term by 1) concentration of softmax to 1/L, 2) the assumption 21 l|s(') = o((l5(")1z)) which implies the dominance of (v, after stage I, and the dominance of (w (.... ). 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 L1 norm, the softmax at the position of the largest noise patch will converge to I, 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, would like to remark that the value matrix in our model setting already acts linear layer while many works Tarzanagh et al. (2023a:b) 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:
Vol. Ls (W)
Σ 2 Co n
1 + 81,21 (W ya14) + (S1,12 + 81,22) E/),,
1=1
where Orr denotes the parameter for the second layer F, at the r-th entry, initialized with j/my. for 1 < r S m. During stage I, it moves at most Tin which is negligible compared to the initialization 1 1/my, and will not affect our analysis for value in stage I. After stage I, the terms (wv..r-jyf. become positive for all 1 < i < and larger than the value signal term (WV.jars 14), which leads to that egn(Vo,,. o,. I.s(W) = -j. Consequently, all 01,- keep increasing while all 1,r keep decreasing. Until then, all 0) 's sign are the same as initialization and will not change during the rest of training. Therefore, 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, and Fu. respectively. When value converges to pattern specified in stage the second layer weights have almost 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.
76

---

Published as a conference paper at ICLR 2025
F.4 EXTENSION TO MULTI-HEAD ATTENTION
In this section, we discuss how to extend our main theoretical results, which focus on a single-head attention layer, multi-head attention layer. We note that theoretical results be easily extended to multi-head attention.
For multi-head attention layer. let the parameters be W Q. WK. WyJ where W q. WK.h K.h E IR and Wv.j.h E Rm for j {#1} and h [H]. Here, the H is the number of attention head, assumed to be a fixed constant. Then, the network can be written as J(W,X) = Fi (W, X) F-1(W, NWW,), X), where Pi (W, X) and F-1(W, X) are defined as:
F(W, = Σ (W)
4=1
F
Σ I Xsoftmax
x
W Wi. K, W Woxx!)).
7hp
Gradients in Multi-Head Attention. Regarding the gradients, the partial derivatives of the single- 0F OF,.A 0F,A. head model outputs with respect to the parameters, i.c., 2W JWKK JW remain unchanged. However, the gradient of the loss with respect to the single-head model outputs, i.e., 8F,,, change. This change, however, is linear and straightforward analyze. Intuitively, the model outputs increase approximately H times, causing the magnitude of the loss derivatives ll to decrease accordingly.
at does
Formally, theory shows that in single-head attention, the loss derivative are close to initialization up to = 477 where 477 is the time when the sign alignment of negative query noise completes. Specifically, for all i E [n], s 4T we have
Oe
1/2 o(1)
Of(w()) X.)
This implies (WTT) X,) = o(1) Then, the H-fold increase in the multi-head model outputs does not alter this result, so the effect of changes in JF, can be neglected.
88
Consequently, the behavior of signals and noise still follows the four-stage dynamics in the single-head case, with the dynamics of all attention heads being the same.
Experiments. In Figure 13, plot the full dynamics of our simplified transformer but with attention heads. We can see that the dynamics of query noise, key noise, query signals, and key signals identicalto those in the single-head model (Figure 18). Additionally, the dynamics of the softmax outputs in each head are consistent. These empirical evidences further support that our theory holds in the multi multi-head attention models.
77

---

Published as a conference paper at ICLR 2025
G MORE DETAILED EXPLANATIONS FOR THE FOUR-STAGE DYNAMICS
<V.EE>
<V
0.6
1.00-
< A
0.75
0.3
0.4
0.30
0.25
0.0
Figure 18: The full training dynamics of two-layer transformers with SignGD. This figure shows the dynamics of value, query, and key signals and noise, as well as softmax outputs and loss across all iterations. (a) Dynamics of value noise, and value signals. (b) Dynamics of query noise, key noise, query signals, and key signals. (c) Dynamics of softmax outputs. (d) Training and test loss curve. Omitted notations are the same Fig. 1.
Phase
Phase II
Phase III
Phase IV
Ti
T3
27:
47
Te
(1))
(v,4)
(55)))
14
(Wes.u)
(Wa..")
4
8izs
Figure 19: Timeline of the four-stage training dynamics. We mark all key points, including Ti, TSON T2. T3. TT. T4, along with their corresponding key behaviors the dynamics: D: Mean value noise increases rapidly, becomes positive, and grows lincarly with t. @: Mcan value signals remain stuck at their initialization. 3: Query and key noise are stuck at initialization. D: Query and key signals stuck at initialization. 5: Noise-signal softmax outputs are stuck at initialization. @: Mean value noise continuously increases. ①: Mean value signals are negligible compared to mean value noise. @ Query noise and key noise align their sign with each other. ②: Sign alignment completes, query noise and key noise continuously increase, and grow linearly with L. @ Query signals leave the initialization, align their sign with the majority of query noise, and grow linearly with t. @: Key signals leave the initialization, align their sign with the minority of query noise, and grow linearly with t. @: Query signals and key signals continuously increase. 3: Noise-signal softmax outputs decrease exponentially. @: Positive key noise continuously increases. Negative key noise flip the update direction, and align with query signals. (:: Noise-signal softmax outputs reach (11) and remain sufficiently small. ( Positive query noise continuously increase; Negative query noise flips update direction and aligns with query signals. @ Negative key noise completes sign alignment, and all key noise continuously increases. (:: Negative query noise completes sign alignment, and all query noise continuously increases. ② All signals and noise continuously increase, and the loss decreases exponentially, Note that by "increase", we refer to increase in magnitude.
78

---

Published as a conference paper at ICLR 2025
H DISCUSSION ON DIFFERENCES FROM REAL- WORLD SETUP
We would like to discuss the following key differences between our theoretical setups and the transformers used in real-world applications:
The gap in the data noise structure.
The data model considered in our study is simplified - we assume Gaussian noise that is i.i.d. across data samples and independent of the true features. All these simplifications inevitably create a gap from real-world datasets. Additionally, our data settings are motivated by image datasets. For language data, since language tokens have more dense semantic information, and since the noise must be discrete, the noise in language data could be more structured compared the image data.
These differences are acknowledged in the Limitation section of our initial manuscript, and make it difficult to compare the noise in our data model and real-world language or image datasets. However, we would like to highlight again that even in this simplified data model, significant challenges present in the theoretical analysis.
The gap in tasks regarding optimization.
The sensitivity of SignGD to noise, demonstrated in our findings, is relative to GD and is observed in a task where both optimizers achieve perfect convergence. In contrast, real-world language and vision tasks using transformers often reveal that GD struggles to optimize effectively, with a significant training loss gap compared to Adam (Zhang et al., 2024c).
This highlights a gap between our task and real-world tasks in terms of optimization. This gap in tasks regarding optimization may also affect the generalization performance and robustness of the optimization methods. Additionally, as GD cannot optimize effectively in real-world transformers, it is difficult to fairly compare the generalization properties of GD and SignGD. Understanding why GD fails to optimize transformers on real-world tasks should be an important prerequisite.
However, we would like to emphasize that, although our task optimizes easily, our work pro- vides a clear example where SignGD, and by extension Adam, fails in learning transformers, and hence is not always more effective than GD.
The reason behind this gap: This gap in the optimization may stem from factors such as simplified data structures and models in our setups. However, since our experiments two-layer and three-layer transformers demonstrate consistent results regarding training dynamics (See Appendix B.6), we suspect this gap more attributable to differences in data.
How to define robustness in the real-world transformers, particularly language models. The concept of robustness in real-world language models requires careful consideration. In practice, when training transformers on language tasks with Adam, training instability and loss spikes are often observed, usually due low-quality data batches (Chowdhery et al., 2023). Given that transformer training typically involves billions of tokens, the generalization performance often closely tied to the training performance. Thus, from this viewpoint, training instability and loss spikes can reasonably be viewed as a form of non-robustness. The message we aim to convey is: To fairly connect theoretically motivated findings to real-world applications, it is crucial to carefully define what robustness means in practical transformers like language models.
In the above, discuss the gap from the perspectives of data noise, optimization complexity, and the definition of robustness. We hope this discussion helps clarify the differences between our theoretical findings and real-world scenarios.
79


Reviews:
[{"review_id": "WBdep4oXzb", "reviewer": "Reviewer_w8BA", "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.", "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.\n\n2. 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)).\n\n3. 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?", "limitations": "", "rating": 8, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 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. \n\n2. 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.\n\n3. The theoretical approach is novel.", "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.\n\n2)  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.\n\n3) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "aMtT5eAGdK", "reviewer": "Reviewer_hqL4", "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.", "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. \n\nFrom 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.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "The paper provides a multi-stage framework to analyze the transformer training using SignGD, making this complex behaviour into small interpretable stages.\n\nBy 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.\n\nThe combination of theoretical proofs and experimental validation strengthens the proposed analysis and its overall message.", "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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BV2EcIKgNh", "reviewer": "Reviewer_YN62", "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.", "questions": "Please see the weaknesses.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 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.", "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.\n\n- 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. \n\n- 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.\n\n**Minor Issues**  \n- 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.\n\n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)