1 Introduction↩︎

Language models (LMs) have shown a remarkable ability to absorb massive amounts of knowledge through large-scale pretraining [1]–[3]. However, knowledge is not static—new information accumulates quickly while old knowledge becomes obsolete. This dynamic nature necessitates frequent model updates, making costly pretraining impractical. A common approach is to fine-tune the model directly on new knowledge. However, recent studies have shown that fine-tuning is brittle: training on long-tail knowledge can lead to unintended disruptions, such as decreased factuality and exacerbated hallucinations [4]–[6]. Furthermore, a fine-tuned model might not properly memorize knowledge, failing to recall or manipulate it effectively [7], [8]. This fragility raises the question of whether fine-tuning can be applied repeatedly as a reliable mechanism for continual knowledge acquisition.

To address this, we investigate the dynamics of memorization in a continual learning setting [9], [10], in which the model acquires new information incrementally, one set at a time. While prior research on continual learning in LMs has focused on general capabilities such as reasoning [11] or broad proxies like language modeling loss over a general corpus [12], we use this framework to study the memorization dynamics of LMs through fine-tuning. We formalize this as continual memorization, in which a model is first trained on a small collection of factoids (factual associations) and must retain this knowledge after training on additional datasets in a subsequent stage. Specifically, we train models to memorize factoid datasets (stage 1) and then evaluate how well these factoids are retained after a second stage of training on a different dataset (Figure 1). We study two stages in our main setting, and provide analysis of more stages in §5.

We conduct extensive experiments to characterize forgetting patterns in continual memorization. First, we find that the most severe forgetting occurs when the second stage of training involves another factoid dataset, regardless of whether the facts overlap with those from stage 1. For example, accuracy on TriviaQA drops from \(100\%\) (after stage 1) to \(39.8\%\) after further training on other factoid datasets, such as LAMA. Forgetting is less pronounced when fine-tuning on non-factoid datasets, such as those involving coding, math, or chat. We also find that long-tail data is the hardest to retain—a randomly generated key-value string is most susceptible to forgetting, with accuracy dropping to \(13\%\) after further training on another factoid dataset. This aligns with recent findings that memorizing knowledge (e.g., factoids) causes greater disruption to other model capabilities, which, in our setting, results in more severe forgetting [4]. Finally, we observe that common experience replay methods, which mix a fraction of data from earlier training stages, fail to prevent forgetting when the second stage involves a factoid dataset, in contrast to their effectiveness in general continual learning settings. For instance, even when mixing in \(10\%\) of the factoids from stage 1, the model fails to recover performance beyond \(60\%\).

Next, we investigate whether data mixing strategies can mitigate forgetting. Through theoretical derivations, we develop intuition that this question may be approached in two ways: 1) teaching the model to protect learned knowledge better in the first stage, or 2) reducing the interference of the second stage by manipulating the data distribution. Based on this hypothesis, we examine a range of data mixing strategies at each stage. Intriguingly, we find that mixing in either generic pretraining data or even random word sequences leads to a considerable reduction in forgetting. We combine both strategies, and refer to this mitigation as REMIX (Random and Generic Data Mixing). Our experiments demonstrate that REMIX is highly effective at helping the model retain learned factoids: in the most severe case, REMIX increases post-stage 2 accuracy from \(13.5\%\) to \(53.2\%\). In comparison, replay can only reach \(41.6\%\) despite using \(10\%\) of the factoids from stage 1. Other common continual learning methods also fall short: weight regularization (EWC) [13] and behavior regularization [14] both lag behind REMIX. These benefits are seen consistently across several choices of factoid and non-factoid tasks in stage 2.

To understand why these mixing strategies help reduce forgetting, we analyze REMIX using Logit Lens [15] and ablation studies. Our analysis suggests that including a broad range of mixed data encourages the model to store facts in relatively earlier layers (compared to the baseline setting), as well as to diversify where it stores the knowledge. This diversification allows it to better protect learned knowledge in subsequent stages of training. In the second stage, jointly learning the mixing data and the stage 2 data helps prevent overfitting to a narrow distribution, alleviating the negative interference on the learned factoids. Finally, we show that more robustly memorized factoids are not only better retained and recalled, but are more easily extracted for manipulation.

We summarize our contributions as follows:

• We formalize the continual memorization setting and demonstrate the fragility of the factoid memorization process in LMs; we further show that it cannot be easily addressed with replay.

• We find that mixing random and generic data (REMIX) in different stages can greatly mitigate forgetting without accessing the factoids from prior stages.

• We find that successful mixing diversifies the layers where the learned knowledge is stored and tends to store it in earlier layers than in models that suffer from forgetting, shedding light on the patterns of robust memorization.

2 Continual Memorization of Factoids↩︎

2.1 Problem Definition↩︎

2.1.0.1 Factoid vs non-factoid datasets.

We define a factoid to be a triplet (subject, relation, object). A dataset \(D \in \mathcal{D}\) in this paper is a set of (prompt, response) pairs. A factoid dataset \(D \in \mathcal{D}_\text{fact} \subset \mathcal{D}\) is a set of factoids formatted as pairs (e.g., “prompt = The <relation> of <subject> is” \(\rightarrow\) response = <object>). If \(D \in \mathcal{D} \setminus \mathcal{D}_\text{fact}\), we call \(D\) a non-factoid dataset. A language model \(\mathcal{M}\): \(p_\theta(y \mid x)\) parameterized by \(\theta\) defines a distribution over response \(y\) given the prompt \(x\). Given a model \(\mathcal{M}\) and dataset \(D\), we denote by \(\mathcal{L}(\theta; D) \in \mathbb{R}^+\) the loss and \(\mathcal{A}(\mathcal{M}; D) \in [0, 1]\) the average exact-match accuracy. We define a factoid \(x\) to be familiar to \(\mathcal{M}\) if \(\mathcal{A}(\mathcal{M}; \{x\}) = 1\) and unfamiliar otherwise. ³ An unfamiliar dataset consists entirely of unfamiliar facts.

