Towards Understanding Continual Factual Knowledge Acquisition of Language Models: From Theory to Algorithm
Abstract
Continual Pre-Training (CPT) is essential for enabling Language Models (LMs) to integrate new factual knowledge without erasing old. While classical CPT techniques like data replay have become the standard paradigm, the mechanisms underlying how LMs acquire and retain facts over time, termed as continual Factual Knowledge Acquisition (cFKA), remain unclear. In this work, we present a theoretical framework that characterizes the training dynamics of cFKA using a single-layer Transformer with linear attention, offering a unified explanation for the behavior of popular CPT methods. Our analysis reveals that regularization-based methods merely adjust the convergence rate of parameters without altering the inherent forgetting tendency, whereas data replay methods shift convergence dynamics and stabilize pretrained knowledge. Building on these insights, we propose a novel generative data replay approach, called Selecting Tokens via attentiOn Contribution (STOC), which identifies influential factual snippets to guide replay generation. Extensive experiments on both synthetic and real-world datasets validate our theoretical findings and demonstrate that STOC effectively enhances cFKA by mitigating catastrophic forgetting.
Lay Summary
AI systems could be highly useful in areas such as medical triage, legal consultation, and mental health support. However, before these systems can be safely deployed, they must acquire reliable knowledge about the specific domain in which they will be used. One way to achieve this is to continue training a Large Language Model like ChatGPT on high-quality domain-specific data. Unfortunately, this process can cause the AI model to lose some of its previously learned general abilities. For example, after such training, the model may make common-sense mistakes, fail to understand user conversations, or forget knowledge it once possessed. This work investigates why this loss happens. We begin with an extramely simple neural network and mathematically analyze its learning process, which helps us understand the mechanisms behind the loss of prior knowledge and capabilities. Based on this understanding, we propose a new training technique called STOC. STOC is designed to help AI models learn specialized domain knowledge while preserving their existing general-purpose abilities. Using the legal domain as an example, we show that STOC allows the model to learn legal knowledge effectively while substantially reducing the collapse of its earlier knowledge and capabilities.