From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics
Authors: Zheng-An Chen
Tao Luo
Venue: NeurIPS 2025 oral ()
Keywords: Two-stage analysis
Training dynamics of transformers
Condensation
Small initialization
Primary area: deep_learning

Abstract: Although transformer-based models have shown exceptional empirical performance, the fundamental principles governing their training dynamics are inadequately characterized beyond configuration-specific studies. Inspired by empirical evidence showing improved reasoning capabilities under small initialization scales in language models, we employ the gradient flow analytical framework established in \cite{zhou2022towards} to systematically investigate linearized Transformer training dynamics. Our theoretical analysis dissects the dynamics of attention modules into two distinct stages. In the first stage, asymmetric weight perturbations from random initialization sustain non-degenerate gradient dynamics in parameter matrices, facilitating systematic escape from small initialization regimes. Subsequently, these matrices undergo condensation, progressively aligning toward the target orientation. In the second stage, the previously static key-query matrices actively participate in training, driving the normalized matrices toward asymptotic rank collapse. This two-stage framework generalizes classical directional convergence results.

From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics
Zheng-An Chen', Tao Luo!2
School of Mathematical Sciences, Shanghai Jiao Tong University
Institute of Natural Sciences, MOF-LSC, CMA-Shanghai, Shanghai Jiao Tong University
Abstract
Although transformer-based models have shown exceptional empirical perfor- mance, the fundamental principles governing their training dynamics are inade- quately characterized beyond configuration-specife studies. Inspired by empiri- cal evidence showing improved reasoning capabilities under small initialization scales in language models, employ the gradient flow analytical framework tablished Whou et al [2021] to systematically investigate lincarized Transformer training dynamics. Our theoretical analysis dissects the dynamics of attention modules into two distinct stages. In the first stage, asymmetric weight perturba- tions from random initialization sustain non-degenerate gradient dynamics in pa- rameter matrices, facilitating systematic escape from small initialization regimes. Subsequently, these matrices undergo condensation, progressively aligning toward the target orientation. In the second stage, the previously static key-query matrices actively participate in training, driving the normalized matrices toward asymptotic rank collapse. This two-stage framework generalizes classical directional conver- gence results.
Introduction
The waasformer-based models Yaswani et all |01 have achieved remarkable breakthroughs in various fields, with the successful application of large language models. However, the theoretical
analysis of the transformer still remains in specific tasks, such as in-context learning settings Brown
etall [2020], Olsson et all [2022]] Bieth all [2023] or single attention block with reparameteriza- tion Tian all [2023]] The use of linear regression tasks Zhang al. [Αيلа]] and Markov chain
tasks Ildi ci all [2122]] has provided highly interpretable theoretical analyses, but crucial question
still remains: Can analyze the characteristics of the training dynamics of transformers indepen-
dently of specific tasks?
Meanwhile, small initialization has been increasingly shown to hold promise in the training process of large models, especially for reasoning tasks. Numerous studies Zhang et al, [202217 20256]. Yao aL [2025] suggest that the implicit regularization effect of small initialization is still effective in large language models. This effectiveness particularly significant the context of modern large models, which are characterized by extreme overparameterization. In these regimes, where explicit regularization techniques like weight decay or dropout may prove insufficient on their own, implicit regularization becomes pivotal. It operates by imposing intrinsic constraints the training dynamics and the resulting parameter space, effectively guiding the model towards solutions with good generalization properties despite the vast hypothesis space. This implicit bias is key under- standing how models with such immense capacity manage to avoid severe overfitting and achieve remarkable performance unseen data.
*Conresponiing author: luotao41 @sjtu.edu.cn.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Motivated by these observations, we propose to investigate the training dynamics of transformers under a small initialization setting. Leveraging the gradient flow theme similarly to Zhou et all [2022], We delineate different training dynamics for outer parameters versus attention parameters We and W Transformers.
We dissect the dynamics of attention modules into two distinct stages. In the first stage, the core attention mechanism, softmax (QKT), remains nearly stagnant, as asymmetric weight perturbations from random initialization drive non-degenerate gradient dynamics in parameter matrices, particu- larly Wv. facilitating escape from small initialization regimes. During this escape, the parameter matrix converges row-wise toward the target orientation, a process we term condensation. We theo- retically prove that condensation guaranteed under small initialization, and experimentally observe that it stabilizes without significant fluctuations.
remains
In the second stage, after the outer parameters, such as Wr, reach a quasi-steady state, the previ- ously static key-query matrices, Wq and Wx. hegin to actively participate in training, driving their collapse. This two-stage framework not only elucidates the training dynamics but also generalizes classical directional convergence results, offering a robust theoretical foundation for Transformer optimization.
To sum up, our contribution can be summarized as follows.
1. Blow-up Dynamics: We prove the blow-up property (Theorem M holds for measure- measure-theoreticall theoretically generic initializations, eliminating reliance on dichotomy assumptions while ensuring model non- degeneracy.
2. Condensation Mechanism: By introducing a condensation condition (Assumption I., we estab- lish theoretical guarantees for condensation emergence (Theorem 2)
3. Key-Query Collapse: After outer parameters stabilize in quasi-steady state (Assumption D), the key-query matrices begin active training, leading to asymptotic rank collapse of the normalized key-query matrices (Theorem B)
4. Experimental evidence:
We validate our hypotheses and theoretical predictions on both syn- thetic and real datasets with one and multi-layer Transformers, consistently observing two-stage dynamics marked by condensation and an eventual rank collapse of the normalized key-query matrices (Figure R)
2 Related Works
Training dynamics of transformer. Given the scale of modern models and the complexity of optimizers, studying the training dynamics of Transformers is challenging problem. Prior works have primarily investigated the optimization dynamics of a single attention layer eu [2021]] Li et allll [2025]] Snell al [ΑΑΣа]] However, these studies mainly focused specific tasks, such as topic structure prediction and translation.
Recently, the dynamics of in-context learning (ICL) has emerged as prominent research area within Transformer dynamics, particularly given ICL's ability to solve novel tasks without parameter up- dates. Many works Mahankali et all [2112]] Zhang et al. [[|zE]] Huang et all [202]]] Collins et al [2024]] have focused on the linear regression setup to theoretically investigate the mechanism of ICL in Annleelayer Transformers, line of work that has also informed algorithmic development Akyüüre et al. [2023), Bai et all [2023), Guo all [2024]]. Another line of research investigates how specific structures within attention emerge during training. notably starting with studies on induc-
tion heads Olsson et all [2022]], Reddy [20224), [Edelman all [2024]] Zhang et al. [2025%], memory
recall mechanisms Biett
[023
Cabannes ef
[024 and even causal structure Nichami
at [2024]]
Despite the sophisticated structure of realistic Transformers, Tian all | proposed novel mathematical framework for analyzing the joint dynamics of MLP and attention blocks and success- fully explained the sparsity of attention score matrices. Meanwhile, Chen [20z0]] provides rigorous proof for the convergence of the ICL linear regression task using gradient flow with suffi- ciently small initialization.
Small initialization and its applications The initialization of neural network significantly affects
its learning outcomes Arora all [201%)] Wilhams et al [M1%) Mei ct all [2018] Lacot et al

---

