Upweighting Easy Samples in Fine-Tuning Mitigates Forgetting

We thank the reviewer for their detailed evaluation and encouraging feedback. We respond to your questions and feedback below:

Performance for larger/different models and larger datasets: We believe that our language experiments are indeed large-scale, and the size of our vision datasets is comparable to prior works such as [A, B]. However, we appreciate your feedback, and in light of that, we performed extra experiments on larger models (ViT-B/16 and CLIP ViT-B/32) and a larger dataset (Food 101). We provide the details and results below:

• ViT-B/16 model on Food101: As requested, we performed experiments with a ViT pre-trained on ImageNet-1k (IN-1K). Also, Food 101 has twice the number of samples (101k samples) as CIFAR (50k samples), the image sizes are very different, and are corrupted with random noise. We follow a similar experimental setup to section 5.1 in the paper. Due to a lack of space, we are omitting the rest of the details here. Here are the results:

FLOW exhibits the same behavior as in our previous experiments in the paper – meaningfully outperforming the baselines.

• CLIP ViT-B/32: Here, we fine-tune the image encoder of a pre-trained (OpenAI) multi-modal CLIP ViT-B/32 model. We use zero-shot ImageNet-1K accuracy as our pre-trained metric and then train a classification head on top of the image encoder of the CLIP model. We follow a similar experimental setup to section 5.1 in the paper, fine-tuning on 7 different downstream tasks. Due to a lack of space, we are omitting the rest of the details here. Please see the results below.

FLOW performs the best by a decent margin.

Distillation paper [1]: Based on the feedback, we compared the distillation-based idea (abbreviated as LwF in [1]) against ours for the new ViT-B/16 experiment on Food101. We used the hyper-parameters for LwF suggested by [1]. We observe that our method FLOW outperforms LwF despite the simplicity; FLOW achieves better performance on the forgetting front, while LwF has higher accuracy on the fine-tuning task.

We believe our method has clear advantages on the efficiency front as well. Note that the LwF method comes with more tunable parameters, additional memory, and evaluation overhead. They have a distillation loss scaling factor $\lambda_o$ and “softening parameter” $T$ that rescales the outputs as $y’_i = (y’_i)^{1/T} / \sum_k^l (y’_k)^{1/T}$ for an $l$-dimensional output vector. In contrast, our method only has $\tau$, for which we use a prescribed value. Second, the logits obtained per sample for the pretrained model body and head must be stored to compute the distillation loss, which is infeasible in the language modeling setting. Our method shows comparable performance with easier tuning and better storage efficiency. Thank you for pointing out this line of work. Combining the distillation idea with our sample selection idea is also an interesting future work.

We hope to have resolved your concerns, and we're happy to discuss further. If you’re satisfied, we sincerely hope you will raise your score!

[A]: Wang et al., 2025. “LiNeS: Post-training Layer Scaling Prevents Forgetting and Enhances Model Merging”

[B]: Wortman et al., 2022. “Robust fine-tuning of zero-shot models”

We thank the reviewer for recognizing the novelties of the simple design of our sample-wise weighting scheme. We appreciate the constructive and encouraging feedback. Please find our responses to your major comments below.

Wider range of fine-tuning tasks, especially for language modeling: While we agree that our language modeling section could benefit from additional fine-tuning domains (e.g., instruction tuning, question answering, safety alignment, etc.), we are unfortunately unable to present any further experiments due to compute and time constraints. More specifically, for example, we don’t have LLM-as-a-judge set up to evaluate instruction tuning.

We appreciate you pointing out several of these fine-tuning domains, but we would like to clarify that we did not intend to focus on this method for specific language tasks. Rather, we have proposed this method oblivious to any particular architectural or data modality, with theoretical intuitions for linear models. However, these are great suggestions for our algorithm’s applications – with potential adaptations – to specific important language tasks. Thanks!

We additionally include results on fine-tuning the image encoder of a pre-trained (OpenAI) multi-modal CLIP ViT-B/32 model. We use zero-shot ImageNet-1K accuracy as our pre-trained metric and then train a classification head on top of the image encoder of the CLIP model. We present their results in the table below. As you see, FLOW performs the best even here by a decent margin.

