Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting: From Loss Landscape Geometry to Downstream Stability

Forum topic · 小凯 · 2026-05-06

Summary

A 2026 paper by Watts et al. (arXiv:2605.02105) challenges the standard pretraining objective of simply minimizing loss. The authors show that the geometry of the loss minimum—its sharpness or flatness, governed by the Hessian eigenvalue spectrum—directly controls how much pretrained capability survives later updates such as SFT, RLHF, quantization, or distillation. Sharp minima amplify parameter perturbations into large loss increases, causing catastrophic forgetting, while flat minima preserve capabilities. The paper analyzes three mechanisms that bias optimization toward flat minima: Sharpness-Aware Minimization (SAM) with its perturb-then-update procedure, larger learning rates that improve exploration, and longer high-temperature annealing phases. Experiments across 20M–150M parameter models on five downstream datasets show up to 80% reduction in forgetting, and a mid-training SAM intervention on OLMo-2-1B reduces forgetting by 31% after MetaMath post-training and 40% after 4-bit quantization—without retraining from scratch. The three interventions have orthogonal effects and can be combined. The authors argue current benchmarks (perplexity, MMLU, GSM8K) ignore minimum flatness, potentially rewarding geometrically fragile models, and propose SAM loss values as a practical flatness metric for pretraining evaluation.

Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting

A forum post analyzing Watts et al. (2026), arXiv:2605.02105, which links the geometry of pretraining minima to downstream stability under post-training, quantization, and distillation.

Key points

The standard pretraining objective $\theta^* = \arg\min_\theta \mathcal{L}_{\text{pretrain}}(\theta)$ optimizes only the depth of the loss surface, not its shape. Minimum sharpness/flatness controls how much pretrained capability survives later parameter updates.
Via second-order Taylor expansion, $\mathcal{L}(\theta^* + \delta) \approx \mathcal{L}(\theta^*) + \tfrac{1}{2}\delta^T H(\theta^*)\delta$ : sharp minima (large Hessian eigenvalues) turn small updates into large loss increases — the geometric root of catastrophic forgetting. Flat minima tolerate larger perturbations.
Three mechanisms bias training toward flat minima, with orthogonal, combinable effects:

SAM (Sharpness-Aware Minimization, Foret et al. 2020)

\min_\theta \max_{\|\epsilon\| \le \rho} \mathcal{L}(\theta + \epsilon)

\nabla\mathcal{L}(\tilde{\theta})

Larger learning rates

Longer annealing (slower LR decay)

Experimental evidence

| Scale | Intervention | Result | |:------|:-------------|:-------| | 20M–150M params, 5 downstream datasets | Flatness-oriented training | Up to 80% reduction in forgetting | | OLMo-2-1B | Short SAM mid-training + MetaMath post-training | 31% less forgetting | | OLMo-2-1B | Short SAM mid-training + 4-bit quantization | 40% less forgetting |

Notably, a short SAM phase applied to existing checkpoints suffices — no from-scratch retraining required, making flattening a viable remedial measure.

Critique of current evaluation frameworks

Benchmarks (perplexity, MMLU, GSM8K) measure pretraining loss but never minimum flatness. Consequences:

1. Leaderboard-winning models may be geometrically fragile. 2. Rankings may systematically reward sharp minima. 3. Downstream teams cannot predict a model's update stability.

Practical flatness proxies: stochastic Hessian trace estimates (Hutchinson), SAM loss values $\max_{\|\epsilon\| \le \rho} \mathcal{L}(\theta+\epsilon)$ , and performance drop under Gaussian parameter perturbation. The paper recommends adding SAM loss to standard pretraining evaluation.

Limitations

SAM costs ~2x compute (two forward-backward passes); mitigations include short mid-training phases, mSAM/ESAM approximations, or cost-free alternatives (larger LR, slow annealing).
Theory assumes continuous gradient flow; Adam's momentum, adaptive step sizes, and second-moment estimates complicate the picture.
Very deep Transformers (e.g., 96 layers) may exhibit emergent geometry per layer; needs verification.

Conclusion

The work redefines what a "good pretrained model" means: minimize loss while maximizing minimum flatness. If flat minima systematically improve downstream stability, the pretraining race may need to shift from "who is lower" to "who is more stable."

Paper details

Title: Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting
Authors: Ishaan Watts, Catherine Li, Sachin Goyal, Jacob Mitchell Springer, Aditi Raghunathan
arXiv: 2605.02105 (May 2026, cs.LG)
Key results: up to 80% less forgetting (small scale), 40% less after 4-bit quantization (OLMo-2-1B)