Continual Memorization of Factoids in Large Language Models

Testing REMIX in a three-stage or four-stage setting could provide better insight into its stability and effectiveness over longer training cycles.

We investigated the three-stage setting (L426 - L431 and Figure 8) and show that REMIX can retain learned factoids for different combinations of stage 2 data.

 

Why forgetting is more pronounced with factoids compared to non-factoids, as well as any observed differences in how REMIX performs on these types?

While we do not have a direct mechanistic explanation for this phenomenon, it corroborates with the findings in 1) the continual learning literature which suggest catastrophic forgetting happens when two tasks are similar and therefore interfere [6, 7], and 2) finetuning on unfamiliar knowledge disrupts the model and causes exacerbated hallucinations [8, 9, 10, 11]. A mechanistic understanding of this phenomenon is an important area for future investigation but is slightly outside of the scope of our paper.

 

Effectiveness of Random vs. Generic Text Mixing: The paper explores both random word sequence mixing and generic pretraining text mixing in REMIX. However, it is not entirely clear whether these two approaches yield similar or differing effects on knowledge retention. Could the authors provide more details on any observed differences in effectiveness between random and generic data mixing?

Both the Random and Generic mixing strategies can be applied at either stage 1 or stage 2 and they yield different results as shown in Table 2. Specifically, mixing Random is only effective in stage 1 and mixing Generic is only effective in stage 2. When combining both it yields the best results. This aligns with our motivation in section 4.1 – mixing uncorrelated data in stage 1 protects the learned factoids and mixing at stage 2 helps reduce interference especially when forgetting is severe and the mixing data aligns with $D_A$.

 

100% Accuracy in Table 1: In Table 1, it is stated that all Stage 1 datasets are trained to 100% accuracy before Stage 2 training. Could the authors clarify how this 100% accuracy is achieved and guaranteed across different datasets? Specifically, were there particular training techniques or criteria used to ensure full memorization of Stage 1 data?

We train the models until reaching a loss below 0.0001 by 5 times before stopping to ensure full convergence. This often entails training many epochs to guarantee perfect accuracy. We provide training details in Appendix B.4 for reference.

 

Has REMIX been tested on other types of tasks, such as generative or dialogue-based tasks?

In our initial exploration, we experimented with generative tasks and found that the forgetting phenomenon is much less pronounced than long-tail factoid data to begin with. This prompted us to further investigate the problem of forgetting specific to factoid datasets. We see the same phenomenon manifested in stage 2 training as well where non-factoid datasets generally have less impact leading to forgetting.

   

Please let us know if there's any more information we can provide to clarify our work, thank you!

References

[1] Kirkpatrick et al., Overcoming Catastrophic Forgetting in Neural Networks. PNAS 2017.

[2] Yang et al., Synthetic Continued Pretraining. Arxiv 2024.

[3] Allen-Zhu and Li, Physics of Language Models: Part 3.1, Knowledge Storage and Extraction. Arxiv 2024.

[4] Allen-Zhu and Li, Physics of Language Models: Part 3.2, Knowledge Manipulation. Arxiv 2024.

[5] Berglund et al., The Reversal Curse: LLMs trained on "A is B" fail to learn "B is A". ICLR 2024.

[6] Farajtabar et al., Orthogonal gradient descent for continual learning. AISTATS 2020.

[7] Bennani et al., Generalisation Guarantees for Continual Learning with Orthogonal Gradient Descent. ICML 2020.

[8] Kang et al., Unfamiliar Finetuning Examples Control How Language Models Hallucinate. Arxiv 2024.

[9] Gekhman et al., Does Fine-Tuning LLMs on New Knowledge Encourage Hallucinations? EMNLP 2024.

[10] Zhang et al., Knowledge Overshadowing Causes Amalgamated Hallucination in Large Language Models. Arxiv 2024.

[11] Ghosal et al., Understanding Finetuning for Factual Knowledge Extraction. ICLR 2024.

[12] Sun et al., Distill and Replay for Continual Language Learning. COLING 2020.

[13] Wang et al., Orthogonal Subspace Learning for Language Model Continual Learning. EMNLP 2023.

