Research map / Architectures and efficiency / Attention and transformers
Eureka-Moments in Transformers: Multi-Step Tasks Reveal Softmax Induced Optimization Problems
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.
Studies why transformers suddenly learn intermediate sub-tasks ("Eureka-moments") on synthetic two-step tasks, traces the failure to Softmax-induced small gradients for W_Q/W_K, and shows NormSoftmax and heat treatment mitigate it.
Problem
Transformers struggle for hundreds of epochs to learn an unsupervised intermediate step in multi-step decision tasks, with training and validation loss saturating before abrupt, sudden improvements; the cause of this optimization failure is unclear.
Approach
Builds controlled synthetic two-step vision and algorithmic tasks where only the second step is supervised, attributes the plateau to local uniform attention in Softmax producing imbalanced, small gradients for W_Q and W_K, and proposes fixes: NormSoftmax (normalizing QK^T by its empirical std) and Heat Treatment (annealing temperature from low to sqrt(d_k)).
Key results
- NormSoftmax with sqrt(d_k) reached a Eureka-ratio of 10/10 with average Eureka-epoch 28.2 versus 7/10 and 207.4 for vanilla ViT on the main dataset.
- On the No position task, NormSoftmax+HT reached average Eureka-epoch 19 with accuracy 90.63, and NormSoftmax 3/4 Eureka-ratio versus 0/4 for ViT.
- In the learning-rate sensitivity study, NormSoftmax achieved Eureka-ratio 36/36 versus 4/36 for vanilla ViT.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace standard attention scaling with NormSoftmax | ViT · Main dataset | Eureka-ratio: 7/10 -> 10/10 | Helped |
| Replace standard attention scaling with NormSoftmax | ViT · No Position Task | Eureka-ratio: 0/4 -> 3/4 | Helped |
| Set attention temperature to 1/3 instead of √d_k | ViT · Main dataset | Eureka-ratio: 7/10 -> 10/10 | Helped |
| Use Heat Treatment, increasing temperature from 1/3 to √d_k | ViT · Main dataset | average Eureka-epoch: 207.43 -> 74.0 | Helped |
| Combine NormSoftmax with Heat Treatment | ViT · Main dataset | average Eureka-epoch: 207.43 -> 17.5 | Helped |
| Increase temperature from 1/3 to √d_k with Heat Treatment | ViT · No Position Task | Eureka-ratio: 1/4 -> 1/4 | No clear effect |
| Use gradient scaling for W_Q, W_K, and W_V | ViT · No Position Task | Eureka-ratio: 0/4 -> 0/4 | No clear effect |
| Use 20 warmup epochs (ablation) | ViT · Main dataset | accuracy: 89.48 -> 87.65 | Hurt |
| Add weight decay 0.5 (ablation) | ViT · Main dataset | Eureka-ratio: 7/10 -> 5/10 | Hurt |
| Use NormSoftmax instead of standard attention scaling | ViT · Learning-rate sensitivity across schedules | Eureka-ratio: 04/36 -> 36/36 | Helped |
| Reduce transformer depth from 7 to 4 (ablation) | ViT · Main dataset, model-scale experiment | Eureka-ratio: 2/4 -> 3/4 | Helped |
| Increase number of heads from 4 to 6 (ablation) | ViT · Main dataset, model-scale experiment | Eureka-ratio: 2/4 -> 3/4 | Helped |
Surprising findings
- Simply rescaling W_QKV gradients to equal norms raised the Eureka-ratio on the main task but completely failed on the No position task (0/4).
- Warmup alone (without changing attention) did not resolve the problem, and increasing weight decay reduced the Eureka-ratio.
Limitations
- Main analysis uses a synthetic two-step vision task; transfer to real data is shown only for RoBERTa and GPT-2 ICL with few seeds.
- Eureka-moments are hard to predict; very long training (3000 epochs) was needed for reliable convergence without the fixes.
Experimental setup
Datasets: MNIST Fashion-MNIST Cifar-10 Omniglot Wikipedia ImageNet-100 based dataset
Benchmarks: main synthetic two-step vision task No position task Cifar task 1 Top if above Same color decision task Color or fashion class Digit grouping reasoning task RoBERTa Wikipedia pretraining GPT-2 Omniglot ICL
Metrics: Eureka-ratio accuracy after Eureka-moment average Eureka-epoch validation accuracy steps to 95% of ViT accuracy
Baselines: ViT ResNet9 ViT+Warmup W_QKV gradient scaling ViT+weight decay
Builds on: ViT for small datasets grokking NormSoftmax
Closest papers by meaning
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.91)
- Towards Understanding How Transformers Learn In-context Through a Representation Learning Lens (NeurIPS 2024 · similarity 0.90)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.89)
- Abrupt Learning in Transformers: A Case Study on Matrix Completion (NeurIPS 2024 · similarity 0.89)
- Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions (ICLR 2024 · similarity 0.89)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.89)
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.