2.1.0.2 Continual memorization.

We now describe the setting of continual memorization, which consists of two or more stages. We describe the setting with two stages below. Let \(D_A \in \mathcal{D}_\text{fact}\) be a factoid dataset, and \(D_B \in \mathcal{D}\) be another dataset (factoid or non-factoid). In the first stage, a pretrained model \(\mathcal{M}_0\) is trained on \(D_A\) until convergence to obtain the trained model \(\mathcal{M}_A\) with near-zero loss \(\mathcal{L}(\theta_A; D_A) \approx 0\) and accuracy \(\mathcal{A}(\mathcal{M}_A; D_A) \approx 1\). In the second stage, \(\mathcal{M}_A\) is further trained on \(D_B\) until convergence. The resulting model \(\mathcal{M}_B\) is evaluated on \(D_A\) to gauge its retention \(\mathcal{A}(\mathcal{M}_B, D_A)\). In this paper, we consider the case where all factoid datasets (in the first as well as second stage—if applicable) are unfamiliar and we refer to them simply as factoid datasets. Typically, one observes \(\mathcal{A}(\mathcal{M}_B, D_A) \ll \mathcal{A}(\mathcal{M}_A, D_A)\) due to catastrophic forgetting. Figure 1 illustrates this setting.

2.2 Constructing Factoid Datasets↩︎

We consider a variety of (unfamiliar) factoid datasets in our experiments. These datasets are either 1) constructed synthetically to ensure that they were not seen by the model \(\mathcal{M}_0\) during pretraining—such as by generating random key-value mappings, or 2) filtered from factoid datasets to remove familiar instances (details in § 9.8). We further describe the specific choice of datasets for the two stages below.

2.2.0.1 Stage 1: Factoid dataset \(D_A\).

Key-Value Recall (KVR): we generate \(2,000\) unique key-value pairs, each containing \(8\) characters from the mix of alphabets and number digits. PopQA: we sample \(2,000\) unfamiliar knowledge triplets from a set of diverse questions and relationships about long-tail entities [16]. TriviaQA: we sample \(2,000\) unfamiliar question-answer pairs from the dataset  [17]. See examples in Figure 1 and §9.9.

2.2.0.2 Stage 2: Dataset \(D_B\).

We explore a wide range of datasets in stage 2 to reflect real-world application scenarios. Specifically, we consider two types of datasets: factoid and non-factoid. We choose this split because we want to see how the effect of stage 2 changes from a knowledge-intensive factoid dataset to, e.g., a general instruction tuning dataset. Additionally, domain-specific knowledge and instruction-tuning data represent two of the most common types of data used for supervised fine-tuning—a fact reflected in our selection of tasks. We explore:

• Factoid datasets: LAMA [1], Entity Questions [18], WebQA [19]. In addition, we also explore adding new (and unfamiliar) examples from the distribution of \(D_A\) (i.e., the same task as in stage 1) referred to as the “In-Domain” (ID) datasets in our results.

• Non-factoid datasets: UltraChat [20], EvolCode [21], APPS [22], GSM8K [23], and MATH [24]. These datasets exemplify common non-factoid datasets used for finetuning: chat, code and math.

2.2.0.3 Training and evaluation.

We use Llama-3-8B [25] and Mistral-7B [26] to initialize \(\mathcal{M}_0\) in our experiments (both are base models). All of our experiments use the Tulu-v2 prompt template [27], i.e., "<user>...<assistant>..." for both stages. We provide training details in §9.8. Our accuracies are computed as Exact String Match and normalized to \([0, 100]\) for all the experiments, as the tasks only need to generate a few tokens. We report averaged accuracy across \(3\) runs.

3 How Do Models Forget Factoids?↩︎

3.1 Understanding the Forgetting Patterns↩︎

We first establish the forgetting patterns in continual memorization by examining which intervening tasks affect the final accuracy most severely when used in the second stage. Table 1 shows the performance degradation of stage 1 tasks after training on stage 2 tasks. We observe that forgetting is most severe when stage 2 is also a factoid dataset, degrading accuracy for KVR to \(13.5\%\), PopQA to \(47.0\%\), and TriviaQA to \(39.8\%\) on average. In fact, with LAMA these accuracies fall to \(2.1\%, 7.7\%\) and \(4.3\%\) respectively—far below the numbers seen with non-factoid datasets. This corroborates findings from the continual learning literature which suggest catastrophic forgetting happens when two tasks are similar and therefore cause interference [28]–[30]. While the overall trend holds, we note caveats in §9.11.

3.2 Replay Does Not Mitigate Forgetting Fully↩︎

Replay-based methods mitigate forgetting by sampling a small portion of data from earlier stages and mixing it with the subsequent dataset during training. Replay from past experience has been a long-established mitigation to prevent forgetting in reinforcement learning research [31] and more recently continual pretraining for LMs. Although replay-based methods have proven helpful for continual learning, we hypothesize that they will be less effective for tasks requiring memorization, as the individual instances are largely unrelated [32], [33]. Figure 2 shows that although replay reduces forgetting across the board, the effectiveness is not uniform. Replay is less effective at preventing forgetting when stage 2 is factoid than non-factoid (full results in §8.2). The experiments suggest that manipulating the training dynamics such as exposing the model to different distributions can affect the model’s ability to recall factoids, even when the replayed factoids are individually independent from other factoids in the same stage.

4 REMIX: Random and Generic Data Mixing↩︎

4.1 Method↩︎