Expanding this method to resource-intensive domains like finance or healthcare could present challenges. If the authors could discuss the trade-offs between added data usage and computational demands [...], it would help assess its feasibility and scalability in high-resource environments.

Scalability is indeed critical for practical adaptations as the more mixing data needed for mitigation, the less scalable and practical it is. We provide analysis on the amount of REMIX data needed for successful forgetting mitigation. We show in L421-L424 the mixing ratio required for different mixing scenarios: most REMIX strategies only need a mixing ratio of 1.0 to be effective, and increasing the mixing ratio hits diminishing returns. We also simulate the scaled-up scenario (resource intensive cases as suggested) by increasing the dataset size from 2000 factoids to 4000 factoids.

KP = Knowledge Pile. We find that the same trend holds for the scaled up case where only the mixing ratio of 1.0 is effective. A promising future direction is to discover the kind of mixing data that can be effective with very small mixing ratios.

 

Exploration of Bidirectional Relational Memory with Atomic Fact Datasets: [...] Could the authors clarify whether such tests were conducted, or suggest if REMIX could potentially extend to this type of bidirectional memorization?

The reviewer rightfully pointed out the importance of transferring the learned knowledge to downstream tasks to make it useful. This is discussed in depth in works such as [2, 3, 4, 5], which highlight the difficulty of manipulating the learned knowledge in downstream tasks. Our work is mainly positioned as the prerequisite before knowledge manipulation – if the knowledge is not retained in the first place, then there’s no chance for it to be recalled and manipulated successfully. While retention does not entail successful manipulation, we aim to understand the dynamics of retention as the first step, which is on its own a challenging task.

Nonetheless, we fully recognize the importance of this question and further conducted evaluations to assess such capabilities of our models. We use three templates to assess the model’s ability to manipulate learned knowledge on the KVR task (since it is guaranteed to have no contamination from pretraining):

Template 1 (reverse recall):
The key of the value DEF is?
Key1: ABC, Value: DEF
Answer: ABC

Template 2 (selective recall):
Here are two keys: ABC and XYZ. What is the value of the first key?
Key1: ABC, Value1: DEF
Key2: XYZ, Value2: GHI
Answer: DEF

Template 3 (recall-then-manipulate):
If the first character in the value of key ABC is changed to X, what is the new value of key?
Key: ABC
Value: DEF
Answer: XEF

We evaluate on the following models: No-Mixing, REMIX (Random / -), and REMIX (Random / Knowledge-Pile).

R = Random Word Sequence. KP = Knowledge Pile. We observe that none of the models can perform Template 1, which corroborates with [8, 9, 10, 11], highlighting the unique challenge of reverse recall in knowledge storage and manipulation. However, we found that REMIX improves other types of knowledge manipulation such as Template 2 and 3 as shown in the table. This is an interesting finding that warrants further study. We appreciate the reviewer’s suggestion and will include this and further findings in the updated version of the paper.

The datasets used in this paper, such as Key-Value Recall and PopQA, are primarily synthetic and consist of isolated factoids, which may not fully reflect the complexity of real-world data. [...] Testing REMIX on more commonly used datasets, such as Wikipedia or open-domain QA datasets (e.g., Natural Questions), could provide a more realistic evaluation of its effectiveness and generalizability.

We totally agree that capturing real-world complexities is important to understand the generalizability of our conclusion. In our investigation, we aim to strike a balance between controllability and the real world complexity. With the Key-Value Recall (KVR) task, we can ensure there’s no data contamination from pretraining despite its simplicity. On the other hand, both PopQA and TriviaQA are sourced from real websites yet maintain the “atomic” nature of each factoid. While PopQA is still template-based, questions in TriviaQA is sourced from real Trivia websites and annotated by humans, which is close to how Natural Question is collected. We aimed to understand the memorization dynamics of these cases from simple/controllable (KVR) to realistic (TriviaQA).

Nonetheless, we fully recognize the importance of using more natural and realistic datasets to strengthen our conclusion suggested by the reviewer. We further investigate the effectiveness of REMIX on the Natural Question dataset as suggested. The results are the following:

