Sharpness-Aware Pretraining Mitigates Catastrophic Forgetting

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.

Large language models are first trained on massive datasets (a phase called pretraining) and then adapted for specific tasks like answering questions or writing code (a phase called post-training). A common assumption is that a stronger pretrained model will always lead to a stronger final model. But we found this isn't always true: adapting a model for a new task can cause it to forget what it learned during pretraining, and a model that looks stronger before post-training can actually perform worse after it. We show that what matters is not just how well a model learns during pretraining, but how sensitive its parameters are to later modifications. By training with methods that flatten the model's loss landscape (making it robust to small parameter changes), we produce models that forget significantly less when adapted for new tasks or compressed for faster deployment. These methods include a technique called Sharpness-Aware Minimization, as well as simpler adjustments like using larger learning rates or shorter training cooldown periods. Our results suggest that pretraining recipes should be optimized not just for raw base model performance, but for how well the model holds up after post-training. This is a shift in how we should think about building language models.