PaperlayerView exact revision on arXiv

Source-linked AI summary

Continual Memorization of Factoids in Language Models

Howard Chen, Jiayi Geng, Adithya Bhaskar, Dan Friedman, Danqi Chen

arXiv:2411.07175v3cs.CL

TL;DR

Fine-tuning can fail to reliably store new facts, motivating study of whether models retain them through repeated training. The paper evaluates continual memorization and finds that REMIX mixing substantially mitigates forgetting, including in severe cases, while robust memorization is associated with earlier and more diversified layer storage.

  • Problem

    Fine-tuning for new knowledge can fail to reliably memorize it, raising the question of whether it supports repeated knowledge acquisition.

    Source: §1 · pp. 1–2

  • Method

    The paper studies continual memorization by training on factoids and then measuring their retention after training on a subsequent dataset, and evaluates REMIX data mixing.

    Source: §1 · pp. 2–3, §2.1 · p. 3

  • Results

    In the most severe case, REMIX raises post-stage 2 accuracy from 13.5% to 53.2%, while replay reaches 41.6% despite using 10% of stage 1 factoids.

    Source: §1 · p. 3

  • Takeaways & Limitations

    REMIX is associated with more diversified fact storage in earlier layers and with easier recall and manipulation of robustly memorized factoids.

    Source: §1 · p. 3

    The theoretical analysis assumes the loss is differentiable and locally well-approximated by a first-order expansion around θA.

    Source: §4.1 · p. 6

Abstract

from arXiv · show

As new knowledge rapidly accumulates, language models (LMs) with pretrained knowledge quickly become obsolete. A common approach to updating LMs is fine-tuning them directly on new knowledge. However, recent studies have shown that fine-tuning for memorization may be ineffective in storing knowledge or may exacerbate hallucinations. In this work, we introduce a setting we call continual memorization, where a model must memorize and retain a set of factoids through multiple stages of fine-tuning on subsequent datasets. We characterized the forgetting patterns through extensive experiments and show that LMs widely suffer from forgetting, especially when needing to memorize factoids in the second stage. We posit that forgetting can be alleviated by modifying training dynamics: (1) protecting the memorization process when learning factoids or (2) reducing interference from subsequent training stages. Intriguingly, we find that mixing randomly generated word sequences or generic data sampled from pretraining corpora at different training stages effectively mitigates forgetting REMIX: Random and Generic Data Mixing). REMIX can recover performance from severe forgetting, outperforming replay methods and other continual learning baselines. We analyze how REMIX influences the learning process and find that robust memorization follows a distinct pattern: the model stores factoids in earlier layers than usual and diversifies the layers that retain them, which results in easier recall and manipulate of the learned factoids.

1 Introduction

The paper frames repeated fine-tuning as a fragile route to retaining new facts and studies this through continual memorization. It finds severe forgetting when later training uses factoids, while REMIX mixing substantially improves retention and changes where facts are stored.

1 Introduction

  • Continual memorization tests whether factoids learned in an initial fine-tuning stage remain after training on a subsequent dataset.The paper focuses on memorization dynamics, rather than general capabilities or broad language-modeling loss.Source: §1 · p. 2
  • TriviaQA accuracy falls from 100% after stage 1 to 39.8% after training on other factoid datasets, and long-tail key-value accuracy drops to 13%.Forgetting is less pronounced after non-factoid training, including coding, math, or chat.Source: §1 · p. 2
  • In the most severe case, REMIX raises post-stage 2 accuracy from 13.5% to 53.2%, while replay reaches 41.6% despite using 10% of stage 1 factoids.REMIX also outperforms weight regularization (EWC) and behavior regularization across several stage 2 task choices.Source: §1 · p. 3
  • REMIX is associated with storing learned facts in earlier layers and diversifying their storage across layers, alongside easier recall and manipulation.The analysis attributes reduced second-stage interference in part to jointly learning mixing data and stage 2 data, which helps prevent overfitting to a narrow distribution.Source: §1 · p. 3
  • REMIX combines random word sequences and generic pretraining data during training stages, without access to factoids from prior stages.The paper presents mixing as protecting learning in the first stage or reducing interference in the second.Source: §1 · p. 3

2 Continual Memorization of Factoids

The section defines continual memorization as retaining unfamiliar factoids after subsequent training, then specifies the factoid and non-factoid datasets and evaluation setup used to study it.

