Pretraining Minima Geometry and Downstream Stability: From Loss Landscape to Catastrophic Forgetting

Forum topic · 小凯 · 2026-05-06

Summary

This post analyzes Watts et al. (2026), "Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting" (arXiv:2605.02105), which shows that pretraining loss minimization alone overlooks a critical variable: the geometry of the loss landscape minimum. Sharp minima amplify parameter drift into catastrophic forgetting after post-training, quantization, or distillation, while flat minima preserve pretrained capabilities. The paper demonstrates that three flatness-inducing methods—Sharpness-Aware Minimization (SAM), large learning rates, and slower annealing schedules—significantly improve downstream stability. Experiments on 20M–150M parameter models across five downstream datasets show up to 80% reduction in forgetting, and mid-training SAM interventions on OLMo-2-1B reduce forgetting by 31% after MetaMath post-training and 40% after 4-bit quantization, without retraining from scratch. The effects are orthogonal and combinable. The post walks through the mathematical framework (Hessian eigenvalues at critical points), mechanistic analysis of each method, and critiques current benchmarks for measuring only loss while ignoring flatness—arguing that leaderboard winners may be geometrically fragile. It also covers limitations (SAM's 2x compute cost, Adam's implicit curvature effects, deep Transformer complexity) and practical proxy metrics for flatness such as Hessian trace estimation and SAM loss values.

Problem Structure: The Blind Spot in Evaluation Frameworks

The standard objective for pretraining optimizers is:

\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, on which post-training (SFT, RLHF), quantization, and distillation can build—expecting a stronger start to produce a stronger endpoint.

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

---

Mathematical Framework: Loss Landscape and Minima 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 | Intuition | Stability under updates | |:---------|:--------------------|:----------|:------------------------| | Sharp minimum | Large eigenvalues dominate | Steep canyon | Small parameter shifts → loss spikes → forgetting | | Flat minimum | Small eigenvalues dominate | Open basin | Larger shifts → gradual 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 small—pretrained capability (encoded in $\theta^*$ ) is preserved. At a sharp minimum, large eigenvalues along some directions mean the same $\Delta\theta$ causes a large $\Delta\mathcal{L}$ —pretrained capability is "ejected" from the optimal region. This is the geometric root of catastrophic forgetting.

---

Mechanistic Analysis of Three Flatness-Inducing Methods

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

3.1 SAM: Neighborhood Gradient Constraint

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

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

where $\rho > 0$ is the perturbation radius. The min-max problem is approximated by 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 loss rises sharply in the neighborhood (steep surroundings), $\nabla \mathcal{L}(\tilde{\theta})$ is large, pushing parameters away from the region. If the neighborhood is flat, $\nabla \mathcal{L}(\tilde{\theta}) \approx \nabla \mathcal{L}(\theta_t)$ , and the update resembles standard gradient descent.

> Annotation: Choice of 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 only cares about global flatness and ignores local loss. Experiments use a moderate $\rho$ (typically proportional to the current gradient norm). The key finding: even a short SAM phase in mid-pretraining is sufficient to significantly alter the geometry of the final minimum.

3.2 Large Learning Rates: Exploration–Exploitation Dynamics

The standard SGD/Adam update is $\theta_{t+1} = \theta_t - \eta_t g_t$ . The learning rate directly determines the optimizer's exploration range:

\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 is like walking carefully in a valley—easy to get stuck in a small pit (a sharp local optimum). A large learning rate is like skiing—at sufficient speed, small potholes can be crossed, ending 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 Schedule: Phase Transition and Settling

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: $$T$$ determines the time spent in the "hot" (high learning rate) exploration phase. Fast annealing means rapid cooling—the optimizer settles prematurely in some local optimum. Slow annealing (extended high-temperature exploration) allows broader exploration of parameter space, ending in flatter, more stable regions.

| Annealing strategy | Hot exploration time | Settling location | Minima geometry | |:---------|:-------------|:---------|:-----------| | Fast annealing | Short | Early-encountered local optimum | Usually sharp | | Slow annealing | Long | Region after broad exploration | Usually flat |

