Continual Memorization of Factoids in Large Language Models

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

2024-11-12

Abstract:Large language models can absorb a massive amount of knowledge through pretraining, but pretraining is inefficient for acquiring long-tailed or specialized facts. Therefore, fine-tuning on specialized or new knowledge that reflects changes in the world has become popular, though it risks disrupting the model's original capabilities. We study this fragility in the context of continual memorization, where the model is trained on a small set of long-tail factoids (factual associations) and must retain these factoids after multiple stages of subsequent training on other datasets. Through extensive experiments, we show that LLMs suffer from forgetting across a wide range of subsequent tasks, and simple replay techniques do not fully prevent forgetting, especially when the factoid datasets are trained in the later stages. We posit that there are two ways to alleviate forgetting: 1) protect the memorization process as the model learns the factoids, or 2) reduce interference from training in later stages. With this insight, we develop an effective mitigation strategy: REMIX (Random and Generic Data Mixing). REMIX prevents forgetting by mixing generic data sampled from pretraining corpora or even randomly generated word sequences during each stage, despite being unrelated to the memorized factoids in the first stage. REMIX can recover performance from severe forgetting, often outperforming replay-based methods that have access to the factoids from the first stage. We then analyze how REMIX alters the learning process and find that successful forgetting prevention is associated with a pattern: the model stores factoids in earlier layers than usual and diversifies the set of layers that store these factoids. The efficacy of REMIX invites further investigation into the underlying dynamics of memorization and forgetting, opening exciting possibilities for future research.

Computation and Language

What problem does this paper attempt to address?

The problem that this paper attempts to solve is to achieve continual memorization of factoids in large - language models (LLMs). Specifically, the research focuses on how, after the model has undergone multi - stage training, it can still retain the initially learned specific long - tail factual data (such as subject - relation - object associations) without forgetting this knowledge. The paper points out that although large - language models can absorb a large amount of knowledge through pre - training, they are inefficient in learning long - tail or specialized facts. Therefore, the fine - tuning method is usually used to make the model learn new or specialized knowledge, but this may undermine the original capabilities of the model. In addition, traditional replay techniques are of limited effectiveness in preventing forgetting, especially when performing poorly on factual data sets trained in later stages.

The paper proposes an effective mitigation strategy - REMIX (Random and Universal data MIXing). By mixing in general - purpose data sampled from the pre - training corpus or randomly generated word sequences at each training stage, it prevents forgetting. Research shows that the REMIX method can significantly improve the performance of the model after severe forgetting, and is generally superior to replay methods that can access the first - stage factual data. Further analysis shows that REMIX effectively protects the learned knowledge by making the model store facts in earlier layers and dispersing the storage locations of these facts.