2.1 Problem Definition

  • A factoid is a subject–relation–object triplet formatted as a prompt–response pair, while non-factoid datasets are other prompt–response datasets.The section defines average exact-match accuracy and calls a factoid familiar when the model's accuracy on it is 1.Source: §2.1 · p. 3
  • Continual memorization trains a pretrained model to convergence on factoid dataset DA, trains it further on dataset DB, then measures retention on DA.The setting considers two or more stages; the stated experiments treat factoid datasets as unfamiliar.Source: §2.1 · p. 3

2.2 Constructing Factoid Datasets

  • Stage 1 uses 2,000 generated key-value pairs or 2,000 sampled unfamiliar PopQA or TriviaQA examples.Datasets are synthetic or filtered to remove facts familiar from pretraining.Source: §2.2 · p. 4
  • Stage 2 includes factoid datasets such as LAMA, Entity Questions, and WebQA, and non-factoid datasets for chat, code, and math.Experiments initialize from Llama-3-8B or Mistral-7B, use the Tulu-v2 prompt template, and report exact-string-match accuracy averaged across three runs.Source: §2.2 · p. 4

3 How Do Models Forget Factoids?

Forgetting is most severe when a second-stage dataset is also factoid, while replay reduces forgetting less in that setting than when stage 2 is non-factoid.

3.1 Understanding the Forgetting Patterns

  • After factoid stage-2 training, KVR, PopQA, and TriviaQA average 13.5%, 47.0%, and 39.8% accuracy, respectively.With LAMA, they fall further to 2.1%, 7.7%, and 4.3%.Source: §3.1 · p. 4, Table 1 · p. 5

3.2 Replay Does Not Mitigate Forgetting Fully

  • Replay reduces forgetting across the board, but is less effective when stage 2 is factoid than when it is non-factoid.Figure 2 reports replay results averaged across stage-2 datasets for four mixing ratios.Source: §3.2 · p. 5, Figure 2 · p. 5

4 REMIX: Random and Generic Data Mixing

REMIX mixes random word sequences or generic text into one or both training stages to alter memorization dynamics. It mitigates forgetting, with outcomes depending on the mixing data and stage.

4.1 Method

  • REMIX mixes random word sequences or generic text, with no overlap with the factoids being memorized in stage 1, into the stage datasets.The paper uses Random Word Sequence and Knowledge Pile as its main mixing datasets.Source: §4.1 · pp. 6–7
  • Stage-1 mixing can reduce opposition between the factoid-task gradients, while stage-2 mixing can align the effective update with the stage-1 gradient.The proposed intuition is that the first strategy protects memorized facts and the second reduces interference from subsequent training.Source: §4.1 · p. 6

4.2 Results

  • At stage 2, Knowledge Pile raises KVR accuracy from 13.5% to 27.8%, PopQA from 47.0% to 86.4%, and TriviaQA from 39.8% to 90.6%.Stage-2 Random Word Sequence mixing instead exacerbates forgetting.Source: §4.2 · p. 7
  • At stage 1, Random Word Sequences improve KVR from 13.5% to 28.8%, PopQA from 47.0% to 53.7%, and TriviaQA from 39.8% to 51.0%.Stage-1 Knowledge Pile mixing hurts performance on these factoid tasks.Source: §4.2 · p. 7
  • On non-factoid stage-2 tasks, stage-1 mixing is more beneficial than stage-2 mixing, though the best data source varies by task.KVR benefits most from stage-1 Random Word Sequence mixing (38.9% → 77.9%); Knowledge Pile helps PopQA (58.4% → 82.2%) and TriviaQA (57.9% → 77.8%).Source: §4.2 · p. 7
  • Combining stage-1 Random Word Sequence mixing with stage-2 Knowledge Pile mixing outperforms either stage's mixing alone.The combination demonstrates that stage-specific mixing strategies can be composed.Source: §4.2 · pp. 7–8
  • On Mistral-7B, REMIX raises KVR accuracy from 15.0% to 43.5% after factoid stage-2 training.Its advantage is less pronounced when No Mixing baselines experience less forgetting.Source: §4.2 · p. 8, Table 3 · p. 8
  • Weight- and behavior-regularization baselines lag behind REMIX by 40% on KVR, 30% on PopQA, and 40% on TriviaQA.The comparison is against factoid datasets at stage 2.Source: §4.2 · p. 8, Table 4 · p. 8

5 Analysis

REMIX improves retention across additional training stages and supports better manipulation of memorized factoids. Its analysis links robust retention to storing facts earlier and across a more diverse set of layers.

