Training Language Models with Memory Augmentation

By Zexuan Zhong, Tao Lei, Danqi Chen. 2022

153 citations recorded by the source corpus.

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.

Subfields: Artificial Intelligence, Large Language Models, Prompting & Reasoning

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Large Language Models

In Prompting & Reasoning

Citations in the corpus

Augmented Language Models: a Survey (2023, background): “Various works augment LMs with a dense retriever by appending the retrieved documents to the current context (Chen et al., 2017; Clark and Gardner, 2017; Lee et al., 2019; Guu et al., 2020; Khandelwal et al., 2020; Lewis et al., 2020; Izacard and Grave, 2020; Zhong et al., 2022; Borgeaud et al., 202”
Retrieval as Attention: End-to-end Learning of Retrieval and Reading within a Single Transformer (2022, background): “Retrieval from external datastore to improve language modeling perplexity has been explored by many works, where additional tokens are retrieved during generation based on contextual representations (Khandelwal et al., 2020; Borgeaud et al., 2021; Wu et al., 2022; Zhong et al., 2022).”
Adaptive Chameleon or Stubborn Sloth: Revealing the Behavior of Large Language Models in Knowledge Conflicts (2023, unknown): “…with up-to-date information, namely tool-augmented (Nakano et al., 2021; Yao et al., 2023; Qin et al., 2023; Schick et al., 2023; Lu et al., 2023) or retrieval-augmented (Guu et al., 2020; Khandelwal et al., 2020; Izacard & Grave, 2021; Borgeaud et al., 2022; Zhong et al., 2022) language models.”
Retrieval-Augmented Generation for AI-Generated Content: A Survey (2024, unknown): “5) Neural Machine Translation : Neural Machine Translation (NMT) is the automated process of translating text from a source language to a target language [161], [242], [243].”
Retrieval-Augmented Generation for Large Language Models: A Survey (2023, background): “The preparation and curation of high-quality training datasets, defining fine-tuning objectives, and providing corresponding computational resources are necessary.”
REPLUG: Retrieval-Augmented Black-Box Language Models (2023, methodology): “For example, Atlas (Izacard et al., 2022b) finetunes an encoder-decoder model jointly with the retriever by modeling documents as latent variables, while RETRO (Borgeaud et al., 2022) changes the decoder-only architecture to incorporate retrieved texts and pretrains the language model from scratch.”
Training Trajectories of Language Models Across Scales (2022, background): “…al., 2021) and BLOOM (Scao et al., 2022), different architectures and objectives, such as retrieval-based language models (Khandelwal et al., 2020; Zhong et al., 2022; Borgeaud et al., 2021) and sparse models (Fedus et al., 2022; Artetxe et al., 2022), the relationship between validation…”
Check Your Facts and Try Again: Improving Large Language Models with External Knowledge and Automated Feedback (2023, unknown): “While there is a growing interest in improving LLMs using external knowledge ( e.g., Ghazvininejad et al., 2017; Guu et al., 2020; Zhong et al., 2022; Gao et al., 2019, 2022), almost all the previously proposed methods require ﬁnetuning the parameters of a LLM, which can be prohibitively expensive…”
SILO Language Models: Isolating Legal Risk In a Nonparametric Datastore (2023, methodology): “Future work can improve k NN-LM, e.g., by training the model to output a nonparametric distribution (Zhong et al., 2022; Lan et al., 2023; Min et al., 2023), by having a vocabulary-specific λ (Huang et al., 2023b), or by modeling λ as a function of the input x (He et al., 2021; Drozdov et al.,…”
LongMemEval: Benchmarking Chat Assistants on Long-Term Interactive Memory (2024, unknown): “A second line of research integrates differentiable memory modules into language models, proposing specialized architectural designs and training strategies to enhance memory capabilities (Weston et al., 2014; Wu et al., 2022; Zhong et al., 2022; Wang et al., 2023).”
Retrieval-based Language Models and Applications (2023, methodology): “The former uses fractions of the full corpus that are carefully designed during joint training (Zhong et al., 2022; de Jong et al., 2022; Min et al., 2022).”
LeanDojo: Theorem Proving with Retrieval-Augmented Language Models (2023, methodology): “Our ReProver is the first retrieval-augmented language model for formal theorem proving, though similar architectures have been studied extensively in NLP [73–81].”
Adapting Language Models to Compress Contexts (2023, background): “A number of approaches have been proposed to infuse external knowledge in the input layer (Guu et al., 2020; Shi et al., 2023), intermediate layers (Borgeaud et al., 2022) or at the output layer (Khandelwal et al., 2020; Zhong et al., 2022).”
In-Context Retrieval-Augmented Language Models (2023, background): “Retrieve and Read Models This family of RALMs creates a clear division between document retrieval and document reading components.”
When Not to Trust Language Models: Investigating Effectiveness and Limitations of Parametric and Non-Parametric Memories (2022, background): “ries in intermediate states or output spaces (Zhong et al., 2022; Khandelwal et al., 2020; Min et al., 2022), which are shown to be effective to overcome LMs’ limited memorization or knowledge updates.”