Demystifying Language Model Forgetting with Low-rank Example Associations

We thank the reviewer for the very thoughtful comments and questions.

W1: The crucial experiment in Table 2, which aims to interpret these associations by checking correlations with similarity measures, was conducted on OLMo-1B only

We thank the reviewer for the question and we agree that the results on OLMo-7B should be added to Table 2. We performed this new set of experiments and we will include the results in the final version of the paper.

To summarize, on OLMo-7B, gradient-based similarity measures show similar or lower correlations with the actual forgetting compared to OLMo-1B.

Nevertheless, we also notice that TF-IDF text similarity correlates positively with forgetting OLMo-7B but negatively in OLMo-1B. To better interpret this, we measure how each component ($\alpha_k\beta_k^T$) in the association matrix $Z$ correlates with the TF-IDF text similarity. We highlight components where Spearman's |$\rho$| > 0.05

The results show that some components are significantly more positively or negatively correlated with the text similarity, and how the components (and their correlations to TF-IDF text similarity) differ between OLMo-1B and OLMo-7B models. We will update the paper with the new results and the analysis.

W2: The paper concludes that similarity measures "correlate poorly" with forgetting, dismissing them as weak predictors. However, this conclusion may be too strong.

We thank the reviewer for pointing out the issue. Yes, the correlations of similarity measures) to forgetting (especially TF-IDF and gradient inner products), despite small, is statistically significantly non-zero. Therefore, they can still be predictors of forgetting. We will improve the preciseness of the writing in the final version of the paper.

W3: A more direct and functional evaluation of forgetting would be to measure performance degradation on a specific, previously fine-tuned task (Task A) after the model learns subsequent new tasks (B, C, etc.). By constructing an association matrix between Task A's examples and the new downstream tasks, one could directly quantify the retention of a specific capability, which is a more intuitive and practical measure of forgetting than perplexity changes across a vast pre-training corpus.

We thank the reviewer for the thoughtful comment. Regarding evaluation on downstream tasks, the most relevant results presented in the paper should be forgetting of OLMo-7B-Instruct on the FLAN subset of Tulu, where we use the exact-match drop metrics (instead of perplexity). The association matrix is visualized in Figure 8 in Appendix, and we have forgetting prediction results (in prediction F1) and forgetting mitigation results (in Token F1 retention) included in Table 3 and Figure 5.

In our new set of experiments, we measured forgetting on the LLM leaderboard task examples. We see the similar low-rank structures and will include the visualization and statistics in the final version of the paper. We agree with the reviewer that the part of functional evaluation of forgetting can be further extended. We will expand our discussions about this line of future work.

W4: The method depends on the access of the upstream data

We acknowledge this limitation in our Limitations section. In fact, we consider this limitation inherited from replay-based algorithms for mitigating forgetting. We still decided to study replay-based methods because they are shown to be scalable and effective when applicable in LLMs.

W5: The best-performing offline mitigation variant doubles the fine-tuning computational cost

We acknowledge the extra computational overhead of the offline variant. We advise future practitioners to apply the online variant proposed in the paper in cost sensitive scenarios.

We thank the reviewer again for recognizing the strength of our paper and the thoughtful comments.

We thank the reviewer for the thoughtful comments and the very helpful suggestions.

W1. Limited insight into why low-rank holds

We thank the reviewer for the thoughtful comment. As we have acknowledged in the Limitations section, we do not have a theoretical explanation of the low-rank structure at this moment. We tried to connect it to first-order approximations of forgetting and low-rank subspaces where weight updates can reside, but the actual forgetting does not correlate strongly (Sec. 4.3). We leave theoretical and mechanistic interpretations of forgetting as future works.

At this stage, our paper contributes to behavioral interpretability of the low-rank structure by studying how various factors affect the low-rank approximation (e.g. model size, knowledge coverage of LMs, etc.) with data-driven experiments.

Q1: What inductive biases do you think induce low-rank structure in forgetting matrices?

We thank the reviewer for this insightful question. We expect that significant study is required to answer the question. We believe mechanistic interpretations on how knowledge is updated and encoded in models (e.g. the suggested reference Chen et al. 2024) is relevant. We will note future lines of work in the updated version of the paper.

Informal intuition about the low-rank structure. one intuition of the low-rank structure is an analogy to human’s forgetting process. Compared to acquiring a new skill, forgetting depends less on the material (training data) that humans learn afterwards, resulting in the low-rank structure. We refer to the interference theory of the human forgetting process for this intuition.

