Abstract: Fine-tuning a pre-trained model on a downstream task often degrades its original capabilities, a phenomenon known as "catastrophic forgetting". This is especially an issue when one does not have access to the data and recipe used to develop the pre-trained model. Under this constraint, most existing methods for mitigating forgetting are inapplicable. To address this challenge, we propose a *sample weighting scheme for the fine-tuning data* solely based on the pre-trained model's losses. Specifically, we upweight the easy samples on which the pre-trained model's loss is low and vice versa to limit the drift from the pre-trained model. Our approach is orthogonal and yet complementary to existing methods; while such methods mostly operate on parameter or gradient space, we concentrate on the sample space. We theoretically analyze the impact of fine-tuning with our method in a linear setting, showing that it stalls learning in a certain subspace, which inhibits overfitting to the target task. We empirically demonstrate the efficacy of our method on both language and vision tasks. As an example, when fine-tuning Gemma 2 2B on MetaMathQA, our method results in only a $0.8$% drop in accuracy on GSM8K (another math dataset) compared to standard fine-tuning, while preserving $5.4$% more accuracy on the pre-training datasets.

Lay Summary:

Pretrained AI models are powerful, but fine-tuning them on new tasks often causes them to forget what they originally learned — a problem known as catastrophic forgetting. This is especially challenging when the original training data is unavailable, as is often the case with publicly released models.To address this, we propose a simple method called FLOW. Unlike traditional approaches focusing on difficult examples during fine-tuning, FLOW does the opposite: it upweights “easy” examples — cases that the model already handles well. This helps preserve the model’s original capabilities while still allowing it to adapt to the new task.FLOW is easy to use, requires no extra parameters, and relies only on a basic calculation using the model’s pre-trained performance. Despite its simplicity, it consistently outperforms standard fine-tuning and other relevant baselines across both vision and language tasks. For example, when fine-tuning for math reasoning, FLOW achieved similar accuracy on new math problems while retaining 5% more accuracy on the original general knowledge tasks than standard fine-tuning.Because FLOW doesn’t require access to original training data, it offers a practical solution for retaining valuable model capabilities while fine-tuning on new tasks.