KP = Knowledge Pile. The trend aligns with the trends for our choice of datasets in the paper – mixing Knowledge Pile at stage 2 leads to best retention of the learned factoids (+12.4% average accuracy over No Mixing), followed by mixing Random Word Sequence at stage 1 + Knowledge Pile at stage 2 (+6.7% average accuracy over No Mixing). We observe that forgetting is generally less severe than the datasets we chose, which might be due to some level of contamination reported in the Llama 3 technical report (Table 15).

 

Unclear Justification for Types of Data Mixing: The paper employs both random word sequences and knowledge-rich text (e.g., Knowledge Pile) as mixed data to prevent forgetting, but it does not provide a clear explanation of why these two disparate types would produce similar effects. [...] However, the choice of these sources appears empirical, lacking theoretical justification or detailed explanation. It remains unclear why certain data sources yield better performance on specific tasks, and this potential variation across tasks is not fully explored.

We motivate our choice of different mixing strategies in section 4.1 (L260 - L267) with the following intuition: REMIX at stage 1 aims to protect the learned factoids by maintaining a good starting weight space for stage 2 training. REMIX at stage 2 aims to reduce the interference of the memorized factoids. With this intuition, we justify the use of the two mixing datasets with mathematical derivation in Appendix A.3 (L866 - L917). Specifically, in stage 1, we need to choose a mixing data that is uncorrelated to $D_A$ and $D_B$, hence the choice of random word sequence. In stage 2, any natural distribution can achieve mitigation when forgetting is already severe; when the dataset aligns with $D_A$ more than $D_B$, the mitigation can be more effective (e.g., Knowledge Pile helps knowledge factoids more than Arxiv Pile in Figure 5).

 

Impact on Performance in New Tasks: While REMIX performs well in retaining early-stage knowledge, the paper does not explore its impact on subsequent new tasks. For instance, it would be useful to know whether REMIX might limit the model's ability to learn these new tasks when introduced for fine-tuning.

When applying mixing at stage 2, we ensure convergence for dataset $D_B$ using the same stopping criterion (loss reaching below 0.0001 by 5 times). We did not observe issues lowering the loss after stage 1 training. Similarly, we observe no issue when performing stage 3 training with mixing data as well (L426-L431 and Figure 8). This suggests that mixing data does not hinder the model's ability to learn new tasks.

We thank the reviewer for the detailed comments, suggestions, and the recognition of the novelty of our methods, the extensiveness of our experiments, and the generalizability of our conclusions. The reviewer pointed out several important aspects that would help strengthen our investigation, which we provide further experiments and explanations in the following passages.

 

[...] the paper lacks a systematic comparison with other common forgetting mitigation techniques, such as Elastic Weight Consolidation (EWC) or Knowledge Distillation.

We provide three more types of baselines to compare with our data mixing method:

• Weight regularization: we use Elastic Weight Consolidation (EWC) [1] and calculate the Fisher score using one backward pass using the current mini-batch for training.
• Behavior regularization: we add the KL between the training model vs the original reference model to the loss. Knowledge distillation can be seen as a type of regularization in the continual learning setting [12].
• Parameter expansion method: we learn separate and none-overlapping LoRA adapters at stage 1 and 2, similar to the IncLoRA model in [13]. We compare these baselines against the No Mixing baseline and REMIX (Random at stage 1 and Knowledge Pile at stage 2). We show results on the datasets that suffer most from forgetting: all factoid datasets and GSM8K from the non-factoid datasets.

We observe that the weight regularization baseline and output regularization baseline can obtain better factoid retention at different tasks but on average lags behind REMIX by a large margin (40%+ on KVR, 30%+ on PopQA, and 20%+ on TriviaQA). In our attempt, the parameter expansion based baseline learns to achieve 100% accuracy at stage 2, but catastrophically forgets at stage 2, achieving close to zero factoid retention.

the idea of mixing generic data into training is not groundbreaking and does not specifically address the unique challenges of factoid memorization

We would like to first point out the important distinctions between our setting and the general continual learning setting where mixing is applied. 1) Mixing often assumes access to distributions of the previous training stages (similar to Replay in Section 3.2 that uses a small percentage of $D_A$). In REMIX, the model does not have access to $D_A$, which renders the problem much more challenging. 2) Mixing random sequence data is unexplored in past literature and its effectiveness is surprising. 3) While mixing generic pretraining data is familiar in settings like continual pretraining, its effectiveness is largely under-explored in the context of factual knowledge retention. 4) Although it appears fairly straightforward, we discovered that such a simple strategy can work extremely well, which we established through mathematical derivation and validated through extensive experimentation.