Despite the shortcomings of replay, we make one key observation: when mixing only \(10\%\) of the factoids used in stage 1, the accuracy increases after learning non-factoid data in stage 2 from no mixing at \(40.1\%\) to \(83.9\%\) (Table 9). This implies the existence of associations that were stored in model weights but could not be retrieved effectively. It is then prudent to ask if these “hidden” associations can be surfaced with a different choice of mixing data. To answer this question, we propose Random and Generic Data Mixing (REMIX), a data mixing strategy that manipulates the memorization dynamics during training to prevent forgetting. The mixing data is sampled from either random word sequences or generic text such as pretraining corpora, which has no overlap with the factoids aiming to memorize in stage 1. Figure 3 illustrates the intuition behind the mixing strategies.

For the purpose of developing intuition, we take the simplification to assume the entire optimization is captured by the one-step gradient update. Let \(\mathcal{L}(\theta; D)\) be the empirical loss on dataset \(D\), and let \(\nabla \mathcal{L}(\theta; D)\) denote its gradient. Starting from \(\theta_0\), stage-1 training on \(D_A\) yields \[\theta_A = \theta_0 - \eta \nabla \mathcal{L}(\theta_0; D_A),\] and subsequent stage-2 training on \(D_B\) yields \[\theta_B = \theta_A - \eta \nabla \mathcal{L}(\theta_A; D_B),\] for learning rate \(\eta > 0\). A first-order Taylor expansion of \(\mathcal{L}(\cdot\,; D_A)\) around \(\theta_A\) gives \[\mathcal{L}(\theta_B; D_A) - \mathcal{L}(\theta_A; D_A) \approx - \eta \, \nabla \mathcal{L}(\theta_A; D_B)^T \, \nabla \mathcal{L}(\theta_A; D_A),\] and catastrophic forgetting on \(D_A\) occurs when \(\nabla \mathcal{L}(\theta_A; D_B)^T \, \nabla \mathcal{L}(\theta_A; D_A) < 0\), i.e., when the gradients induced by \(D_B\) point in a direction that increases the loss on \(D_A\). REMIX changes the learning dynamics by mixing \(D_M\) in the first stage and/or \(D_M'\) in the second stage to prevent forgetting. We state the condition for REMIX to mitigate forgetting in the following.

At stage 1, the mixing data can teach the model to diversify where to store the knowledge, resulting in a better starting position in the parameter space for stage 2 training (smaller angle between \(\nabla \mathcal{L}(\theta; D_A)\) and \(\nabla \mathcal{L}(\theta; D_B)\)), achieving better protection of the memorized factoids. At stage 2, the mixing data can rotate the direction of \(\nabla \mathcal{L}(\theta; D_B)\) to align with \(\nabla \mathcal{L}(\theta; D_A)\), thus reducing the interference on the memorized factoids from stage 2 training; if the two gradients are in extreme opposing directions, it becomes easier for the mixing data to align them. We provide detailed derivations to concretize the intuition in §8.3. Based on the above insight, we posit: 1) mixing at stage 1 mitigates forgetting most when the mixing data is unrelated to both \(D_A\) and \(D_B\), and 2) mixing at stage 2 is most effective if the forgetting is severe, and is more effective when \(D_M'\) aligns with \(D_A\).

4.1.0.1 REMIX datasets \(D_M\).

We explore three data sources for generic data mixing: 1) Knowledge Pile [34], 2) Arxiv Pile [35], and 3) Fineweb [36]. We construct the Random Word Sequence data by collecting a set of uniformly sampled \(50\) random word sequences from the NLTK Word Corpus [37]. We check and ensure no overlap between the factoid data and the mixing data (see details in §9.10). When applying REMIX, we add the mixing data directly to \(D_A\) in stage 1 and \(D_B\) in stage 2. Therefore, the model trains on more data at each stage with mixing. We use Random Word Sequence and Knowledge Pile as the main datasets in the following experiments and later show that other mixing datasets show similar trends. We use \(D_A:D_M = 1:2\) and \(D_B:D_M = 1:2\) for the main experiments.

4.2 Results↩︎

4.2.0.1 Factoid tasks.

Table 2 shows the results of factoid tasks with Llama-3-8B. We observe that mixing Random Word Sequences prevents forgetting across the board, improving average accuracy for all \(D_A\), improving Key-Value Recall (\(13.5\% \rightarrow 28.8\%\)), PopQA (\(47.0\% \rightarrow 53.7\%\)), and TriviaQA (\(39.8\% \rightarrow 51.0\%\)). On the other hand, mixing Knowledge Pile at stage 1 hurts the performance. Mixing at stage 2 shows an opposite trend. We observe drastically better performance with mixing Knowledge Pile, improving the average accuracy for Key-Value Recall (\(13.5\% \rightarrow 27.8\%\)), PopQA (\(47.0\% \rightarrow 86.4\%\)), and TriviaQA (\(39.8\% \rightarrow 90.6\%\)). In contrast, mixing Random Word Sequence at stage 2 exacerbates forgetting. The results align with our prediction that stage 1 mixing relies on data that is unrelated to either \(D_A\) or \(D_B\), while stage 2 mixing benefits most when forgetting is severe and the mixing data aligns with \(D_A\).

4.2.0.2 Non-factoid tasks.

Figure 2 shows that the model exhibits consistent results after training on non-factoid data at stage 2. We observe that stage 1 mixing is more beneficial than stage 2 mixing across the board. However, the best mixing data varies for different \(D_A\). KVR benefits most from mixing Random Word Sequence at stage 1 (\(38.9\% \rightarrow 77.9\%\)), while Knowledge Pile benefits most on PopQA (\(58.4\% \rightarrow 82.2\%\)) and TriviaQA (\(57.9\% \rightarrow 77.8\%\)).

4.2.0.3 Applying mixing at both stages.

Based on the observation that mixing with Random Word Sequence at stage 1 and mixing Knowledge Pile at stage 2 individually benefit memorization intensive tasks, we examine if the two stages can be combined. Figure 2 shows that the combination outperforms individual stage mixing, demonstrating the possibility of composing mixing strategies. We also provide stage 2 task performance in §9.2.

