💬 NLP

Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting

This paper demonstrates that biasing pretraining optimization toward flatter minima through methods like Sharpness-Aware Minimization, large learning rates, or shortened annealing periods significantly mitigates catastrophic forgetting and improves downstream performance across various model sizes and post-training tasks.

Original authors: Ishaan Watts, Catherine Li, Sachin Goyal, Jacob Mitchell Springer, Aditi Raghunathan

Published 2026-05-06

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ishaan Watts, Catherine Li, Sachin Goyal, Jacob Mitchell Springer, Aditi Raghunathan

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are training a master chef (the AI model) by having them taste millions of different dishes (pretraining). The goal is to make them the best possible chef before you send them to work in a specific restaurant kitchen (post-training/fine-tuning).

For a long time, researchers thought the only thing that mattered was making the chef taste the best during that initial training phase. They assumed that if the chef started out as a "perfect" generalist, they would automatically be a great specialist later, no matter what changes happened in the kitchen.

The Problem: The "Rigid" Chef
This paper argues that assumption is wrong. It turns out that some chefs are trained to be so rigidly perfect at the general tasting that they break easily when asked to do something new.

Think of it like a ball sitting in a valley.

The Standard Approach (AdamW): This trains the chef to sit in a very deep, narrow, sharp valley. They are very good at staying exactly where they are. But if you nudge them slightly (like asking them to learn a new recipe or compressing their memory to fit in a smaller pocket), they roll out of the valley immediately. They "forget" everything they knew.
The Paper's Solution (Sharpness-Aware): This trains the chef to sit in a wide, flat valley. They might not be at the absolute deepest point of the valley, but they are surrounded by flat ground. If you nudge them, they don't fall out; they just roll a little bit and stay safe. They are "robust."

What the Researchers Did
The team tested three ways to train these "flat valley" chefs instead of "sharp valley" chefs:

The "Shake" Method (SAM): They used a special technique called Sharpness-Aware Minimization. Imagine that while the chef is learning, you gently shake the table. If the chef can still taste the food correctly while the table is shaking, they are learning to be stable. This forces the model to find those wide, flat valleys.
The "Fast Lane" Method (High Learning Rates): They made the chef learn faster and more aggressively at the start. This is like running a race; if you run fast enough, you naturally avoid getting stuck in tiny, deep holes in the road.
The "Quick Stop" Method (Shorter Annealing): Usually, training ends by slowly slowing the chef down. The researchers found that stopping the slowdown early (keeping the pace up for longer) helped the chef stay in that stable, flat zone.

The Results
They tested this on models of different sizes (from small to medium-large) and found:

Less Forgetting: When these "flat valley" models were later asked to learn specific tasks (like math or coding) or had their memory compressed (quantization) to run faster, they forgot much less of their original knowledge.
The "Mid-Training" Hack: You don't need to use this special "shaking" method for the entire training process, which is expensive. They found that just using it for the last 10% of training (the "annealing" phase) gave almost all the benefits for a tiny fraction of the cost.
Real-World Proof: They tested this on a large 1-billion-parameter model (OLMo-2-1B). Even though the model trained with this method started out slightly "weaker" at general tasks, it became 31% better at remembering its skills after learning math, and 40% better at keeping its skills after being compressed for efficiency.

The Big Takeaway
The paper concludes that we shouldn't just train AI to be the "best" at the starting line. We should train it to be flexible and stable so it doesn't break when we change its environment later.

It's the difference between training a gymnast to hold a perfect, rigid pose (which breaks if the floor tilts) versus training them to have a strong, balanced center of gravity (which lets them adapt to any new move without falling). By focusing on this "balance" (flatness) rather than just "perfection" (sharpness), we get AI models that are much better at learning new things without forgetting the old ones.

Technical Summary: Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting

Problem Statement

Current pretraining optimization strategies are typically tuned to maximize base-model quality, measured by pretraining loss or benchmark performance, under the assumption that a stronger base model yields a stronger model after subsequent modifications like fine-tuning (post-training) or quantization. However, recent work indicates that this assumption can fail: extending pretraining beyond a certain point may improve pretraining loss while degrading downstream performance after post-training.

The core issue identified is that base-model evaluation ignores the stability of a model's capabilities under parameter updates. Models that are highly sensitive to these updates suffer from "catastrophic forgetting," losing pretrained abilities regardless of their initial strength. This sensitivity is linked to the geometry of the loss landscape; specifically, models converging to "sharp" minima are more susceptible to performance degradation when parameters are perturbed during downstream tasks.

Methodology

The authors investigate three pretraining optimization approaches designed to bias the model toward flatter minima, thereby reducing sensitivity to local parameter perturbations:

