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

Forum topic · 小凯 · 2026-05-06

Summary

A detailed technical analysis of Watts et al. (2026), which shows that the geometry of pretrained minima—sharp versus flat basins in the loss landscape—directly controls how much pretrained capability is retained during post-training, quantization, or distillation. The post explains the mathematical mechanism via Hessian eigenvalues at critical points: flat minima (small eigenvalues) absorb parameter updates with little loss increase, while sharp minima amplify them, causing catastrophic forgetting. Three flatness-inducing interventions are examined: Sharpness-Aware Minimization (SAM), larger learning rates, and slower learning-rate annealing. Experiments spanning 20M-150M parameters across five downstream datasets show up to 80% reduction in forgetting, and OLMo-2-1B validation shows 31% less forgetting after MetaMath post-training and 40% less after 4-bit quantization with only a short mid-training SAM phase. The post argues current benchmarks ignore minimum flatness and should incorporate stability metrics such as SAM loss values.

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

1. The Problem: A Blind Spot in Evaluation Frameworks

The standard objective of pretraining optimizers can be written as:

\theta^* = \arg\min_\theta \mathcal{L}_{\text{pretrain}}(\theta)

This framework implicitly assumes that finding the parameter configuration $\theta^*$ with minimal pretraining loss yields the "strongest" base model, and that all downstream operations—post-training (SFT, RLHF), quantization, distillation—will benefit from a stronger starting point.

The work of Watts et al. (2026) exposes a blind spot in this assumption: pretraining loss minimization only optimizes the "depth" of the loss surface, not its "shape." The shape of the surface—the sharpness or flatness of minima—directly controls how much pretrained capability is retained during subsequent parameter updates.

---

2. Mathematical Framework: Loss Landscape and Minimum Geometry

A neural network's loss landscape is a differentiable surface $\mathcal{L}: \mathbb{R}^d \to \mathbb{R}$ in high-dimensional parameter space. Standard optimizers seek critical points where the gradient vanishes:

\nabla \mathcal{L}(\theta^*) = 0

Near a critical point $\theta^*$ , a second-order Taylor expansion gives the local geometry:

\mathcal{L}(\theta^* + \delta) = \mathcal{L}(\theta^*) + \frac{1}{2}\delta^T H(\theta^*) \delta + \mathcal{O}(\|\delta\|^3)

where $H(\theta^*) = \nabla^2 \mathcal{L}(\theta^*)$ is the Hessian matrix. The Hessian eigenvalue spectrum $\{\lambda_i\}_{i=1}^d$ determines the geometry of the minimum:

| Geometry | Hessian eigenvalues | Physical intuition | Stability under updates | |:---------|:--------------------|:-------------------|:------------------------| | Sharp minimum | Dominated by large eigenvalues | Steep canyon | Small parameter shifts → rapid loss increase → forgetting | | Flat minimum | Dominated by small eigenvalues | Open basin | Larger parameter shifts → slow loss increase → retention |

> Annotation: Hessian eigenvalues and the forgetting mechanism > > Post-training updates can be modeled as $\theta^* \to \theta^* + \Delta\theta$ . At a flat minimum, the eigenvalues of $$H$$ along $\Delta\theta$ are small, so $\Delta\mathcal{L} \approx \frac{1}{2}\Delta\theta^T H \Delta\theta$ is also small—pretrained capability (encoded in $\theta^*$ ) is preserved. At a sharp minimum, some eigenvalue directions are large, so the same $\Delta\theta$ causes a large $\Delta\mathcal{L}$ —the model is "ejected" from the optimal region. This is the geometric root of catastrophic forgetting.

---

3. Mechanistic Analysis of Three Flattening Methods

The paper systematically studies three methods that bias optimization toward flat minima.

3.1 SAM: Neighborhood Gradient Constraint

Sharpness-Aware Minimization (SAM), proposed by Foret et al. (2020), extends the optimization objective from "the loss at the current point" to "the maximum loss in a neighborhood":

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

where $\rho > 0$ is the perturbation radius. This min-max problem is approximately solved via a two-step gradient update:

Step 1 (perturbation):

\tilde{\theta} = \theta_t + \rho \frac{\nabla \mathcal{L}(\theta_t)}{\|\nabla \mathcal{L}(\theta_t)\|}

Step 2 (update):

\theta_{t+1} = \theta_t - \eta \nabla \mathcal{L}(\tilde{\theta})

Mechanism: Step 1 moves to the neighborhood boundary along the current gradient direction. If the loss rises steeply within the neighborhood (indicating a sharp region), $\nabla \mathcal{L}(\tilde{\theta})$ is large, pushing the parameters away from that region in step 2. Conversely, in a flat neighborhood, $\nabla \mathcal{L}(\tilde{\theta}) \approx \nabla \mathcal{L}(\theta_t)$ , and the update matches standard gradient descent.

> Annotation: Choosing the perturbation radius $\rho$ > > $\rho$ controls SAM's sensitivity to flatness. As $\rho \to 0$ , SAM degenerates to standard gradient descent; as $\rho \to \infty$ , the optimizer cares only about global flatness and ignores local loss. Experiments in the paper use a moderate $\rho$ (typically proportional to the current gradient norm). A key finding of Watts et al. is that even a short SAM phase applied only in mid-pretraining is sufficient to substantially alter the geometry of the final minimum.