Along with the empirical efficacy of REMIX, we provide intuition (section 4.1) and derivations (Appendix A.3) to justify the choice of random data (help protecting the learned knowledge) and the generic data (help reduce the interference of the stage 2 data), which we hope supplement the existing data mixing methods such as replay.

 

The authors only explore experience replay as the baseline approaches, whereas there exists other methods in literature that can mitigate forgetting during continued pretraining (parameter expansion-based methods, regularization methods, etc.

We provide three more types of baselines to compare with our data mixing method:

• Weight regularization: we use Elastic Weight Consolidation (EWC) and calculate the Fisher score using one backward pass using the current mini-batch for training [9].
• Behavior regularization: we add the KL between the training model vs the original reference model to the loss [10].
• Parameter expansion method: we learn separate and none-overlapping LoRA adapters at stage 1 and 2, similar to the IncLoRA model in [11].

We compare these baselines against the No Mixing baseline and REMIX (Random at stage 1 and Knowledge Pile at stage 2). We show results on the datasets that suffer most from forgetting: all factoid datasets and GSM8K from the non-factoid datasets.

We observe that the weight regularization baseline and output regularization baseline can obtain better factoid retention at different tasks but on average lags behind REMIX by a large margin (40%+ on KVR, 30%+ on PopQA, and 20%+ on TriviaQA). In our attempt, the parameter expansion based baseline learns to achieve 100% accuracy at stage 2, but catastrophically forgets at stage 2, achieving close to zero factoid retention.

We thank the reviewer for recognizing the rigor in our experimentation and the clearness in our writing. We also appreciate that the reviewer pointed out several aspects that would benefit from further clarification and experimentation to solidify the arguments in our investigation. We provide the explanations and further evidence in the following passages.

 

The paper does not convincingly explain why memorizing long-tail knowledge in the form of factoids is important in practical applications. [...] If we are only concerned about factoids, why use LLMs in the first place? Why not just use traditional knowledge-based systems? The authors should show how memorizing factoids leads to downstream applications, such as utilizing the information from the factoids on tasks that specifically require LLMs.

We agree that the choice between parametric vs non-parametric representation for factoid knowledge has a long standing tension in practical scenarios. However, we position our work more as part of the larger body of works that aim to understand the knowledge acquisition dynamics of language models through finetuning and the unintended risk ([1, 2, 3, 4]), which also evaluate on knowledge datasets such as PopQA and TriviaQA.

The reviewer rightfully pointed out the importance of transferring the learned knowledge to downstream tasks to make it useful. This is discussed in depth in works such as [5, 6, 7, 8] which highlight the difficulty of manipulating the learned knowledge in downstream tasks. Our work is positioned as the prerequisite before knowledge manipulation – if the knowledge is not retained in the first place, then there’s no chance for it to be recalled and manipulated successfully. While retention does not entail successful manipulation, we aim to understand the dynamics of retention as the first step, which is on its own a challenging task. We appreciate this profound point and will improve our framing to reflect this emphasis and make the distinction clearer.

Nonetheless, we fully recognize the importance of this question and further conducted evaluations to assess such capabilities of our models. We use three templates to assess the model’s ability to manipulate learned knowledge on the KVR task (since it is guaranteed to have no contamination from pretraining):

Template 1 (reverse recall):
The key of the value DEF is?
Key1: ABC, Value: DEF
Answer: ABC

Template 2 (selective recall):
Here are two keys: ABC and XYZ. What is the value of the first key?
Key1: ABC, Value1: DEF
Key2: XYZ, Value2: GHI
Answer: DEF

Template 3 (recall-then-manipulate):
If the first character in the value of key ABC is changed to X, what is the new value of key?
Key: ABC
Value: DEF
Answer: XEF

We evaluate on the following models: No-Mixing, REMIX (Random / -), and REMIX (Random / Knowledge-Pile).

R = Random Word Sequence. KP = Knowledge Pile. We observe that none of the models can perform Template 1, which corroborates with [7, 8], highlighting the unique challenge of reverse recall in knowledge storage and manipulation. However, we found that REMIX improves other types of knowledge manipulation such as Template 2 and 3 as shown in the table. This is an interesting finding that warrants further study. We appreciate the reviewer’s suggestion and will include this and further findings in the updated version of the paper.

Please let us know if there's any more information we can provide to clarify our work, thank you!

References

[1] Kang et al., Unfamiliar Finetuning Examples Control How Language Models Hallucinate. Arxiv 2024.

[2] Gekhman et al., Does Fine-Tuning LLMs on New Knowledge Encourage Hallucinations? EMNLP 2024.

[3] Zhang et al., Knowledge Overshadowing Causes Amalgamated Hallucination in Large Language Models. Arxiv 2024.

[4] Ghosal et al., Understanding Finetuning for Factual Knowledge Extraction. ICLR 2024.

[5] Yang et al., Synthetic Continued Pretraining. Arxiv 2024.

[6] Allen-Zhu and Li, Physics of Language Models: Part 3.1, Knowledge Storage and Extraction. Arxiv 2024.

[7] Allen-Zhu and Li, Physics of Language Models: Part 3.2, Knowledge Manipulation. Arxiv 2024.

[8] Berglund et al., The Reversal Curse: LLMs trained on "A is B" fail to learn "B is A". ICLR 2024.

[9] Kirkpatrick et al., Overcoming Catastrophic Forgetting in Neural Networks. PNAS 2017.

[10] Sun et al., Distill and Replay for Continual Language Learning. COLING 2020.

[11] Wang et al., Orthogonal Subspace Learning for Language Model Continual Learning. EMNLP 2023.

We thank the reviewer for the comments and suggestions. We appreciate your positive comments on the clarity and simplicity of our problem setup and the proposed solutions.

There are some other methods to reduce the interference among the datasets of stage 1 and stage 2. For example, the method needs to compare with another baseline, i.e. "mixing of Data A and Data B"

Thank you for pointing out the need for more baselines in order to strengthen our conclusions.

We would like to clarify that since we focus on understanding how the knowledge of D_A is retained after another stage of training, stage 2 shouldn't have access to D_A (including mixing D_A and D_B).

Nonetheless, we recognize that more baselines can strengthen our conclusions. We provide three more types of representative baselines to compare with our data mixing method:

• Weight regularization: we use Elastic Weight Consolidation (EWC) and calculate the Fisher score using one backward pass using the current mini-batch for training [1].
• Behavior regularization: we add the KL between the training model vs the original reference model to the loss [8].
• Parameter expansion method: we learn separate and none-overlapping LoRA adapters at stage 1 and 2, similar to the IncLoRA model in [9].

We compare these baselines against the No Mixing baseline and REMIX (Random at stage 1 and Knowledge Pile at stage 2). We show results on the datasets that suffer most from forgetting: all factoid datasets and GSM8K from the non-factoid datasets.

We observe that the weight regularization baseline and output regularization baseline can obtain better factoid retention at different tasks but on average lags behind REMIX by a large margin (40%+ on KVR, 30%+ on PopQA, and 20%+ on TriviaQA). In our attempt, the parameter expansion based baseline learns to achieve 100% accuracy at stage 2, but catastrophically forgets at stage 2, achieving close to zero factoid retention.

For the Replay approach, what if we use a ratio = 1.0?

Similar to the previous point, a ratio of 1.0 means using the entire $D_A$ together with $D_B$ for training at stage 2, which is in conflict with the continual multi-stage setting. Replay is not meant to be a baseline it uses D_A but rather as a motivational study that inspires REMIX. In the factoid memorization task, using $D_A$ is essentially cheating, so it should be understood as a different settings (therefore 10% reported in section 3.2 seems adequate for the motivational purpose).

Table 1, the degradation is more severe for Stage 2 is also a factoid dataset. Do you have any explanation? Also, there is a big drop when using GSK8k. It will be very insightful to understand the interplays of the datasets.

While it is hard to obtain a direct mechanistic explanation for this phenomenon, it corroborates with the findings in 1) the continual learning literature which suggest catastrophic forgetting happens when two tasks are similar and therefore interfere [2, 3], and 2) finetuning on unfamiliar knowledge disrupts the model and causes exacerbated hallucinations [4, 5, 6, 7]. A mechanistic understanding of this phenomenon is an important area for future investigation but is slightly outside of the scope of our paper.