4.2.0.4 Mistral results.

We report REMIX results for Mistral in Table 3 (full results in §9.3). For KVR, REMIX can successfully prevent forgetting and improve performance after stage 2 training on factoid data (\(15.0\% \rightarrow 43.5\%\)). REMIX’s advantages are less pronounced since the No Mixing baselines are not affected by forgetting too severely.

4.2.0.5 Comparison to other baselines.

We compare with three other representative baselines against REMIX in Table 4: 1) weight regularization [13], 2) behavior regularization [14], and 3) parameter expansion [38]. We use Elastic Weight Consolidation (EWC) for weight regularization and calculate the Fisher score using one backward pass using the current mini-batch for training. For behavior regularization, we add the KL between the training model vs the original reference model to the loss. We provide the parameter expansion based baseline using LoRA adaptors [39]. We randomly select parameters to train and do not explicitly avoid overlapping of the stage 1 and stage 2 tunable parameters. We observe that the weight regularization baseline and behavior regularization baselines lag behind REMIX by a large margin (\(40\%\) on KVR, \(30\%\) on PopQA, and \(40\%\) on TriviaQA).

5 Analysis↩︎

5.0.0.1 Robust memorization learns factoids in earlier layers.

We use Logit Lens [15] to decode the top \(10\) tokens from the representations at each layer using the output embedding. We record the layer index of the first occurrence of the correct token, referred to as layer of first occurrence (LoF). LoF is then normalized by the total number of occurrences. This measure indicates how early the correct token first appears. In Figure 4 (left), we compare 1) No Mixing, 2) Random / K-Pile which successfully prevents forgetting, and 3) K-Pile / None which suffers from forgetting for KVR. We notice two main differences between the two runs – first, the successful run moves the knowledge to an earlier layer, whereas the unsuccessful one does not change where the factoids are stored. The successful run also diversifies the set of layers that are used. We aggregate the mean and standard deviation of LoF over 100 examples in Figure 4 (right). The results corroborate our intuition: the model protects the factoids from interference when the knowledge is stored earlier (lower mean) and diversified (larger STD) in the layers.

5.0.0.2 REMIX enables better knowledge manipulation.

Recent work has shown that manipulating learned knowledge is challenging especially during fine-tuning  [7], [8]. We design two templates to evaluate: 1) Selective Recall and 2) Recall & Manipulate. For selective recall, the model that has memorized the factoids “X: A, Y: B” needs to answer the question “Here are two keys: X and Y. What is the value of the first key?” with A. For recall & manipulate, the model that has memorized the factoid “XYZ: ABC” needs to answer the question “If the first character in the value of key XYZ is changed to Q, what is the new value of key?” with “QBC”. We show in Table 5 that even though knowledge manipulation remains extremely hard for fine-tuning, REMIX still enables better manipulation of learned knowledge than no mixing, especially on the selective recall template.

5.0.0.3 Can REMIX go beyond two stages?

We test REMIX after more training stages to assess its effectiveness beyond the main two-stage setting. Figure 5 shows the accuracy of the Key-Value Recall task when trained on the combination of WebQA, EntityQA, MATH, and UltraChat. We observe a severe degradation when the two consecutive stages are both memorization-intensive. When the two following data are both factoid tasks, the No Mixing baseline is able to retain \(37.0\%\) accuracy. In contrast, REMIX can largely enhance the model’s ability to retain knowledge, and is robust after two stages of training, leading at least \(30\%\) accuracy above the baseline across the board.

5.0.0.4 Other ablation and analysis.

We provide extensive ablations and analysis on the effects of different mixing data, mixing data lengths, ablation of mixing ratio, and REMIX’s impact on downstream tasks in §9.7.

6 Related Work↩︎

6.0.0.1 Continual learning.

Continual learning has been the subject of investigation since early research on connectionist models, which identified catastrophic forgetting as a fundamental challenge [9], [10]. Many methods have been proposed for mitigating forgetting in continual learning. The simplest approach involves maintaining a memory of examples from previous tasks and replaying them during subsequent training [40]–[42]. Other methods involve regularization techniques that preserve important weights [13], [43] or reduce the divergence between model predictions [44]. One group of methods project the gradient for a new task to be orthogonal to the gradients from previous tasks, with the aim of reducing interference between tasks [28], [45]. A number of studies have attempted to characterize the relationship between task similarity and forgetting, empirically and theoretically [46]–[48]. In this paper, we restrict the class of approaches to those that do not change model weights, e.g., via regularization.

6.0.0.2 Memorization and forgetting in LMs.

In the context of LMs, many prior works have investigated the factors that influence memorization during pre-training [16], [49]–[51]. In particular, prior work has observed that instruction tuning can lead to some degradation on general NLP tasks, which has been called an “alignment tax” [52], [53]. [52] find that this alignment tax can be partly mitigated by mixing pre-training data into the alignment data, and [11] find that LMs forget less when the instruction-tuning data is more diverse.  [54] find that fine-tuning LMs leads to bigger performance degradation on tasks that are more similar to the fine-tuning task (as measured by likelihood under the learned fine-tuning distribution). See [55] and [56] for more extensive surveys of continual learning in the context of LMs.

6.0.0.3 Fine-tuning on unfamiliar facts.

Our work builds on several recent observations about the effect of fine-tuning an LM on unfamiliar facts. [4] find that fine-tuning LMs on unfamiliar examples (questions that the LM cannot answer correctly via few-shot prompting) leads the model to “hallucinate” plausible-sounding but incorrect answers to unfamiliar test examples. Similarly, [5] find that unknown examples take longer to learn, and learning unknown examples leads to more hallucination. These studies highlight the difficulty of encoding new facts into a model during fine-tuning. Instead of directly learning the facts, [57] and [58] study the setting where the facts are embedded in the corpora and need to be learned continually. [59] propose to address this challenge by generating synthetic data for continual pretraining. This approach can be motivated by mechanistic studies [7], [8], which have found that knowledge extraction is possible only when information appears in diverse forms in the training data (e.g. paraphrases), which leads models to encode information more effectively for later extraction.