W2: The definition of forgetting

Thank you for the thoughtful comment. We choose to examine forgetting with log perplexity increase or exact match drop because (1) they are connected to theoretical approximations of forgetting with first-order techniques (Sec. 3.3), allowing proper comparison and (2) we hoped to study forgetting on upstream examples where we are sure the models have seen before.

We agree with the reviewer that forgetting defined in other ways can have different and interesting practical implications. This includes the "shallow forgetting" (i.e. forgetting that is recovered when twisting the prompts), and the positive side of the forgetting the reviewer suggested in the comment. We agree that including the discussions of these works can position our work paper and inspire future analysis of forgetting under these definitions. We thank the reviewer for the helpful references and we will discuss them in the related works section.

Q2-1: How stable is the low-rank approximation with respect to noisy or adversarial fine-tuning?

We thank the reviewer for the suggestion. In our new on-going experiments, we perform supervised fine-tuning over noisy outputs or rejected responses in the FLAN subset of the UltraFeedback dataset and will include the results in the final version of the paper.

Q2-2: Does PEFT (parameter efficient finetuning) such as LoRA change the conclusion?

Our empirical results presented in Appendix H and Table 9 shows LoRA improves low-rank approximation of the forgetting, indicating simpler associations between learned tasks and forgotten examples. Therefore, we do not expect LoRA to change the conclusion about the low-rank property.

Q3: Could this method be used during training to halt updates early when forgetting risk is high?

Thank you for bringing up this extension. Yes, we believe that by predicting forgetting during training allows us to efficiently determine early stopping of the training process. This should be computationally more efficient than evaluating actual forgetting over all upstream examples.

We thank the reviewer again for recognizing the contribution of our paper and the helpful comments.

We thank the reviewer for the thoughtful comments. We respond to the weaknesses and the questions suggested by the reviewer below.

W1: The experimental setting is overly idealized. A critical limitation is that the method assumes access to the previous training data, which is often unavailable in practice due to expired licenses or unreleased datasets.

Thank you for pointing out the limitation. Although we agree with the reviewer about the inapplicability of our approach under data unavailability (inherited from replay-based methods), we decided to study replay-based methods because they are shown to be scalable and effective when applicable, which probably explains why they are still actively studied for LLMs [1,2]. This setup applies to, for example, developers that own their model and the training data. We will better clarify the applicability scope of the proposed method following the reviewer’s comment. Currently, this limitation is discussed in the Limitations section in Appendix A.

W2: The mitigation method’s performance is not consistently superior, for example predict+replay performance is close to random replay in Figure 5(c)

We understand the concern from the reviewer regarding small numerical performance gaps like in Figure 5(c). At this stage, our main goal is to analyze patterns in forgetting and showcase the benefit of predicting forgetting. We leave more effective forgetting mitigation algorithms that utilize the predicted forgetting as future works.

Are the improvements consistent? Yes - we measured statistical significance of the improvements with paired t-tests in all setups in Figure 5, and achieved p<0.05 or 0.01 in most of the setups. To elaborate, a small p-value in this test suggests over all 20 tasks, our method brings consistent (but probably small) improvements.

W3 / Q2: Question about using log perplexity over the upstream training data as an appropriate metric of forgetting. Why not use task-specific metrics, for example, context-aware accuracy or answer correctness for QA datasets? What justifies log perplexity as the primary metric?

We thank the reviewer for the thoughtful question. The reason why we chose log perplexity over upstream data as the main metrics is many-fold.

First, log perplexity bridges theoretical approximations and empirical observations of forgetting, allowing proper comparison (in Sec. 3.3). Log perplexity, equivalent to the training loss of language modeling, is often approximated with first-order techniques (gradient x gradient, gradient x weight differences) in prior works [3,4], assuming small weight updates. These measures can be, in theory, taken as predictors of forgetting; which, however, do not correlate strongly with empirical forgetting under practical (longer) fine-tuning runs in our experiments (Table 2 in Sec 3.3.). To compare with such theoretical approximations, we decided to measure forgetting with log perplexity increase.

Second, we try to measure forgetting over examples that the model has surely been trained on, which is the upstream pre-training or instruction-tuning data. We also dissect forgetting from confounders such as generalization effects on unseen tasks.

We will significantly improve the clarity about the choice of log perplexity metrics in the final version of the paper.

In addition, we understand the limitations of using log-perplexity as the only metrics and that performance drop measured on benchmark dataset is also very interesting.

