Meta-Learning Online Adaptation of Language Models (2023)

Abstract: Large language models encode impressively broad world knowledge in their parameters. However, the knowledge in static language models falls out of date, limiting the model’s effective “shelf life.” While online fine-tuning can reduce this degradation, we find that naively fine-tuning on a stream of documents leads to a low level of information uptake. We hypothesize that online fine-tuning does not sufficiently attend to important information. That is, the gradient signal from important tokens representing factual information is drowned out by the gradient from inherently noisy tokens, suggesting that a dynamic, context-aware learning rate may be beneficial. We therefore propose learning which tokens to upweight. We meta-train a small, autoregressive model to reweight the language modeling loss for each token during online fine-tuning, with the objective of maximizing the out-of-date base question-answering model’s ability to answer questions about a document after a single weighted gradient step. We call this approach Context-aware Meta-learned Loss Scaling (CaMeLS). Across three different distributions of documents, our experiments find that CaMeLS provides substantially improved information uptake on streams of thousands of documents compared with standard fine-tuning and baseline heuristics for reweighting token losses.

## introduction large language models learn impressively broad world knowledge through large-scale unsupervised pre-training, which they can leverage for a wide variety of downstream tasks @xcite @xcite . however, large language models are typically static artifacts, and as the world changes, the knowledge encoded in their parameters becomes stale. while retrieval-augmented models are one approach to * equal contribution. figure 1 : the proposed method camels learns to rescale the per-token online loss, sparsifying the fine-tuning gradients to emphasize informative timesteps. the middle row shows the weights output by camels. the top and bottom rows show raw and weighted per-token gradient norms, respectively. mitigating the staleness issue, even very large language models often fail to correctly update their memorized predictions when presented with counterfactual retrieved information @xcite @xcite . moreover, purely parametric language models are uniquely suited for edge computing due to their compact size (relative to a large retrieval index) and simplicity of inference @xcite . recent work has thus considered variants of online fine-tuning on a stream of documents to efficiently perform direct updates to the knowledge inside of a large language model @xcite . ideally, we could simply fine-tune a language model on an online stream of documents, and the information contained in those documents would be readily available for the model to use in a variety of downstream tasks such as answering questions about the information in the documents. unfortunately, we find that in this online adaptation setting, fine-tuning with a well-tuned learning rate leads to a nearly negligible improvement in a question-figure 2 : we study the setting of a language model being adapted unsupervised (without annotation of important tokens) on an online stream of documents and being later evaluated on queries (e.g., questions) about those documents. downstream inputs are not provided during the adaptation phase, requiring the model to integrate as much information as possible about the documents. answering model's ability to answer questions relating to the stream of documents. we hypothesize that naive fine-tuning is not effective in the online adaptation setting because the negative log likelihood (nll) loss does not accurately reflect the importance of a token. that is, tokens containing important factual information may receive relatively small nll loss and therefore a small fine-tuning gradient. for example, consider the nll of the word rishi and the word reports in the phrase the uk prime minister is rishi sunak. reports suggest . . . for a slightly out-of-date language model. because rishi sunak was a well-known politician before becoming prime minister, a model may place reasonably high probability mass on his name (even if other completions are higher probability). on the other hand, 'reports' will invariably receive low probability, because the distribution over the first word in a sentence is unavoidably high entropy. this hypothesis suggests that we can improve upon online adaptation by only fine tuning on a subset of tokens which are most likely to lead to useful updates. one natural approach to identify such factual tokens is through salient spans @xcite . another common technique used to weight words it via tf-idf scores @xcite . we find that fine-tuning while using these heuristics does improve information uptake. however, it is unclear if such heuristic choices are optimal. as an alternative, we explore a method for learning a per-token importance weights corresponding to the utility of fine-tuning on that token. however, such utility is difficult to define, and even with a suitable definition, dense per-token annotations of utility are extremely time-consuming to collect. we thus select a definition of utility that enables using distant supervision of the utility of each token: a high utility token is one whose fine-tuning gradient improves a question-answering model's ability to answer questions about the contents of the surrounding document. using this notion of a token's utility for online learning, we propose context-aware meta-learned loss scaling (camels), an approach to online adaptation that meta-trains an importance weighting model to identify such tokens in a document. given a dataset of documents and queries, we use a meta-learning loss to train our weighting model: first, in an 'inner loop,' we update a base model (a proxy for the model we will update at test time) using the gradient of nll of the document, weighted by the outputs of the importance weighting model. next, in the 'outer loop', the loss is computed by evaluating the updated base model's performance on the corresponding query. this outer loss is used to updated the parameters of the importance weighting model. during online fine-tuning on a stream of documents, we simply re-weight the online loss using the importance-weighting model's output. although the process used to train camels uses a proxy model (i.e., a stand-in for the model we will update at test time), one might hope that the importance of tokens would be independent of the model used for inner loop updates; to a significant degree, we intuit that the importance of a token should be an innate trait of underlying text. indeed, we find that the meta-learned importance weights generalize across models; for each dataset, we metatrain our importance weighting model once using distilgpt-2 @xcite as the base model and successfully use these weighting model without modification to update gpt-j 6b @xcite . across three online adaptation benchmarks based on streams of news and wikipedia articles, camels substantially improves knowledge acquisition over naive fine-tuning as well as salient span and tf-idf based baselines. ## related work adapting to new data or task distributions is typically studied in the context of continual or lifelong learning @xcite . continual learning in deep networks involves the challenge of simultaneously avoiding catastrophic forgetting @xcite , the process under which a neural network's performance on old tasks or data is dramatically degraded by the process of learning new information, while maintaining plasticity @xcite , or the ability to adapt to the latest change in the data distribution, even after many changes have already been experienced. while most work in continual learning considers sequences of supervised data @xcite @xcite , some work also studies continual few-shot @xcite or unsupervised learning @xcite , which is closer to the setting in this paper. however, these works typically focus on streams of visual data. dynamic, or streaming, language models were first considered in the context of n-gram language models, combining a cache of recently-used words to update the predictive probabilities of a tri-gram model @xcite @xcite . later work describes online embased algorithms for efficiently updating n-gram models @xcite . other studies investigate the evolution of decontextualized word embeddings over as a result of temporal shifts in the use of language @xcite or the use of vector memories to store recent information when training recurrent neural networks online @xcite . more recently, several studies have explored methods for updating large neural language models, typically through online fine-tuning on a stream of documents @xcite with architectural constraints @xcite or explicit conditioning on time @xcite used as strategies to reduce forgetting of old information. @xcite use meta-learning to reduce the compute requirements of online fine-tuning. however, recent work suggests that while increasing the size of language models may largely mitigate the problem of forgetting old information @xcite , improving the efficiency of acquisition of new knowledge is still a challenge, and this problem is therefore the focus of the present work. other methods for dynamically updating the knowledge in parametric language models develop specialized techniques, called model editors, designed to make targeted edits to individual facts @xcite @xcite or behaviors @xcite . however, model editors assume access to annotations of the tokens or facts that must be updated; in this work, we study the problem of learning which tokens in an unlabeled sequence of documents are important. ## meta-learning improved online adaptation of large language models given an out-of-date language model and a stream of recent documents, we aim to update the model such that it effectively answers typical queries about the documents in the stream. by focusing only on retaining knowledge relevant to the 'typical' queries, we avoid the need to completely memorize the documents, making the problem tractable. we study question-answering (qa) models specifically, as the question-answer format makes assessing a model's knowledge straightforward. in this section, we formalize this problem setting and then describe an approach to this setting, context-aware meta-learned loss scaling. ## experiments after outlining datasets and experimental details, we present several experiments aimed at understanding camels's behavior in unsupervised online adaptation. section 4.3 studies the extent to which camels's importance weights improve knowledge retention in online adaptation. section 4.4 qualitatively and quantitatively explores the weights themselves, suggesting several ablations of camels that we explore in section 4.5. section 4.6 evaluates the cross-dataset generalization of camels weights, and finally we examine the forgetting and plasticity dynamics of camels within the document stream in section 4.7. ## discussion while large language models are powerful, keeping them up-to-date remains a challenge. in this paper, we consider the unsupervised online language model adaptation setting, in which a language model's knowledge must be updated using a stream of documents, without annotations of key facts or information. finding that naive online fine-tuning provides little retention of knowledge from the document stream, we propose context-aware metalearned loss scaling (camels), a meta-learning algorithm that learns an importance weighting model to reweight the per-token loss of the online data stream. camels leverages side information of the form of paired documents and knowledge queries about those documents to identify which tokens in the documents are most likely to be informative for answering downstream queries. empirically, we find that the importance weighting model learned by camels consistently improves knowledge retention across three datasets of documents and questions. crucially, we find that camels's importance weighting model generalizes across outdated language models and datasets, meaning that an importance weighting model can be trained once on a small proxy language model (such as distilgpt-2) and then be immediately used to improve online adaptation of much larger models, like gpt-j 6b. this transferrability of camels's weighting model significantly increases its practical utility.