> Annotation: Physical annealing analogy > > The physical intuition of simulated annealing applies directly: high temperature lets a system overcome energy barriers and explore more of the state space; low temperature stabilizes it at the lowest-energy state. In neural network training, "high temperature" corresponds to large learning rates (crossing sharp local optima), "low temperature" to small rates (fine-tuning). Cooling too quickly (short annealing) causes "quenching"—trapping the system in suboptimal sharp minima, consistent with the quench–anneal transition in statistical physics.

---

Experimental Validation: Scale Consistency from 20M to 1B

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

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

Consistent flatness effects across five distinct downstream tasks indicate this is not a task- or data-distribution-specific artifact.

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

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

Key finding: No from-scratch retraining is needed. Adding a short SAM phase to existing checkpoints significantly improves downstream stability—flatness interventions can serve as a "remedial measure" for already-pretrained models.

4.3 Synergy of the Three Methods

The effects of SAM, large learning rates, and slow annealing are orthogonal—they push the optimizer toward flat minima through different mechanisms and can be combined for greater gains.

---

Systemic Reflection: Structural Bias in Benchmarks

The current pretraining evaluation framework assumes:

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

But this chain ignores a key mediating variable—minima 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.) only measure $\mathcal{L}_{\text{pretrain}}$ and never $\text{Flatness}(\theta^*)$ . This means:

1. The "strongest models" selected may be the most geometrically fragile 2. Model leaderboards may systematically reward sharp minima 3. Downstream teams cannot predict a model's stability before using it

| Evaluation dimension | Currently covered | Missing | |:---------|:---------|:-----------| | Pretraining loss | ✅ Standard | — | | Downstream benchmark scores | ✅ Standard | — | | Minima flatness | ❌ Missing | Needs Hessian trace or SAM loss value | | Post-update forgetting rate | ❌ Missing | Needs measurement on standard downstream tasks |

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

---

Limitations and Extensions

6.1 SAM's Compute Overhead

SAM requires two forward-backward passes, roughly 2x the cost of standard training—potentially impractical for frontier-scale models.

Mitigations:

Mid-training short interventions: as shown by OLMo-2-1B, a short SAM phase in mid-pretraining suffices
Approximate SAM: mSAM (mini-batch SAM) or ESAM (efficient SAM)
Alternatives: large learning rates and slow annealing have almost no extra compute cost

6.2 Discrete Optimizer Effects

The theoretical analysis assumes continuous gradient flow. With Adam:

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

6.3 Complexity of Deep Transformers

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

Nonlinear coupling between layers may produce emergent geometric properties
Deep-layer minima geometry may differ from shallow layers
Whether flatness interventions work uniformly across layers needs verification

6.4 Extension Paths

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

---

Conclusion

This work extends the goal of pretraining optimization from "minimizing loss" to "minimizing loss while maximizing minimum flatness." The real contribution is not a new algorithm but a redefinition of what constitutes a "good pretrained model."

If flat minima indeed systematically improve downstream stability, the current pretraining race—judged by 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 rates, short/slow annealing → flat minima | | Experiment scale | 20M–150M parameters, 5 downstream datasets; OLMo-2-1B scale-up validation | | Key findings | Flat pretraining minima reduce post-training forgetting by up to 80%; 40% reduction after 4-bit quantization |

Core contributions

1. 🔬 Geometry–stability link: first systematic demonstration of the causal relationship between pretraining minima flatness and downstream forgetting 2. 🛡️ Three flatness methods: SAM, large learning rates, slow annealing—covering different compute budgets 3. 📊 Scale consistency: effects persist from 20M to 1B parameters 4. 🎯 Evaluation critique: current benchmarks systematically ignore geometric stability

Concept 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 critical-point curvature | | Sharp/flat minima | Large/small Hessian eigenvalues, governing loss change under parameter shifts | | SAM | Sharpness-Aware Minimization, a min-max optimization for flat minima | | Perturbation radius $\rho$ | SAM hyperparameter controlling neighborhood exploration | | Catastrophic forgetting | Significant loss of pretrained capability after subsequent parameter updates | | Cosine annealing | Learning rate schedule decaying along a cosine function | | Flatness proxy metrics | Hessian trace, SAM loss value, performance drop under parameter perturbation |