6.0.0.4 Model editing and unlearning.

Our work is also related to a line of research aimed at explicitly modifying facts that are encoded in an LLM—for example, to update information about entities to reflect changes in the world [60]–[63]. Studies have shown that these methods can update individual facts, but do not lead to consistent changes about all of the implications of these updates [64], [65]. A related line of work has investigated whether specific information can be deliberately removed from neural networks [66], [67]. Our focus in this paper is on introducing new information while retaining existing knowledge, rather than modifying or erasing existing knowledge.

7 Conclusion↩︎

In this paper, we formalize finetuning a language model with factual knowledge in the continual memorization framework. In contrast to continual learning, which focuses on general capabilities, we focus on the specific challenges inherent to finetuning to memorize knowledge. Through careful experiments, we establish that finetuning on factoid data causes the most severe forgetting on the memorized factoids from previous stages of finetuning. We then evaluate experience replay methods that are often used in continual learning and find that they do not satisfactorily revive forgotten factoids. To address the issue of forgetting, we propose a surprisingly effective strategy REMIX. By mixing random word sequences or generic pretraining data into different stages of training, REMIX outperforms replay-based methods and other baselines in our experiments despite not using any factoids from the original set in its mixing process. Finally, we analyze REMIX using Logit Lens and ablation studies to find that it teaches the model to change where it stores facts—moving it to earlier layers or diversifying the knowledge storage location. Studying the continual memorization problem opens up many new directions for future research. For example, future work may explore REMIX and similar approaches to ensure that safety-tuning is not easily undone by further finetuning. Its efficacy poses interesting questions about the dynamics of memorization in language models, which we are excited to see investigated in future work.

Acknowledgments↩︎

We thank Alexander Wettig, Howard Yen, Tianyu Gao, Mengzhou Xia, and Yihe Dong for their valuable feedback on the manuscript. This research is supported by the National Science Foundation (IIS-2211779), a Sloan Research Fellowship, and Cisco Research.

8 Derivations for Forgetting, Replay, and REMIX↩︎

8.1 Forgetting in Continual Memorization↩︎

We give a formulation of when forgetting happens and how random and generic data mixing (REMIX) can mitigate forgetting.