For GSM8K, we hypothesize that the special format of its training data (e.g., extensive use of angle brackets) might contribute to forgetting as the model picks up such irregularities very quickly.

Please let us know if there's any more information we can provide to clarify our work, thank you!

Reference:

[1] Kirkpatrick et al., Overcoming Catastrophic Forgetting in Neural Networks. PANS 2017.

[2] Farajtabar et al., Orthogonal gradient descent for continual learning. AISTATS 2020.

[3] Bennani et al., Generalisation Guarantees for Continual Learning with Orthogonal Gradient Descent. ICML 2020.

[4] Kang et al., Unfamiliar Finetuning Examples Control How Language Models Hallucinate. Arxiv 2024.

[5] Gekhman et al., Does Fine-Tuning LLMs on New Knowledge Encourage Hallucinations? EMNLP 2024.

[6] Zhang et al., Knowledge Overshadowing Causes Amalgamated Hallucination in Large Language Models. Arxiv 2024.

[7] Ghosal et al., Understanding Finetuning for Factual Knowledge Extraction. ICLR 2024.

[8] Sun et al., Distill and Replay for Continual Language Learning. COLING 2020.

[9] Wang et al., Orthogonal Subspace Learning for Language Model Continual Learning. EMNLP 2023.

Dear Reviewer,