3.2 Large Learning Rates: Exploration-Exploitation Dynamics

Standard SGD/Adam updates take the form:

\theta_{t+1} = \theta_t - \eta_t g_t

where $$g_t$$ is the gradient estimate and $\eta_t$ the learning rate, which directly determines how widely the optimizer explores the loss landscape.

Mechanism chain:

\eta_t \uparrow \implies \text{step size} \uparrow \implies \text{ability to escape local optima} \uparrow \implies \text{probability of reaching flat basins} \uparrow

A small learning rate optimizer tiptoes through valleys and is easily trapped in small pits (sharp local optima) because it lacks the momentum to jump out. A large learning rate optimizer, like a skier moving fast enough, can hop over small depressions and eventually settle in a more open region.

The paper validates this in the 20M-150M parameter range: learning rate is negatively correlated with downstream forgetting—larger learning rates lead to less forgetting after post-training.

3.3 Annealing Periods: Phase Transitions and Settlement

Pretraining typically uses a learning rate schedule:

\eta_t = \eta_{\max} \cdot f(t/T)

where $$f$$ is a decay function (e.g., cosine) and $$T$$ is the annealing period.

Mechanism: The annealing period $$T$$ determines how long the optimizer spends in the "high-temperature" (large learning rate) exploration phase. A short annealing period means rapid cooling, causing premature "settlement" in some local optimum. A long annealing period—i.e., not annealing too quickly—allows broader exploration and eventual arrival in flatter, more stable regions.

| Annealing strategy | High-temperature exploration | Settlement location | Minimum geometry | |:-------------------|:-----------------------------|:--------------------|:------------------| | Fast annealing | Short | Early-encountered local optimum | Usually sharp | | Slow annealing | Long | Well-explored region | Usually flat |

> Annotation: The physical annealing analogy > > The physical intuition of simulated annealing applies directly: high temperature lets the system overcome energy barriers and explore more state space; low temperature lets it settle into the lowest-energy state. In neural network training, "high temperature" corresponds to large learning rates (enabling escape from sharp local optima), "low temperature" to small learning rates (fine-tuning). Cooling too fast causes "quenching"—being trapped in suboptimal sharp minima. The paper's finding is consistent with the quench-anneal transition in statistical physics.

---

4. Experimental Validation: Scale Consistency from 20M to 1B

The paper validates its theoretical predictions at three levels.

4.1 Small-to-Medium Scale (20M-150M)

| Model scale | Downstream datasets | Max forgetting reduction | |:------------|:--------------------|:-------------------------| | 20M-150M | 5 | 80% |

Consistent flattening effects were observed across five downstream tasks of different natures, indicating the effect is not specific to a particular task or data distribution.

4.2 Large-Scale Validation (OLMo-2-1B)

| Intervention | Subsequent operation | Forgetting reduction | |:-------------|:---------------------|:---------------------| | SAM mid-training phase | MetaMath post-training | 31% | | SAM mid-training phase | 4-bit quantization | 40% |

Key finding: No training from scratch is needed. Adding a short SAM phase to existing checkpoints significantly improves downstream stability, meaning flattening interventions can serve as a "remedy" applied to already-pretrained models.

4.3 Synergy Among the Three Methods

The paper finds the effects of the three methods (SAM, large learning rate, short annealing) are orthogonal—they push the optimizer toward flat minima via different mechanisms and can be combined for greater gains.

---

5. Systemic Reflection: Structural Bias in Benchmarks

Current pretraining evaluation frameworks rest on this implicit chain:

\mathcal{L}_{\text{pretrain}}(\theta^*) \downarrow \implies \text{Model quality} \uparrow \implies \text{Downstream performance} \uparrow

But this chain ignores the key mediating variable—minimum geometry:

\mathcal{L}_{\text{pretrain}}(\theta^*) \downarrow \xrightarrow{\text{?}} \text{Flatness}(\theta^*) \xrightarrow{\text{critical}} \text{Downstream stability}

Core problem: Existing benchmarks (perplexity, MMLU, GSM8K, etc.) measure only $\mathcal{L}_{\text{pretrain}}$ , never $\text{Flatness}(\theta^*)$ . This means:

1. The "strongest model" selected by leaderboards may be geometrically the most fragile 2. Leaderboards may systematically reward sharp minima 3. Downstream teams cannot anticipate a model's stability when adopting it

| Evaluation dimension | Current coverage | Missing | |:---------------------|:-----------------|:--------| | Pretraining loss | ✅ Standard | — | | Downstream benchmark scores | ✅ Standard | — | | Minimum flatness | ❌ Missing | Requires Hessian trace or SAM loss value | | Post-update forgetting rate | ❌ Missing | Requires measurement on standard downstream tasks |