We aim to analyze how mixing data during training affects memorization. Assume access to the mixing dataset \(D_M\) while learning either \(D_A\) or \(D_B\) – training on \(D_A' = D_A \cup D_M\) at stage 1 and converges to \(\theta_A'\) or \(D_B' = D_B \cup D_M\) at stage 2 and converges to \(\theta_B'\). Our goal is to examine under what condition does the following occur: \[\begin{align} \mathcal{L}(\theta_B; \mathcal{D}_A) > \mathcal{L}(\theta_B'; \mathcal{D}_A), \end{align}\] which means that through mixing, the final model \(\theta_B'\) achieves a lower loss under \(D_A\) than \(\theta_B\).

We can track the progression of the model with the following stages: \[\begin{align} \theta_A &= \theta_0 - \eta \nabla \mathcal{L}(\theta_0; \mathcal{D}_A) ~~~~ \text{(Stage 1; no mixing)} \\ \theta_B &= \theta_A - \eta \nabla \mathcal{L}(\theta_A; \mathcal{D}_B) ~~~~ \text{(Stage 2; no mixing)} \end{align}\]

Note that this is a simplification of the actual optimization process as the local one-step gradient may point in a different direction from the final parameter difference (\(\theta_A - \theta_0\)). We use \(\nabla \mathcal{L}(\theta; D)\) to represent the conceptual overall direction for model \(\theta\) to point to the low loss region of data \(D\). The goal can be expressed as the difference: \[\begin{align} \Delta &= \mathcal{L}(\theta_B; \mathcal{D}_A) - \mathcal{L}(\theta_B'; \mathcal{D}_A) \\ &= \Big( \mathcal{L}(\theta_A; D_A) + (\theta_B - \theta_A)^T \nabla \mathcal{L}(\theta_A; D_A) + \underbrace{~~~~~~R_1~~~~~~}_{\text{Higher-Order Terms}} \Big) \\ &~~~~ - \Big( \mathcal{L}(\theta_A'; D_A) + (\theta_B' - \theta_A')^T \nabla \mathcal{L}(\theta_A'; D_A) + \underbrace{~~~~~~R_2~~~~~~}_{\text{Higher-Order Terms}} \Big) \\ &= \Big( \mathcal{L}(\theta_A; D_A) - \eta \nabla \mathcal{L}(\theta_A; D_B) ^T \nabla \mathcal{L}(\theta_A; D_A) \Big) \\ &~~~~ - \Big( \mathcal{L}(\theta_A'; D_A) - \eta \nabla \mathcal{L}(\theta_A'; D_B \cup D_M)^T \nabla \mathcal{L}(\theta_A'; D_A) \Big) + (R_1 - R_2) \\ &= \underbrace{ \mathcal{L}(\theta_A; D_A) - \mathcal{L}(\theta_A'; D_A) }_{\Delta_1} \\ &~~~~ + \underbrace{ \eta \Big( \nabla \mathcal{L}(\theta_A'; D_B \cup D_M)^T \nabla \mathcal{L}(\theta_A'; D_A) - \nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L}(\theta_A; D_A) \Big) }_{\Delta_2} \\ &~~~~ + \underbrace{(R_1 - R_2)}_{\Delta_3} \end{align}\]

We assume that the first two terms \(\Delta_1, \Delta_2\) are the main source contributing to forgetting and ignore the higher-order terms.

8.2 Replay↩︎

In the replay scenario, the mixing data \(D_M\) is a subset of \(D_A\). We denote the \(r\%\) subset of \(D_A\) as \(D_A^r\). With \(D_M = D_A^r\), we can assert that \(\Delta_1 \approx 0\) since the converged model should obtain the same loss under \(D_A\) and \(D_A \cup D_A^r\). The second term \(\Delta_2 = \nabla \mathcal{L}(\theta_A'; D_B \cup D_A^r)^T \nabla \mathcal{L}(\theta_A'; D_A) - \nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L}(\theta_A; D_A) > 0\). \[\begin{align} \Delta_2 &= \nabla \mathcal{L}(\theta_A'; D_B \cup D_A^r)^T \nabla \mathcal{L}(\theta_A'; D_A) - \nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L}(\theta_A; D_A) \\ &\approx \Big( \nabla \mathcal{L}(\theta_A'; D_B \cup D_A^r) - \nabla \mathcal{L}(\theta_A; D_B) \Big)^T \nabla \mathcal{L}(\theta_A; D_A) \\ &> 0 \end{align}\]

8.3 REMIX↩︎

8.3.0.1 Mixing at stage 1: \(D_A' = D_A \cup D_M\).

\(\Delta_1 \approx 0\) due to convergence in either no mixing or mixing training scenarios. We turn to analyzing \(\Delta_2\). The term \(\nabla \mathcal{L}(\theta_A'; D_A) \approx \nabla \mathcal{L}(\theta_A; D_A) + H_A(\theta_A' - \theta_A)\) and \(\nabla \mathcal{L}(\theta_A'; D_B) \approx \nabla \mathcal{L}(\theta_A; D_B) + H_B(\theta_A' - \theta_A)\), where \(H_A\) is the Hessian of \(\nabla \mathcal{L}(\theta; D_A)\) at \(\theta=\theta_A\), and \(H_B\) is the Hessian of \(\nabla \mathcal{L}(\theta; D_B)\) at \(\theta=\theta_B\). With mixing at stage 1, we have \(\theta_A' = \theta_0 - \eta \nabla \mathcal{L}(\theta_0; D_A \cup D_M)\), which gives us \(\theta_A' - \theta_A = \eta (\nabla \mathcal{L}(\theta_0; D_A) - \nabla \mathcal{L}(\theta_0; D_A \cup D_M)) = -\eta \nabla \mathcal{L}(\theta_0; D_M)\). \[\begin{align} \Delta_2 &= \eta \Big( \nabla \mathcal{L}(\theta_A'; D_B)^T \nabla \mathcal{L}(\theta_A'; D_A) - \nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L}(\theta_A; D_A) \Big) \\ &= \eta \bigg( \Big( \nabla \mathcal{L}(\theta_A; D_B) + H_B(\theta_A' - \theta_A) \Big)^T \Big(\nabla \mathcal{L}(\theta_A; D_A) + H_A(\theta_A' - \theta_A) \Big) \\ &~~~~~~ - \nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L}(\theta_A; D_A) \bigg) \\ &= \eta \bigg( \Big( \nabla \mathcal{L}(\theta_A; D_B) + H_B (-\eta \nabla \mathcal{L}(\theta_0; D_M)) \Big)^T \Big(\nabla \mathcal{L}(\theta_A; D_A) + H_A (-\eta \nabla \mathcal{L}(\theta_0; D_M)) \Big) \\ &~~~~~~ - \nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L}(\theta_A; D_A) \bigg) \\ &= -\eta^2 \nabla \mathcal{L}(\theta_A; D_B)^T H_A \nabla \mathcal{L}(\theta_0; D_M) - \eta^2 \nabla \mathcal{L}(\theta_A; D_A)^T H_B \nabla \mathcal{L}(\theta_0; D_M) \\ &~~~~~~ + \eta^3 \nabla \mathcal{L}(\theta_0; D_M)^T H_B H_A \nabla \mathcal{L}(\theta_0; D_M) \end{align}\]

We analyze the three terms under the assumption that \(H_A\), \(H_B\), and \(H_B H_A\) are positive semi-definite. If the distributions for \(D_M\) and \(D_B\) are uncorrelated, then in expectation \(\mathbb{E}[ \nabla \mathcal{L}(\theta_A; D_B)^T H_A \nabla \mathcal{L}(\theta_0; D_M) ] = 0\). The same holds for \(D_M\) and \(D_A\). And the last term will be positive, contributing to \(\Delta_2\) and thus mitigate forgetting. Note that the norm \(\lvert \lvert \nabla \mathcal{L}(\theta_0; D_M) \rvert \rvert\) and the eigenvalues of the Hessians \(H_A\) and \(H_B\) are not bounded, which may be large and compensate for the leading \(\eta^3\). If we assume that mixing \(D_M\) does not drift the parameters away too far, making \(\lvert\lvert \nabla \mathcal{L}(\theta_A'; D_B) - \nabla \mathcal{L}(\theta_A; D_B)^ \rvert \rvert_2^2 < L_1\), and \(\lvert\lvert \nabla \mathcal{L}(\theta_A'; D_A) - \nabla \mathcal{L}(\theta_A; D_A) \rvert \rvert_2^2 < L_2\), where \(L_1, L_2\in \mathbb{R}\), we can expect the contribution to the \(\Delta_2\) term comes from the change in the angle.

8.3.0.2 Mixing at stage 2: \(D_B' = D_B \cup D_M\).

With no mixing in stage 1, we have \(A' = A\). Therefore, the first term \(\Delta_1 = \mathcal{L}(\theta_A; D_A) - \mathcal{L}(\theta_A'; D_A) = 0\) since \(D_A' = D_A\). We can also express: \[\begin{align} \Delta_2 &= \eta \Big( \nabla \mathcal{L}(\theta_A; D_B \cup D_M)^T \nabla \mathcal{L}(\theta_A; D_A) - \nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L}(\theta_A; D_A) \Big) \\ &= \eta \Big( \beta_1 \nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L} (\theta_A; D_A) + \beta_2 \nabla \mathcal{L}(\theta_A; D_M)^T \nabla \mathcal{L} (\theta_A; D_A) \\ &~~~~ - \nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L} (\theta_A; D_A) \Big) \\ &= \eta \Big( \beta_1 \nabla \mathcal{L}(\theta_A; D_M) - (1-\beta_2) \nabla \mathcal{L}(\theta_A; D_B) \Big)^T \nabla \mathcal{L} (\theta_A; D_A), \end{align}\] where \(\beta_1, \beta_2 \in [0, 1]\).

Consequently, the condition for forgetting mitigation requires \(\nabla \mathcal{L}(\theta_A; D_M)^T \nabla \mathcal{L} (\theta_A; D_A) > \frac{1-\beta_2}{\beta_1} \nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L} (\theta_A; D_A)\). This condition posits that mixing data can reduce forgetting as long as it aligns with the original data \(D_A\) more than \(D_B\). When \(D_A\) and \(D_B\) are already pointing in drastically opposite directions, making the term \(\nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L} (\theta_A; D_A)\) negative, the mixing has a higher chance to lower \(\Delta_2\). On the other hand, if \(\nabla \mathcal{L}(\theta_A; D_B)^T \nabla \mathcal{L} (\theta_A; D_A)\) is positive, it is harder for mixing to mitigate forgetting.

9 Supplementary Results↩︎

9.1 Main Results with Standard Deviation↩︎

We provide our main results with error bars over \(3\) runs in Figure 6 where stage 2 training uses factoid datasets and Figure 7 where stage 2 training uses non-factoid datasets.

9.2 Stage 2 Performance↩︎

For factoid datasets, the goal is full memorization of the trained examples. For non-factoid datasets, we also train to near perfect training accuracy (the overfitting regime) since we aim to assess the maximum disruption that training can cause. We show the corresponding accuracy in Table 6 of the main paper (for non-factoid data, we only show the ones where accuracy can be calculated). We show the representative strategies: No Mixing and Random / -. Results show near perfect accuracy for \(D_B\), indicating that learning does not hinder performance. The only exception is KVR (No Mixing)— this further highlights the benefit of REMIX in facilitating learning. With REMIX, all training reaches over \(95\%\) accuracy.

We also provide test accuracies in Table 7 for the non-factoid datasets where separate test sets are available. Note that for factoid datasets, each example is an isolated fact to be memorized exactly, therefore the notion of generalization does not apply. We observe that for KVR, REMIX improves generalization noticeably across all tasks. For PopQA and TriviaQA, REMIX’s generalization ability is close to No Mixing (within 2 point range).

We use only \(2000\) examples for all datasets during training and deliberately overfit on to induce maximal forgetting on, so the test performance level is expected. We would like to emphasize that overfitting on non-factoid is necessary for the purpose of our goal to induce the forgetting pattern in Table 2, which allows us to stress test retention of \(D_A\), or otherwise the forgetting is much less pronounced for non-factoid to begin with.

9.3 Mistral Results↩︎

We report the complete results of Mistral-7B-v0.3 in Table 8.

9.4 Replay Results↩︎

We report the full replay results in Table 9. Even though replay reduces more forgetting across the board, especially when we increase the ratio \(r\), the replay-based method does not effectively mitigate forgetting in the factoid knowledge dataset.

9.5 Forgetting in Familiar Factoid Instances↩︎

We also investigate whether REMIX can retain the memorization of familiar factoid instances after directly fine-tuning on both factoid and non-factoid data in stage 2. After fine-tuning in stage 2, we evaluated the familiar instances from the factoid dataset \(D_A\). The evaluation results for Llama-3-8B are shown in Table 10. We observe that mixing Knowledge-Pile, Arxiv-Pile, and FineWeb with factoid data in stage 2 helps mitigate the forgetting of familiar factoid instances for both Llama-3-8B and Mistral-7B-v0.3, aligning with the results in Figure 6.

9.6 Probing Results↩︎

9.7 Ablations↩︎

9.7.0.1 Ablating mixing data length.

Figure 10 shows the effect of sequence length when using Random Word Sequences and Knowledge Pile for mixing. We observe that longer Random Word Sequences hurt the performance, highlighting the risk of incorporating wildly out of distribution data. On the other hand, Knowledge Pile also saturates after \(50\) words, indicating the limits of the generic data. The ablation also affirms that the role of the mixing data serves as a way to manipulate the memorization dynamics as opposed to providing extra information.

9.7.0.2 Effect of mixing ratio.

We show in Figure 11 the model’s KVR performance under varying mixing ratio across all stage 2 tasks. We observe that stage 2 mixing is particularly sensitive to the increase of mixing ratio. On the other hand, stage 1 mixing enjoys less decrease or even increase in performance as the mixing ratio goes up, suggesting a different memorization dynamics than stage 1.

9.7.0.3 Impact on downstream performance.

Intuitively, adding random word sequences might risk disrupting capabilities in other domains. We evaluate the model’s performance on MMLU [68] shown in Table 11. We observe REMIX maintains better performance across the board compared to no mixing.

9.7.0.4 Effect of different mixing data.

We investigate how the choice of the mixing data impacts the results for factoid-tasks. Figure 12 shows no difference between Knowledge Pile and other generic mixing data such as ArXiv Pile and FineWeb. This affirms that the effectiveness of REMIX does not rely on Knowledge Pile’s potential distributional overlap with memorization-intensive tasks.

9.8 Training Details↩︎

For all experiments with Llama-3-8B, we average the results over three seeds and use a learning rate of 5e-5. For all experiments with Mistral-7B-v0.3, we use a learning rate of 1e-5. For experiments measuring forgetting of familiar factoid datasets, we use a batch size of \(128\). For the rest of the experiments, we set the batch size to \(32\). Additionally, different stopping conditions are applied for the different factoid datasets: for the KVR task, we use a fixed number of epochs (\(20\)), while for other factoid tasks, training stops when the loss drops below \(0.0001\). We provide our training prompt in §9.9.

9.9 Dataset Examples and Prompts↩︎

We provide examples of the stage 1 factoid datasets \(D_A\) and the mixing datasets \(D_M\). Since stage 2 non-factoid datasets are standard instruction tuning datasets, we omit these examples in the following sections.

9.9.1 Factoid dataset (\(D_A\)) Examples↩︎

1. Key-Value Recall
   Input:  The value of key e6395973 is?
Target: 8219acf2

2. PopQA
   Input:  Question: What is New Lands’s author? The answer is:
Target: Charles Fort

3. TriviaQA
   Input:  Question: Which city does David Soul come from? The answer is:
Target: Chicago

9.9.2 Mixing data (\(D_M\)) Examples↩︎

1. Knowledge-Pile
   Input: Complete the following partial passage: Processing hyperspectral images allows you to decode images and recognize objects in the scene on the base of analysis of spectrums. In some problems, information about the spectra may not be sufficient. In this case, visualization of data sets may use, for object recognition, by use additional non-formalized external attributes
Target: (for example, indicating the relative position of objects). Target visualization is a visualization adapted to a specific task of application. The method discussed in this chapter uses a way to visualize a measure of similarity to the sample. As a result of the transformation, the hyperspectral (multichannel) image is converted [...]

2. Arxiv-Pile
   Input: Complete the following partial passage: --- abstract: ’The purpose of this article is to study the problem of finding sharp lower bounds for the norm of the product of polynomials in the ultraproducts of Banach spaces\((X_i)_{\mathfrak U}\). We show that, under certain hypotheses, there is a strong relation between this problem and the same
Target: problem for the spaces\(X_i\).’ address: ’IMAS-CONICET’ author: - Jorge Tomás Rodríguez title: On the norm of products of polynomials on ultraproducts of Banach spaces --- Introduction ============ In this article we study the factor problem in the context of ultraproducts of Banach spaces. This problem can be stated as [...]

3. FineWeb
   Input: Complete the following partial passage: *sigh* Fundamentalist community, let me pass on some advice to you I learned from the atheistic community: If you have set yourself on fire, do not run. Okay? Okay?? Please? Look, D, you had two months to say to Harvard in private emails, "I’m sorry, I shouldn’t have been using
Target: that animation in my paid presentations. I wont use it again. I really do like ’Inner Life’, though, and would love to use it in classroom presentations, from the BioVisions site, if that is acceptable." I sat here, for two months, waiting for that to happen, anything to happen, and [...]

4. Random Word Sequence
   Input: Memorize the following random-string passage: pliosaur bismuth assertoric decentralization emerse redemonstrate sleepwaker Coracias thirstland Stercorariinae Cytherean autobolide pergamentaceous ophthalmodynamometer tensify tarefitch educement wime cockneity holotype spreng justiciary unseparate ascogonial chirimen Styphelia emotivity heller hystazarin unthinkable Corinth vicianose incommunicative sorcerous lineograph dochmiacal heresiographer interrenal anes mercal embryogenic swoon diptote funniness unwreathed contection rhapsodical infolding colorature multifurcate
Target: pliosaur bismuth assertoric decentralization emerse redemonstrate sleepwaker Coracias thirstland Stercorariinae Cytherean autobolide pergamentaceous ophthalmodynamometer tensify tarefitch educement wime cockneity holotype spreng justiciary unseparate ascogonial chirimen Styphelia emotivity heller hystazarin unthinkable Corinth vicianose incommunicative sorcerous lineograph dochmiacal heresiographer interrenal anes mercal embryogenic swoon diptote funniness unwreathed contection rhapsodical infolding colorature multifurcate

9.10 Dataset Details↩︎

We examine the strict overlap of knowledge entities between PopQA, TriviaQA, and the generic data used for mixing. By extracting knowledge entity pairs from the questions and target answers, we calculate the exact overlap between these pairs. The overlap percentage among PopQA, TriviaQA, and the generic data is less than 1.3%. Note that the entity overlap ratio calculated is an overestimate. For example, parsing the instance “Question: Behind Russia, what is the second largest country in Europe? Answer: Ukraine” results in entities like “Russia”, “Europe”, and “Ukraine”. It’s easy to find in generic corpora a text that mentions all the entities, yet the relation “the second largest" is not present. This is different from the replay data which is the exact instance that contains this relation. If we apply strict entity parsing rules and only extract tags such as PERSON, PRODUCT, ORG, then the overlap is close to zero (\(\sim0.1\%\)).

9.11 Caveats for Interpreting Table 1↩︎

We provide a more nuanced reading of Table 1. First, the factoid and non-factoid \(D_B\)’s differ along more dimensions beyond factoid vs. non-factoid, and we do not control for all the possible variables. Second, individual datasets do not always follow the aggregate trend. For example, GSM8K (non-factoid) induces severe forgetting on PopQA (19.0%) and TriviaQA (9.4%), comparable to or exceeding some factoid \(D_B\)’s. Conversely, WebQA (factoid) causes relatively mild forgetting on KVR (33.8%) and TriviaQA (68.6%). The overall pattern is broadly consistent with findings from the continual learning literature suggesting that catastrophic forgetting tends to be more pronounced when two tasks are similar, potentially causing greater interference [28]–[30]. The trend warrants further controlled investigation to disentangle the effect of task type from other confounding factors.

References↩︎