We'd like to send a gentle reminder that we have submitted the rebuttal to address your comments. We sincerely appreciate your feedback and are happy to address any additional questions you may have during this discussion period.

We thank you again for taking the time to review our work.

We thank the reviewer for the positive comments, especially recognition of the importance of the problem, the novelty of our proposed mitigation method, and the extensiveness of our experiments. We also appreciate the reviewer for the constructive and actionable suggestions.

 

Section 3.2 (on replay) is lacking detail in comparison to other sections, especially as it justifies the use of REMIX compared to other replay methods. In particular, I could not find which of the two LLMs was used to measure the effect of replay methods.

Thank you for pointing this out. The model used for replay was the Llama 3 8B model. We did not directly compare replay and REMIX due to the difference in their settings – replay assumes access to stage 1 dataset $D_A$ and REMIX does not. Therefore, we use replay mainly as inspiration for REMIX as opposed to a competitive baseline. For clarity as suggested, we will update the REMIX section to refer and draw appropriate comparison to the replay method.

 

It would have been interesting to investigate whether this results in any significant degradation of other model capabilities (e.g. fluency) compared to the basic two-stage training process. I understand that this paper specifically focuses on factoid memorization and contains many experiments already, but this could be mentioned as future work.

We find models trained with REMIX actually contain slightly better fluency compared to other models mainly because the models without mixing suffer from overfitting more severely. However, we totally agree a more comprehensive analysis on the side effects when using REMIX will be very valuable. We will aim to include such an analysis in the updated version.

 

[...] vary either the model or the dataset's size to evaluate the link between model capacity and the efficacy of REMIX/replay techniques.

We provide an extra result on increasing the $D_A$ dataset size from 2000 factoids to 4000 factoids ($D_B$ size maintains to be 2000).

KP = Knowledge Pile. We find that the effectiveness of REMIX holds for the scaled up dataset size case. We also find that there seems to be generally less forgetting in the No Mixing case, which might be an interesting phenomenon to further study. With the computation budgets at hand we defer the model scaling experiments to future research.

 

Have the authors considered/tried combining REMIX with classic replay techniques? This seems like a natural next step to know whether the use of both methods leads to even better results.

In our investigation, we aim to fully separate the two settings: 1) allowing access to $D_A$ at later stages, and 2) strict continual learning setting where $D_A$ cannot be used in stage 2. We use the second setting throughout since using $D_A$ can be seen as a form of cheating because they are the exact factoids to memorize. Therefore, Replay (section 3.2) was only meant to motivate REMIX instead of treated as a fully comparable baseline.

Please let us know if there's any more information we can provide to clarify our work, thank you!