Research map / Architectures and efficiency / Attention and transformers
What Happens During the Loss Plateau? Understanding Abrupt Learning in Transformers
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.
Shows that abrupt learning in shallow Transformers on algorithmic tasks coincides with a partial solution, repetition bias, and representation collapse, with slow attention-map learning as the bottleneck, and that these degenerate patterns appear in early pretraining of Pythia and OLMo-2.
Problem
Understanding why Transformers exhibit an extended loss plateau followed by sudden improvement (abrupt learning) when trained on algorithmic tasks, and what hidden progress occurs during the plateau.
Approach
Identifies common degenerate patterns during the plateau (partial solutions, output repetition bias, and post-attention representation collapse) and shows via training-time attention interventions that slow learning of the optimal attention map is a key bottleneck driving plateau length and these patterns.
Key results
- Attention progress measure rises gradually from near 0 to about 0.8 during the plateau, before the sudden loss drop.
- Repetition frequency grows to about 0.8 within roughly 50 steps and average hidden-state cosine similarity reaches about 0.95 during the plateau in the MWS setup.
- Pythia checkpoints show average cosine similarity above 0.9 within a few training steps, and OLMo-2-7B's earliest checkpoint (step 150) has similarity about 0.93.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| bias attention toward optimal positions from the start (c=10) (ablation) | 1-layer, 1-head Transformer · MWS | peak representation cosine similarity: ≈ 0.95 -> ≈ 0.6 | Helped |
| bias attention toward optimal positions (c>1) (ablation) | 1-layer, 1-head Transformer · MWS | convergence speed | Helped |
| bias attention away from optimal positions (c<1) (ablation) | 1-layer, 1-head Transformer · MWS | convergence speed | Hurt |
| initialize K,Q to final values and retrain subsequent layers (ablation) | 1-layer, 1-head Transformer · MWS | convergence time | Helped |
| also initialize V,O to final values in addition to K,Q (ablation) | 1-layer, 1-head Transformer · MWS | average representation cosine similarity: ≈ 0.45 -> approx $0.15$ | Helped |
| fix MLP or embeddings together with LM head to final values (ablation) | 1-layer, 1-head Transformer · MWS | loss plateau and degeneracy | No clear effect |
| train on repetitive target sequences where every output copies the first input | Transformer · REPEAT1 | loss plateau | Helped |
| use smaller initialization scale (α<1) (ablation) | 1-layer, 1-head Transformer · MWS | loss plateau duration | Hurt |
| use initialization scale α=10 (ablation) | 1-layer, 1-head Transformer · MWS | peak representation cosine similarity: ≈ 0.95 -> ≈ 0.6 | Helped |
| use Muon with Adam for embeddings, LM head, and biases instead of Adam alone (ablation) | 1-layer, 1-head Transformer · MWS | plateau duration: ∼ 40 steps | Helped |
Surprising findings
- Representation collapse and repetition bias arise after a few training steps rather than at initialization, contrasting with rank collapse in deep softmax Transformers.
- Fixing optimal attention weights at initialization largely removes the plateau, while fixing optimal MLP or embeddings does not.
Limitations
- Main analysis uses 1-2 layer linear-attention Transformers on simple synthetic tasks.
- The mechanism behind why repetition and collapse arise is not rigorously explained; the 'search' hypothesis is untested theoretically.
Experimental setup
Datasets: synthetic algorithmic tasks (moving-window-sum, prefix sum, multi-digit addition, permutation, histogram, reverse, copy) AI2 ARC-Easy AI2 ARC-Challenge GSM8K
Benchmarks: AI2 ARC-Easy AI2 ARC-Challenge GSM8K
Metrics: accuracy loss repetition frequency sequence entropy average pairwise cosine similarity attention progress measure first-token accuracy
Baselines: softmax attention Transformer SGD Muon
Builds on: minGPT Pythia OLMo-2 Muon
Closest papers by meaning
- Abrupt Learning in Transformers: A Case Study on Matrix Completion (NeurIPS 2024 · similarity 0.91)
- The emergence of sparse attention: impact of data distribution and benefits of repetition (NeurIPS 2025 · similarity 0.88)
- From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics (NeurIPS 2025 · similarity 0.88)
- Eureka-Moments in Transformers: Multi-Step Tasks Reveal Softmax Induced Optimization Problems (ICML 2024 · similarity 0.88)
- Breaking through the learning plateaus of in-context learning in Transformer (ICML 2024 · similarity 0.88)
- Transformers Converge to Invariant Algorithmic Cores (NeurIPS 2026 · 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.