Title: From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics
Venue: NeurIPS 2025
Accepted: yes
Keywords: Two-stage analysis

Training dynamics of transformers

Condensation

Small initialization
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.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nS1 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.\n\nQUESTIONS:\n1. 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.\n\nLIMITATIONS:\nI 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?).", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 2, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-06-27"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n- 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?\n\nLIMITATIONS:\nBoth 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.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \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.\n\nQUESTIONS:\n1. 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)\n\nLIMITATIONS:\nYes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-07-01"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: \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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 2, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-07-02"}