Training Language Models with Memory Augmentation (2022)

Abstract: Recent work has improved language models (LMs) remarkably by equipping them with a non-parametric memory component. However, most existing approaches only introduce mem-ories at testing time or represent them using a separately trained encoder, resulting in suboptimal training of the language model. In this work, we present TRIME, a novel yet simple training approach designed for training LMs with memory augmentation. Our approach uses a training objective that directly takes in-batch examples as accessible memory. We also present new methods for memory construction and data batching, which are used for adapting to different sets of memories—local, long-term, and external memory—at testing time. We evaluate TRIME on multiple language modeling and machine translation benchmarks and show that it is able to achieve significant improvements across all the settings. Concretely, TRIME reduces the perplexity from 18.70 to 15.37 on WIKITEXT-103, by effectively leveraging a large memory set from the training corpus. Compared to standard LM training, TRIME adds negligible computational overhead and is compatible with different neural architectures, making it a versatile solution for training memory-augmented LMs.

## introduction memory augmentation has become a remarkable approach to enhance language modeling performance without significantly increasing the amount of parameters and computation. by accessing memory units such as a neural cache of recent inputs @xcite and an external look-up table @xcite , a memory-augmented language model (lm) enjoys increased memorization capacity and sets * tl currently works at google research. the collaboration was initialized before tl joined google. 1 our code and pre-trained models are publicly available at https://github.com/princeton-nlp/trime . new state-of-the-art records in various language modeling benchmarks. a major limitation of existing approaches, however, is that the memory units are either introduced at testing time @xcite or taken from a separately trained model @xcite . as a consequence, they are not directly optimized during the training process, resulting in a missed opportunity to achieve even stronger results. in this paper, we pioneer and present a novel yet simple training approach trime (training with in-batch memories) foot_0 , that is well-suited for memory augmentation in language modeling. our approach makes two major departures compared to standard language model training: training objective inspired by contrastive representation learning, we propose a training objective that directly leverages in-batch examples as accessible memory (figure 1 ). our training ob-jective is closely connected to neural cache models @xcite and nearest-neighbor language models @xcite , where the next-token probabilities are calculated by comparing encoder outputs against static token embeddings and memory representations. however, previous work only considers incorporating memories at testing time, while we do for both training and testing. in-batch memory construction with this training objective in mind, the key challenge is how to construct memories effectively during training while keeping it efficient. we identify three types of memories that can be leveraged at testing time and have been explored in the literature: (a) local memory denotes the words that appear in the recent past and are modeled using attention @xcite ; (b) long-term memory foot_1 denotes longrange context from the same document but cannot be directly accessed due to the limit of input length; (c) external memory is used to store the entire training set or any additional corpus @xcite . to better leverage these memories at testing time, we devise new data batching strategies to improve the construction of training memories ( §4). by packing consecutive segments from the same document in one training batch, our model can access long-term memories beyond the attention context. we pack segments from other documents that have high lexical overlap as a proxy to all external memory units. importantly, these working memories are generated on the fly during training, allowing us to back-propagate to all memory representations. we instantiate trime in three models by considering different sets of training and testing memories ( training testing memory memory vanilla lm none none cont. cache none m local or m long in summary, we propose a simple approach trime for optimizing language models with memory augmentation and demonstrate consistent and significant gains in multiple experimental settings. our approach only uses memories at the final prediction step, and hence adds little computational overhead and can be combined with different model architectures such as recurrent networks and other attention variants @xcite @xcite . we hope that our work can encourage the research community to think about better training objectives for language models, given their significant societal impacts @xcite @xcite . ## training with in-batch memories in this section, we propose a new training approach trime for language model training. compared to standard language model training, our training objective assumes a set of training memories m @xmath0, x j )}. we differentiate training memories from testing memories, as they are constructed on the fly during training and may deviate from the testing memories used during inference. importantly, the training memories are constructed from the same training batch, which enables backpropagating the training signal to the current hidden representation as well as all the memory representations. we will discuss how to construct training memories in the next section ( §4) and only discuss the training objective in a general form. our training objective is illustrated in figure 1 . given a memory set m and a context c, trime defines the next-token probability distribution as: here, f θ (c) is the output representation of a transformer model and e w is the token embedding. g θ (⋅) denotes the representations that can be used to compute similarity between c and all the contexts c j in the memory m train . it is possible to simply take g θ = f θ ; however, we find that taking g θ to be the input of the final feed-forward layer in transformer works better, which is consistent with the observation in @xcite . in addition, sim(⋅, ⋅) is a similarity function and we found using the scaled dot-product sim(q, k) = q⋅k √ d @xcite leads to stable training and better performance in our preliminary experiments. this training objective can be viewed as a contrastive loss @xcite : for a contexttarget pair (c, w * ), the goal is to align the query representation f θ (c) (and g θ (c)) with the static token representation e w * , and contextualized representations that share the same next token i.e., g θ (c j ) for x @xmath1. our objective handles rare words nicely-if w * does not appear in the training memory, the objective will fall back to aligning f θ (c) with only the word embedding e w * . similar to the vanilla training loss (eq. 1), our trime loss is optimized to minimize the negative log-likelihood of next token w * and all the parameters θ and e w are updated during training. our training objective is also inspired by the success of contrastive learning in dense retrieval @xcite . as we will show in §6, it can help improve retrieving contexts that share the same next token effectively when the set of testing memories is large. our objective is also closely connected to the objective used in @xcite ; @xcite , which linearly interpolates the distribution of standard language modeling, and a distribution defined by cache/external datastore, e.g., p (w c) = (1-λ)p lm (w c)+λp knn (w c). our work differs from previous works that we use this objective during training (and testing), while they only used it at testing time-the key is how to construct training memories that we will elaborate next. foot_4 4 adaption to different memories inference we are interested in incorporating the three types of memories defined in §2.2 and their combinations at testing time. the testing objective is basically the same as the training objective (eq. 5) except that we take testing memories as a combination of m local , m long and m ext . as m ext can be very large, we approximate it by retrieving the top-k closest terms to g θ (c). we tune a temperature term τ to adjust the weight of the memory component (see appendix a for details). notation throughout this section, we use l to denote segment length, b to denote the total number of segments used in the one training batch, and m to denote the number of consecutive segments from each document in the batch. correspondingly, each batch will contain b ≈ b m different documents. l, b and m are hyper-parameters that we will choose for training, and will vary as we consider different memories during inference. a key challenge is that the testing memories can be very large (e.g., m long ∼ 10 4 and m ext ∼ 10 8 in our experiments) and it is computationally infeasible to keep training memories the same as testing memories. in the following, we will discuss three ways of constructing training memories and data batching, aiming to reduce the discrepancy between training and testing. along the way, we will also present three major model instantiations: trimelm, trimelm long , trimelm ext ( ## analysis we conduct ablation studies and analysis to further understand individual components of our approach. due to the limited computation budget, some experiments on wikitext-103 are conducted with a small 7m transformer model (8 layers, hidden dimension 128) in this section and the trends are generally similar for smaller models (see appendix d and appendix f for details). ## conclusion in this work, we propose trime, a training approach for language modeling. we present three model instantiations trimelm, trimelm long , trimelm ext : through carefully-designed data batching and memory construction during training, we show that our models can leverage long-range contexts and external memory effectively at testing time. our approach adds little computational overhead and does not modify model architectures, making it compatible with other neural models and techniques. for future work, we are interested in training trime with large language models and other text generation tasks.