Training Language Models with Memory Augmentation

Proceedings of the 2022 Conference on Empirical Methods in Natural Language Processing, pages 5657 - 5673 December 7-11, 2022 ©2022 Association for Computational Linguistics Training Language Models with Memory Augmentation Zexuan Zhong† Tao Lei∗ Danqi Chen† †Princeton University {zzhong, danqic}@cs.princeton.edu, taole@google.com 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 subop- timal 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 adapt- ing to different sets of memories—local, long- term, and external memory—at testing time. We evaluate TRIME on multiple language mod- eling and machine translation benchmarks and show that it is able to achieve significant im- provements across all the settings. Concretely, TRIME reduces the perplexity from 18.70 to 15.37 on WIKITEXT-103, by effectively lever- aging a large memory set from the training corpus. Compared to standard LM training, TRIME adds negligible computational over- head and is compatible with different neural architectures, making it a versatile solution for training memory-augmented LMs.1 1 Introduction Memory augmentation has become a remarkable approach to enhance language modeling perfor- mance without significantly increasing the amount of parameters and computation. By accessing memory units such as a neural cache of recent inputs (Merity et al., 2017; Grave et al., 2017b) and an external look-up table (Khandelwal et al., 2020), a memory-augmented language model (LM) enjoys increased memorization capacity and sets ∗TL currently works at Google Research. The collabora- tion was initialized before TL joined Google. 1Our code and pre-trained models are publicly available at https://github.com/princeton-nlp/TRIME. Forward pass encoder Jobs became CEO of _ and Apple first … similarity … works at Microsoft … Jobs became CEO … returned to Apple … moves to Apple Target token’s embedding Other token embeddings Positive in-batch memory Negative in-batch memory prediction (target: “Apple”) … Microsoft color |V| token embeddings In-batch memories Back-propagation Figure 1: An illustration of our training objective. Our objective aligns the hidden representation with both to- ken embeddings and a set of in-batch contextualized representations that are constructed during training. new state-of-the-art records in various language modeling benchmarks. A major limitation of existing approaches, how- ever, is that the memory units are either intro- duced at testing time (Grave et al., 2017b,a; Khan- delwal et al., 2020) or taken from a separately trained model (Yogatama et al., 2021). As a con- sequence, they are not directly optimized during the training process, resulting in a missed oppor- tunity to achieve even stronger results. In this pa- per, we pioneer and present a novel yet simple training approach TRIME (Training with In-batch Memories)2, that is well-suited for memory aug- mentation in language modeling. …

Content truncated to 500 words