Recent work has improved language models (LMs)  remarkably by equipping them with a non-parametric memory component.
However, most existing approaches only introduce memories at testing time or represent them using a separately trained encoder, resulting in suboptimal training of the language model.
In this work, we presentTrime, 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 evaluateTrimeon multiple language modeling and machine translation benchmarks and show that it is able to achieve significant improvements across all the settings. Concretely,Trimereduces the perplexity from 18.70 to 15.37 onWikiText-103, by effectively leveraging a large memory set from the training corpus. Compared to standard LM training,Trimeadds negligible computational overhead and is compatible with different neural architectures, making it a versatile solution for training memory-augmented LMs.111Our code and pre-trained models are publicly available athttps://github.com/princeton-nlp/TRIME.