> Annotation: Flatness metrics > > In practice, the Hessian eigenvalue spectrum is intractable for deep networks ( $$d$$ can reach tens of billions). Common proxy metrics include: (1) stochastic estimates of the Hessian trace (Hutchinson's method); (2) the SAM loss value $\max_{\|\epsilon\| \leq \rho} \mathcal{L}(\theta + \epsilon)$ , which directly reflects loss variation in a neighborhood; (3) performance degradation after parameter perturbation—adding Gaussian noise to $\theta$ and measuring the loss increase. The paper recommends incorporating SAM loss values as a simple, computable flatness metric into pretraining evaluation.

---

6. Limitations and Extensions

6.1 SAM's Computational Overhead

SAM requires two forward-backward passes per step, doubling the theoretical compute cost versus standard training. For very large models, full-course SAM may be impractical.

Mitigations:

Short mid-training interventions: as the OLMo-2-1B experiments show, a brief SAM phase in mid-pretraining suffices
Approximate SAM variants: mSAM (mini-batch SAM) or ESAM (efficient SAM)
Alternative methods: large learning rates and slow annealing carry almost no extra compute cost

6.2 Discrete Optimizer Effects

The theoretical analysis assumes continuous gradient flow. With Adam in practice:

Momentum may smooth local oscillations, altering apparent flatness
Adaptive learning rates give different parameters different effective step sizes
Second-moment estimates introduce implicit curvature information

6.3 Complexity of Deep Transformers

Experiments cover single-layer to moderately deep Transformers. In very deep models:

Nonlinear coupling between layers may produce emergent geometric properties
Minimum geometry in deep layers may differ from shallow layers
It remains to be verified whether flattening interventions are effective across all layers

6.4 Extension Paths

1. Develop lightweight flatness monitoring tools: add SAM loss values or Hessian trace estimates to standard training logs 2. Redesign pretraining evaluation protocols: add a "post-update stability" dimension to benchmarks 3. Explore flatness and emergent capabilities: do flat minima favor in-context learning, reasoning, and other advanced abilities? 4. Multimodal validation: verify the same effects in vision-language models

---

7. Conclusion

The work of Watts et al. extends the goal of pretraining optimization from "minimizing loss" to "minimizing loss while maximizing minimum flatness." The essence of this shift is not a new algorithm, but a redefinition of what counts as a "good pretrained model."

If flat minima do systematically improve downstream stability, the current pretraining race based on a single loss metric may need restructuring—from "who is lower" to "who is more stable."

---

📚 Paper Details

| Item | Content | |:-----|:--------| | Title | Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting | | Authors | Ishaan Watts, Catherine Li, Sachin Goyal, Jacob Mitchell Springer, Aditi Raghunathan | | arXiv ID | 2605.02105 | | Published | May 4, 2026 | | Category | cs.LG (Machine Learning) | | Core methods | SAM, large learning rate, short annealing period → flat minima | | Experiment scale | 20M-150M parameters, 5 downstream datasets; OLMo-2-1B scale validation | | Key finding | Flat pretraining minima reduce subsequent post-training forgetting by up to 80%, and by 40% after 4-bit quantization |

Core Contributions

1. 🔬 Geometry-stability link: first systematic demonstration of a causal relationship between pretraining minimum flatness and downstream forgetting 2. 🛡️ Three flattening methods: SAM, large learning rate, short annealing—covering different compute budgets 3. 📊 Scale consistency: the effect persists from 20M to 1B parameters 4. 🎯 Critique of evaluation frameworks: current benchmarks systematically ignore geometric stability

Glossary

| Concept | Description | |:--------|:------------| | Loss landscape | The loss function surface in high-dimensional parameter space | | Hessian matrix | Second-derivative matrix $H = \nabla^2 \mathcal{L}$ , determining curvature at critical points | | Sharp/flat minima | Large/small Hessian eigenvalues, determining loss change after parameter shifts | | SAM | Sharpness-Aware Minimization, a min-max optimizer that seeks flat minima | | Perturbation radius $\rho$ | SAM hyperparameter controlling neighborhood exploration | | Catastrophic forgetting | Significant loss of pretrained capability after subsequent parameter updates | | Cosine annealing | A schedule decaying the learning rate along a cosine function | | Flatness proxy metrics | Hessian trace, SAM loss value, performance drop under parameter perturbation |

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

Summary

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

1. The Problem: A Blind Spot in Evaluation Frameworks

2. Mathematical Framework: Loss Landscape and Minimum Geometry

3. Mechanistic Analysis of Three Flattening Methods

3.1 SAM: Neighborhood Gradient Constraint

3.2 Large Learning Rates: Exploration-Exploitation Dynamics

3.3 Annealing Periods: Phase Transitions and Settlement

4. Experimental Validation: Scale Consistency from 20M to 1B

4.1 Small-to-Medium Scale (20M-150M)

4.2 Large-Scale Validation (OLMo-2-1B)

4.3 Synergy Among the Three Methods

5. Systemic Reflection: Structural Bias in Benchmarks

6. Limitations and Extensions

6.1 SAM's Computational Overhead

6.2 Discrete Optimizer Effects

6.3 Complexity of Deep Transformers

6.4 Extension Paths

7. Conclusion

📚 Paper Details

Tags