Research map / Architectures and efficiency / Attention and transformers
Abrupt Learning in Transformers: A Case Study on Matrix Completion
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.
Trains a 4-layer BERT on masked low-rank matrix completion and interprets the sudden loss drop as an algorithmic shift from copying the input to computing missing entries.
Problem
Understanding why Transformer training loss plateaus and then suddenly drops, using low-rank matrix completion as a controlled masked-language-modeling-style task.
Approach
Casts low-rank matrix completion as masked language modeling over tokenized matrix entries, trains BERT online, and uses interventions (token replacement, uniform attention ablation, weight switching, activation patching, position permutation) plus probing to show the sudden drop coincides with a change from copying to completing masked entries.
Key results
- BERT reaches a final MSE of about 4e-3 on 7x7 rank-2 matrix completion, lower than nuclear norm minimization.
- The loss drops sharply around step 15000 after a plateau, with the pre-shift model copying inputs and predicting ~0 at masked entries.
- Uniformly ablating all attention heads post-shift raises masked MSE from 0.0128 to 0.2183 while observed MSE barely changes, indicating attention drives missing-entry prediction.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| train BERT to complete masked low-rank matrices | 4-layer, 8-head BERT · 7×7 rank-2 matrix completion | MSE | Helped |
| replace nuclear norm minimization with trained BERT | 7×7 rank-2 matrix completion · matrix completion at various mask levels | MSE | Helped |
| uniformly ablate all attention heads (ablation) | pre-shift BERT · 7×7 rank-2 matrices | L_mask: 0.2236 -> 0.2236 | No clear effect |
| replace pre-shift attention weights with post-shift weights (ablation) | pre-shift BERT · 7×7 rank-2 matrices | L_mask: 0.2246 | No clear effect |
| uniformly ablate all attention heads (ablation) | post-shift BERT · 7×7 rank-2 matrices | L_mask: 0.0128 -> 0.2183 | Hurt |
| replace post-shift attention weights with pre-shift weights (ablation) | post-shift BERT · 7×7 rank-2 matrices | L_obs: 9.44e-5 -> 9.5e-4 | Hurt |
| randomly permute positional embeddings (ablation) | post-shift BERT · 7×7 rank-2 matrices | L_mask: 0.01428 -> 0.5687 | Hurt |
| ablate row-head group (2,1), (3,4), (4,8) (ablation) | post-shift BERT · 7×7 rank-2 matrices | L: $0.0035$ -> $0.0073$ | Hurt |
| ablate observed-copy heads (4,3), (4,4) (ablation) | post-shift BERT · 7×7 rank-2 matrices | L: $0.0045$ -> $0.0079$ | Hurt |
| ablate mask-ignore heads (ablation) | post-shift BERT · 7×7 rank-2 matrices | L: $0.0043$ -> $0.057$ | Hurt |
| ablate longest-contiguous-column heads (2,5), (2,7), (2,8) (ablation) | post-shift BERT · 7×7 rank-2 matrices | L: $0.0049$ -> $0.0112$ | Hurt |
| ablate input-processing heads (1,1–2), (1,5–8) (ablation) | post-shift BERT · 7×7 rank-2 matrices | L: $0.0048$ -> $0.0314$ | Hurt |
Surprising findings
- Post-shift BERT outperforms nuclear norm minimization in MSE while having a larger nuclear norm, so it does not implement that algorithm.
- Pre-shift attention weights and heads are largely inconsequential; copying is done without attention.
Limitations
- Experiments are restricted to small matrices (largest 15x15), and the authors note scaling would likely need modifications.
- The mechanism for computing missing entries is not formalized as an explicit algorithm.
Experimental setup
Datasets: synthetic low-rank matrices (Unif[-1,1] factors)
Benchmarks: synthetic 7x7 rank-2 matrix completion
Metrics: MSE overall MSE on observed entries MSE on masked entries
Baselines: nuclear norm minimization regularized MSE minimization
Builds on: BERT nuclear norm minimization grokking
Closest papers by meaning
- What Happens During the Loss Plateau? Understanding Abrupt Learning in Transformers (NeurIPS 2025 · similarity 0.91)
- Eureka-Moments in Transformers: Multi-Step Tasks Reveal Softmax Induced Optimization Problems (ICML 2024 · similarity 0.89)
- From Condensation to Rank Collapse: A Two-Stage Analysis of Transformer Training Dynamics (NeurIPS 2025 · similarity 0.88)
- Training Dynamics of Transformers to Recognize Word Co-occurrence via Gradient Flow Analysis (NeurIPS 2024 · similarity 0.88)
- From Attention to Activation: Unraveling the Enigmas of Large Language Models (ICLR 2025 · similarity 0.87)
- On the Role of Attention Masks and LayerNorm in Transformers (NeurIPS 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.