Research map / Architectures and efficiency / Attention and transformers
Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis
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 a one-layer softmax transformer (W_K, W_Q, W_V, and a linear MLP trained from random init) on detecting co-occurrence of two target words, proving a two-phase dynamic and near-zero training loss.
Problem
Understanding how transformers provably learn via gradient-based training without simplifying assumptions (weight reparameterization, attention linearization, special initialization, lazy regime), using a word co-occurrence detection task.
Approach
Treats joint gradient-flow training of three attention matrices and a linear MLP from random initialization as a coupled dynamical system over inner-product variables, showing a two-phase process (MLP alignment, then joint attention-MLP margin growth) and a novel "automatic balancing of gradients" property that makes losses of different sample types decrease at similar rates.
Key results
- Phase 1: the linear MLP aligns with target signals within T_1 = Õ(1/m) while attention stays near initialization; all training samples are correctly classified but training loss remains Θ(1).
- Phase 2: attention and MLP jointly enlarge the classification margin, and the training loss converges to at most 1/poly(m) (both training and generalization).
- Synthetic experiments (60 samples, m=128, 256 neurons, GD lr 0.01, 30000 epochs) qualitatively verify the two-phase dynamics and attention-score trends.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Jointly evolve the attention matrices and linear MLP after initial MLP alignment | One-layer transformer trained by gradient flow · Synthetic co-occurrence task | training loss | Helped |
| Align the linear MLP with the target words while attention remains nearly unchanged | One-layer transformer trained by gradient flow · Synthetic co-occurrence task, Phase 1 | training loss | Mixed |
| Train the full transformer from random initialization | One-layer transformer trained by gradient flow · Synthetic co-occurrence task | generalization loss | Helped |
Surprising findings
- Losses for samples with only one target signal first increase during Phase 1 before decreasing as the MLP aligns negatively with the common token.
Limitations
- Analysis restricted to a one-layer, single-head transformer with a linear MLP and gradient flow rather than discrete gradient descent.
- Data model assumes target and irrelevant tokens appear once and irrelevant words are all distinct, with orthonormal embeddings.
- Experiments are small synthetic settings only.
Experimental setup
Datasets: synthetic co-occurrence data (theoretical data distribution)
Metrics: training loss attention score correlation MLP output G(mu)
Builds on: gradient flow analysis of transformers
Closest papers by meaning
- From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics (NeurIPS 2025 · similarity 0.90)
- Transformers Learn Nonlinear Features In Context: Nonconvex Mean-field Dynamics on the Attention Landscape (ICML 2024 · similarity 0.90)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.89)
- How Transformers Learn Regular Language Recognition: A Theoretical Study on Training Dynamics and Implicit Bias (ICML 2025 · similarity 0.88)
- Abrupt Learning in Transformers: A Case Study on Matrix Completion (NeurIPS 2024 · similarity 0.88)
- In-context Convergence of Transformers (ICML 2024 · similarity 0.87)
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.