|010] Rotskoft and Vanden-Eing Eijnden [018], Zhang et all [2020] Small initialization is common
setting investigated the study of neural network optimization dynamics, which different with the Neural Tangent Kernel (NTK) perspective in infinitely wide networks. For linear model, Di
and galssry [201]] establish matrix alignment results theoretically. For nonlinear model, Zhou et al [202]] found that small initialization can similarly promote parameter condensation, thereby
reducing model complexity. Theoretically, et [202T]] Chen al [20221]] Zhou of all [202%]
Kumar and Haupt [2072]] have deepened the understanding of this phenomenon. The recent survey
article all [2025] systematically synthesizes empirical and theoretical findings.
The implicit bias induced by small initialization has also been discussed in the context of linear
regression Saxe all [2013]] Min at [[02T]], Marre at [112]]] and matrix factorization tasks Li
at [2013] Arora ct all [2019a], Stöger and Soltanolkotabi [202]] Soltanolktabl et [2023]] ER
atall [2024]] More recently, many researchers have adopted small initialization settings to simplify
the analysis of training dynamics in more complex models. From theoretical perspective, Zhang
et |0225 applied small initialization to ICL. tasks to analyze the behavior of linear attention.
al | considered the training dynamics of the embedding space under small initialization
using a synthetic dataset designed for reasoning and memorization. From an applied perspective, Zhang al [2019], Huang al. [2020]] Zhu et all [2021]] highlighted the importance of initialization in Transformers, while Bachlechner
[2021]] combined zero- initialization with residual blocks in Transformers. Some research Zhang al [20224]]] Yao et all [2026] shows that small initialization helps Transformers learn the reasoning aspects of data rather than just memorization, a principle
already applied in realistic LLM training (Yn et al [2025]].
3 Preliminaries
3.1 Basic Notations
First, we introduce some notations that will be used in the rest of this paper. Let n and dm be the number of samples and the width of hidden layers, respectively. Let [r]] denote the set of integers from 1 to n. Denote vector L norm as Il- and matrix Frobenius norm as ||-||r. Let (.. represent standard inner product between two vectors. For a vector v, denote its k-th entry as Uk. For matrix A. denote the element in the k-th row and k'thh column as Are. And denote k-th row as A, and kh column as A Unless otherwise specified, summation 'E' is performed over the network width.
3.2 Classification Task
For decision tasks, the network produces scalar output fx( R. The predicted class assignment is determined by the sign of the output. The dataset is denoted by {(X,, s, w where X; E Rexdm stands for input sequence in which represents the sequence length and dm represents the hidden dimension, and y E {+1} stands for label. For a loss function e: R R,, we define the empirical risk as L(0) 1 Σ E (X.))
Binary classification:
For probabilistic tasks, the network outputs logit vectors fa(X) E Rd that parameterize categorical distribution via the softmax transformation Py = il xx:0) = (x)(fe(X):)
Multi-class Classification:
Eie, exp(f)(X);) where d, denotes the vocabulary size. For cross-entropy loss, we define the em-
pirical risk as ((0) = - Σ log P(y y4 X):
3.3 Condensation and Rank Collapse
We formalize the two geometric phenomena that will recur throughout our analysis.
Definition 1 (Condensation). Let (() be matrix with rows Wa (l) (or columns (*))). We say W condenses to direction v if, as
We(t) [Tll?
+1 every index k with [t)ll) # 0
(equivalently, the same holds columnwise).
Condensation is a directional notion and implies rank-1 collapse when a unique direction emerges. Rank collapse is a spectral notion and allows k > 1 when multiple top singular directions are tied.
3

---

Definition (Asymptotic rank collapse). Let (t) be a matrix. We say W exhibits rank-k collapse the limit
(
W 00 = lim t-I
[Wttl
exists and rank(W.>)) S
Theoretical Results
4.1 Problem Formulation
To analyze condensation phenomenon in transformers, we begin by formulating the problem. Specif- ically, we consider the following one-layer transformer model:
Definition (One-layer transformer). Let X E Rexdm be an input sequence of length s with model dimension dm. The Transformer function fe Raxx - R is defined by the composition of attention and forward operations:
Jo( (x FFN(AAU1(X))
(All(())w!)) w
(1)
The atlention sublayer Atin : Bexdd - Rexdm is computed as:
XW XWv.
XX) softmax
(2)
Vom
where parameter matrices satisfy Wo, WK, Wv, W E RRddxdd xdm and W E Rdd, The activa- tion function
+ R. is tanh.
We use one-layer transformer fe to solve binary classification tasks and take the last dimension of the output fo (X)), the output. So the empirical risk be minimized is given by
n Σ (()..((X.).). n2
(
For simplicity of presentation, employ the exponential loss function E(q) e which is com- monly used in the analysis of classification tasks Lyu and Li [2020]] The analysis can be readily extended to other loss functions such as the logistic loss.
The model parameters are initialized with Gaussian distributions scaled by small perturbation parameter e:
w k P]] NO0,
EE ~ =3,, WO WK WV.K (0,
(4)
where < 1 controls initialization magnitude. To analyze training dynamics, adopt the gradient flow (GF) framework- continuous-tima limit of gradient descent. Given the small initialization scale, derive effective dynamics through a perturbative expansion of the empirical risk (()) in powers of e.
First, we normalize parameters by absorbing the initialization scale:
W = =-111 Will = - -wil,, Wq = -Wo W
Wy=
Performing Taylor expansion of (()) about = 0 yields the leading-order asymptotic form:
((
Σ e²
Σ
W. Will Wo
+ o(E3)
(5)
2n
This expansion induces simplified gradient dynamics characterized by the following proposition. Proposition 1 (Effective training dynamics), Given a binary dataset {Xx,, y we define con- densation direction a and rescaled time coordinate as follows:
Σm yi
- X
T:=
Σ i=
1/1 Σ
Xij
1.
(6)
E
Vi
Σ
X1,)
1=1

---

Then, normalized parameters 0 follow leading-orde order dynamics after rescaling:
d0 Va (wWy
df
Proposition M reveals hierarchical learning mechanism: During initial training phases, the fully- connected layers w W and value projection matrix Wv exhibit substantial updates, while the query/key matrices We and WK in the self-attention module remain quasi-static.
For subsequent analysis, we define the projection of Wy onto v W. == AWy (omitting bar notation for simplicity) and introduce the energy functional:
E Wn w
(8)
The effective dynamics can thus be interpreted gradient ascent on this energy landscape.
4.2
Blow Up Dynamics
We first elucidate why Transformers with small initialization can successfully train and eventually escape the small initialization regime. This phenomenon emerges from the interplay between two fundamental mechanisms:
1. Effective dynamics driving:
The parameter evolution governed by the effective dynamics ex- hibits remarkable symmetry, manifested through strict conservation laws that preserve key quan- tities during training.
Random normal initialization: While degenerate cases theoretically exist under Gaussian ini- tialization, they occur with vanishing probability (measure zero in parameter space). Conse- quently, the dynamics almost surely demonstrate non-degenerate characteristics, ensuring stable training trajectories.
To preserve dynamical symmetry, we invoke the following proposition following the approach es- tablished in prior works J1 and Telgarsky [
Proposition 2 (Conservation laws). Under the gradient flow dynamics prescribed by system Eq. (ZA) the following system of conservation laws emerges:
d dt
W v.k
Σ
(
(wll) RE
=0
and
dt (w! Σ
d
(w!?)) 2
(will) 0.
(9)
We now analyze the non-symmetric property arising from Gaussian random initialization, with par- ticular focus on the degeneracy mechanism. Crucially, we establish that degeneracy exclusively occurs when initialization violates the following non-degenerate initialization:
Definition 4 (Non-degenerate initialization). Lei 0 (Wv, Will, w (21) denote parameters initial- ized from a Gaussian distribution. The initialization is called non-degenerate if W 112 + |W and
2
W.
1 + min(l]
| 12 (I 11 I
+
(10)
Having clarified the definition non-degenerate initialization, we present the following theorem that reveals the non-degeneracy property of effective training dynamics.
Theorem 1 (Blow-up in finite time). Let the parameters be initialized randomly as above from Gaussian distribution. Then, almost surely, the initialization is non-degenerate in the sense of Definition A and the effective training dynamics Eq. ( blows up in finite time. That is, there exists T O such that
lim E(t) = +00 67
Proof sketch. We prove finitetime blow-up via Riccati-type differential inequality for the energy E()). Full technical details are provided in Appendix A.
(1) Superlinear growth. direct computation gives (t) 2 3E(1)4/3, hence O,E(t) < and
1
E(I) 2
(11)
(E(0) 1/3 1)²
5

---

For E(0) > this yields T - E(0)-1/3 (2) Negative initial energy. If ((() 0, then
E(t) Z ((-E(0)) t)s
(12)
so E' is increasing and cannot remain negative indefinitely. Assuming E(t) I 0 for all t leads to contradictions with (i) standard continuation at finite T", or (ii) monotone limits at T oc. reducing to the borderline case E(t) t O.
(3) Borderline exclusion. In the regime 7 = and E()) t structural identities and conservation give (t)ll (t) + 0 and E(t) - O. Under the non non-degenerate initialization (Def. HD.
this forces contradiction, since the limiting E must be strictly positive.
4.3 Condensation Dynamics
We have proved that energy and parameters norm will blow up almost surely. It implies the effective dynamics drive parameters escape small initialization area in finite time. The next question is how the effective dynamics affects the emergence of condensation and whether there exist observables to help us characterize condensation.
We propose a condition of condensation and verify its effectiveness using experimental and theoret- ical methods. In particular, we theoretically prove that the solution of the effective dynamics has specific properties, which is to some extent sufficiency argument. The necessity argument is quite difficult theory. But experimental results provide us a strong implication that this condition maybe also necessary.
Assumption 1 (Condensation condition). The parameters satisfy the condensation condition at time t. That
For each index E [dml) w! W. w 1) and Wo,, W. wii)
For each pair i.j E [dm]. (w:!²) w
[1]. w will)) 0, and (We,, w!!i W., wr!l) O.
This hypothesis can be verified experimentally in See. formalize the statement of the culminating theorem as follows:
Then, based on Assumption M.
Theorem (Condensation). Under Assumption i the effective dynamical system governed by Eq. (Z)) drives the parameter matrix Wv undergo condensation in the sense of Definition
This section gives a highlevel proof sketch; full details appear in Appendix
Proof sketch. We establish finitetime directional convergence (condensation) via geometric propa- gation and twosided energy control.
(1) Geometric consistency and alignment dynamics. Under Assumption Proposition g shows that once the alignment condition holds at some to < T it propagates throughout (lo. T* Propo- sition & yields structural dichotomy of the columns of Wv into condensing class C1 and uniformly bounded class C2. Propositions R and @ further establish dynamical alignment between ( and its time derivative, ensuring coherence of the evolving direction.
(2) Singularity structure and condensation. Proposition E supplies an energy upper bound which, combined with the lower bound in Eq. (ID), furnishes bilateral estimate on E(I). A telescoping- integral argument then proves that the condensing indices dominate in finite time, completing the proof of condensation via Theorem
4.4 Key-Query Dynamics
Following the initial training stage, the parameter matrices Wv. Will and W22] exhibit substantial growth in magnitude, effectively escaping the amall-initialization regime. In contrast, the key-query matrices We and Wx demonstrate remarkable stability scale. This separation phenomenon fundamentally governed by the effective dynamics of the learning system.

---

A pivotal question arises: Under what conditions do the quy-query matrices become dynamically ac- tivated, thereby enabling the attention mechanism to exert its structural influence? We hypothesize that during early training, Wv. W and W 2] converge almost vanish, temporarily stabilizing in this dormant state. The following analysis provides mech- anistic insights into this dynamical freezing phenomenon. The final activation function omitted from our analysis. This is justified because the layer's pre-activations consistently operate within the linear regime of the function. Furthermore, empirical results confirm that its inclusion does not alter the model's learning dynamics (refer to B.Z). Now empirical loss (()) has the following decomposition:
critical point where Wq and WK
((0) ~ (1)(0)22,(0),
(13)
where (i,((0) =exp (-V [Ejn Xi, Wv W.
(exp0)
(0) =exp (exp))exp(( { -y: [Σi= X., WoWkk
+ Ei
X. W
X., Wy wiel))
}
ddm
Based on the above discussion, we now formalize the following assumption. Assumption 2 (Dynamics separation stage). After the breakdown of the effective dynamics in Eq. (Z)) let 8 denote a small purameter. The gradient flow subsequently enters a stage characterized by:
1. Criticality conditions: The outer parameters (Wv, W W(z) converge to quasi-stationan configuration such that Ww, = Wwu = V we L1 = ( where Li = E, f1,i-
Key-query stunting: The attention parameters remain small, satisfying Weeg WK.jl = (6)), until their norms [[Wall and |W&| exceed critical scale.
To facilitate empirical validation, we introduce modified version of the basic equivalence of the first part of Assumption D denoted as Assumption P Assumption 2 The outer parameters (Wy. W w(R) reach a quasi-stationary state whose directions vary negligibly over time, i.e., d Wy de (WWIT) at WOT aE
d will
w
PJ 0, and the
W
loss evolution satisfies dC O(82) for all i. dt
Since we assume the small parameter d is still relatively small, we get the following Proposition which illustrate evident dynamics separation and the leading order dynamics of key-query matrices. Proposition 3 (Effective dynamics during dynamics separation stage). Under Assumption D or As- sumption the empirical risk L(O) exhibits the following properties:
1. Dynamics separation The gradients of the empirical risk with respect to Wv. Will and W are of order (82), while the gradients with respect to the query matrix Wq and key WK are of order (5))
2. Key-query dynamics Treating Wv. W and W as fixed due to dynamics separation, the leuding-order dynamics of key-query matrices are given by
dWo = FWK. dWk df
F
(14)
dt
where F is defined as follows
F
1
M. M. W W Wy
Σx,,
X,
1
(15)
n. dm
8
Since the dynamics governing Wa and Wx form linear ordinary differential equation system in this context, we can rigorously establish the subsequent conclusions.
Theorem 3 (Asymptotic rank collapse). Given the key-query dynamics governed by Eq. (LA), the normalized key and query matrices exhibit rank collapse as Definition A Specifically, when F possesses unique largest singular value, both normalized matrices asymptotically become rank 1.
The detailed proofs of Proposition B and Theorem B can be found in the Appendix A
7

---

Experimental Results
In this section, we first demonstrate the phenomena cohesion and rank collapse using synthetic data and confirm the assumptions required for our theoretical analysis of the one-layer Transformer model. We then present experiments on natural language processing tasks demonstrate the gener- ality of our theoretical findings with respect to various datasets and network architectures.
5.1 Synthetic Dataset
We employ the concept of the anchor function Zhang et al. [2024]] to construct synthetic dataset that simulates simplified language modeling scenario. The model is one-layer Transformer with tanh activation, trained using cross-entropy loss and the Adam optimizer. Further experimental settings are detailed in Appendix B.l
5.1.1 Phenomenon: Condensation and Rank Collapse
To dissect the learning dynamics, we visualize the training process through three complementary lenses: the cosine similarity of parameters (Calculation method refer Sec. B. the relative change of norms, and the effective rank weight matrices. As shown in Figure these analyses collectively reveal distinct three-stage training trajectory, which we characterize as Condensation, Key-Query Rank Collapse, and the further training.
1977
(8)
X10->
143
Stage
5.2
Stage
20
hustae haining
19:1 1M
15
4.8
5
3.0
4.6
Average
0.5
4.2
0.0
10000
Figure 1: (a) Evolution of cosine similarity matrices for outer and attention parameters. The training process is partitioned into three stages: Condensation (Stage 1), Key- Bey-Query rank collapse (Stage 2), and further training stage. Stage transitions are identified by plateaus in the loss curve and structural shifts in these matrices. (b) The relative change of norms between attention and outer parameters. The gray dashed line marks the onset of Stage 2, where updates to the attention param- eters begin to dominate, (c) Evolution of the effective rank for both parameter groups, tracking the change in their intrinsic dimensionality throughout training.
The training process begins with rapid decrease in loss, driven almost exclusively by the outer- layer parameters since the relative change of the outer parameters far exceed those of the attention parameters (Ib) during this initial phase. This intense optimization leads the condensation phenomenon, where the initially random outer parameters organize into low-rank configuration. This is visually evident from the emergence of block structures their cosine similarity matrices (Fig ([) and is quantified by monotonic and significant decrease in their effective rank (Fig. II(c)). Throughout this stage, the attention parameters remain largely static and unstructured.
Following the initial phase, the training loss enters a prolonged plateau. This signals critical transition in the learning dynamics, marked by the gray dashed line in Fib. At this stage, clear dynamics separation occurs: the updates to the outer parameters subside, and the attention parameters become the primary focus of optimization. This empirical observation validates our theoretical framework, particularly Proposition B. As the changes in the outer parameters become slower (supporting Assumption 2), the attention parameters begin learn their specialized roles. This is characterized by rank collapse, confirmed visually by the sudden formation of structure in their similarity matrices and quantitatively by a precipitous drop their effective rank (Fig. IIKa), IIC))..
8

---

5.1.2 Experimental Validation of Key Assumptions
To ground our theoretical analysis in the observed dynamics, provide direct empirical val- idation for the key assumptions that underpin our framework: Assumption II and Assumption
Wo We
Wy W
(b)
(c)
15.0
Satisfied Proportion
Direction similarty
0.9
10.0
7.5
0.4
5.0
Condition
0.2
0.5
Condition 2
0.0
2.5
0.0
3000
Step
Figure 2: (a) Proportion satisfied conditions in Assumption II, measured as dm initions of Ai and A2 refer to Sec. [B.I). (b) Similarity between singular vectors of two adja- cent time steps. For example, let V, and E++1 V441 be the singular value decomposi- tions of parameter matrix W, and W The similarity is defined as 1 cos(u): u(+1) (or dm
[A and (Def- de
di i=
cos( vf, 1))1 (c) Frobenius norms of parameter groups.
First, we examine the condensation condition. Figure Daa) plots the proportion of satisfied conditions in Assumption M1 The proportion rapidly approaches 1 within the first 200 training steps, confirming that the outer parameters quickly converge to a state where this assumption holds.
Next, we validate the assumption of dynamics separation. As discussed in the previous section, our observation that the gradual change of outer parameters during Stage and flat loss curve (often means critical point has appeared) already provide strong qualitative support for the first part of Assumption D. To analyze this more rigorously, we examine its empirical variant, Assumption Pa This assumption points that the direction of parameters remains unchanged and the leading-order loss changes very slowly.
Figure (ibb shows the cosine similarity between the singular vectors of the outer parameter matrices adjacent time steps. The similarity for all outer parameters remains extremely close after the first stage. This indicates that the subspace spanned by these parameters is highly stable, meaning their directional structure is effectively frozen. This stability, combined with the flat loss curve observed in Stage 2, provides compelling evidence for Assumption 2X Figure (ic) validates the
scale separation implied by the assumption. It shows that by the onset of Stage 2, the Frobenius norms of the outer parameters have grown significantly, while the norms of the attention parameters remain small and close to their initialization values. This confirms the expected scale difference between the two parameter groups, where outer parameters are O(1) and attention parameters are (6).
5.2 Real Task
We further validate our theoretical predictions on a real-world language modeling benchmark, Wiki- Text MMeity al. [2017]. Unlike the synthetic setup, where anchor functions explicitly defined, WikiText provides natural linguistic dependencies and high distributional variability. This allows us test whether the proposed two-stage dynamics, early condensation of outer parameters followed by attention-drive rank collapse, persist in realistic Transformer training. In this setting, we employ two-layer transformer with GeLU activation and residual connections. To keep the consistency of architecture and focused on the core dynamics, layer normalization is omitted. Further experimental settings are provided in Appendix B.3.
As shown in Figure the two-layer Transformer on WikiText exhibits the same stage-wiss wise dynamics observed in the synthetic experiments. During the initial phase, the outer parame- ters (Wv, Will, w Ply in both layers undergo rapid condensation, while the attention weights

---

Stage
Stage
10.2
Purther Training
10.0
922
M42
9.8
2030
4000
6000
01000 1.0000
12020
Stop
Figure 3: Evolution of cosine similarity between parameter of the two-layer transformer WikiText dataset. Training dynamics also show a similar three-phase characteristic. Superscripts are used to indicate parameters of different layers, and subscripts indicate different parameters within a layer. For example, W represents the value matrix of the first layer.
(Wo: Wx) remain largely unchanged. As training proceeds and the loss enters plateau, at- tention parameters begin to evolve, displaying sharp rank collapse that reorganizes internal repre- sentations.
This empirical observation confirms that the separation between outer-parameter condensation and attention-driven rank reduction not an artifact of the synthetic dataset also emerges naturally in real-world text modeling. The consistent appearance of this two-stage dynamic across both syn- thetic and natural settings suggests that implicit regularization, first through low-rank condensation and then through targeted attention adaptation, may serve general mechanism underlying the emergence structured representations in Transformer models.
Discussion
6.1
Conclusion
This work advances the theoretical understanding of transformer training dynamics by establishing two-stage analytical framework. Through gradient flow analysis, show that small initialization helps models escape degenerate regions via asymmetric weight updates, leading to condensation of parameter matrices toward task-relevant directions. In the subsequent stage, the key-query matrices undergo a coordinated collapse that further refines the learned representations. Together, these re- sults clarify the mechanisms underlying the condensation and rank collapse phenomena, providing a principled foundation for future studies on Transformer optimization and generalization.
6.2 Limitations
While this work provides valuable theoretical insights, its most significant constraint stems from analyzing exclusively binary classification scenarios: a simplification dictated by technical barriers in gradient flow analysis. This narrow scope inherently precludes insights into transformers' dynam- ics in practical multi-class classification or interactions between multiple prediction targets and attention mechanisms likely emerge. Though focused theoretical simplification is methodologically justified, extending this framework to broader problem domains remains critical for unifying theory with real-world transformer optimization. Fu- ture work should prioritize overcoming these technical limitations to theoretically verify whether conclusions hold true beyond binary settings.
learning contexts, where complex
10

---

Acknowledgments and Disclosure of Funding
This work sponsored by the National Key R&D Program China Grant No. 2022YFA1008200 (T. L.). We also thank Shanghai Institute for Mathematics and Interdisciplinary Sciences (SIMIS) for their financial support. This research was funded by SIMIS under grant number SIMIS-ID-2UU ST. The authors grateful for the resources and facilities provided by SIMIS, which essential for the completion of this work. We thank Pengxiao Lin for insightful discussions and support and encouragement to the authors.
References
Ekin Akyürek, Dale Schuurmans, Jacob Andreas, Tengyu Ma, and Denny Zhou. What learning algorithm is in-context learning? investigations with linear models. In The Eleventh International Conference on Learning Representations, 2023, URL https://openrevieww net/ net/ wwwpee///orm/m/// orum?id=
1gUX4H8yN411
Sanjeev Arora, Nadav Cohen, Wei Hu, and Yuping Luo. Implicit regularization in deep matrix factorization. Advances in Neural Information Processing Systems, 32, 2019a.
Sanjeev Arora, Simon S Du, Wei Hu, Zhiyuan Li, Russ R Salakhutdinov, and Ruosong Wang. On exact computation with an infinitely wide neural net. Advances in neural information processing systems, 32, 2019b.
Thomas Bachlechner, Bodhisattwa Prasad Majumder, Henry Mao, Gary Cottrell, and Julian McAuley. Rezero all you need: fast convergence at large depth. In Cassio Campos and Mar- loes H. Maathuis, editors, Proceedings of the Thirty-Seventh Conference Uncertainty in Arti- ficial Intelligence, volume 161 of Proceedings of Machine Learning Research, pages 1352-1361. PMLR, 27-30 Jul 2021. URL https://proceedingg mmr...mmm..... ptap. html.
Yu Bai, Fan Chen, Huan Wang, Caiming Xiong, and Song Mei. Transformers as statisticians: Prov- able in-context learning with in-context algorithm selection. In Thirty-sevnnth Conference Neural Information Processing Systems, 2023. URL https::/ http:://pppaarrrrreeee//////maoo LiMSquuVgy.
Zhiwei Bai, Jiajie Zhao, and Yaoyu Zhang. Connectivity shapes implicit regularization in matrix fac- torization models for matrix completion. In The Thirty-eighth Annual Conference on Neural Infor-
mation Processing Systems, 2024. URL https://openreviiet http:///ppperrriooo...tterrrrrra wtt::oorummieeegggiiiitaaa
Alberto Bietti, Vivien Cabannes, Diane Bouchacourt, Herve Jegou, and Leon Bottou. Birth of transformer: A memory viewpoint. In seirth- Conference on Neural Information Process- ing Systems, 2023. URL http:://pppaarrrrerr....//tttttt
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, 33:1877-1901, 2020.
Vivien Cabannes, Berfin Simsek, and Alberto Bietti. Learning associative memories with gradient descent, 2024. URL https://arxiss org/ 2t/. 18724.
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 M.XXv:2402..94442 2024a.
Zheng-An Chen, Yuqing Li, Tao Luo, Zhangchen Zhou, and Zhi-Qin John Xu. Phase diagram of initial condensation for two-layer neural networks. CSIAM Transactions Applied Mathematics, 5(3)-448-514 2024b. ISSN 2708-0579. doi: https://doi.
org/10. 21t.22//esiaammmm-00002-22-2---0
URI. httpps///ggoooaaacccccccarrrrrra

---

Liam Collins, Advait Parulekar, Aryan Mokhtari, Sujay Sanghavi, and Sanjay Shakkottai. In-context learning with transformers: Softmax attention adapts function lipschitzness. In A. Glober- son, I.. Mackey, D. Belgrave, A. Fan, U. Paquet, J. Tomezak, and C. Zhang, editors, Advances in Neural Information Processing Systems, volume 37, pages 92638-92696, Curran Associates, Inc., 2024. URL attps://proceedd/ meerrii neurips. cctp.cccpppprr files/ htless/aaarr2022/E1122 665p5227000666666622222aa/eee/e2
pdf
Benjamin L.. Edelman, Ezra Edelman, Surbhi Goel, Eran Malach, and Nikolaos Tsilivis. The evo- lution of statistical induction heads: In-context learning markov chains, 2024. URL https:
//arxiv.o sarxiv.org// abs/2402. 11004.
Tianyu Guo, Wei Hu, Song Mei, Huan Wang. Caiming Xiong, Silvio Savarese, and Yu Bai. How do transformers learn in-context beyond simple functions? case study on learning with repre- sentations. In The Twelfth International Conference on Learning Representations, 2024. URL
https://openreviies http://www....ee.ooooo/d//////es
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer opti- mization through better initialization. In International Conference on Machine Learning, pages 4475 4483. PMLR, 2020.
Yu Huang, Yuan Cheng, and Yingbin Liang. In-context convergence of transformers, 2023. URL 05229
M Emrullah Ildiz, Yixiao Huang, Yingcong Li, Ankit Singh Rawat, and Samet Oymak. From self- attention to markov models: Unveiling the dynamics of generative transformers. arXiv preprint 133552 2024.
Arthur Jacot, Franck Gabriel, and Clément Hongler. Neural tangent kernel: Convergence and gener- alization in neural networks. Advances in neural information processing systems, 31, 2018.
Ziwei Ji and Matus Telgarsky. Gradient descent aligns the layers of deep linear networks. In In- lernational Conference on Learning Representations, 2019. URL https://openreview.ntt
Forum/id=HJj kg30gkX
Akshay Kumar and Jarvis Haupt. Early directional convergence in deep homogeneous neural net- works for small initializations. arXiv preprint arXiv:2403.08121, 2024.
Yuanzhi Li, Tengyu Ma, and Hongyang Zhang. Algorithmic regularization in over-parameterized matrix sensing and neural networks with quadratic activations. In Conference On Learning The- ory, pages 2-47, PMLR, 2018.
Yuchen Li, Yuanzhi Li, and Andrej Risteski. How do transformers learn topic structure: Towards a mechanistic understanding, 2023. URL https://aarxvv..g////////// 04244
Haoye Lu, Yongyi Mao, and Amiya Nayak. On the dynamics of training attention models. In In- ternational Conference on Learning Representations, 2021. URL https://openrrvvier net/
Tao Luo, Zhi-Qin John Xu, Zheng Ma, and Yaoyu Zhang. Phase diagram for two-layer relu neural networks at infinite-win -with limit. The Journal of Machine Learning Research, 22(1):3327-333 2021.
Kaifeng Lyu and Jian Li. Gradient descent maximizes the margin of homogeneous neural networks. In International Conference on Learning Representations, 2020. URL. https://openreviev. met//foru/iddSJeLLgg...
Arvind V. Mahankali, Tatsunori Hashimoto, and Tengyu Ma. One step of gradient descent is prov- ably the optimal in-context learner with one layer of linear self-attention. In The Twelfth In- ternational Conference on Learning Representations, 2024. URL https://openrevier net/ forum/1d=8ppfuuo
Song Mei, Andrea Montanari, and Phan-Minh Nguyen. A mean field view of the landscape of two- layer neural networks. Proceedings of the National Academy of Sciences, 1766 2018.
12

---

Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. Pointer sentinel mix- ture models. In International Conference on Learning Representations, 2017. URL https:
V/openreview.
Hancheng Min, Salma Tarmoun, Rene Vidal, and Enrique Mallada. On the explicit role of initializa- tion on the convergence and implicit bias of overparametrized linear networks. In Marina Meila and Tong Zhang. editors, Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pages 7760-7768. PMLR, 18-24 Jul 2021. URL
mmr press/v139/ http//vv19/mmmmmm/m html
Eshaan Nichani, Alex Damian, and Jason D. Lee. How transformers learn causal structure with gradient descent. In Forty-first International Conference on Machine Learning. 2024. URL
https://openreviest met//forumice JNK4mmmm1
Catherine Olsson, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, Amanda Askell, Yuntao Bai, Anna Chen, Tom Conerly, Dawn Drain, Deep Ganguli, Zac Hatfield-Dodds, Danny Hernandez, Scott Johnston, Andy Jones, Jackson Kernion, Liane Lovitt, Kamal Ndousse, Dario Amodei, Tom Brown, Jack Clark, Jared Kaplan, Sam McCandlish, and Chris Olah. In-context learning and induction heads, 2022. URL sttps://arx// org/
abs abs/2209. 11895
Gautam Reddy. The mechanistic basis of data dependence and abrupt learning in an in-context classification task. In The Twelfth International Conference on Learning Representations, 2024.
URL https::/epenreeviok wtw.ppparrrrrei....oeooorrraaaaa ott.://
Grant Rotskoff and Eric Gunden-Eijnden Parameters as interacting particles: long time convergence and asymptotic error scaling of neural networks. Advances in neural information processing systems, 31, 2018
Andrew M Saxe, James L McClelland, and Surya Ganguli. Exact solutions the nonlinear dynam- ics of learning in deep linear neural networks. arXiv preprint arXiv:1312.6120 2013.
Charles Burton Snell, Ruiqi Zhong, Dan Klein, and Jacob Steinhardt. Approximating how single head attention learns. ArXiv, abs/2103.07601, 2021. URL httpss//ppp..emmtt.ccch///
org/CorpusID 232232786
Mahdi Soltanolkotabi, Dominik Stöger, and Changzhi Xie. Implicit balancing and regularization: Generalization and convergence guarantees for overparameterized asymmetric matrix sensing. In The Thirty Sixth Annual Conference on Learning Theory, pages 5140-5142. PMLR, 2023.
Dominik Stöger and Mahdi Soltanolkotabi. Small random initialization is akin to spectral learning: Optimization and generalization guarantees for overparametterized low-rank matrix reconstruc- tion. Advances in Neural Information Processing Systems, 01--23331-28443, 2021.
Yuandong Tian, Yiping Wang, Beidi Chen, and Simon Du. Scan and snap: Understanding training dynamics and token composition in l-layer transformer. Advances in neural information process- ing systems, 36:71911-71947 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 The Twelfth In- ternational Conference on Learning Representations, 2024. URL attps://openrews
net/
Iorum 1Or=m/Id=LDJJKGNYYC
Aditya Vardhan Varre, Maria-Luiza Vladarean, Loucas Vivien, and Nicolas Flammarion. On the spectral bias of two-layer linear networks. In Thirty-seventh Conference on Neural Informa- fion Processing Systems, 2023. URL https://openreviee
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
Francis Williams, Matthew Trager, Daniele Panozzo, Claudio Silva, Denis Zorin, and Joan Bruna. Gradient dynamics of shallow univariate relu networks. Advances in neural information process- ing systems, 32, 2019.
13

---

Zhi- Qin John Xu, Yaoyu Zhang, and Zhangehen Zhou, An overview of condensation phenomenon in deep learning. arXiv preprint (54::22040004444, 2025.
Junjie Yao, Zhongwang Zhang, and Zhi-Qin John Xu. An analysis for reasoning bias language models with small initialization, 2025. URL https://arxiv wwwpxii..rgggba///.oo
org/ wwwp/b/////....o..////////// 04375.
Yichun Yin, Wenyong Huang. Kaikai Song, Yehui Tang, Xueyu Wu, Wei Guo, Peng Guo, Yaoyuan Wang, Xiaojun Meng, Yasheng Wang, Dong Li, Can Chen, Dandan Tu, Yin Li, Fisher Yu, Ruiming Tang, Yunhe Wang, Baojun Wang. Bin Wang, Bo Wang. Boxiao Liu, Changzheng Zhang, Duyu Tang, Fei Mi, Hui Jin, Jiansheng Wei, Jiarui Qin, Jinpeng Li, Jun Zhao, Liqun Deng, Lin Li, Minghui Xu, Naifu Zhang, Nianzu Zheng, Qiang Li, Rongju Ruan, Shengjun Cheng. Tianyu Guo, Wei He, Wei Li, Weiwen Liu, Wulong Liu, Xinyi Dai, Yonghan Dong, Yu Pan, Yue Li, Yufei Wang, Yujun Li, Yunsheng Ni, Zhe Liu, Zhenhe Zhang, and Zhicheng Liu. Pangu ultra: Pushing the limits of dense large language models on ascend NPUs, 2025. URL https: //arxv. org/ abs/ 07866
Biao Zhang, Ivan Titov, and Rico Sennrich. Improving deep transformer with depth-scaled initial- ization and merged attention. In Kentaro Inui, Jing Jiang. Vincent Ng. and Xiaojun Wan, edi- tors, Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNIPP) pages 898-909, Hong Kong, China, November 2019. Association for Computational Linguistics. doi: 10.18653/v 1D19-1088 URL
Ruiqi Zhang. Spencer Frei, and Peter L Bartlett. Trained transformers learn linear models in-context. Journal of Machine Learning Research, 25(49): 1-55, 2024a.
Yaoyu Zhang, Zhi-Qin John Xu, Tao Luo, and Zheng Ma. A type of generalization error induced by initialization in deep neural networks. In Mathematical and Scientific Machine Learning, pages 144-164. PMLR, 2020.
Yedi Zhang, Aaditya K. Singh, Peter E. Latham, and Andrew Saxe. Training dynamics of in-context learning in linear attention, 2025a. URL
16265.
Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang, Yaoyu Zhang, and Zhi-Qin John Xu. Initialization is critical to whether transformers fit composite functions by inference or memorizing, 2024b. URL https::/t//attt //tpvv/aaxvvvv/v/ abgg/bbs/4405 05409.
Zhongwang Zhang. Zhiwei Wang, Junjie Yao, Zhangchen Zhou, Xiaolong Li, Zhi-Qin John Xu, et al. Anchor function: type of benchmark functions for studying language models. arXiv preprint 24Xiv:2401. 08309,, 2024c.
Zhongwang Zhang, Pengxiao Lin, Zhiwei Wang, Yaoyu Zhang, and Zhi-Qin John Xu. Complexity control facilitates reasoning- based compositional generalization in transformers. arXiv preprint arXiv:2501.08537 2025b.
Hanxu Zhou, Zhou Qixuan, Tao I Luo, Yaoyu Zhang, and Zhi-Qin Xu. Towards understanding the condensation of neural networks at initial training. Advances in Neural Information Processing Systems, 35:2184-2196, 2022.
Zhangchen Zhou, Hanxu Zhou, Yuqing Li, and Zhi-Qin John Xu. Understanding the initial conden- sation of convolutional neural networks. arXiv preprint Xivv:2005.00947 2023.
Chen Zhu, Renkun Ni, Zheng Xu, Kezhi Kong, W Ronny Huang, and Tom Goldstein. Gradinit: Learning initialize neural networks for stable and efficient training. Advances in Neural Infor- mation Processing Systems, 34:16410-16422, 2011.
14

---

NeurIPS Paper Checklist
The checklist designed to encourage best practices for responsible machine learning research, ad- dressing issues of reproducibility, transparency. research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.
Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:
You should answer [Yes], [No] or [NA].
[NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.
Please provide short (12 sentence) justification right after your answer (even for NA).
The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.
The reviewers of your paper will be asked to use the checklist as one of the factors in their evalu- is perfectly acceptable answer "[No]
ation. While "[Yes] " generally preferable to "[No] provided a proper justification is given (e.g.. "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] or "[NA] is not grounds for rejection. While the questions are phrased binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to question, in the justification please point to the section(s) where related material for the question can be found.
IMPORTANT, please:
Delete this instruction block, but keep the section heading "NeurIPS Paper Check- list".
Keep the checklist subsection headings, questions/answers and guidelines below.
Do not modify the questions and only use the provided macros for your answers.
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We elaborate setups and contribution in the abstract and introduction, especially in the last paragraph of the introduction.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
fine to include aspirational goals motivation as long it is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations of the work performed by the authors? Answer: [Yes]
15

---

Justification:
The limitations could be found at Sec 5
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The au- thors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should articulated.
The authors should reflect on the factors that influence the performance of the ap- proach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-tex system might not be used reliably to provide closed captions for online lectures because fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach ad- dress problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification:
We provide the definitions, assumptions and proofs at Sec. and Appendix
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theo- rems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be comple- mented by formal proof' provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclu- sions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: We show the experiment setup in Sec.
Guidelines:
16

---

The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make it possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate the research performed.
While NeurIPS does not require releasing code, the conference does require all sub- missions to provide some reasonable avenue for reproducibilityy which may depend on the nature of the contribution. For example
(a) the contribution primarily new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be a way to access this model for reproducing the results or way to re- produce the model (e.g., with an open-souree dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case au- thors welcome deseribe the partieular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification:
The code is provided in the supplementary materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so No is acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:)
/ /maps..c/public/guudee Sooeedeuee//oSSSummiioa issiisssioniii Pnolicy) for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
17

---

Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g.. data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
We show the experiment setup in Sec.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined other appropri- ate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We show the error bar in Fig, for the anchor function experiments and Fig. B for WikiText task.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the eтror bar is the standard deviation or the standard error of the mean.
It OK to report I-sigma error bars, but one should state it. The authors should prefer- ably report 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide the information of compute resources in Appendix. C
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
18

---

The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification:
The research conducted in the paper conform with the NeurIPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g.. if there special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper diseuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This work is phenomenological study, therefore, there is socictal impact of the work performed.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses
(e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact spe- cific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitiga- tion strategies (e.g., gated release of models, providing defenses addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system lears from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (c.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: Justification: This work is a phenomenological study, therefore, this work poses such risks.
19

---

Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by re- quiring that users adhere to usage guidelines restrictions access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: All the assets mentioned in paper is open-sourced and properly cited.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the pack- age should be provided. For popular datasets, paperswithcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the li- cense of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documenta- tion provided alongside the assets?
Answer:
[NA]
Justification:
The paper does not release new assets.
Guidelines:
The answer NA means that the paper does release new assets.
Researchers should communicate the details the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the pa- per include the full text of instructions given to participants and screenshots, if applicable, well as details about compensation (if any)?
20

---

Answer: [NA]
Justification:
The paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contri- bution of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, tion, or other labor should be paid at least the minimum wage the country of the data collector.
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: These type of studys are not involved this paper.
Guidelines:
The answer NA
that the paper does involve crowdsoureing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equiva- lent) may be required for any human subjects research. you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, originality of the research, declaration is not required.
Question:
Answer: [Yes]
Justification: LLMs are not used for any core methods in this paper. Therefore, we are not required to make such a declaration.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non standard components.
Please refer to our LLM policy http:///eoaurrrooeo/o/o///eerees for what should or should not be described.
21

---

A Theory Details
A.1
Theory Details for Blow up Dynamics
A.1.1 Proof for Proposition
Proof. Taking advantage of the inherent symmetry of the system, the proof focuses on analyzing the coupled dynamics of W and w
d ( ww Σ T
Wo,
W RK
W
2 d di
Σ
W
[1]
Kike
k'
k'
This finished the proof of the first two equations. We also derive the relation between energy and the evolution of parameters
W
W v.k
2
W b.s
W W
2E
dt
RK
k,k'
This is just the third equation.
A.1.2 Proof for Theorem 1
Proof. Since use Gaussian random initialization, the initialization satisfies Definition ~ almost surely. Therefore, we establish our results under the assumptions specified in Definition E By local Lipshcitz condition the right hand side of dynamical system Eq. (α)) it has solution for t E (0, 7*)) where 7 is maximum existence time of solution and can be infinity. Taking derivative of E, we obtain
d W.wii w
E dt
W.12
| 1 + W. W
(16)
The inequality of arithmetic and geometric means leads to
É > 3(1 W. 1
13 w 12 w 13)}
2 (.W.
3E
This implies that energy F increase monotonically. If F.(0) O,
dt
Integrating both sides of the inequality yields lower bound for the energy E
E(I) 2 (E(0)-3 t
(17)
Thus, in the case where E(0) > the dynamical system explodes before T E(00
In the case where E(O)= S O, we consider - EE) instead, and obtain
E(1) 2 + (-E(0))) 3ja"
(18)
We claim that there exists some time to > 0, such that E(to) > 0. This claim can be proved by contradiction.
Suppose that F(() S for all 0 < 1 < T Recall that É(()) 2 0 throughout this interval. The boundedness of W.. Will, and 2 together with the monotonicity of energy E, implies that E(T") lim,-7- E(t) exists and satisfies -oo E(T") < 0. We now consider different cases separately.
(i) The case of T < +oo. The solutions can be extended to time larger than T since || |.,l12 |w(ill | are bounded due to the conservation law. This contradicts the definition of T
22

---

(ii) The case of T = +o0 and E(T*) < 0. That is lim,+++00 E(t) < 0. However, this contradicts Eq. (TS).
(iii) The case of T = +o0 and F(T*) = 0. That is lim (1) = 0. We prove this case in three steps.
Step I: We show that lim W. (t) (0) Since E(t) S 0 for all t, the quantities IW lž, || || are monotonically decreasing. However, each of them is bounded below by zero, and hence they all converge to finite limits and remain uniformly bounded.
Moreover, note that E > 0 for all t. If lim inf E(t) lim;->c E(t) = 0. Therefore, it must hold that
0, it contradicts the fact that
lim W J.W.O1B W 2 (1)) 0. +00
This implies that either lim,- [. (t) = or lim [ 12 (t)1 = We can also obtain
lim (t) = since Wlll W.Wa
Step 2: We show lim E(t)
Without loss generality, we assume that W.(t)lB = is similar. That is [|W(())l2 > (t)(t]l 0.
lim; || (t)l = The case lim;
2- By conservation law, we have lim,
Considering the second derivative of W kr we obtain
W 12 A
Σ
I1 W
2
w
[1] kk
+ W o w
Thus lim (t) = O since lim (t) = 0. Note that E(t) 2 0 and
lim E()) = 0. Recall that lim inf, E(I) = 0 and É are bounded. And we also have
1 < M. We claim that lim || (2. (1)11 = O, which implies
lim E(t) 0.
(19)
Since lim [WP = O and lim [+) 12 = 0. Using Taylor expansion, we have for some 4 E [t, 11
w & (t 1) w wPP() w # (t) W 2 (w), Vk.
which implies lim W. (t)ll2 Therefore, the assertion holds.
Step 3: We show that Eq. obtain
(1)) contradicts the condition in Definition EL By direct calculation, we
d
d
W.
1
(
W w wi))
dt
dt
w w 22] + w W (4 2
=
w w W w 3 w
21 ! 11]T w [1]
w [2]
W. w. w w
w
w
W. will w
+ 2 [1]
[1] W 22] w + 2 [I] (1 (1 W!
2E w 2 w
and
d
11
dt Ww. w " 11' w!)
dt
-W.w!" W [1]T W. + W. w w'l'!
W,W W will' WT - W,, w w!
= =WI2I w' 111 w WI' WJ + W, W W.W.W WI w W!
+W. W. w W. WJ W, wi w wia
2EE W. 2W, P1
23

---

Therefore,
d
W
11
W 2 [
(0)1 (0]3))
dt Integrating both sides the equality, we obtain
dt
lim -1
[ W
(0)
0| (0) (0) (w.))))
1-00
However, according to Definition @ and the fact that lim
| (t)l 0. we have that
lim [[ (t) 0.
Based on Eq. (IG)) we have
lim E(t) = lim
(t)(1 + | 2 ( + (
t100
t-000
lim
W, (t) +
1-000
contradicts with Eq.
(T9) which claims limy-+00 Ett = 0. This completes the proof.
A.2 Theory Details for Condensation
In this section, we prove the main theorems which characterize the condensation. In retrospect of the proof of Theorem M. Eq. (IA) provides lower bound that leads the presence of explosion. However, this inequality leaves the precise growth rate of energy E undetermined. The key idea here is that Assumption can give us an upper limit on how fast the energy can grow. Once we understand this growth rate, we can then move forward with proving the main theorems.
We begin our proof by the following proposition.
Proposition 4 (induction). Consider dynamical system Eq. (Z)) If Assumption D holds at some time to with < T", then Assumption will hold & (ta, T")
Proof. First, we consider the second condition in Assumption II By direct calculation, we have d dt
(ww/²)² + (W!el)
w 2 witls wfFlwitt)) (w) W, w(lls)
& W W
W (1 will (yl.
(m//p))
W. wlltt) + (w.
+
w
2 w 2 will
w
(w!2))
and
d di (Wo. wi We!l = (Wo, w/!l w.ll//leee wlll
( +
We, W. v, w!!i w/Ih
W2 v.j
(Wo, w/ll
W2
W w.. Wo., w w/lu))
W i,i
By Assumption MI we know the above equations are larger than O. So (w. wile w 2 will) and ( (W., W w/ll)) will be monotonically increasing since to Calculating the derivative of left hand side of first condition, we have
d w 22 W. W.Wll
W 0,
W.. v W.,
willd + (w!!))² I
dt
and
d W. v. 1111
W/
w2 o. +
dm Σ
W. v W w. w("" + (W..,) ||
di
j=1
24

---

Hence, w W. Wil and W u W;" W will also increase monotonically since lo. Therefore the condensation condition will hold until T".
Next, analyze the angle relation between W and its derivative W For the simplicity of proof and description, we adopt standardized notation to represent angles between distinct vectors throughout the ensuing discussion.
Definition 5. Let (t) denote the angle between the vectors W. ( (1) and W We.g (t. (t), and (t(t) denote the angle between the vectors (t and W ( (t). Let (t(() denote the angle between (t) and We,a(t) (i (t), while ((t) denote the angle between ( and W (t). In subsequent expressions, the variable will be omitted unless there is specific emphasis on the temporal change of angles.
We divide the entries vector Wi into two classes according to whether their limit is finite.
Proposition Suppose that Assumption holds. Consider the effective dynamics Eq. (.)) The indices [dm] can be partitioned into two disjoint classes, denoted by C1 = (ir.... 4 # Ø and C2 [dm]] \ C1. The partition satisfies the following properties:
(i) For each [ml, the limits of W v, exist. In particular,
00,
i C1,
lim Wo,
(20)
W 0,6
E
(t), (ii) The angle Eig between the vectors W. (t) (t) and W (t) (t), as defined in Defini- tion B fulfills the condition:
lim cos Exj 1, for C1.
(21)
(iii) The following limits exist
WW W
| I
lim TT
lim TT
1.
(22)
Proof. 1. First, we find that for every index the (W v. increases monotonically. So their limits exist. We define the index set of the parameters that tend to infinity as C1 and the others as C2. Based on Theorem M and conservation laws, we know that Ci * Ø. Property 1 is automatically satisfied due to our partition.
2. We introduce new variables
p (W,, w!!! W., w''l),
9 W 0. v.j
W
According to the proof of Proposition H1 we find that
dp
(w., w wl)) W +
1 W., n W., w! w!liT)
dt
W
+ ( w w(ll)
W2,
W u.j o. w." w
(23)
W2
W w/!h w w/W + W. W.,
w w(ll)) (1+2)
Thanks to Eq. (.), we obtain
dq (W.,) w!'llw + W., w/! W.,
(24)
25

---

Combining Eq. (2:3) and Eq. (224). we obtain
dp
1 p
(25)
dq
2 q
2
Let p/g. Note that dp d we get
dg
(ug)
du dg
+ Combining this with the right hand of Eq. (25)
dg
2du
(26)
q
The Eq. (26) can be solved explicitly.
In [t)] - In |g(to)] = - 2 2n (u(t) - 1| + 2 (u(to) - 1/.
(27)
For i,j E C₁, we have lim ul) 1 since tends to infinite as / tends to T
By definition,
Ww w Wo w IW: 1 | 4 ll2 W W.sl
cos Eij-
(28)
9
W 0,
W v.j
Using the conservation laws, have
|w - |2 W (0).
For i E C1, we have
w [1] 12
lim
1.
(29)
W2 0,
Combining Equations (BX) and (29)), we get
lim cos Eij = 1, i,j Ec. 8-7
(30)
This finishes the proof statement (ii).
3. Finally, calculate the norm
W22l [22. By definition, obtain
w
W o. w W., W
We divide the into three parts due to the boundedness of entries of W.. 9-
[1]
Σ Σ
W n.
W
Ei +2 Σ Σ
We, W. IW; 2 1all ///l 11 cos Eij
W
W
iEC, JEC1
#EC1 JEC2
Σ
W o. W
w
||2 cos
iECs JECs
[|: ll2 liw/llla
cos cos -1 -1)
1)
Σ
Σ
W W
Σ
W W2
Wo,l
[Wal
iEC, iEC1
AECIJEC
Σ iEC, JEC,
W. W W
| la W
cos Eij + Σ Σ
W., W., W cos Eij-
IEC, JEC,
| W...
|w!l Wo
Since lim
cos Eij = 1, we have
(w W/(2)
lim t
(Eec W2) = 1.
(31)
Based on statement (i), we obtain
W2 e. Σieo
lim *
1.
(32)
W, 1
26

---

Combining Equations (.I)) (12) and conservation law, we have
(3
[2]
2l = 1.
lim
lim 6
(33)
W 1
This finishes the proof of statement (ii)..
Proposition E describes the angle between W., w!ll and w.,, W.J. Since W a linear combi- nation of W., W!!l,, we immediately have following corollary.
Corollary 1. Suppose that Assumption 0 holds. Consider the effective dynamics Eq. ( and recall index class defined Proposition B
The angle V/s between the vectors W and W., W 11 as
defined in Definition B satisfies:
lim cos v/ =1, i t-7
(34)
Proof. By definition,
W
W. w
cos م/ا
W e.
W ||
Recall the definition of Eij, the above equation can be reformulated as follows:
-==1 W
m
W
[1]
ll2 cos Eij
cos 1/4
W.. W W//12
m
[ W cos Enj
|W[2] [2]
Σ jeC, W
w
cos Eiz + Ejec, (Wo.)) |
122 cos Eij
W
| 1|2
According to Equations and we have lim,-2 cos the = 1. This completes the proof.
So far, we have characterized some properties of W. w!') [1] which is component of W We have shown that some of them will have the same direction when é tends to T However, it is not enough for our seek for upper bound for energy E. Luckily, based Corollary M. we can analyze the angle between W and W., W!! which provides an upper bound. Before this, we give the following proposition. The subsequent proposition demonstrates an extension of statement 1 of Assumption MI going beyond the condition of W.. W being greater than zero to include
additional anglc-related information.
Proposition 6. Suppose that Assumption holds. Consider the effective dynamics Eq. (1 and recall index class defined in Proposition & There exists constants Ti E (la, T*) and O E [0, 4) such that for each index E C1, the follow inequality holds:
cos Pi 2 cos 11, EE (T,
Proof. It is sufficient prove the statement for any fixed E Ci due to the finiteness of [C1]. In Proposition El we have shown that (W., will| w(l) > 0 for t E (to, which implies
cos p > O, E (to, T.). Hence we can focus on its square, i.e., cos2 (w Ulya tion of cos2 44 is
(35)
By direct calculation, the deriva-
2
(W W w
[2] w. + (
1 w 1
1 W
W
2
2 [2] w (1
2
(
+
27

---

We can rewrite the numerator
2 W 11 wil w!ll)
W W
11)
1 (wi2 will) (wiz), wlll)]
+ 2| W wP W., w!!l)
W 2
" ||2 (wl2, w w/!1)3]
The second term of above expression obviously greater than zero by inequality of arithmetic and geometric means. Also we can rewrite the first term
W (wll), W.,, w (wiel W.,, w wlll (wlz), W.,, w!") (W wlll)
W 0,2
We find that the first two factors W W Wo, (.!' of above expression are positive. Ac- cording Definition E. the difference term can be reformulated
(1 12 W D, W U1) 12 (w W.. W.. W. (WW w
W W | W.. n. W. 11 2 s/s - cos cos yp:)
Note that lim,-7 CO8 t/:
1 for i Cj. So for every 0, there exists d such that
1 - = < cos , - 1, E (T - 6, )")
S - and 0, arccos(1 - <). Then have either cos yo 2 cos t E (f,, I*), or there exists E T*) such that cos sp; I cos fg. then it will increase monotonically until it goes up to 0. No matter in which case, we can find 0, E [0, 2) such that cos pi 2 cos 0g. Let Zi maxiec, & and θ maxiec, 0,. Thus, for each i C1. the following inequality holds
cos pi 2 cos 11 E (T,
This completes the proof.
In fact, Proposition & provides the angular relationship between W and its derivative W We summarize it as follows.
Proposition Suppose that Assumption Z holds. Consider the effective dynamics Eq. (Z) There exists T2 E (Ti, T*) and O2 E [0, 7 such that
cos Ç 2 cos O2, LE (T2,T*).
(36)
Moreover, recall the definition of energy E, the following inequality holds
= (w w 2 cos co W ||] || 1|,, LE T").
(37)
Proof. By definition of Ç, we obtain
cos Ç (W W (2))
EM W w!!! w(l)
W W
W 1
According to Proposition we have
Σieo,
W w lla
cos ( 2 cos O
tE (TiTT".
||
Since we have shown that lim Σiec, W., w.w lla
according to the proof of Proposition B
there exists T2 E (Ti, T and O₂ [0, 5 such that
cos 2 cos Oz, LE (T2,7"))
Recall the definition of energy F, we have
E (w wlll > (cos O2) W This completes the proof.
W | LE (T2, T").
28

---

Now we prove the proposition with all preparations above.
Proposition 8 (energy upper bound and blow up estimate). Suppose that Assumption M holds. Con- sider the effective dynamics Eq. (Z)) There exist T3 (T2, T and C 2 1 such that the following upper bound of Energy E holds
E(t) S
TEStteT.
(38)
(E(s)- (E( C(t C(t-
Moreover, the blow up time T is bounded below by
2 E(t) TT
(39)
Proof. First, by calculating the derivative of energy E, we have
E W +
w 2
W
We rewrite the derivative of energy E as
W 11
W 12
w
E = |
+
|13
|| (21
12
W 11
W W
|wel
1
W (2)
12
W 12
According to conservation laws and Equation (200 there exists Ta T2 such that
W
2
W 11
4.
W [2]
W
22
Then we have
E
w
(3]
11 T3.
Based Proposition D. we have
|
21
121 1/
1|2 S
cos O2 E, Vt > T3.
Thus obtain
< 4(
E.
Vt > T3.
cos 2
We denote 4( (o0s,)) 3 as Co. Then we have Co 2 4 and
E F
Opposite to proof of Theorem II we obtain
F $ Co. V> Tz.
Co C. Thus C > I and we have We denote d
E(t) < C(t- 3:
<<<<t<rr
(40)
Hence, for each time < T", the time of blow up is bounded below by
T > t E(t)
(41)
C
This completes the proof.
29

---

Now we begin the proof for Theorem
Proof. just prove the case for W., v.i > And the case for W., < 0 follows by similar argument.
Because the derivative of W is Wiil W we only need to show that
W dt +00.
(42)
I 2 [W1-
Since lim
due to the statement (iii) of Proposition B there exists T4 T3 such
that [ - 2 w 12 which implies
E = W 2 w(2) & || |2]l 11 - 2 wl
(43)
The idea is we can find a infinite division of (T,TT*) such that the integral of W w on each sub-interval larger than positive constant which is an independent constant. Then we consider the integral Sr W dield. By direct calculation of the derivative of W22), we have
(W 1 Wie)= W.. W 21
w + Wi!!!
Σ W., W.Jl
Integrating both sides of the equality, we have
ww!"/w(')>( = wll))) + I" W.., WR. w w!"l
Σ Wo. w/") w!""
ds
2 (T,)(Ta) / w wlidss
Note that Eq.
( implies
1 / E3 (s)ds
W W ds 2
2 |
2 2 Jn (E((t)
ds
b2))²
1
1
1
2 E(t,) -(t - t1) E(t,) 3
Thus the integral satisfies
pt
wi 2 Wo,;(Ta)
=
E(t,) dl
2 11
E(t,) (t (t-
12 4
(Ta)(Ta))
In(E(tt)) -(t t2)) In((((t))
E(t,)
= 20
According to Theorem choose to - t1
Thus obtain
J" w willt (TA)
1
In
(44)
2C
2C
Then we introduce auxiliary function
1
f(t) In
- 1-t
In(1
- L.
We have f(0) and f(t) > Then we have 1-1 20 such sub-intervals, the proof of the theorem is completed.
Since there are infinitely many
30

---

A.3 Theory details for Key-Query Dynamics
This appendix provides the detailed derivations for Proposition & and Theorem assuming linear activation function and holding to Assumption @ (or its empirical variant, Assumption 25).. Our approach begins with a standard asymptotic analysis decompose the loss function, as shown in Eq. (TE).
Starting from the definition of the empirical risk, we have:
(())
X W.W.
X., Wv W w wrel ))
exp
vdm
exp
X.,, W...W
1=1 exp
vdl
(45)
Σ exp
1
Xi, WOW
1/1
+
12
Vmm
5=1
1
X W WOWTXY
+ +0(8%) x x.s wwwll )
The third equality results from applying a Taylor expansion to the softmax function, which is justi- fied by Assumption D. The higher-order term, O ), is subsequently omitted as it does not affect the reading-order training dynamics. Given the property of the exponential loss function, E(g) c e the empirical loss can be decomposed as follows:
X X W. w(wll )
22
((0)
exp
-//
Σ j=1
i=1
x WoW
X WOW ddm
x.,, Wr Wi.Wwill
exp
V/
s
dmm
1=1
(1)(0)2,,(0),,
1=1
(46)
where (0) and (0) correspond to the two exponential factors in the preceding expression.
A.3.1 Proof for Proposition B under Assumption P
Proof. The proof is structured in two parts. First, demonstrate the separation of dynamics by showing that the gradients with respect to different sets of weights have different orders magni- tude. Second, we derive the specific dynamics for the key and query matrices.
Dynamics Separation: By symmetry, we will detail the calculations for the partial derivatives with respect to Wy and Wq. The derivations for the other weight matrices follow similar procedure.
We begin by computing the partial derivative of the loss with respect to Wy. Applying the product rule to the decomposed loss from Eq. (HG), we obtain:
OL
Ø OWv
1
Σ i=1
(.)
dWv
@LI, OWv
aL2,
C2,t + L1.8
(47)
OWv
(1+ + Li, ..x))
7
OWy
31

---

The third equality holds based on Assumption A which implies that 22, = 1 + 0 and its derivative 0Lgı is also of order (6))). Furthermore, Assumption states that the leading-order term of the loss, 1,, independent of the attention mechanism weights at initialization. Consequently, the term bwy of order (8")).
aWv
BC is dominated by terms 0C is of order O This implies that the entire gradient bWv
Next, we consider the partial derivative with respect to Wq:
ac
n
Σ
(1,
aWq
aWQ
(48)
1
8L2. dWa
(0)
n
since Li is independent of We. Based Assumption D. the term 0C2s is of order O()), which DW
establishes that the overall gradient 8C
is also of order (6)).
aWa
Key-Query Dynamics: The principle of dynamics separation, established above, shows that the gra- dients with respect to Wº and Wx (order (6)) are significantly larger than those for Wv. Will, and W (order O (()2)). Therefore, during the initial phase of training, the dynamics are dominated by the updates to We and WK- We can thus analyze their leading-order dynamics by treating the other weight matrices effectively constant.
To derive these dynamics, we employ matrix calculus with differentials. For a sealar function f(X) of a matrix variable X, the differential given by df (( 8x
dx). We apply this to the argument of the exponential in L2,5, which denote as (0). The differential of A; with respect to Wq is:
x., Wy ))
X W
WOWIX bj
X
W WOW
d
V/
dy
82
ddm
1=1
X (dW
XT
XT
X
Wow
sydm
1=1
tr
Us
WK
bj
X
Σ
X
Wy Pl
(49)
By identifying the coefficient of dW2 from the trace form, we obtain the gradient. Consequently, after neglecting higher-order terms, the leading-order dynamics for We under the gradient flow dWa 8C are given by: dt oW
] WK.
dW di
Σ VrCiii (Wi)
Σ
X1,
Xij
1 Σ X
ns dm
(50)
A symmetric argument yields the corresponding dynamics for W K. This completes the proof.
D
32

---

A.4
Proof for Proposition 3 under Assumption 25
We use another scheme estimate the gradient. Consider Proof. Dynamies Separation:
Oc
1 (1,(0)
dWv
dwv
n
2 + L1,t 2C2.
dwv
oWv
(51)
1
Σ
OLI.
(1+0 + L1.5 +1.5
Owv
ULi,i 0
owv
ac Take the inner product of both sides of the equation with aw we have
OC
Oc
0L1, Oc
+0(8)).
(52)
OWv
OWv
OWv
dWv
However, note
that
d
1Li, OL
n
0Li,
Oc
0L1,; Oc
C
dt
OW dwv
aw
ww OW
(53)
Based on Assumption PA one.
we have ac awv
(6²)). The rest of the proof is similar the previous
Bey-Query Dynamics:
This part is almost the same. Just need to note that the condition ~ gives the same F after normalization.
d
W.
w
wizl WFT
de
W
de
A.5 Proof for Theorem B
Proof. Based on the leading-order dynamics key-query matrices, we prove the theorem for Wo and the technique for Wx is similar. Differentiating the dynamics again, we obtain:
d²
Wq FF
(54)
dt2
Let the singular value decomposition of F be F UEVT, then the dynamics can be rewritten
W Wo = UE'UWW
(55)
de
Let We U'U we have
d²
We E'
(56)
dr
The evolutions of entries are
(o,ú(() = C1,ig gx + +C2ig e Ast
(57)
As result,
We
rank
lim (+0
k
(58)
WQllF
where k is the multiplicity of the largest singular value. This result finishes the proof.
33

---

B Experimental Details
In this section, we present more experimental details to supplement the main text.
B.1 Experimental Setting of Synthetic Dataset
We introduce the dataset construction method and training hyperparameters used to train the syn- thetic dataset.
First, give some calculation methods for experimental pictures.
Satisfaction rate. We denote the conditions in Assumption M follows:
A = [dm]
w! W. Wil > O, W., w 22 > 0)=
A₂ = (i,5) [dm]
[dm]
(w P2 Ww!ll Uil
w w > 0, (W W.w.will wills, W w > 0}
Cosine similarity. To visualize the internal structure of a weight matrix W, we generate a heatmap of its reordered row-wise cosine similarity matrix. The procedure as follows: first, the row-wise cosine similarity matrix S is computed, where each entry Sty = cos W,) measures the similarity between row vectors w; and w
To reveal underlying block structures, we then employ a spectral reordering technique. This in- volves finding the principal eigenvector Mmax (the corresponding to the largest eigenvalue) of the similarity matrix S. The sorted order of this eigenvector's 's components, P argsort( Vinax). provides permutation index. By applying this permutation to both the and columns of S, group highly correlated row vectors together, making low-rank patterns visually apparent in the final heatmap.
Dataset Construction. We construct synthetic datasets using the concept of the anchor func- tion Zhang al, [20224]] which enables controlled simulation of linguistic relationships. Let the set of prompt anchors be A {a E N+ amin S a S Qmax and the set of keys be Z {2 N+
Smin 2 Smax), ). where A and Z are disjoint, i.c., An A Z
We define an anchor function F(X): : N
where X (I1, T2,..., ........,))..................).. ) a sequence of length Each sequence contains exactly one anchor token E A among the first 1 positions, and the function outputs the token immediately following the anchor, shifted by a:
xxs)
F(....... = I++1 +a, where 2, = a.
(59)
In our experiments, we set A {1, 2, 3). (5..... 100). and 10. To introduce synonymy among anchors, we modify the mapping
F(II,...,..) I++1 (a mod 2), where x; a,
(60)
so that anchors {1, 2) and (3, 4} produce equivalent outputs, mimicking synonymous relationships observed in natural language.
Model and training hyperparameters. Our model is decoder-only Transformer with a single layer and single attention head. The architecture follows the standard GPT design, consisting of a multi-head self-attention block and a position-wise feed-forward network. The Tanh activation function is used the feed-forward network. A key aspect our experimental setup that both the token embedding layer and the final output projection layer are fixed and not updated during training. This allows us isolate the learning dynamics exclusively within the Transformer's attention and feed-forward weights.
All trainable weights in the model are initialized from a normal distribution with mean of 0. The standard deviation for different components is set based on the model dimension dmodel as a 0.855 The loss is computed only on the prediction of the last token the sequence.
The model was trained for 30 epochs using the Adam optimizer. We employed a learning rate scheduler that combines a gradual warmup phase for the first epochs followed by a cosine anneal- ing schedule. The specific hyperparameters are detailed in Table M.
34

---

Table Model and Training Hyperparameters
Parameter
Value
Parameter
Value
Model Architecture
Training Settings
Vocabulary Size
201
Optimizer
AdamW
Model Dimension (dmode)
640
Batch Size
1000
Dim. (dn)
1280
Epochs
30
Key/Value Dim. (dz, de)
640
Weight Decay
0.0
Number of Layers
Gradient Clipping
1.0
Number of Heads
AdamW 81,82
0.9, 0.999
Activation Function
Tanh
Learning Rate Scheduler (Warmup CosineAnnealing))
Initial LR Immi
10-3
Warmup Epochs
Warmup Multiplier 15.0
Cosine Annealing Tmax
200
B.2 Synthetic Dataset without activation function
In this section, show similar result for model without activation function, as supplement to the synthetic data experiments. This shows that is reasonable to ignore activation our analysis
(W12)
5:25
Stape
Siage 2
20
Further Training
1M (Wis
4.50
6.55
0.5
4.00
0.0
6000
LO000 12000
4000
Stop
Step
Stop / 100
Figure 4: (a) Evolution of cosine similarity matrices for outer and attention parameters. The training process is partitioned into three stages: Condensation (Stage 1), Key-Query rank collapse (Stage 2), and further training stage. Stage transitions are identified by plateaus in the loss curve and structural shifts in these matrices. (b) The relative change of norms between attention and outer parameters. The gray dashed line marks the onset of Stage 2, where updates to the attention param- eters begin dominate. (c) Evolution of the effective rank for both parameter groups, tracking the change in their intrinsic dimensionality throughout training.
Wo
(b)
(c)
1.0
1.0
15.0
Satisfied Propertion
Direction similarity
12.5
0.8
Average 1
10.0
0.6
7.5
0.4
5.0
Condition Condition
0.2
2.5
0.5
0.0
0.0
200
400
10
20
30
40
1000
2000
3000
4000
Step 100
Figure 5: (a) Proportion of satisfied conditions Assumption II measured as du. Similarity between singular vectors of two adjacent time steps. (c) Frobenius norms of parameter groups.
|Ail and d2
A (b)
35

---

B.3 Experimental Setting of Real Task
To validate that our theoretical insights generalize beyond simplified settings, we conducted exper- iments on the WikiText dataset, a standard benchmark for language modeling. This experimental setup intentionally incorporates more complex and commonly used architectural features.
Dataset and Task. We use the Wiki'ext dataset, which consists of high-quality articles from Wikipedia. The task is next-token prediction, where the model is trained predict the next word in sequence. Consistent with our synthetic experiments, the training objective is calculated exclusively based on the prediction loss for the final token of each input sequence. The sequence length is set to 2048.
Model and Training Hyperparameters. We 2-layer decoder-only Transformer. To test the robustness of our findings, this model's architecture includes standard components that were abstracted away in the synthetic setup. Specifically. it incorporates residual connections after both the self-attention and feed-forward sub-layers, and it utilizes the GeLU activation function in the feed-forward network. This more realistic configuration allows to demonstrate that our theory holds even in the presence such non-linearities and standard architectural features.
All model weights are initialized from a normal distribution with a standard deviation of dmdid The model was trained for 5 epochs using the AdamW optimizer with an initial learning rate of 2 10-4 4, which was managed by cosine decay schedule with warmup phase. The detailed hyperparameters for this experiment are listed in Table 2
Table 2: Model and Training Hyperparameters for WikiText
Parameter
Value
Parameter
Value
Model Architecture
Training Settings
Vocabulary Size
31,999
Dataset
WikiText
Model Dimension (dmode))
64
Sequence Length Optimizer Batch Size
2048
Feed-Forward Dim. (dr)
800
A
Key/Value Dim. (ds, do)
64
500
Number of Layers
2
Epochs
5
Number Heads
Learning Rate
2 x 10-4
Activation Function
GeLU 2048
AdamW $1,82
0.9, 0.999
Pos. Emb. Length
Weight Decay
0.0
Gradient Clipping
1.0
C
Experiments Compute Resources
The experiments were conducted on server with the following configuration:
48 AMD EPYC 7352 24-Core Processors, each with 512KB of cache
2B1GB of total system memory
8 NVIDIA Geffrce RTX 4080 GPUs with 16GB of video memory each
The experiments were run using Ubuntu 22.04 LTS operating system
36


Reviews:
[{"review_id": "R8kCllRVrn", "reviewer": "Reviewer_S1TU", "summary": "The paper investigates the training dynamics of transformers under a small initialization setting from the perspective of the gradient flow.  Specifically, the dynamics of attention modules are dissected into two stages: a) the first stage is called condensation stage, in which the core attention module (i.e., key-query matrices) is nearly stagnant (i.e., quasi-static) without significant fluctuations; b) the second stage, the quasi-static key-query matrices in the previous stage start to participate in training, leading to rank collapse. The paper provides rigorous analysis from the gradient flow to interprete the blow up dynamics and the condensation dynamics, and also show a so-called \"dynamics separation\" dynamics in the second stage.", "questions": "1. While the paper provides a rigorous and lengthy analysis for the so-called blow up dynamics, it depends an non-degenerate initialization condition. Whether the conditions are met in practice, or in what probability these condition are true in practice? Is there any theoretical evaluation (from the perspective of high dimension probability) or empirical evidence? \n\n2. The theoretical analysis for the so-called condensation dynamics also depends upon Assumption 1. So the questions are:  what is the probability when these conditions are met? Are there any theoretical evaluation or emprical evidence to support these conditions?\n\n3. For the dynamics separation stage, it is shown that the gradients of the empirical risk with respect to W_v, W^[1], W^[2] are very smaller, and the gradients  of the empirical risk with respect to W_Q and W_K are at the order of \\delta---which characterizes the dynamics separation phenomenon. Have the theoretical results in Proposition 3 been supposed already? What is the essential theoretical contribution in the analysis for the second stage?\n\n4. Are there any theoretical or empirical evidence to support the so-called dynamics separation assumptions in Assumption 2?", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths: \n1. The paper is well-motivated and the literature review is comprehensive.  \n \n2. The paper provides a rigorous analysis for two sort of dynamics in the first stage: a) the blow up dynamics, which interprete why Transformers with small random initialization can eventually escape from the small initialization, and b) the condensation dynamics, which interprete the condensation phenomenon in training Transformers. \n\n3. It also provides the key-query dynamics in the second stage under a so-called \"dynamics separation\" assumption to show that the gradients of the empirical risk with respect to W_v, W^[1], W^[2] are very smaller, and the gradients  of the empirical risk with respect to W_Q and W_K are at the order of \\delta---characterizing the dynamics separation phenomenon. \n\n\nWeaknesses: \n1. The paper provides a rigorous and lengthy analysis for the so-called blow up dynamics, which depends a non-degenerate initialization condition. It is not clear whether the conditions are met in practice. Or what is the probability of the condition is true in practice? Is there any theoretical evaluation (from the perspective of high dimension probability) or empirical evidence? \n\nSome symbols in the proofs provided in the appendix are mess, not consistent with the main part. \n\n2. For the so-called condensation dynamics, the theoretical analysis is also depend upon the Assumption 1 (Condensation condition). The reviewer is wondering that in what probability these condition is met, and would like to see any theoretical evaluation or emprical evidence to support the conditions in Assumption 1. \n\n3. For the dynamics separation stage, it is shown that the gradients of the empirical risk with respect to W_v, W^[1], W^[2] are very smaller, and the gradients  of the empirical risk with respect to W_Q and W_K are at the order of \\delta---which characterizes the dynamics separation phenomenon. However, it seems that the theoretical results in Proposition 3 have been supposed, no need to further proof and repeatation.  Thus, the reviewer can hardly see the substantial theoretical contribution in the analysis for the second stage. \n\n4. Moreover, are there any theoretical or empirical evidence to support the so-called dynamics separation assumptions in Assumption 2? It is not difficult to provide. But, the reivewer did not find any empirical evidence to support the conditions in Assumption 2.", "was_revised": true, "final_justification": "The reviewer would like to recommend a weak accept for the paper, due to well-motivated and clear presentation, comprehensive literature review and rigorous analysis. The concerns of the reviewers listed in the weaknesses have been resolved in the rebuttal. Thus, the reviewer changes the rating to positive.", "rebuttal": "We deeply appreciate your rigorous feedback, which has helped us significantly strengthen the theoretical grounding of our work.​​  We respectfully ask if you might be willing to reconsider the negative overall evaluation, in light of your positive scores on quality, clarity, significance, and originality. Below are point-by-point responses to your concerns.\n\n## Response to Question 1: Practical Probability of Non-Degenerate Initialization  \n> *\"What is the probability that non-degenerate initialization conditions hold in practice?\"*\n\n**Our Response:** We thank the reviewer for this insightful question. Our non-degenerate condition holds with probability 1 under standard initialization schemes (i.e., almost surely). To address this rigorously, we have added a formal proof in Appendix A of our revised manuscript, which we summarize below.\n\n### Theoretical Justification (Probability Measure Perspective)\n\n**Proof Strategy:**  \nOur proof strategy is to demonstrate that the complementary (failure) event has a probability of zero. The non-degenerate condition, as specified in Definition 3, fails if either of the equalities it seeks to avoid occurs.\n\n#### Part 1: Proof for Definition 3.1\nThe first condition for failure is $\\\\|\\boldsymbol{W}^{[2]}(0)\\\\|\\_2 = \\\\|\\boldsymbol{W}\\_{\\boldsymbol{v}}(0)\\\\|\\_2$. This is equivalent to the event $\\\\|\\boldsymbol{W}^{[2]}(0)\\\\|\\_2^2 - \\\\|\\boldsymbol{W}\\_{\\boldsymbol{v}}(0)\\\\|\\_2^2 = 0$.\n\nLet's define $D\\_1 := \\\\|\\mathbf{W}^{[2]}(0)\\\\|\\_2^2 - \\\\|\\mathbf{W}\\_{\\mathbf{v}}(0)\\\\|\\_2^2$. The weights in $\\boldsymbol{W}^{[2]}(0)$ and $\\boldsymbol{W}\\_{\\boldsymbol{v}}(0)$ are initialized by sampling from Gaussian distributions. As a result, $D\\_1$ is a continuous random variable, as it is a sum and difference of squared random variables drawn from continuous distributions. For any continuous random variable, the probability of it taking on a specific value is zero.\n\nTherefore, $\\mathbb{P}(D\\_1 = 0) = 0$. This means the failure event almost surely does not occur, and consequently, Definition 3.1 is satisfied with probability 1.\n\n#### Part 2: Proof for Definition 3.2\nThe second condition (Case 1) for failure, assuming $\\\\|\\boldsymbol{W}^{[2]}(0)\\\\|\\_2 > \\\\|\\boldsymbol{W}\\_{\\boldsymbol{v}}(0)\\\\|\\_2$ is $\\\\|\\dot{\\mathbf{W}}\\_{\\mathbf{v}}(0)\\\\|\\_2^2 - \\\\|\\dot{\\mathbf{W}}^{[2]}(0)\\\\|\\_2^2 - \\\\|\\mathbf{W}^{[2]}(0)\\\\|\\_2^2 D\\_1 =0$.\n\nLet's define $D\\_2 := \\\\|\\dot{\\mathbf{W}}\\_{\\mathbf{v}}(0)\\\\|\\_2^2 - \\\\|\\dot{\\mathbf{W}}^{[2]}(0)\\\\|\\_2^2 - \\\\|\\mathbf{W}^{[2]}(0)\\\\|\\_2^2 D\\_1$. The continuous distribution characteristic of standard initialization schemes renders $D\\_2$ absolutely continuous. This property necessarily implies $\\mathbb{P}(D\\_2 = 0) = 0$ almost surely. Analogous reasoning applies to Case 2 ($\\\\|\\mathbf{W}^{[2]}(0)\\\\|\\_2 < \\\\|\\mathbf{W}\\_{\\mathbf{v}}(0)\\\\|\\_2$).\n\n## Response to Question 2: Probability of Condensation Condition (Assumption 1)\n> *\"In what probability is the condensation condition (Assumption 1) met? Please provide theoretical evaluation or empirical evidence.\"*\n\n**Our Response:** Empirical measurements demonstrate ​​Assumption 1 holds with probability 1 during the early condensation phase​​ (≤0.3 epochs), satisfying the preconditions for our theoretical analysis.\n\n### Empirical Validation of Assumption 1\n\n#### 1. Empirical Validation Framework\n\nTo quantitatively evaluate Assumption 1, we define \"satisfaction rates\" that measure the proportion of neuron pairs satisfying each of its two conditions during training.\n\n   - **Assumption 1.1 satisfied rate** = $\\frac{|A\\_1|}{d\\_m}$ where \n     $A\\_1 = \\left \\\\{ i \\in [d\\_m] \\mid \\boldsymbol{W}^{[2]}\\_i \\boldsymbol{W}\\_{\\mathbf{v}} \\boldsymbol{W}^{[1],i} >0 \\right\\\\}$\n   - **Assumption 1.2 satisfied rate** = $\\frac{|A\\_2|}{d\\_m^2}$ where \n     $A\\_2 = \\left\\\\{ (i,j) \\in [d\\_m] \\\\times [d\\_m] \\mid \\langle \\boldsymbol{W}^{[2]}\\_i \\boldsymbol{W}^{[1],i}, \\boldsymbol{W}^{[2]}\\_j \\boldsymbol{W}^{[1],j} \\rangle > 0 \\right\\\\}$\n\n#### 2. Key Results: Rapid Convergence to Full Satisfaction\nQuantitative measurements across initial training epochs reveal:\n\n| Metric / Epoch         | 0     | 0.022 | 0.056 | 0.09 | >0.3 |\n|-------------------------|-------|-------|-------|------|------|\n| Assumption 1.1 satisfied rate | 0.525 | 0.52  | 0.65  | 0.95 | 1.0  |\n| Assumption 1.2 satisfied rate | 0.50  | 0.51  | 0.70  | 0.95 | 1.0  |\n\nThe satisfaction rates start near 50% at initialization, as expected for random vectors. They then rapidly converge, surpassing 95% by 0.1 epochs and reaching a stable 100% by 0.3 epochs. This state of full satisfaction persists throughout the condensation phase we analyze.\n\n### Conclusion\nThis empirical evidence confirms that Assumption 1 is not merely a theoretical convenience but a condition that is reliably and quickly met in practice. This validates the foundational premise of our condensation analysis, directly addressing the reviewer's concern.\n\n## Response to Question 3: Theoretical Contribution of Dynamics Separation Stage\n> *\"The theoretical results in Proposition 3 appear self-contained without needing further proof or repetition.\"*\n\n**Our Response:** We thank the reviewer for this sharp observation. We respectfully clarify that Proposition 3 is a non-trivial consequence of Assumption 2, not a restatement. Our core contribution is to prove how the state described in the assumption leads to specific, separated gradient behaviors in the full, coupled system.\n\n### 1. What is Assumed vs. What is Proven\n- Assumption 2 (The Precondition): We assume that after the initial condensation stage, the system enters a specific state where:\n  - The linear subsystem (involving $\\boldsymbol{W}\\_V, \\boldsymbol{W}^{[1]}, \\boldsymbol{W}^{[2]}$) has effectively reached a critical point.\n  - The attention score dynamics (involving $\\boldsymbol{W}\\_Q, \\boldsymbol{W}\\_K$) remain quasi-static.\n- Proposition 3 (The Proven Consequence): We then prove that if the system is in this state, the gradients of the total empirical loss $\\mathcal{L}(\\boldsymbol{\\theta})$ with respect to the different parameter groups will separate and scale differently. Specifically, we prove that $\\nabla_{\\boldsymbol{W}\\_V, \\boldsymbol{W}^{[1]}, \\boldsymbol{W}^{[2]}} \\mathcal{L} \\rightarrow \\mathcal{O}(\\delta^2)$ and $\\nabla_{\\boldsymbol{W}\\_Q, \\boldsymbol{W}\\_K} \\mathcal{L} \\rightarrow \\mathcal{O}(\\delta)$.\n\n### 2. The Non-Trivial Derivation: From Assumption to Proposition\nThe key to our proof is the novel decomposition of the loss function, which is motivated by our multi-stage dynamics framework:\n$$\n\\mathcal{L}(\\boldsymbol{\\theta}) \\approx \\frac{1}{n}\\sum\\_{i=1}^n \\underbrace{\\mathcal{L}\\_{1,i}(\\boldsymbol{\\theta})}\\_{\\text{linear dynamics}} \\cdot \\underbrace{\\mathcal{L}\\_{2,i}(\\boldsymbol{\\theta})}\\_{\\text{attention dynamics}}\n$$\nThis decomposition allows us to analyze the total gradient using the chain rule. The gradient for $\\boldsymbol{W}\\_V$, for instance, is not assumed to be small. Instead, we derive its structure:\n\n\\begin{align}\n\\frac{\\partial \\mathcal{L}}{\\partial \\boldsymbol{W}\\_V} \n&= \\frac{1}{n}\\sum\\_{i=1}^n \\left( \\frac{\\partial \\mathcal{L}\\_{1,i}}{\\partial \\boldsymbol{W}\\_V} \\mathcal{L}\\_{2,i} + \\mathcal{L}\\_{1,i} \\frac{\\partial \\mathcal{L}\\_{2,i}}{\\partial \\boldsymbol{W}\\_V} \\right) \\\\\\\\\n&= \\frac{1}{n}\\sum\\_{i=1}^n \\frac{\\partial \\mathcal{L}\\_{1,i}}{\\partial \\boldsymbol{W}\\_V} (1 + \\mathcal{O}(\\delta^2)) + \\mathcal{O}(\\delta^2) \\quad \\text{(by Stage 1 condensation)} \\\\\\\\\n&= \\underbrace{0}\\_{\\text{critical point}} + \\mathcal{O}(\\delta^2) \\quad \\text{(via Assumption 2.2)}\n\\end{align}\n\n### Conclusion\nOur significant theoretical contribution is thus two-fold: we propose a novel dynamics separation framework, and we provide the rigorous proof (Proposition 3) that technically validates it by connecting the assumed system state (Assumption 2) to the concrete, observable gradient dynamics.\n\n## Response to Question 4: Validation of Dynamics Separation Assumptions\n> *\"Are there theoretical or empirical evidence supporting Assumption 2? No empirical evidence was found in the paper.\"*\n\n**Our Response:** We provide direct empirical validation of Assumption 2 through quantitative measurement of parameter dynamics during training. Our analysis reveals clear separation behavior consistent with the theoretical assumptions.\n\n### Empirical Validation Method\nWe quantify the dynamical separation phenomenon using the **relative change rate** metric:\n$$\n\\text{relative change rate} = \\frac{\\\\|\\boldsymbol{W}\\_{t+1} - \\boldsymbol{W}\\_{t}\\\\|\\_F}{\\\\|\\boldsymbol{W}\\_{t}\\\\|\\_F}\n$$\nThis measures the relative magnitude of parameter updates between training steps.\n\n### Key Results: Separation at Epoch 13\nMeasurements at the critical transition point (epoch 13) demonstrate:\n\n| epoch                | 0          | 2          | 10         | 13         | 14         | 16         | 25         |\n|----------------------|------------|------------|------------|------------|------------|------------|------------|\n| $W\\_Q$                 | 8.3e-06  | 1.3e-06  | 9.4e-07  | 8.2e-03  | 5.0e-04  | 1.4e-06  | 6.8e-05  |\n| $W\\_K$                 | 8.0e-06  | 1.2e-06  | 1.0e-06  | 7.8e-03  | 4.3e-04  | 8.5e-07  | 1.2e-05  |\n| $W\\_V$                 | 1.2e-03  | 3.3e-05  | 1.4e-04  | 5.3e-05  | 7.5e-04  | 1.9e-06  | 5.5e-05  |\n| $W^{[1]}$ | 1.1e-03  | 3.1e-05  | 2.5e-04  | 6.4e-05  | 1.0e-04  | 6.7e-05  | 5.3e-05  |\n| $W^{[2]}$| 1.9e-03  | 3.3e-04  | 1.1e-04  | 9.6e-05  | 8.5e-05  | 4.0e-05  | 4.5e-05  |\n\n**Critical observations**:\n1. **Attention parameters ($\\boldsymbol{W}\\_Q,\\boldsymbol{W}\\_K$)** exhibit large update magnitudes $\\approx 1e-2$\n2. **Linear path parameters ($\\boldsymbol{W}\\_V,\\boldsymbol{W}^{[1]},\\boldsymbol{W}^{[2]}$)** show minimal updates ($\\le 1e-4$)\n\n### Conclusion\nThe significant difference in relative change rates provides direct empirical confirmation of the dynamics separation phenomenon postulated in Assumption 2, validating the foundation of our Stage 2 analysis."}, {"review_id": "wO06AuaNd7", "reviewer": "Reviewer_7v7g", "summary": "This paper aims to rigorously investigate gradient flow dynamics of Transformers with small initialization via dynamics of individual weight matrices in the model. Their main results are based on the ‘condensation’ phenomenon wherein the rows and columns of weight matrices converge in direction (upto sign) to the same vector, in stages starting with W_V and MLP matrices followed by Key, Query matrices. This phenomenon is empirically demonstrated via training on a toy dataset.", "questions": "1. To understand if condensation occurs more generally, is it possible to show that the same empirical results hold (for any of the below)\n\n- for Transformers with residual connections,\n- beyond the toy dataset used,\n- for MLP activation functions other than tanh like ReLU/GELU, and\n- for Transformers with more than one layer?\n\n2. Section 4:\n- What is the exact initialization used for this experiment?\n- I could not understand the motivation for the constructed dataset: could the authors elaborate what language task it is simulating? \n- Could the authors clarify how the cosine similarity matrices of Fig 1 are calculated? Also, what are the numerical values on y-axis in Fig 1?\n- It looks like $W^{[1]}$ is no longer in the condensation state at the end of training in Fig 1 - is this expected?\n- Is there a chance that training for longer time beyond the epochs reported leads to a sudden change in the model, somewhat similar to the grokking phenomenon [1]? \n\n3. Eq 10/ Line 193 : The product $vW_V$ cannot be computed since $v$ is a vector; did the authors mean $v^\\top W_V$? Also, by definition 1 $W_V \\in \\mathbb{R}^{d_m \\times d_m}$, while $v \\in \\mathbb{R}^d$ from Eq. (8); is $d=d_m$?\n\n4. Line 289: \"matrices $W_Q$ and $W_K$ demonstrate remarkable stability in scale\" - is it possible to show that this holds empirically? (since this does not follow from cosine similarity maps in Figure 1)", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 3, "strengths_and_weaknesses": "Strengths: \n\nRigorously analyzes stagewise optimization dynamics of Transformer using weight matrix dynamics, supported by experiments on a simple dataset and model.\n\nWeaknesses:\n\n1. The data setup seems to be very specific, and it is not clear if it is applicable to natural language data (or any other dataset except the one in the paper). Unless condensation is consistently observed for at least some other datasets, it is difficult to assess its general applicability.\n\n2. The Transformer architecture is unconventional: there are no residual connections for the Attention / MLP output, and activation function for MLPs is tanh instead of the commonly used ReLU/GELU.", "was_revised": true, "final_justification": "My main concern about this work were limited applicability beyond the toy experimental setup in the original version. The authors have addressed that via experiments on wikitext-103 data with a decoder-only Transformer with GELU activation and residual connection. They have also clarified some questions I had about measuring cosine similarity / experimental setup etc. and a question on experimentally verifying stability in scales of weight matrices. Hence I've raised my score to 5.", "rebuttal": "We sincerely appreciate your insightful comments, which have greatly contributed to improving the quality of our work. Below is a point-by-point response to your feedback.\n\n## Weakness 1\n>\"Is condensation consistently observed beyond synthetic datasets?\"\n\n**Response:** \nWe confirm the ​​universality of condensation dynamics​​ through new experiments on ​​Wikitext-103​​ (natural language corpus) using a standard 2-layer decoder-only transformer with GeLU activation and residual connections. This demonstrates the phenomenon’s independence from specific data distributions.\n\n### Experimental Results\n\n1. Condensation Rate\n\nWe define the following aggregation index to reflect our results. According to the definition 1 of the article. A neuron is condensed if its alignment with the dominant direction exceeds:\n$$\n \\left| \\left\\langle \\frac{{W}\\_k}{\\| {W}\\_k \\|}, {v} \\right\\rangle \\right| \\ge 0.95\n$$\nThe condensation rate is defined as the number of condensed neurons divided by the total number of neurons.\n\n2. Key Result\nWe plot the condensation rate of the first layer of decoder.\n\n|Condensation Rate/Epochs|0.025|0.075|0.125|0.15|0.2|\n|-|-|-|-|-|-|\n|$W\\_Q$|0.11|0.11|0.77|0.99|0.99|\n|$W\\_V$|0.11|0.9|0.91|0.84|0.76|\n|$W^{[1]}$|0.1|0.91|0.97|0.98|0.96|\n\nOn Wikitext-103, all key matrices rapidly condense (rates approach 1.0), confirming the phenomenon's generality. The observed temporal separation, where linear parameters condense earlier than attention parameters, also provides strong support for our multi-stage dynamics theory.\n\n## Weakness 2\n>\"Unconventional architecture: no residual connections, tanh instead of ReLU/GELU\"\n\n**Response:** We address Weakness 2 from two perspectives: experimental validation and theoretical justification.\n\n### 1. Experimental Setting\n\nAs demonstrated in our response to Weakness 1, experiments using a 2-layer, 1-head decoder-only Transformer on Wikitext-103—which employs ​​GeLU activations​​ and ​​residual connections​​—confirm that our core findings generalize to architectures with modern activation functions and residual pathways. It shows the generality of our results to multi-layer gelu activation functions.\n\n### 2. Theoretical Development\n\n**On the versatility of activation functions**\n\nOur theoretical results apply to all activation functions that satisfy the following conditions:\n$$\n\\sigma(0) = 0 ~ \\text{and} ~ \\sigma'(0) \\neq 0.\n$$\nThis condition encompasses GeLU (validated experimentally) and excludes only functions with undefined or zero derivatives at the origin (e.g., ReLU). \n\nStage 1 dynamics remain consistent​​ under this condition. The empirical loss expansion retains the same asymptotic form:\n$$\n    \\mathcal{L} ({\\theta})= \\frac{1}{n} \\sum\\_{i=1}^n \\left[ 1- \\varepsilon^3 \\left( \\sum\\_{j=1}^s \\frac{1}{s} y\\_i X\\_{i,j} \\bar{{W}}\\_V \\bar{{W}}^{[1]} \\bar{{W}}^{[2]} \\right) + o(\\varepsilon^3) \\right],\n$$\nyielding identical leading-order dynamics:\n$$\n    \\frac{\\mathrm{d} \\bar{{\\theta}}}{\\mathrm{d} \\bar{t}} = \\nabla\\_{\\bar{{\\theta}}} \\left( {v}^{\\intercal} \\bar{{W}}\\_V \\bar{{W}}^{[1]} \\bar{{W}}^{[2]} \\right)\n$$\nSo we get the same result for general activation function in the first stage of training. \n\nStage 2 analysis​​ relies on the assumption that after Stage 1 convergence to a linear critical point, higher-order perturbations remain negligible. This holds for all σ satisfying the above condition, preserving our theoretical conclusions.\n\n**Discussion about residual connections**\n\nWhile residual connections warrant deeper exploration, we provide preliminary theoretical insights. Consider a single-layer Transformer with residuals:\n$$\nf\\_{{\\theta}} (X) = (X + Attn (X) + FFN(Attn(X))) W\\_{\\text{proj}}. \n$$\nThe empirical loss expands as:\n\\begin{equation*}\n    \\mathcal{L} ({\\theta})= \\frac{1}{n} \\sum\\_{i=1}^n \\left[ 1- \\varepsilon \\left(y\\_i X\\_{i,s} \\bar{{W}}\\_{proj} \\right) + o(\\varepsilon) \\right].\n\\end{equation*}\nTherefore, the first stage of analysis will be much simpler, since the leading order dynamics will be a linear ode governing ${W}\\_{proj}$. Next, following the same principles as our second-stage analysis, we assume that the model reaches the critical point of the linear model dominated by $W\\_{proj}$ and decompose the empirical loss as follows:\n$$\n\\mathcal{L} ({\\theta}) \\approx \\frac{1}{n} \\sum\\_{i=1}^n \\mathcal{L}\\_{1,i}({\\theta}) \\mathcal{L}\\_{2,i} ({\\theta}) \\mathcal{L}\\_{3,i} ({\\theta}) \n$$\nwhere \n$$\n\\begin{aligned}\n\\mathcal{L}\\_{1,i} &= e^{-y\\_i X\\_{i,s} W\\_{proj}} \\\\\\\\\n\\mathcal{L}\\_{2,i} &= e^{-y\\_i Attn(X\\_i)\\_s W\\_{proj}} \\\\\\\\\n\\mathcal{L}\\_{3,i} &= e^{-y\\_i FFN(X\\_{i,s}+Attn(X\\_i)\\_s) W\\_{proj}} \n\\end{aligned}\n$$\nBy making reasonable assumptions, we can analyze the dominant dynamics at the moment. At this moment, the change of $W\\_V$ will be most significant, rather than $W\\_Q$ and $W\\_K$. After all linear model changes stop, $W\\_Q$ and $W\\_K$ in the attention module begin to change. This ​​sequential learning phenomenon​​ aligns with the module-specific update ordering observed in our original work. Future work​​ will formalize this analysis for multi-layer architectures, but current results suggest residual connections preserve the core mechanistic principles we describe.\n\n## Response to Question 1:\nTo address the reviewer's concern about the generalizability of our findings on condensation, we provide the following clarifications. We have indeed confirmed that the condensation phenomenon holds across the various dimensions queried by the reviewer. These results are empirically and theoretically demonstrated in our responses to Weaknesses 1 and 2, as summarized below:\n\n- Residuals & Multi-Layer Models: Addressed by our extended theoretical analysis (Weakness 2) and empirically verified with a 2-layer model on Wikitext-103 (Weakness 1).\n- Alternative Activations (GELU/ReLU): Covered by our general theoretical proof and validated with GELU experiments (Weakness 1).\n- Beyond Toy Datasets: Confirmed via large-scale experiments on Wikitext-103 (Weakness 1, Table 1).\n\n## Response to Question 2:\nWe appreciate these technical inquiries and address them sequentially below.\n\n> What is the exact initialization used for this experiment?\n\nGiven any trainable parameter matrix $W \\in \\mathbb{R}^{d\\_1 \\times d\\_2}$, where $d\\_1$ and $d\\_2$ denote the input and output dimensions, respectively, its elements are initialized according to a normal distribution:\n\n$$\nW\\_{i,j} \\sim \\mathcal{N}\\left( 0,  d\\_{1}^{-2\\gamma} \\right),\n$$\n\nwhere $\\gamma$ is the initialization rate. Specifically, the initialization scale decreases as $\\gamma$ increases. Note that $\\gamma = 0.5$. In our experiments, we take $\\gamma = 0.8$.\n\n> Could the authors elaborate what language task it is simulating?\n\nOur constructed dataset is designed to simulate algorithmic reasoning tasks, similar to those found in works like [arxiv:2306.00802, 2401.08309]. The primary goal is to create a controlled environment that isolates a model's ability to learn fundamental structural relationships between tokens, a core capability required for more complex language understanding.\n\n> How the cosine similarity matrices of Fig 1 are calculated? What are the numerical values on y-axis in Fig 1?\n\n- Calculation Method: For parameter matrix ${W} \\in \\mathbb{R}^{n \\times d}$:\n  - Row-wise normalize: $\\tilde{{W}}\\_i = \\frac{{W}\\_i}{\\\\| {W}\\_i\\\\|\\_2}$\n  - Compute $\\tilde{{W}} \\tilde{{W}}^{\\intercal}$\n\n- Y-axis Values: The numerical values on y-axis in Fig 1 is about 4-4.5 which corresponds to relative early stage of training.\n\n> It looks like $W^{[1]}$ is no longer in the condensation state at the end of training in Fig 1 - is this expected?\n\nYes, this is entirely expected and a key part of our findings. Our results show that training proceeds in stages:\n- Stage 1 (epoch 0- 15): ${W}\\_V$, ${W}^{[1]}$ and ${W}^{[2]}$ first condense into a low-complexity state to capture the task's simplest patterns.\n- Stage 2 (epoch 15-50): Other components, like ${W}\\_Q$ and  ${W}\\_K$, begin to condense.\n- Furthur Stage (epoch 50+) To achieve a lower loss and fit the task more accurately, the model must develop more complex and specialized representations. This requires the parameters to differentiate and move out of the condensed state, as observed for $W^{[1]}$.\n\n> Is there a chance that training for longer time beyond the epochs reported leads to a sudden change in the model, somewhat similar to the grokking phenomenon?\n\nThis is an insightful question. While we did not observe the specific delayed generalization dynamic that defines grokking, the staged learning process we identify could be related. The transition from a condensed, low-complexity state to a differentiated, high-complexity one might be a foundational mechanism that precedes phenomena like grokking. We see this as a very promising avenue for future investigation.\n\n## Response to Question 3:\nWe confirm\n\n- Eq 10 contains a notation error. The correct term is ​${v}^{\\intercal} \\bar{{W}}\\_V$.\n- Proposition 1 contains a notation error. The correct term is $X\\_i \\in \\mathbb{R}^{s \\times d\\_m}$.\n\n## Response to Question 4:\nThis is an excellent point. We agree that cosine similarity does not capture matrix scale, and we appreciate the opportunity to provide direct empirical evidence.\n\nTo demonstrate the different scaling dynamics, we report the Frobenius norm of $W\\_Q$ and $W^{[2]}$ at key training epochs below.\n\n|Frobenius Norm/Epoch|0|13|15|19|25|\n|-|-|-|-|-|-|\n| $W_Q$ |1.1|1.1|11.0|19.4|47|\n| $W^{[2]}$ |3.5|12.8|13.2|18.8|20.0|\n\nAs the data shows, the norm of $W\\_Q$ is remarkably stable during the first training phase (e.g., from epoch 0 to 13). The subsequent sharp increase in its norm precisely marks the transition into the second learning phase we described. This behavior contrasts sharply with $W^{[2]}$ whose norm grows significantly within the first phase. This empirically confirms our claim that different model components exhibit distinct dynamic behaviors."}, {"review_id": "erUwYShfm3", "reviewer": "Reviewer_DNga", "summary": "The authors provide an in depth analysis on the training dynamics of a single-layer transformer model with small initialization trained on a binary classification task with exponential loss. Their theoretical contributions can be separated into three key parts/stages:\n1. By analyzing the linearization of the loss, they show that the loss is initially guided by the gradient of an energy functional that depends only on W^V, W^[1], and W^[2]. This first shows that the initial stages of optimization optimize only these three matrices while leaving W^Q, and W^K primarily static. They additionally show that this energy functional diverges to infinity in finite time (almost surely) and use this as a means to show that the norms of the three matrices are blowing up and capable of escaping the small initialization.\n2. Using the dynamics dictated in the previous part, they use this to establish a condensation condition under which W^V experiences condensation.\n3. Lastly, the authors leverage the dynamics to motivate assumptions that would lead to rank collapse in the matrices W^Q, and W^K. Notably, these assumptions assert that the second stage in training emerges after the W^V, W^[1], and W^[2] have been sufficiently optimized where the W^Q and W^K have also shrunk in norm (Key-Query stunting). During this second stage, the W^V, W^[1], and W^[2] are assumed to be constant and the W^Q, W^K matrices evolve according to a linear ODE. The extent of the rank-collapse hence depends on the multiplicity of the top singular value of the linear ODE.", "questions": "- The experimental setup is also lacking a bit of clarity. What is the objective? Is it to approximate the anchor function? If so, why would it make sense to use the cross-entropy loss?\n- For the experiments, it is unclear to me what the cosine similarity is taken over. The caption says that it’s taken across the neural network layers but the model is only a single layer. It would be great if an equation is added describing exactly what the visualizations are showing.\n- Does Assumption 2 hold in the experiments provided in Section 4? Would it not be possible to measure and show that the key-query stunting and that the criticality conditions hold?\n- Similarly, with Theorem 3, could a soft measure of rank be used to show that rank collapse emerges with training in the experimental setting provided in Section 4?\n- Could the measurements in Section 4 be performed on at scale transformers to show that the stages emerge in actual training and that the results are not just limited to the synthetic setting?", "limitations": "Both theoretical and empirical analysis is performed in synthetic settings. While the theoretical limitations is adequately addressed in their Limitations section and reasonable, the paper would strongly benefit including additional at-scale empirical experiments that would provide evidence for the behaviour that their paper is showing.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths:\n- Provides a solid theoretical analysis on the training dynamics of a single layer transformer model that could be used to explain why rank collapse emerges in practise.\n- Separates and groups the training of a transformer block into two different stages which is backed by both empirical experiments and a theoretical derivation.\n\nWeaknesses:\n- $m$ is not defined prior to definition 1.\n- W^[1] and W^[2] are being referenced prior to being properly defined and cannot be as easily deduced as W^Q, W^K, W^V.\n- Structurally, the placement of the experimental results (Section 4) seems out of place to me. There is little to no motivation provided prior to the presentation of the experiments so it is unclear to me what the purpose of the experiments are when they are introduced. It is only until after the theoretical exposition does it become clear what the experiments are trying to show (i.e. the two different stages that are governed by the training dynamics).\n- Section 5.3 is referencing experiments/empirical evidence that don’t seem to be included in the paper.\n- All of the work done in the paper (including empirical results) are performed on miniature/synthetic settings and additionally, many of the theoretical assumptions that can be verified empirically in the synthetic setting in Section 4 are not verified empirically. While it is reasonable to perform the theoretical analysis in the synthetic setting, the paper would benefit by including empirical evidence that the behaviour dictated by their training dynamics emerge in larger models/in practise as well. Currently, it is unclear if any of the results in the paper actually generalize.", "was_revised": true, "final_justification": "The authors have addressed all the concerns I raised, and I have raised the score to accept in response.", "rebuttal": "We would like to express our sincerest gratitude for your thorough and highly constructive review.\n\nWe were very encouraged by your positive assessment of our paper's quality, clarity, significance, and originality. We respectfully ask if you might be willing to reconsider the negative overall evaluation, in light of your positive scores on these individual aspects. We hope that following detailed responses have fully addressed your concerns and now better align the paper's overall contribution with the strengths you identified.\n\n## Response to Weaknesses\nWe sincerely appreciate these structural observations and will implement the following revisions:\n\n> $m$ is not defined prior to definition 1.\n\nWe will reposition ​​Definition 1​​ to immediately follow Section 3.1.\n\n> $W^{[1]}$ and $W^{[2]}$ are being referenced prior to being properly defined and cannot be as easily deduced as $W\\_Q$, $W\\_K$ and $W_V$.\n\nWe will adjust the statement in L37 to  “We delineate different training dynamics for parameter matrices in Transformers.” and adjust the statement in section 4.\n\n> The placement of the experimental results (Section 4) is inappropriate.\n\nWe will relocate experiments to follow Section 5 (Theoretical Analysis) and add new introductory paragraph:  \"We now empirically validate our theoretical findings regarding stage-dependent dynamics (Section 5). These experiments specifically test two predictions: (1) initial condensation phenomenon, and (2) bifurcation of training dynamics into distinct stages.\"\n\n> Section 5.3 is referencing experiments/empirical evidence that don’t seem to be included in the paper.\n\nWe apologize for the lack of experimental evidence and add the following:\n\n### Empirical Validation of Assumption 1\n\n**Empirical Validation Framework**\n\nWe quantitatively evaluate both conditions in Assumption 1 using custom-defined satisfaction metrics which reflects how many neurons satisfy the hypothesis:\n\n   - **Assumption 1.1 satisfied rate** = $\\frac{|A\\_1|}{d\\_m}$ where \n     $A\\_1 = \\left \\\\{ i \\in [d\\_m] \\mid \\boldsymbol{W}^{[2]}\\_i \\boldsymbol{W}\\_{\\mathbf{v}} \\boldsymbol{W}^{[1],i} >0 \\right\\\\}$\n   - **Assumption 1.2 satisfied rate** = $\\frac{|A\\_2|}{d\\_m^2}$ where \n     $A\\_2 = \\left\\\\{ (i,j) \\in [d\\_m] \\\\times [d\\_m] \\mid \\langle \\boldsymbol{W}^{[2]}\\_i \\boldsymbol{W}^{[1],i}, \\boldsymbol{W}^{[2]}\\_j \\boldsymbol{W}^{[1],j} \\rangle > 0 \\right\\\\}$\n\n#### 2. Key Results: Rapid Convergence to Full Satisfaction\nQuantitative measurements across initial training epochs reveal:\n\n| Metric / Epoch         | 0     | 0.022 | 0.056 | 0.09 | >0.3 |\n|-------------------------|-------|-------|-------|------|------|\n| Assumption 1.1 satisfied rate | 0.525 | 0.52  | 0.65  | 0.95 | 1.0  |\n| Assumption 1.2 satisfied rate | 0.50  | 0.51  | 0.70  | 0.95 | 1.0  |\n\nThe satisfaction rates start near 50% at initialization, as expected for random vectors. They then rapidly converge, surpassing 95% by 0.1 epochs and reaching a stable 100% by 0.3 epochs. This state of full satisfaction persists throughout the condensation phase we analyze.\n\n### Conclusion\nThis empirical evidence confirms that Assumption 1 is not merely a theoretical convenience but a condition that is reliably and quickly met in practice. This validates the foundational premise of our condensation analysis, directly addressing the reviewer's concern.\n\n> The lack of theoretical hypothesis verification and larger model experiments.\n\nWe will verify the theoretical hypothesis and provide experimental evidence for a larger transformer. \n\n### 1.Validation of Dynamics Separation Assumptions\n\nWe provide direct empirical validation of Assumption 2 through quantitative measurement of parameter dynamics during training. Our analysis reveals clear separation behavior consistent with the theoretical assumptions.\n\n**Empirical Validation Method**\n\nWe quantify the dynamical separation phenomenon using the **relative change rate** metric:\n$$\n\\text{relative change rate} = \\frac{\\\\|\\boldsymbol{W}\\_{t+1} - \\boldsymbol{W}\\_{t}\\\\|}{\\\\|\\boldsymbol{W}\\_{t}\\\\|}\n$$\nThis measures the relative magnitude of parameter updates between training steps.\n\n**Key Results: Separation at Epoch 13**\n\nMeasurements at the critical transition point (epoch 13) demonstrate:\n\n| relative change rate / epoch | 0 | 2 | 10 | 13 | 14 | 16 | 25 |\n|-|-|-|-|-|-|-|-|\n| $W_Q$ | 8.3e-06  | 1.3e-06  | 9.4e-07  | 8.2e-03  | 5.0e-04  | 1.4e-06  | 6.8e-05  |\n| $W_K$ | 8.0e-06  | 1.2e-06  | 1.0e-06  | 7.8e-03  | 4.3e-04  | 8.5e-07  | 1.2e-05  |\n| $W_V$ | 1.2e-03  | 3.3e-05  | 1.4e-04  | 5.3e-05  | 7.5e-04  | 1.9e-06  | 5.5e-05  |\n| $W^{[1]}$ | 1.1e-03  | 3.1e-05  | 2.5e-04  | 6.4e-05  | 1.0e-04  | 6.7e-05  | 5.3e-05  |\n| $W^{[2]}$| 1.9e-03  | 3.3e-04  | 1.1e-04  | 9.6e-05  | 8.5e-05  | 4.0e-05  | 4.5e-05  |\n\n**Critical observations**:\n1. **Attention parameters ($\\boldsymbol{W}_Q,\\boldsymbol{W}_K$)** exhibit large update magnitudes $\\approx 1e-2$\n2. **Linear path parameters ($\\boldsymbol{W}_V,\\boldsymbol{W}^{[1]},\\boldsymbol{W}^{[2]}$)** show minimal updates ($\\le 1e-4$)\n\nThis empirical evidence directly supports Assumption 2 which validates the core claim that $\\boldsymbol{W}_Q,\\boldsymbol{W}_K$ evolve while $\\boldsymbol{W}_V,\\boldsymbol{W}^{[1]},\\boldsymbol{W}^{[2]}$ remain quasi-static.\n\n### 2. Larger transformer experiments\n\nWe confirm the ​​universality of condensation dynamics​​ through new experiments on ​​Wikitext-103​​ (natural language corpus) using a standard 2-layer decoder-only transformer with GeLU activation and residual connections. This demonstrates the phenomenon’s independence from specific data distributions.\n\n### Experimental Results\n\n**Condensation Rate**:\n\nWe define the following aggregation index to reflect our results. According to the definition 1 of the article. A neuron is condensed if its alignment with the dominant direction exceeds:\n$$\n \\left| \\left\\langle \\frac{\\boldsymbol{W}\\_k}{\\| \\boldsymbol{W}\\_k \\|}, \\boldsymbol{v} \\right\\rangle \\right| \\ge 0.95\n$$\nThe condensation rate is defined as the number of condensed neurons divided by the total number of neurons.\n\n**Key Result**\nWe plot the condensation rate of the first layer of decoder.\n\n| Condensation Rate / Epochs | 0.025 | 0.075 | 0.125 | 0.15 | 0.2 |\n|-------------------------|-------|-------|-------|------|------|\n| $\\boldsymbol{W}\\_Q$ | 0.11 | 0.11  | 0.77  | 0.99 | 0.99  |\n| $\\boldsymbol{W}\\_V$ | 0.11 | 0.9  | 0.91  | 0.84 | 0.76  |\n| $\\boldsymbol{W}^{[1]}$ | 0.10  | 0.91  | 0.97  | 0.98 | 0.96  |\n\nExperimental results show that $\\boldsymbol{W}\\_V$ $W^{[1]}$ parameters converge quickly (at epoch 0.075), while $\\boldsymbol{W}\\_Q$ converges significantly later than the former (at epoch 0.15). This shows the applicability of our results on real language tasks.\n\n## Response to Question 1\n\nYes, the anchor function is formulated as a classification task, following common practice in prior works (e.g., [arxiv:2306.00802, 2401.08309]). We use cross-entropy loss to align with standard transformer training and better approximate downstream language tasks.\n\n## Response to Question 2\nThank you for pointing this out, and we apologize for the confusion caused by the caption. You are correct that the model consists of a single layer. Our intention was not to present cosine similarities across layers, but rather to show the cosine similarities of different parameter matrices.\n\nWe clarify the computation of cosine similarity of the prarameter matrix. Given a parameter matrix $\\boldsymbol{W} \\in \\mathbb{R}^{n \\times d}$, we compute the cosine similarity matrix as follows:\n\n- Normalize each row: $\\tilde{\\boldsymbol{W}}_i = \\frac{\\boldsymbol{W}_i}{|\\boldsymbol{W}_i|_2}$\n\n- Compute the cosine similarity matrix: $\\tilde{\\boldsymbol{W}} \\tilde{\\boldsymbol{W}}^\\top \\in \\mathbb{R}^{n \\times n}$\n\nThis results in a matrix where each entry represents the cosine similarity between a pair of rows in $\\boldsymbol{W}$, i.e., between two parameter vectors.\n\n## Response to Question 3  \n\nYes, Assumption 2 is verified in our experiments. As shown in our response to Weakness 5, the key-query stunting and criticality conditions are satisfied around epoch 13.\n\n## Response to Question 4\nYes, we address this by introducing the concept of effective rank as a soft measure of matrix rank. Specifically, for a parameter matrix $\\boldsymbol{M}$ with singular values $\\sigma_1 \\ge \\sigma_2 \\ge \\dots \\ge \\sigma_n$, we define:\n$$\n\\text{erank}(\\boldsymbol{M}) = \\sum_{i=1}^n \\frac{\\sigma_i}{\\sigma_1}\n$$\nThis quantity captures how concentrated the spectrum is, with lower values indicating stronger rank collapse. We compute the effective rank of several parameter matrices across different training epochs, as shown below:\n| Effective rank / Epochs | 0.025 | 0.075 | 0.125 | 0.15 | 0.2 |\n|-|-|-|-|-|-|\n| $W_Q$ | 6.8 | 6.8  | 3.36  | 1.28 | 1.26  |\n| $W_V$ | 6.8 | 2.54  | 1.92  | 2.22 | 2.28  |\n| $W^{[1]}$ | 7.9  | 3.38  | 2.38  | 1.94 | 1.98  |\n\nWe observe a sharp drop in effective rank beginning around epoch 13, which corresponds to the second stage of training. This provides clear evidence that rank collapse emerges during this phase, in line with the theoretical predictions of Theorem 3.\n\n## Response to Question 5\n\nWe have included additional experiments on the WikiText dataset, which show that the same training stages also emerge in real-world settings."}, {"review_id": "z06jSey4xh", "reviewer": "Reviewer_t9Wk", "summary": "The authors present a theoretical analysis of linearized transformer training dynamics under small initialization using a gradient flow framework. They identify a two-stage learning process: first, outer-layer weights (MLP, value matrices) escape the small initialization regime and undergo a directional alignment process called condensation. Then, the key and query matrices become active and exhibit rank collapse. The analysis is formalized for a one-layer transformer trained on a binary classification task. The authors prove that condensation and rank collapse emerge generically from non-degenerate random initialization. They support their theory with a synthetic experiment that shows the evolution of the weight matrices during training for a single task.", "questions": "1. See W1.\n2. In L.28 you say that modern models are extremely overparameterized. Can you give a reference that underscores this points for language models?\n3. Same for using dropout in language models.\n4. There are some non-standard terms that are used without clear definition or citation, e.g. rank-collapse (L.58), dichotomy assumptions (L.53), could you define them?\n5. The related work part explains other work in dynamics, but it would be useful to include further details on how the work mentioned relates to your contributions, other than also being about dynamics. For example, it is unclear to me why ICL is mentioned so much, but I might be missing the connection. In addition L.106 should not go into the Related Work section, as you are defining new novel terms.\n6. Paragraph L.113 I have a hard time understanding. How do the assymetries come into play that make your definition advantageous. Especially as later on you introduce some assumptions on dynamical symmetries?\n7. This is probably my missing background, but what are the assumptions you need to make on X and y for binary classification? Do they need to scale in a certain way? Can it be any task?\n8. Where is the dot{w} in L.218/219 defined?\n9. L.175 claims the analysis can easily be extended to other loss functions. All loss functions? Only logistic loss? Does (7) change then?\n10. You omit the activation function in L.296 for the second part of the analysis. Why is this ok? Does it change the dynamics in your example if you leave it out?\n11. Where does the index in $\\mathcal L$ come from in (15)?\n12. In the experimental section, what was the criterion for outlining phases 1 and 2? Since you consider the blow up in weights, a less quantitative approach could be to measure precisely those quantities. in addition, it would be great to see the weights before and after the given phase, to see if the change is then aprupt or not.\n13. Your experiment is executed once. How robust is the timing of the stage, and the analysis you have for different seeds?\n\nSmall things\n- You could consider using commands to cite that do not include the author names when not directly used in the text. Right now it is slightly difficult to parse at times because it is so bloated.\n- L37 you use some parameters without truly defining them, consider refering to the definition and naming them outer/inner parameters already here (L.57). Without the context what follows this is difficult to understand.", "limitations": "I understand that linearlization is a limited setting, but still worth exploring on its own right. However, it is unclear to see how the assumptions that happen before training phase 2 would translate to larger networks and more complex tasks, I have a hard time imagining and analogy which would result in a deeper motivation (small initializations help big models, but do they follow two stage learning?).", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "S1 The authors analyse a linearized transformer’s dynamics for small initializations, showing that there are two phases, which is an advance over previous work.\n\nS2 The idea of condensating amtrices seems to be interesting, but could be motivated a bit more clearly.\n\nS3 The two phases are shown in an example task.\n\nW1 In the introduction L.22, the authors claim that it is desirable to have dynamical analysis that holds independently of the specific task. However, I find it difficult to understand where in the manuscript this question is adressed, e.g. the task in section 4 only describes a single instance, and the assumptions on the input data in the theoretical analysis are not spelled out in the main.", "was_revised": true, "final_justification": "I updated my score for borderline acceptance (see my reply to the authors). But probably I was not the ideal person to do the review, hence my hesitation.", "rebuttal": "We sincerely thank you for your detailed and insightful feedback, which helps us significantly improve the clarity and rigor of our paper. We address each of your points below.\n\n## Response to Question 1\nOur theoretical analysis does not rely on strong data assumptions—only the mild condition\n\n\\begin{equation*}\n    \\boldsymbol{v} := \\frac{\\sum_{i=1}^n y_i \\left(\\sum_{j=1}^s X_{i,j}\\right)}{\\left\\| \\sum_{i=1}^n y_i \\left(\\sum_{j=1}^s X_{i,j}\\right) \\right\\|_2} \\neq 0.\n\\end{equation*}\nwhich holds in very general settings. Thus, the analysis is largely task-independent. To empirically demonstrate this generality beyond our synthetic task, we conducted new experiments on the Wikitext-103 dataset using a standard 2-layer Transformer with GeLU activation and residual connections. To visualize the result, we define the following aggregation index to reflect our results. According to the definition 1 of the article. A neuron is condensed if its alignment with the dominant direction exceeds:\n$$\n \\left| \\left\\langle \\frac{\\boldsymbol{W}\\_k}{\\| \\boldsymbol{W}\\_k \\|}, \\boldsymbol{v} \\right\\rangle \\right| \\ge 0.95\n$$\nThe condensation rate is defined as follows:\n$$\n\\text{Condensation Rate} = \\frac{\\text{\\\\# condensed neurons}}{\\text{total neurons}}\n$$\n\nWe plot the condensation rate of the first layer of decoder.\n\n| Condensation Rate / Epochs | 0.025 | 0.075 | 0.125 | 0.15 | 0.2 |\n|-------------------------|-------|-------|-------|------|------|\n| $\\boldsymbol{W}\\_Q$ | 0.11 | 0.11  | 0.77  | 0.99 | 0.99  |\n| $\\boldsymbol{W}\\_V$ | 0.11 | 0.9  | 0.91  | 0.84 | 0.76  |\n| $\\boldsymbol{W}^{[1]}$ | 0.10  | 0.91  | 0.97  | 0.98 | 0.96  |\n\nOur findings on this 2-layer Transformer confirm that the same two-stage learning dynamics emerge, directly addressing the concern about whether our results generalize to larger networks and more complex tasks. Specifically, we observed that the $W_V$ and MLP weights condense first, followed by the $W_Q$ and $W_K$ matrices, consistent with our theory.\n\n\n## Response to Question 2 \nBy \"extremely overparameterized,\" we refer to the empirical observation that modern LLMs have far more parameters than needed to fit their training data. This is supported by scaling laws, as shown by Kaplan et al. [1, Section 3.2], where larger models consistently yield better performance despite having orders of magnitude more parameters than data points.\n\n## Response to Question 3\nDropout remains an important regularization technique in modern language models. For example, the Qwen-7B/72B technical reports [3] explicitly state in Section 2.3 that a dropout rate of 0.1 is consistently applied after both attention and feed-forward layers.\n\n## Response to Question 4\nRank-collapse refers to the degeneration of parameter or attention matrices into low-rank states. In our work, this specifically describes the collapse of the $W_Q$ and $W_K$ matrices. Similar phenomena have been discussed in prior works such as [2].\nRegarding the dichotomy assumptions, its absence means that Theorem 1 holds almost everywhere, rather than implying the dynamics necessarily degenerates or diverges. In other words, Theorem 1’s conclusions apply broadly without requiring this assumption.\n\n## Response to Question 5\nWe emphasize in-context learning (ICL) because recent studies in ICL provide well-developed dynamical frameworks under the sequential prediction paradigm, which inspire task-agnostic analysis. This perspective informs our approach and allows our dynamical analysis to remain independent of specific tasks, which we view as a key contribution. We agree that Line 106 introduces novel terminology and will relocate it outside the Related Work section in the revision.\n\n## Response to Question 6\nIn the linearized model, our dynamics (10) exhibits a high degree of symmetry, which induces useful conservation laws. However, purely symmetric setups can lead to degeneracies—for example, certain measurements may vanish due to exact cancellations. To avoid this, we rely on initial asymmetries to break such degeneracies and make the conserved quantities informative. As training progresses, the dynamics tend to restore symmetry through parameter evolution. This motivates our introduction of the cohesion condition to ensure that the system evolves toward a symmetric structure in a controlled way.\n\n## Response to Question 7\nAs mentioned in our response to Question 1, we do not require any specific assumptions on the dataset. Our analysis is general and can be applied to various tasks, including binary classification, without constraints on the scaling of $X$ and $y$. This flexibility is one of the strengths of our approach.\n\n## Response to Question 8\nThe notation $\\dot{\\boldsymbol{W}}$ is shorthand for the time derivative of the matrix, i.e., $\\frac{\\mathrm{d}}{\\mathrm{d} t} \\boldsymbol{W}$. We will clarify this in the manuscript to avoid confusion.\n\n## Response to Question 9\nOur analysis can be extended to all loss functions for which the condition\n$$\nl'(0) < 0.\n$$\nholds. In such cases, equation (7) will remain unchanged.\n\n## Response to Question 10\nAlthough we omit the activation function in the second part of the analysis, this is acceptable because, at this stage, the entries for $W_V$, $W_1$, and $W_2$ remain small. For activation functions like $\\tanh$, the absence of a second-order term in its Taylor expansion means the model remains approximately linear with respect to $W_V$, $W_1$, and $W_2$. Therefore, omitting the activation function simplifies the analysis without affecting the dynamics, as the higher-order terms of $\\tanh$ are negligible during this stage.\n\nTo further address the reviewer's concern, we conducted additional experiments by removing the activation function on synthetic dataset. The condensation rate, defined as the ratio of condensed neurons to total neurons, showed that the dynamics remained consistent with the original experiment, with condensation occurring at similar time points. The minor misalignment is due to adjustments made to the initialization for training stability.\n\n| Condensation Rate / Epochs | 0 | 15 | 20 | 25 |  \n|-------------------------|-------|-------|-------|------|\n| $W_Q$ | 0.04 | 0.04  |  0.05  | 0.84 |  0.97 |\n| $W^{[2]}$ | 0.02  | 0.83  |  0.85  |  0.81 | 0.83  |\n\n## Response to Question 11\nThe index in equation (15) corresponds to the sample index, which is defined in Section 3 of the paper.\n\n## Response to Question 12 \nThank you for your insightful comment. We define the first and second stages based on the degree of parameter aggregation and the change of norm. Specifically, the transition between stages is marked by the following:\n- In the first stage, we observe minimal change in the Frobenius norm of $W_Q$, while the norm of $W^{[2]}$ increases significantly.\n- The second stage begins when the norm of $W_Q$ starts to increase, which coincides with a more substantial change in model behavior.\nTo visualize this, we provide the Frobenius norms of $W_Q$ and $W^{[2]}$ at various epochs in the table below. As shown, the norm of $W_Q$ remains static until epoch 13, where it begins to grow abruptly by an order of magnitude, marking a clear and sudden transition to the second phase.\n\n| Frobenius Norm / Epoch         | 0     | 13 | 15 | 19 | 25 |\n|-------------------------|-------|-------|-------|------|------|\n| $W_Q$ | 1.1 | 1.1  | 11.0  | 19.4 | 47  |\n| $W^{[2]}$ | 3.5  | 12.8  | 13.2  | 18.8 | 20.0  |\n\n## Response to Question 13\nTo evaluate the robustness of the stage timing and the analysis across different seeds, we've conducted five additional runs with seeds 1843, 1984, 2000, 2025, and 2048. The timing of the stage transitions has been measured by tracking the condensation rate of $W_Q$. The results demonstrate a high consistency in the transition between stages, even across different initializations.\n\nThe table below shows the condensation rates for  $W_Q$ at different epochs, with the corresponding mean and standard deviation:\n| Condensation Rate / Epoch         | 10 | 11 | 12 | 13 | 14 |\n|-------------------------|-------|-------|-------|------|------|\n| $W_Q$ | $0.04 \\pm 0.0004 $  | $0.05 \\pm 0.007$  | $0.38 \\pm 0.26$ | $0.87 \\pm 0.03$ | $0.93 \\pm 0.01$  |\n\nAs seen in the table, the condensation rate shows a clear progression across epochs, with minimal variation between runs. The standard deviations indicate that the transitions are stable, and there is a high level of reproducibility in the dynamics of the system for different seeds.\n\nThis consistency reinforces the robustness of the stage timing in the training process.\n\n\n## Reference\n[1] Kaplan, Jared, et al. \"Scaling laws for neural language models.\" arXiv preprint arXiv:2001.08361 (2020).\n\n[2] Noci, Lorenzo, et al. \"Signal propagation in transformers: Theoretical perspectives and the role of rank collapse.\" Advances in Neural Information Processing Systems 35 (2022): 27198-27211.\n\n[3] Bai, Jinze, et al. \"Qwen technical report.\" arXiv preprint arXiv:2309.16609 (2023)."}]

Author rebuttal:

Meta-review: This is an exceptional manuscript that is strongly supported by the reviewers as well as the AC.  It is a strong contender for an oral presentation.

The manuscript develops the gradient flow training dynamics of a simple, foundational, transformer architecture. It proves that there are two stages to the training dynamics, first the MLP matrices undergo an alignment phenomenon where the rows and columns approach prescribed vectors.  Following this the key and query matrices begin to train and suffer rank collapse, seemingly towards rank one.  The theoretical analysis is backed-up by empirical studies of small problems where the phenomenon is most clearly observed.  This is an interesting result in that it connects with both gradient flow analysis, and the phenomenon of rank-collapse in attention mechanisms; both of which are under theoretical investigation.  

The reviewers were very supportive with only modest concerns, such as the relatively small scale of the experiments.  The authors engaged well with the reviewers and conducted excellent experiments that addressed as many concerns as could reasonably be done in the time allotted.

The only weakness of the manuscript is that it doesn't include very large experiments, and that it employs gradient flow which while interesting theoretically isn't necessarily advisable in practice due to the very small initialisation of the weights.  All that said, the manuscript has a foundational theoretical focus and such compromises are often necessary when developing this kind of theory.

This is an exceptional manuscript that certainly should be accepted and is a contender for an oral talk.  If I were attending NeurIPS this is exactly the kind of results I would want to hear as an oral presentation.
Decision: Accept (oral)