We follow a similar experimental setup to section 5.1 in the paper, fine-tuning on 7 different downstream tasks. Due to a lack of space, we are omitting the rest of the details here.

Relationship to prior works on factuality (Ghosal et al., 2024; Gekhman et al., 2024) and the effect of training with unfamiliar/unpopular facts: This is an interesting perspective regarding our weighting scheme and on what basis the samples could be “easy” or “hard”. We came up with FLOW with a more general perspective and problem scope in mind, without a specific consideration for language modeling and factuality. However, we do see the connection. Thank you for bringing it to our attention.

Within the context of factuality and LLMs, for a fine-tuning task on some factual knowledge-based dataset, which has samples that are not properly represented in the pre-training distribution, the respective samples could be considered hard, and vice versa. In a data-oblivious setting, as we do not have access to pre-training data, one way (probably the only principled way) would be to rank hardness based on pre-trained losses. Intuitively, such “hard” samples would have high losses, and our algorithm would give them smaller weights. Given the apparent connections, FLOW should intuitively behave in parallel with the findings of (Ghosal et al., 2024; Gekhman et al., 2024). We will include these discussions in the next version of our paper.

Thanks for the positive assessment of our work! We address your questions below.

Ablation for $\tau$: In the table below, we show the pre-training and fine-tuning accuracies as a function of $\tau$ when the model is ResNet-50, the pre-training dataset is ImageNet-1K, and the fine-tuning dataset is Caltech 101. The last column represents the percentile of pre-training losses, which we use as the value of $\tau$.

In our experiments, we used $\tau = 50^\text{th}$ percentile. While smaller $\tau$ works better here, in general, it seems very difficult to “guess” the optimal value of $\tau$ without any access to the pre-training data. But this is an interesting point which we will investigate in the future. Thanks for bringing this up!

Regarding Machine Unlearning (MU): Thanks for pointing out papers on MU! [1, 2] fall within the more prominent category of weighted updates in the parameter space. As we discussed in our paper, there are several ideas in the parameter space for mitigating forgetting. Our approach is orthogonal to such ideas as we perform weighting in the sample space. One could extend the ideas of [1, 2] for forgetting and apply them with our sample-weighting approach (as we did with existing baselines in Section 6.2). Similarly, as you pointed out, one could extend our sample-weighting scheme for MU. These are great future directions! However, it is worth mentioning that while MU and forgetting sound similar, they are not the same in spirit; so it may not be straightforward to extend techniques of one into another. In MU, we deliberately (and ideally, provably) want to induce “forgetting” on some samples, whereas in our context, forgetting is an undesirable side effect that we want to avoid as much as possible. However, we will discuss the connections between MU and [1,2] in the next version of our paper.

Comparison between adaptive and fixed weighing scenarios and also from multi-task finetuning experiments: These are great points! We opted for the fixed weighting scheme due to two reasons: (a) a principled adaptive scheme in the absence of pre-training data was not clear to us (specifically, in the way we derived Alg. 1 in Section 4), and (b) to keep the fine-tuning procedure light-weight. The multi-task finetuning setting may require a bit more thought, especially when the tasks appear sequentially – specifically, how the weights are chosen for each new task. This is left for future work.

Questions For Authors:

1. Potentially, yes! As we discussed above, a combination of our approach and the approach of [1] may lead to better performance.

2. This is a great direction! Potentially, our sample weighting scheme could be extended to the machine unlearning problem.

3. If there are out-of-distribution outliers with large pre-training losses, then our weighting scheme would assign them low weights and so the influence of such outliers would be curtailed. Hence, our algorithm handles such bad samples. On the other hand, if by outliers, you mean in-distribution samples with large pre-training losses, then it seems difficult to do well on such samples in general if we wish to also do reasonably well on the pre-training data, especially without any access to the pre-training data. The premise of our work is that such samples should be considered detrimental to the pre-training performance as they’d make the model drift away from the pre-training weights.

