Research map / Optimization / Implicit bias and generalization
Grokking at the Edge of Numerical Stability
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.
Argues that grokking without regularization is blocked by floating-point absorption errors in softmax cross-entropy (Softmax Collapse), driven by gradient alignment with a logit-scaling direction (naive loss minimization); proposes StableMax and orthogonal-gradient optimizers that induce grokking without weight decay.
Problem
Grokking (delayed generalization after overfitting) is usually observed only with regularization such as weight decay, and it is unclear why generalization is delayed and why regularization is needed, especially with cross-entropy loss and deeper models.
Approach
Without regularization, training beyond 100% train accuracy drives gradients toward a logit-scaling (naive loss minimization) direction; logits grow until float absorption errors in the softmax sum zero out gradients (Softmax Collapse). Remedies are a numerically stable replacement for softmax (StableMax) and an optimizer (⊥Grad) that projects out the weight-direction component of the gradient.
Key results
- StableMax with cross-entropy and no regularization yields grokking on modular addition, sparse parity and MNIST at 40% training data where standard SCE gets stuck at random-level performance.
- ⊥AdamW/⊥SGD generalize quickly without an initial overfitting phase; ⊥SGD reaches 100% test accuracy within 400 iterations on modular addition.
- In realistic settings, ⊥Grad with CIFAR100 reaches 62.69% vs 59.98% for softmax CE; on ImageNet-1k, StableMax CE drops to 65.87% vs 69.33% for softmax CE.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Replace Softmax CE with StableMax CE | ResNet18 · CIFAR10 | accuracy: 87.17% ± 0.2 -> 87.01% ± 0.2 | No clear effect |
| Replace Softmax CE with StableMax CE | ResNet18 · CIFAR100 | accuracy: 59.98% ± 0.4 -> 60.63% ± 0.4 | Helped |
| Replace Softmax CE with StableMax CE | ResNet18 · ImageNet-1k | accuracy: 69.33% ± 0.04 -> 65.87% ± 0.22 | Hurt |
| Replace Softmax CE with StableMax CE | GPT2-Small · WikiText-103 | Top-5 accuracy: 60.48% ± 0.04 -> 51.85% ± 0.47 | Hurt |
| Replace SGD with ⊥SGD | ResNet18 · CIFAR10 | accuracy: 87.17% ± 0.2 -> 87.22% ± 0.2 | No clear effect |
| Replace SGD with ⊥SGD | ResNet18 · CIFAR100 | accuracy: 59.98% ± 0.4 -> 62.69% ± 0.1 | Helped |
| Replace SGD with ⊥SGD | ResNet18 · ImageNet-1k | accuracy: 69.33% ± 0.04 -> 68.95% ± 0.03 | No clear effect |
| Replace AdamW with ⊥Grad | GPT2-Small · WikiText-103 | Top-5 accuracy: 60.48% ± 0.04 -> 59.64% ± 0.04 | No clear effect |
| Replace Softmax with StableMax in the attention mechanism (ablation) | GPT2-Small · WikiText-103 | Top-5 accuracy: 60.48% ± 0.04 -> 58.52% ± 0.04 | Hurt |
| Use StableMax instead of Softmax CE | MLP and transformer grokking tasks without regularization · Modular arithmetic, Sparse Parity, and MNIST | test accuracy | Helped |
| Remove the gradient component parallel to the weights using ⊥Grad | Grokking models · Modular arithmetic | test accuracy | Helped |
| Increase Softmax calculation precision from float32 to float64 (ablation) | MLP · Modular arithmetic | test accuracy | Helped |
Surprising findings
- Gradient alignment with the weight direction reaches cosine similarity up to 0.9 for output layers after overfitting, so logit scaling dominates updates.
- A re-implementation of log_softmax outside official CUDA kernels kept learning past the point where loss is exactly zero, outperforming the official PyTorch implementation on grokking tasks.
Limitations
- Analysis of naive loss minimization focuses on homogeneous or approximately homogeneous models; models with bias and skip connections are only quasi-homogeneous and lack formal characterization.
- StableMax is not uniformly better in realistic settings: on ImageNet-1k and WikiText-103 it underperforms softmax CE.
- Explanation of why weight decay produces the sudden grokking dynamics (effective learning rate) is left as an open question.
Experimental setup
Datasets: modular arithmetic mod 113 sparse parity MNIST WikiText-103 CIFAR10 CIFAR100 ImageNet-1k
Benchmarks: modular arithmetic mod 113 sparse parity MNIST CIFAR10 CIFAR100 ImageNet-1k WikiText-103
Metrics: train accuracy test accuracy test loss accuracy (top-5 for WikiText-103) fraction of softmax collapse samples weight norm
Baselines: AdamW with SCE AdamW with weight decay SGD with weight decay float64 softmax label smoothing temperature scaling
Builds on: AdamW SGD Softplus grokking (Power et al.) lazy training
Closest papers by meaning
- Grokking Beyond the Euclidean Norm of Model Parameters (ICML 2025 · similarity 0.93)
- Deep Networks Always Grok and Here is Why (ICML 2024 · similarity 0.92)
- Grokking at the Edge of Linear Separability (ICML 2025 · similarity 0.92)
- Let Me Grok for You: Accelerating Grokking via Embedding Transfer from a Weaker Model (ICLR 2025 · similarity 0.91)
- Grokking as the transition from lazy to rich training dynamics (ICLR 2024 · similarity 0.91)
- Flatness is Necessary, Neural Collapse is Not: Rethinking Generalization via Grokking (NeurIPS 2025 · similarity 0.91)
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.