We would like to highlight results in the paper that utilize metrics other than perplexity in the current version of the paper. For instruction-tuned model OLMo-7B-Instruct, we measured binary exact-match drop on the FLAN subset of Tulu V2 as the metrics of forgetting. The dataset involves QA pairs of moderately diverse tasks, and the exact match metrics is close to the accuracy metrics the reviewer asked for. We visualized the association matrix (Figure 8) and analyzed its low-rank patterns (Figure 3(a) - OLMo-7B Inst F1). We measured the performance of predicting binary exact drop and showed statistically significant improvement of targeted replay in Figure 5(b). In summary, the results are very consistent with forgetting measured with log-perplexity.

Current results & future works about forgetting on benchmark datasets. Currently, we showed and concluded that targeted replay does not improve LLM leaderboard tasks performance (Table 7, Appendix F). We leave more effective algorithms as future works. We will include example-level forgetting statistics on LLM leaderboard benchmark datasets and analyze their low-rank approximations in the final version of the paper.

Q1: Extension to MoE and other model architectures

We thank the reviewer for the suggestion. We have tried to make our empirical analysis as comprehensive as possible given computational resources available. We understand MoE models are getting increasingly important but at this stage we have to leave it as future works.

Q3: Does the replay approach impact the learning of the new task? Replaying a mini-batch of upstream examples every few training steps could interfere with the training dynamics on the new task.

We thank the reviewer for the insightful question. We have experiments included in Appendix that can address the reviewer’s question. Per the reviewer’s comment, we will move more discussions to the main body of the paper.

In Appendix F and Table 6, we showed that targeted or random replay can achieve comparable or improved perplexity on the validation sets of the fine-tuned tasks. In some setups, the improvements are statistically significant, which we will annotate in the updated paper.

In Appendix H and Table 10, we measured the correlations between forgetting under different replay strategies. We see replay can change the patterns of forgetting compared to no replay, but the change is mild.

Q4: If the previous data is already known, why is matrix-based prediction necessary? One could directly measure actual forgetting using the original data.

We thank the reviewer for the question. Although we assume access to the upstream data, it is unknown how much the examples suffer from forgetting after fine-tuning on a new task. This can be expensive to acquire, as measuring forgetting requires running inference over the possibly massive upstream data. Our motivation of predicting forgetting with matrix completion instead of directly measuring them is to avoid this heavy computation, by leveraging the empirical low-rank property of the matrix. We compared computational costs of different replay strategies in Table 4. As the cost analysis suggests, we expect further increased benefit of predicting forgetting in larger-scale setups.

An analogy we made in the paper is collaborative filtering in recommender systems, which saves cost by inferring the missing user-item interactions by leveraging the low-rank property of the data matrix, instead of directly querying the user.

We thank the reviewer again for the thoughtful comments, and we will update the paper accordingly.

References

[1] Ibrahim et al. Simple and Scalable Strategies to Continually Pre-train Large Language Models, 2024

[2] Que et al. D-CPT Law: Domain-specific Continual Pre-Training Scaling Law for Large Language Models, NeurIPS’ 24

[3] Lepez-Paz et al. Gradient episodic memory for continual learning, NeurIPS’ 17

[4] Doan et al. A Theoretical Analysis of Catastrophic Forgetting through the NTK Overlap Matrix, AISTATS ‘21

Dear reviewer J9oc,

We hope to hear back from you as the end of the author response period is coming close. We are happy to take further questions. If the questions are addressed, we will be happy to see the reviewer re-evaluate the quality and significance of this work.

Thank you.

We thank the reviewer for the very thoughtful comments and questions.

W1/Q2: Some empirical results are hard to interpret (e.g. the matrix plot in the Appendix). Interactive websites to allow users to explore the results with more granularity can be helpful; Explicit examples in the appendix of data pairs/tasks where forgetting is high and where forgetting is low can be included.

We thank the reviewer for the suggestions. We have already started building an interactive online explorer of the matrix that also shows contents of the data pairs. We will improve the details and clarity of the figures and captions in the final version of the paper.

W2: How much more efficient the use of predict-then-replay result is in practice? Is this method and approach something we should all be using, or just something scientifically interesting?

Thank you for the thoughtful question. We consider that the predict-then-replay approach can potentially be extended to practical setups, and that the computational cost benefit can signify. As our cost analysis in Sec. 4.3 and Table 6 shows, the computational cost benefit of predict-then-replay increases when the upstream data is large, and when the same model is repetitively fine-tuned for different tasks.