4. Answered at the top.

5. We believe our language experiments are on reasonably large-scale models and datasets. As for vision, please see our response to Reviewer 1DND, where we show additional results with (i) a ViT-B/16 model on Food101 which is much a harder dataset with 101k images, and (ii) CLIP ViT-B/32 (OpenAI) which is pre-trained with the contrastive loss but we fine-tune the image encoder with a cross-entropy loss.

6. The premise of this work is that there is some kind of alignment between the pre-training and fine-tuning tasks; otherwise, it is very hard to find a model that does well on both tasks. Moreover, we assume that the pre-training losses are partially reflective of the “hardness” of the fine-tuning examples; otherwise, in the absence of any pre-training data, this seems like a very hard problem. In the case of a pre-trained generative model, the extension should be straightforward. Consider BERT, which is trained using masked language modeling. We would do linear probing of the [CLS] token’s output embedding, followed by applying Algorithm 2 (Appendix B).

Thanks for the positive assessment of our work! We address your questions below.

Number of fine-tuning epochs: For the language experiments, this was chosen based on existing literature recommendations for fine-tuning epochs. Following the observation of [1], we trained for 2 epochs due to compute constraints. For vision, we selected the number of epochs to achieve a fine-tuning accuracy comparable to that of standard fine-tuning on specific models and tasks, aligning with the performance reported in [2].

Theoretical Claims: Our theory setting is a typical linear setting, commonly used for analysis purposes. However, as mentioned in footnote 2, our insights carry over to wide neural networks following the dynamics of linear models under gradient descent (Lee et al., 2019). For a general choice of $\widetilde{\Sigma}$, we derived the weighted covariance matrix $\widetilde{\Sigma}'$ impacting the dynamics of FLOW in Appendix E (see eq. 32). Note that $\widetilde{\Sigma}'$ depends on $\widetilde{\Sigma}$ and $e = \theta_{\ast} - \widetilde{\theta}_{\ast}$; the latter is the difference between the optimal solutions of the pre-training and fine-tuning tasks.

Unfortunately, as we explained in Remark E.1, it is hard in general to characterize the eigen-spectrum of $\widetilde{\Sigma}'$. But informally, more is the alignment between $e$ ($= \theta_{\ast} - \widetilde{\theta}_{\ast}$) and the eigenvectors of $\widetilde{\Sigma}$, slower the convergence along a direction that is aligned to $e$, inhibiting overfitting to the fine-tuning data. Formalizing this intuition is left for future work.

Questions For Authors:

1. Please see Fig. 3 (in Appendix I) for the distribution of sample weights of Gemma 2 2B, when we train with sequence-wise weights (results in the main paper) as well as token-wise weights (which we did as an ablation in the Appendix). Note that the distribution is Gaussian-like in the sequence-wise case, while it is concentrated in the token-wise case; as a result, the latter didn’t work very well. We’ll also produce similar plots for our vision experiments in the future.

2. Indeed, the samples with high pre-training losses have lower accuracy when using FLOW compared to standard fine-tuning (FT).

In fact, this behavior is in line with the premise of our approach, namely, sacrificing performance on the hard examples from the fine-tuning data to maintain performance on the pre-training data.

3. If the fine-tuning examples have high losses with the pre-trained model ($\theta_{\ast}$), it probably indicates that there is no solution in the vicinity of $\theta_{\ast}$ for which the performance on the fine-tuning data is good. Without any access to the pre-training data, all we can hope to do is to remain "close" to $\theta_{\ast}$ to have decent performance on the pre-training data. In that case, all the methods that we are aware of should struggle to achieve good performance jointly on pre-training and fine-tuning data. Regarding the distribution of the weights, the value of the pre-training losses doesn’t really matter; what matters is their relative spread because the temperature $\tau$ is chosen as the median of the pre-training losses.

[1]: Biderman et al., 2024. “LoRA Learns Less and Forgets Less.”

[2]: Wightman et al., 2021. “Resnet strikes back: An improved training procedure in timm.”