5 Analysis

  • Successful mixing moves factoid information to earlier layers and diversifies the layers storing it, patterns associated with better protection from interference.Figure 4 probes correct-token occurrence across layers and summarizes layer of first occurrence over 100 examples; lower mean and higher standard deviation correlate with better performance.Source: §5 · p. 9, Figure 4 · p. 9
  • REMIX improves knowledge manipulation over No Mixing, especially on the selective-recall template, although manipulation remains difficult for fine-tuning.The templates test selecting a value by key and changing the first character of a memorized value.Source: §5 · p. 9, Table 5 · p. 9
  • After two further training stages, REMIX retains at least 30% more Key-Value Recall accuracy than the No Mixing baseline across the tested sequences.When both subsequent datasets are factoid tasks, No Mixing retains 37.0% accuracy.Source: §5 · p. 10, Figure 5 · p. 10

6 Related Work

Related work spans continual-learning methods, studies of memorization and forgetting in language models, fine-tuning on unfamiliar facts, and model editing. This paper focuses on adding new information while retaining existing knowledge.

Continual learning

  • Continual-learning approaches include replaying earlier examples, regularizing weights or predictions, and projecting new-task gradients to reduce interference.Prior work also studies how task similarity relates to forgetting.Source: §6 · p. 10

Memorization and forgetting in LMs

  • Language-model studies examine pretraining memorization, alignment tax, instruction-data diversity, and how fine-tuning similarity relates to degradation.Mixing pretraining data has partly mitigated alignment tax, while more diverse instruction-tuning data has been associated with less forgetting.Source: §6 · pp. 10–11

Fine-tuning on unfamiliar facts

  • Fine-tuning on unfamiliar examples has been linked to hallucination and slower learning, while other work studies continual learning from facts embedded in corpora or synthetic data.Mechanistic studies report that diverse forms such as paraphrases support more effective later extraction.Source: §6 · p. 11

Model editing and unlearning

  • Model editing can update individual facts without consistently changing all their implications, whereas unlearning research studies deliberate information removal.This paper instead addresses introducing information while retaining existing knowledge.Source: §6 · p. 11

7 Conclusion

The paper formalizes continual memorization and finds that factoid training causes severe forgetting, while replay does not adequately restore forgotten facts. It proposes REMIX, which outperforms replay-based methods and other baselines in the reported experiments, and analyzes its effects on knowledge storage.

7 Conclusion

  • Continual memorization formalizes fine-tuning a language model to acquire factual knowledge over successive stages.The framework focuses on memorization challenges rather than general capabilities.Source: §7 · p. 11
  • Factoid training produces the most severe forgetting of previously memorized factoids, and replay methods do not satisfactorily revive them.Source: §7 · p. 11
  • Mixing random word sequences or generic pretraining data with training stages makes REMIX outperform replay-based methods and other baselines without mixing in original factoids.The analysis finds that REMIX shifts storage toward earlier layers or diversifies where knowledge is stored.Source: §7 · p. 11

A Derivations for Forgetting, Replay, and REMIX

The appendix derives conditions for forgetting and for mitigating it through replay or REMIX, then reports main results with variation across runs. Its derivations rely on simplifying assumptions about optimization and the dominant terms.

A.1 Forgetting in Continual Memorization

  • The forgetting analysis simplifies optimization to conceptual gradient directions and treats the first two terms, Δ1 and Δ2, as the main contributors while ignoring higher-order terms.The local one-step gradient may differ from the final parameter change.Source: §A.1 · p. 16

A.2 Replay

  • In the replay scenario, the mixing dataset DM is a subset of the earlier dataset DA; the derivation assumes convergence makes Δ1 approximately zero.Source: §A.2 · p. 16

A.3 REMIX

  • REMIX can mitigate forgetting when mixing gradients align with DA more than DB, and stage mixing may alter gradient alignment or the parameter starting position.The analysis assumes positive semi-definite Hessians and uncorrelated mixing and task distributions for one expectation argument.Source: §A.3 · p. 17

B.1 Main Results with Standard Deviation

B.2 Stage 2 Performance

The section reports how stage-2 training affects performance and retention, including comparisons with replay and tests on familiar factoid instances.