Q1: I assume all models except OLMo-7B-Instruct are base models

Yes, we referred to non-instruction-tuned models as base models. We will improve the clarity in the writing.

Q3: Why is the matrix for OLMo-7B-Instruct much more different than the matrices for the other models in Appendix Figure 6?

We thank the reviewer for noticing this pattern. We believe it is mainly attributed to the differences in the data used in OLMo-7B-Instruct experiments. To elaborate, we evaluated forgetting over the instruction-tuning data, rather than the pretraining corpora in other non-instruction-tuned model experiments. As there are increased chances of task or format similarity between the fine-tuned tasks and the instruction-tuning data (compared to pretraining data), forgetting (log perplexity increase) is sometimes negative, shown in blue pixels. Besides, the set of fine-tuning tasks (T_1..N, y-axis labels of Figure 6) are also different from non-instruction-tuned models, as we make sure the fine-tuning tasks are unseen by the models in the instruction-tuning phase.

Q4: “We empirically notice that the associations between learned tasks and forgotten upstream examples are more complicated in more recent and capable LLMs” - is this simply because more of the finetuned datasets are more likely to be included in OLMo pretraining?

We thank the reviewer for this insightful question. For open-data LLMs like OLMo and OLMo2, the source of the upstream data is released and splitted into dedicated domains. As far as we know, the fine-tuning datasets are not (intentionally) included in the pretraining phase (although a small portion may slip into the commoncrawl partition of the corpora). Therefore, we do not consider that the complexity of the pattern can be solely attributed to the training data inclusion.

Besides, although our Appendix J suggests that inclusion of exact examples of the learned task in the upstream examples can lead to more complicated patterns of forgetting, we further notice that including the data from the same distribution (instead of the exact examples) results in the same patterns of forgetting. Therefore, we hypothesize that complexity depends on the overlap in the “distribution” or “knowledge” level, instead of the exact examples. We consider that understanding what attributes to the complexity of the patterns can require extensive study, which we leave as the future work.

Q5: Why does OLMo2 have zero forgetting in Figure 3(d)?

Thank you for noticing this pattern. Although forgetting appears on the zero line, the forgetting is actually non zero. We present additional statistics below.

We will include these statistics and the visualization of the association matrix of OLMo2-7B in the final version. We further note the low forgetting does not mean the OLMo2 has lost plasticity for updates. This is supported by the drop of perplexity on the fine-tuned tasks in a level similar to OLMo models.

Nevertheless, we notice that OLMo2 models are much more resistant to forgetting when learning the new tasks. This can depend on many factors, such as flat loss minima (so that the upstream data perplexity is resistant to any updates in weights, often affected by choice of optimizers, regularization, dropout etc. during pre-training) [1], and whether the new task data is considered significantly out-of-distribution by the model (so that learning new task updates model parameters in a direction that causes forgetting). There are also empirical studies showing how forgetting depends on pre-trained models such as [2,3].

[1] Mirzadeh et al. Understanding the Role of Training Regimes in Continual Learning, NeurIPS 2021

[2] Scialom et al. Fine-tuned Language Models are Continual Learners, EMNLP 2022

[3] Luo et al. An Empirical Study of Catastrophic Forgetting in Large Language Models During Continual Fine-tuning

Q6: Why do you use all-distilroberta-v1 as your embedding model? This seems quite outdated

We thank the reviewer for the suggestion. We will update Table 2 and Table 3 with the results of OpenAI text embedding models in the final version of the paper.

Q7: What exactly is the “ground truth forgetting” in Figure 5?

We are sorry about the clarity issue. The dash line “ground truth forgetting” here means the replay examples are upweighted according to the actual forgetting, instead of the predicted forgetting.

This result is included as a reference for predict-then-replaying approaches; computing the actual forgetting is expensive in practice, as it requires running LLM inference over the entire upstream data.

Q8: What are the pros and cons between online and offline MF?

The advantage of online MF is computation efficiency. As summarized in Table 4, online MF does not incur extra fine-tuning. However, the prediction of forgetting is less accurate as it is based on forgetting under a fraction of total fine-tuning steps. This is also evidenced by the subpar forgetting mitigation performance compared to the offline variant in Figure 5 in setups such as OLMo-1B.

Minor writing and presentation issues

We thank the reviewer for spotting the issues and the suggestions to enhance the figure presentation. We will update the paper accordingly.

We thank the reviewer again for recognizing the contribution of our paper and the insightful questions.