Research map / Large language models / Language, safety and interpretability
Avoiding spurious sharpness minimization broadens applicability of SAM
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 SAM underperforms AdamW in language modeling because its sharpness reduction is dominated by the logit path, and proposes Functional SAM (perturbing only the Jacobian) plus Adam preconditioning, which improves C4 language modeling across model sizes.
Problem
Sharpness Aware Minimization (SAM) improves generalization in vision but degrades performance in language modeling, even with twice the compute of AdamW. The paper asks why and how to make sharpness minimization work for NLP.
Approach
Decomposes the SAM sharpness gradient into a logit path (via GGN) and a functional path (via functional Hessian), finds the logit path dominates in language modeling, and proposes Functional SAM, which perturbs only the network Jacobian, plus preconditioning the perturbation with Adam's second-moment matrix.
Key results
- At 23.9M params (Chinchilla setting), preconditioned Functional SAM reaches eval loss 3.63 vs AdamW 3.69 and SAM 3.71.
- At 1208M params, preconditioned Functional SAM achieves 2.61 eval loss vs AdamW 2.63, while SAM and preconditioned SAM yield NaNs.
- In the 2M prototype, plain SAM (3.91) is worse than AdamW (3.90), while preconditioned Functional SAM reaches 3.86.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use funcSAM with AdamW preconditioning instead of AdamW | Nanodo Transformer · C4, 23.9M–1208M, fixed-length training | eval loss: 3.57 -> 3.53; 3.45 -> 3.41; 3.28 -> 3.25; 3.08 -> 3.05 | Helped |
| Use funcSAM with AdamW preconditioning instead of AdamW | Nanodo Transformer · C4, 23.9M–1208M, Chinchilla-style training | eval loss: 3.69 -> 3.63; 3.45 -> 3.41; 3.12 -> 3.10; 2.63 -> 2.61 | Helped |
| Use SAM instead of AdamW | Nanodo Transformer · C4 language modeling | eval loss: 3.57 -> 3.59; 3.45 -> 3.46; 3.28 -> 3.29 | Hurt |
| Add preconditioning to funcSAM (ablation) | funcSAM · C4, 2M, fixed-length | eval loss: 3.88 -> 3.86 | Helped |
| Add preconditioning to SAM (ablation) | SAM · C4, 2M, fixed-length | eval loss: 3.91 -> 3.89 | Helped |
| Replace GeLU with ReLU (ablation) | Preconditioned funcSAM · C4, 2M | eval loss: 3.8614 -> 3.8777 | Mixed |
| Increase the SAM perturbation radius beyond the best small radius (ablation) | SAM · C4 language modeling | eval loss | Hurt |
| Warm up the perturbation radius (ablation) | Preconditioned SAM and SAM · C4, 1208M, fixed-length | eval loss: 3.07; 3.08 | Helped |
Surprising findings
- SAM performs worse than AdamW on language modeling despite using twice the compute per step.
- SAM attains the lowest Hessian maximum eigenvalue among methods even though its validation loss is worse than AdamW, suggesting spurious sharpness reduction via logit manipulation.
Limitations
- Functional SAM still costs twice the gradient computation of AdamW per step.
- Perturbation radius tuned only coarsely for the largest models; optimal rho transfer across scales unresolved.
- Preconditioning's interplay with Functional SAM is not fully understood; theoretical argument is a random-matrix heuristic.
Experimental setup
Datasets: C4 ImageNet-1K JFT ImageNet-21K
Benchmarks: C4 validation loss
Metrics: evaluation loss maximum Hessian eigenvalue Hessian trace GGN trace one-shot pruning degradation
Baselines: AdamW SAM preconditioned SAM
Builds on: Sharpness Aware Minimization Adam Nanodo Chinchilla scaling laws Fisher SAM
Closest papers by meaning
- Sharpness-Aware Minimization: General Analysis and Improved Rates (ICLR 2025 · similarity 0.90)
- CR-SAM: Curvature Regularized Sharpness-Aware Minimization (AAAI 2024 · similarity 0.90)
- Friendly Sharpness-Aware Minimization (CVPR 2024 · similarity 0.89)
- Implicit Regularization of Sharpness-Aware Minimization for Scale-Invariant Problems (NeurIPS 2024 · similarity 0.89)
- Sharpness-Aware Minimization Enhances Feature Quality via Balanced Learning (ICLR 2024 · similarity 0.89)
- Sharpness-Aware Minimization Efficiently Selects Flatter Minima Late In Training (ICLR 2025 · similarity 0.88)
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.