Research map / Architectures and efficiency / RNNs and normalization
From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Analyzes gradient-flow training of one-layer (and empirically two-layer) transformers under small initialization, showing a two-stage process: condensation of value/FFN outer parameters followed by asymptotic rank collapse of normalized key-query matrices.
Problem
Training dynamics of transformers are poorly characterized beyond configuration-specific studies; it is unclear how small initialization shapes attention parameter dynamics in binary classification.
Approach
Uses a small-initialization perturbative expansion of the loss to derive effective dynamics, proves finite-time blow-up and condensation of W_V under a condensation condition, and then shows that after outer parameters reach a quasi-steady state, the key-query matrices obey linear ODEs whose normalized solutions collapse to rank one.
Key results
- Theory proves finite-time blow-up of the energy E(t) under almost surely non-degenerate Gaussian initialization.
- Experiments show the condensation condition proportion rapidly approaches 1 within the first 200 training steps on the synthetic task.
- The normalized key-query matrices exhibit rank collapse (sharp drop in effective rank) after the outer parameters reach a quasi-steady state, on both synthetic and WikiText two-layer experiments.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| include the final activation function (ablation) | Transformer training dynamics · Transformer training dynamics | learning dynamics | No clear effect |
Limitations
- Analysis restricted to binary classification; multi-class and sequence-to-sequence settings are not covered.
- Key steps rely on assumptions (condensation condition, dynamics separation) that are verified empirically rather than proven, and the necessity of the condensation condition is unproven.
- Final activation omitted and layer normalization removed in the real-task experiments, moving away from standard transformers.
Experimental setup
Datasets: synthetic anchor-function dataset WikiText
Benchmarks: WikiText
Metrics: cosine similarity of parameters effective rank relative change of norms training loss
Builds on: small initialization condensation (Zhou et al. 2022) anchor functions
Closest papers by meaning
- The underlying structures of self-attention: symmetry, directionality, and emergent dynamics in Transformer training (ICML 2025 · similarity 0.91)
- Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis (NeurIPS 2024 · similarity 0.90)
- An Analysis for Reasoning Bias of Language Models with Small Initialization (ICML 2025 · similarity 0.89)
- Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape (ICML 2024 · similarity 0.89)
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.89)
- What Happens During the Loss Plateau? Understanding Abrupt Learning in Transformers (NeurIPS 2025 · similarity 0.88)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.