Gradient Localization Improves Lifelong Pretraining of Language Models (2024)

Abstract: Large Language Models (LLMs) trained on web-scale text corpora have been shown to capture world knowledge in their parameters. However, the mechanism by which language models store different types of knowledge is poorly understood. In this work, we examine two types of knowledge relating to temporally sensitive entities and demonstrate that each type is localized to different sets of parameters within the LLMs. We hypothesize that the lack of consideration of the locality of knowledge in existing continual learning methods contributes to both: the failed uptake of new information, and catastrophic forgetting of previously learned information. We observe that sequences containing references to updated and newly mentioned entities exhibit larger gradient norms in a subset of layers. We demonstrate that targeting parameter updates to these relevant layers can improve the performance of continually pretraining on language containing temporal drift.

## introduction pretraining over diverse datasets has been shown to encode world knowledge in the parameters of large language models (llms) @xcite @xcite from massive static web-scale datasets. however, these models are normally trained on large static text corpora which do not reflect changes in world knowledge or language usage that occur after the initial data collection. in practice language models are deployed in dynamic real-world settings, and their learned knowledge becomes stale over time; the temporal degradation can be evaluated according to intrinsic measures such as perplexity, or extrinsic downstream performance (e.g. question answering) @xcite @xcite @xcite . incrementally training language models on streams of data has been explored as a method to mitigate temporal performance degradation without incurring the heavy computational and environmental costs of retraining models on large pretraining corpora @xcite @xcite . however, naive online finetuning on these datastreams is known to: induce hallucinations in model generations @xcite , fail to uptake new information @xcite , and catastrophically forget previously learned information @xcite . to address these problems, recent work has applied continual learning and online learning methods to adapting large language models to streams of documents @xcite @xcite as one potential solution, continual pretraining has been shown to improve performance when training on a sequence of natural language domains @xcite , but these methods often fail to acquire new knowledge @xcite . while continual learning methods have been shown to mitigate temporal degradation on the task-level, the mechanisms by which neural language models store and update information are not well understood. in this work, we consider a real-world use case of continual language learning setting, that of temporal language drift, and probe the performance of language models on two types of entity relationships which exhibit temporal degradation: (1) acquisition of information about new entities, and (2) updating relationships between existing entities. we hypothesize that the poor performance of existing continual learning methods on these tasks can be in part attributed to a misalignment in the autoregressive language modeling pretraining objective and the ideal parameter updates required to acquire new information or update existing knowledge. as an indicator of this misalignment, we examine models' gradient updates computed on knowledge intensive salient entity spans and compare them with those seen instandard continual pretraining, and observe that the gradient norms observe high values in distinct groups of layers based on the type of entity relationship presented in the sequence (see fig. 1 ). based on these observations, we propose new methods for aligning the updates steps during continual pretraining which better align with the parameter updates with these layers with high gradient norms. through empirical study, we show that the observed characteristic gradient patterns occur across autoregressive, transformer language models of various of sizes; and we demonstrate the efficacy of our proposed method through performance improvements on knowledge probing tasks when applied on top of existing continual learning methods in pretraining. ## knowledge probing using salient span we probe language models using the task of salient span prediction, which has previously shown success as a pretraining objective for knowledgeintensive tasks such as closed-book question answering @xcite . in salient span prediction, a model is provided with a sequence and tasked with completing a masked slot corresponding to a named entity or noun phrase. specifically, we examine language models on probing tasks for temporal entity knowledge in which the masked sequence corresponds: (1) to an update or change to an existing temporally sensitive entities; (2) to a mention of emerging new entities that were not previously seen during pretraining. ## gradient localized continual pretraining ideally, naive pretraining of a language model on a changing stream of data would be sufficient to update a model to capture the relevant changes in knowledge. however, recent work has demonstrated that current methods for continual learning often suffer from both catastrophic forgetting and a failure to uptake new knowledge even when it is directly contained in the training corpus @xcite . we hypothesize that one cause of failed transfer is due to a misalignment of the gradients from the nll objective function with the desired update based on the information content of the data observed during continual pretraining. from our observations from §3, we hypothesize that the acquisition of entity knowledge can be improved by amplifying updates to the layers that are relevant to the learning of salient entity spans. to identify relevant layers, we compute the relative gradient norm for each layer i as: the ratio between the gradient norm ∇i in the layer i w.r.t. randomly sampled data from the continual pretraining data stream, and data sampled from the validation set of the templama diagnostic dataset: we propose two methods to improve knowledge uptake by aligning gradient updates during continual pretraining. for relevant salient spans from the templama diagnostic dataset, we construct a profile of the relative gradient norms with respect to the gradients for randomly sampled pretraining sequences. we then adjust the learning rates for layers in this profile to increase the updates to layers with large relative gradient norms. we refer to our methods as traced gradient layers (tgl). ## training and dataset details to perform domain adaptive pretraining, we sample and preprocess a snapshot of wikipedia from january 2019 using wikiextractor. for continual pretraining, we follow the methodology of @xcite to collect snapshots of wikipedia from each of the subsequent years until 2022 and filter each corpus to contain the edits to wikipedia made in the intervening year, consisting of new articles and sentences within existing articles that were edited between succeeding snapshots. ## conclusion in this work, we proposed traced gradient layers (tgl) a method for identifying relevant layers to target during continual pretraining of language models. we observe that our proposed approach improve language model performance on tasks probing for entity and relational knowledge; without the need for fine-grained annotations.