Sharpness-Aware Minimization (SAM): An explicit optimization method that penalizes loss curvature. SAM seeks parameters $\theta$ that minimize the maximum loss within a neighborhood of radius $\rho$ , effectively finding minima that remain low-loss under small perturbations.
Increased Peak Learning Rates: An implicit approach based on the "Edge of Stability" phenomenon, where larger learning rates are hypothesized to implicitly cap the maximum eigenvalue of the Hessian, thereby reducing loss sharpness.
Shortened Learning Rate Annealing: Modifying the learning rate schedule to reduce the duration of the decay phase at the end of training, motivated by the finding that late-stage dynamics significantly influence the final geometry of the solution.

The study evaluates these interventions across model sizes ranging from 20M to 150M parameters and validates them at scale on a 1B parameter model (OLMo-2-1B). The evaluation focuses on two primary tradeoffs:

Learning-Forgetting Tradeoff: The balance between acquiring new capabilities (fine-tuning loss) and retaining pretrained capabilities (pretraining loss after fine-tuning).
Compression-Forgetting Tradeoff: The degradation of capabilities when the model is compressed via quantization (e.g., 4-bit).

The authors also analyze the Hessian to verify if these methods actually reduce directional sharpness along the fine-tuning direction and whether the second-order Taylor approximation of loss holds for fine-tuning perturbations.

Key Contributions and Results

1. Improved Learning-Forgetting Tradeoff

Across over 80 pretraining runs and 3,500 fine-tuning experiments, the authors find that sharpness-aware interventions consistently improve the learning-forgetting tradeoff:

SAM Performance: SAM-pretrained models achieve a superior Pareto frontier compared to standard AdamW. For example, on the StarCoder dataset, SAM reduces forgetting by 80% at matched fine-tuning loss.
Scaling Effects: The advantage of SAM over AdamW grows as the pretraining token budget increases. While AdamW checkpoints become increasingly sensitive with more training (catastrophic overtraining), SAM mitigates this, showing the strongest benefits in high-token regimes.
Model Size: The benefits persist and even grow with model scale (20M to 150M parameters) when controlling for token-per-parameter ratios.

2. Robustness Beyond Fine-Tuning

The robustness conferred by sharpness-aware pretraining extends to other downstream perturbations:

Quantization: SAM-trained models exhibit significantly less performance degradation under 4-bit quantization. In high-token budget regimes, SAM yields 2–3× less quantization-induced loss increase compared to the baseline.
Gaussian Perturbations: Models trained with SAM are less sensitive to isotropic Gaussian noise added to weights, indicating a more robust local geometry.

3. Implicit Methods and Practical Scalability

Learning Rate Dynamics: Both higher peak learning rates and shorter annealing periods improve the learning-forgetting tradeoff, even when they result in slightly worse base-model pretraining loss.
Mid-Training Intervention: To address the computational cost of SAM (which roughly doubles per-step compute), the authors propose applying SAM only during the learning rate annealing phase.
- In experiments on OLMo-2-1B, applying SAM only during the final 10% of training (mid-training) reduced forgetting by 31% after MetaMath post-training and 40% after 4-bit quantization, compared to the standard baseline. This approach captures most of the benefit with negligible additional compute.

4. Hessian Analysis

The paper validates its theoretical intuition by analyzing the Hessian:

Quadratic Approximation: The second-order Taylor expansion of the loss serves as a reasonable upper bound for forgetting, particularly for SAM-trained models.
Directional Sharpness: Both SAM and larger peak learning rates explicitly reduce the directional sharpness ( $\Delta_{FT}^\top H \Delta_{FT} / \|\Delta_{FT}\|^2$ ) along the fine-tuning direction. The reduction in forgetting is driven by this flattening of the loss landscape rather than a reduction in the step size of fine-tuning.

Significance and Claims

The paper claims that parameter sensitivity to post-training updates is a critical, yet overlooked, link between pretraining dynamics and downstream performance. The authors argue that the standard practice of optimizing solely for base-model quality is insufficient.

Key claims include:

Redefining Pretraining Goals: Pretraining recipes should be selected for both base-model quality and low sensitivity to downstream parameter shifts.
Optimization for the Pipeline: Language model pretraining should be optimized for the full model-development pipeline (including fine-tuning, alignment, and compression) rather than for the base checkpoint in isolation.
Practical Impact: Practitioners can capture significant robustness benefits by deploying sharpness-aware techniques (like SAM) selectively during late training (annealing) rather than throughout the entire pretraining process, making the approach scalable.

The authors conclude that treating adaptability and low sensitivity as first-class evaluation criteria can make pretrained models stronger starting points for future modifications, even if the base model itself appears slightly weaker by traditional pretraining loss metrics.