B.2 Stage 2 Performance

  • With REMIX, all stage-2 training reaches over 95% accuracy, while KVR is the exception under No Mixing.The experiments train DB datasets to full convergence to induce maximal forgetting in DA.Source: §B.2 · p. 19, Table 6 · p. 20
  • For KVR and TQA, REMIX improves generalization over the no-mixing baseline; PopQA and TriviaQA remain within 2 points of No Mixing.The table caption identifies training and test accuracy for DB datasets after stage-2 training.Source: §B.2 · pp. 19–20, Table 7 · p. 20
  • The experiments use 2,000 examples and deliberately overfit non-factoid datasets to stress-test retention; without this, their forgetting is much less pronounced.Source: §B.2 · p. 20

B.4 Replay Results

  • Replay reduces forgetting more across the board as its ratio increases, but does not effectively mitigate forgetting in the factoid knowledge dataset.Source: §B.4 · p. 20

B.5 Forgetting in Familiar Factoid Instances

  • Mixing Knowledge Pile, Arxiv-Pile, or FineWeb in stage 2 helps retain familiar factoid instances for both Llama-3-8B and Mistral-7B-v0.3.The Llama-3-8B results are shown in Table 10.Source: §B.5 · p. 20

B.6 Probing Results

The supplied captions identify a Mistral-7B-v0.3 comparison and a replay evaluation, but do not report probing outcomes.

B.6 Probing Results

  • Table 8 compares No Mixing with REMIX using random word sequences at stage 1 and Knowledge Pile at stage 2 for Mistral-7B-v0.3.Source: Table 8 · p. 21
  • Table 9 reports replay accuracy on DA across four replay ratios after stage-2 training on unfamiliar factoid and non-factoid datasets.Source: Table 9 · p. 21

B.7 Ablations

The ablations show that mixing-data length, stage, and ratio matter, while several generic data sources support REMIX; training settings and example formats are also specified.

B.7 Ablations

  • Longer random-word sequences hurt performance, while Knowledge Pile saturates after 50 words, suggesting mixing changes memorization dynamics rather than adding information.Figure 10 plots Key-Value Recall accuracy against sequence length for stage-1 random sequences and stage-2 Knowledge Pile.Source: §B.7 · p. 21, Figure 10 · p. 23
  • Increasing the mixing ratio particularly affects stage-2 mixing, whereas stage-1 mixing decreases less or can improve performance.Figure 11 plots Key-Value Recall accuracy against mixing ratio for both stages.Source: §B.7 · p. 21, Figure 11 · p. 23
  • REMIX maintains better MMLU performance across the board than No Mixing.Table 11 compares No Mixing with REMIX using random word sequences at stage 1 and Knowledge Pile at stage 2.Source: Table 11 · p. 22, §B.7 · p. 22
  • Knowledge Pile, ArXiv Pile, and FineWeb show no difference on factoid tasks, so REMIX’s effectiveness does not depend on Knowledge Pile overlap.Figure 12 compares these generic mixing sources on KVR.Source: §B.7 · pp. 22–23, Figure 12 · p. 24

B.8 Training Details

  • Experiments average Llama-3-8B results over three seeds; KVR uses 20 epochs, while other factoid tasks stop when loss falls below 0.0001.Mistral-7B-v0.3 uses a 1e-5 learning rate; batch size is 128 for familiar-factoid forgetting experiments and 32 otherwise.Source: §B.8 · p. 24

B.9 Dataset Examples and Prompts

  • The dataset examples include stage-1 factoid data and mixing data; standard instruction-tuning examples for stage-2 non-factoid data are omitted.Source: §B.9 · p. 24

B.9.1 Factoid dataset (DA) Examples

  • Factoid examples show Key-Value Recall prompts and question-answer pairs, with a target object as the response.Examples include an artificial key-value pair and factual questions about an author and a city.Source: §B.9.1 · p. 24

B.9.2 Mixing data (DM) Examples

  • Mixing-data examples include passages from technical, academic, and informal text, plus a random-string passage used as a memorization target.Source: §B.9.2 · pp. 24–25

B.10 Dataset Details

The dataset analysis finds little strict entity overlap between generic mixing data and factoid tasks, while the Table 1 pattern has important dataset-level caveats.

B.10 Dataset Details

  • Entity overlap among PopQA, TriviaQA, and generic mixing data is below 1.3%, an overestimate; strict parsing of entity types brings it close to 0.1%.The example question yields entity mentions even when the relation itself is absent from generic-corpus text.Source: §B.10 · pp. 25–26

B.11 Caveats for Interpreting Table 1

  • Factoid versus non-factoid status is confounded with other dataset differences, and individual datasets can depart from the aggregate forgetting pattern.The authors call for further controlled investigation to disentangle task type from confounders.Source: §B.11 · p. 26
Loading 2411.07175v3…