Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting
For practitioners of pretraining and fine-tuning, this work identifies a simple, scalable intervention to preserve pretrained capabilities during post-training and quantization.
The paper shows that pretraining optimizers biased toward flatter minima, such as Sharpness-Aware Minimization (SAM), large learning rates, and shortened annealing, reduce catastrophic forgetting in downstream tasks by up to 80% across model sizes from 20M to 150M parameters, and by 31-40% for a 1B model.
Pretraining optimizers are tuned to produce the strongest possible base model, on the assumption that a stronger starting point yields a stronger model after subsequent changes like post-training and quantization. This overlooks the geometry of the base model which controls how much of the base model's capabilities survive subsequent parameter updates. We study three pretraining optimization approaches that bias optimization toward flatter minima: Sharpness-Aware Minimization (SAM), large learning rates, and shortened learning rate annealing periods. Across model sizes ranging from 20M to 150M parameters, we find that these interventions consistently improve downstream performance after post-training on five common datasets with up to 80% less forgetting. These principles hold at scale: a short SAM mid-training phase applied to an existing OLMo-2-1B checkpoint reduces forgetting by 31% after MetaMath post-training and by 40% after 4-bit quantization.