PaperlayerView exact revision on arXiv

Source-linked AI summary

Training Language Models with Memory Augmentation

Zexuan Zhong, Tao Lei, Danqi Chen

arXiv:2205.12674v3cs.CLcs.LG

TL;DR

Memory-augmented language models often introduce memories only at testing time or obtain them from separately trained encoders, limiting how the language model trains with memory. TRIME instead uses in-batch examples as training memory and adds batching strategies for local, long-term, and external memories, yielding gains across language modeling and machine translation benchmarks. The approach adds little computational overhead and remains compatible with different neural architectures, although large external memories can make inference 10× slower and machine-translation evaluation was limited to a small dataset.

  • Problem

    Existing memory-augmented language models usually introduce memory only at testing time or use representations from separately trained models, leaving memory units outside direct language-model optimization.

  • Method

    TRIME trains language models with an objective that uses in-batch examples as memory, supported by batching and construction strategies for local, long-term, and external memories.

  • Results

    TRIME produces consistent gains across language-modeling and machine-translation settings, including WIKITEXT-103 improvements from 18.70 to 17.76 and from 16.23 to 15.41 against kNN-LM.

  • Takeaways & Limitations

    TRIME offers a simple, low-overhead memory-training approach that does not modify model architectures and can leverage long-range context and external memory at testing time.

  • Takeaways & Limitations

    Large external memories reduce inference efficiency, making the model 10× slower in one example, while machine-translation evaluation used only the small IWSLT’14 dataset.

Abstract

from arXiv · show

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.

1 Introduction

TRIME addresses the limited training of memory-augmented language models by using in-batch examples as memory during optimization. It combines this objective with memory construction and batching strategies that support local, long-term, and external memories at testing time.

  • Motivation: Existing methods typically add memories only at testing time or obtain them from separately trained encoders, leaving memory representations outside direct language-model optimization.TRIME is introduced to address this training gap.
  • Method: TRIME directly leverages in-batch examples as accessible memory, extending objectives related to neural-cache and nearest-neighbor language models to training as well as testing.The training objective aligns representations with token embeddings and contextualized in-batch representations.
  • Method: New batching strategies construct training memories for local, long-term, and external settings, including consecutive same-document segments and high-overlap segments from other documents.These strategies aim to make training memories useful for the different memories available at inference.
  • Results: A 247M Transformer trained with TRIME improves WIKITEXT-103 perplexity from 18.70 to 17.76 with negligible overhead, without long-term or external memory.This setup is presented as a simple replacement for vanilla language-model training.
  • Results: Training with consecutive segments enables context lengths of 15k-25k tokens while remaining at least competitive with memory-aware Transformer architectures.The approach is described as conceptually simpler and computationally cheaper than those architectural modifications.
  • Results: TRIME reduces WIKITEXT-103 perplexity from 16.23 to 15.41 relative to kNN-LM and also improves machine translation over kNN-MT on IWSLT’14 De-En.The model incorporates other segments in the same batch as memories and performs better with a large testing datastore.
  • Summary: Across multiple experimental settings, TRIME reports consistent gains, adds little computational overhead, and remains compatible with recurrent and attention-based architectures.The approach uses memories at the final prediction step rather than modifying the model architecture.

2 Preliminaries

The paper represents language-model memories as context-target pairs and distinguishes local, long-term, and external sources. These memories can be aggregated through similarity-weighted next-token probabilities, while long-range access is constrained by Transformer context length.

  • Language modeling: Standard language modeling maps the preceding token sequence to a vector representation and optimizes parameters by minimizing the ground-truth token’s negative log-likelihood.The output embedding represents each candidate token in the vocabulary.
  • Memory representation: Memory augmentation represents stored context-target pairs whose contributions can be aggregated into next-token probabilities using hidden-representation similarity.The framework formalizes local, long-term, and external memory types.
  • Local memory: Local memory consists of preceding tokens in the same input and is accessed through recent context.Although local cache memory has been considered less effective for Transformers at testing time, the paper reports benefits when it is included during training.
  • Long-term memory: Long-term memory supplies long-range context from the same document that self-attention cannot directly access within a limited input segment.Documents are divided into consecutive segments to represent this memory.
  • Relation to prior architectures: TRIME does not modify the underlying model architecture and is compatible with architectures that use hidden representations from previous segments.Its memory mechanism is applied through the prediction objective rather than architectural changes.
  • External memory: External memory stores context-target pairs from a large corpus, potentially the training corpus or a domain-specific corpus used under domain shift.Because external sets can be orders of magnitude larger, accessing them requires approximate nearest-neighbor search.

3 Training with In-batch Memories

TRIME trains language models with in-batch examples as memory, constructing training memories from the same batch so gradients can reach memory representations. Its objective aligns context representations with target-token embeddings and contexts sharing the same next token.

  • TRIME replaces standard language-model training with an objective that uses examples from the same batch as training memories, enabling backpropagation through those memories.Training memories are constructed on the fly and may differ from testing memories.
  • The objective combines a Transformer output representation with similarity-based contributions from memory contexts associated with the target token.The model uses fθ(c), token embeddings E_w, context representations gθ(c), and a similarity function.
  • The method uses scaled dot-product similarity because preliminary experiments found it more stable and effective.The similarity is defined as sim(q,k) = q·k / √d.
  • TRIME treats the training objective as contrastive learning that aligns a query with the target-token embedding and contextualized representations sharing the same next token.When the target token is absent from training memory, the objective falls back to the word-embedding term.
  • Unlike prior cache and nearest-neighbor language models that use memory objectives only at testing time, TRIME applies the objective during both training and testing.The approach is connected to interpolated cache and datastore objectives but differs in when the objective is optimized.

4 Adaption to Different Memories

TRIME adapts its training memories and batching to local, long-term, and external memories at inference. It uses unchanged architecture for long contexts and BM25-based batching for large external datastores, with reported gains and competitive performance.

  • Because testing memories can be very large, TRIME uses distinct training-memory constructions and three model variants to reduce the training–testing discrepancy.The experiments describe testing-memory sizes around |M_long| ∼ 10^4 and |M_ext| ∼ 10^8.
  • 4.1 Local Memory: For local memory, TRIMELM changes only the training objective, adding negligible computational overhead relative to encoding the segment itself.Its memory consists of previous tokens within the same segment.
  • 4.2 Long-term Memory: TRIMELMlong reaches competitive performance with prior long-context memory methods while using a simpler, cheaper architecture that propagates gradients to previous-segment memories.It stores only last-layer outputs rather than modifying self-attention or retaining every layer’s hidden representations.
  • 4.3 External Memory: For external memory, TRIME batches segments with high BM25 lexical overlap to approximate nearest-neighbor relationships in a large datastore.The method can exclude local memory during training with probability p, with p = 90% reported as best for large external-memory evaluation.

5 Experiments

TRIME improves language modeling across local, long-term, and external-memory settings, while extending to domain adaptation and machine translation. Its gains include lower perplexity, competitive long-context results, and better translation quality with negligible overhead.

  • Local memory: TRIMELM improves the 247M Transformer’s WIKITEXT-103 perplexity from 18.70 to 17.76 without computational overhead compared with vanilla language modeling.A continuous cache reaches 18.26, while TRIMELM reaches 17.76 using local memory during training and testing.
  • Long-term memory: TRIMELMlong improves perplexity from 25.87 to 22.66 on WIKITEXT-103 and from 1.16 to 1.05 on ENWIK8 by leveraging long-range contexts.The model performs best with 15,000 tokens on WIKITEXT-103 and 24,576 tokens on ENWIK8, despite shorter training segments.
  • External and cross-task evaluation: TRIMELMext reduces WIKITEXT-103 perplexity from 16.23 to 15.41 versus kNN-LM and also significantly improves over kNN-MT on IWSLT’14 De-En.The external-memory model is evaluated with a large datastore, while the translation result is reported as a significant improvement over kNN-MT.
  • Domain adaptation: TRIME improves domain adaptation on BOOKSCORPUS without retraining or fine-tuning, with TRIMELMext outperforming kNN-LM when using target-domain external memory.Models trained on WIKITEXT-103 are evaluated on BOOKSCORPUS, and target-domain datastore construction improves performance under domain shift.
  • Machine translation: TRIME improves the vanilla Transformer by 1.15 BLEU on IWSLT’14 De-En and outperforms kNN-MT.For this sentence-level task, the adapted model does not use local or long-term memory.

6 Analysis

The analysis examines how batching and memory construction affect TRIME, showing benefits from consecutive or high-overlap segments and effective use of long-term context. Additional ablations identify BM25 batching and back-propagation into memory representations as important components.

  • Memory construction: Packing consecutive or high-BM25 segments into training batches and constructing memories properly improves performance with long-range and external memories.The comparison covers TRIMELM, TRIMELMlong, and TRIMELMext across different testing-memory sets.
  • Leveraging long-range contexts: TRIMELM effectively handles more than 10k tokens of long-term context, with larger gains over continuous cache as more memory is incorporated.Figure 3 compares TRIMELMlong with continuous cache on WIKITEXT-103 and ENWIK8 development sets.
  • Additional analysis: BM25 batching and back-propagation to update memory representations are important for TRIME, while local memory remains effective across segment lengths.Additional ablations also report that TRIMELMext has higher top-K retrieval accuracy than kNN-LM with the external memory set.

7 Related Work

TRIME differs from prior memory-augmented language models by training memory representations through the end loss, preserving the base architecture, and unifying local, long-term, and external memories.

  • Memory-augmented language models: Unlike prior approaches, TRIME updates memory representations through back-propagation from the end loss rather than relying on testing-time or separately trained memories.The framework applies the same training perspective across multiple memory types.
  • Memory-augmented language models: TRIME does not modify the base architecture and can consider different memory types within a unified framework.This distinguishes it from approaches that alter Transformer architectures or use separately constructed representations.
  • Transformers for long inputs: Prior long-context methods leverage previous-segment representations by modifying Transformer architectures, whereas TRIME is compatible with those neural architectures.The paper presents TRIME as orthogonal to architecture-level approaches.

8 Conclusion

TRIME trains language models with in-batch memories and uses tailored batching and memory construction to exploit local, long-range, and external memories at test time. It adds little computational overhead, preserves architectural compatibility, and is demonstrated on language modeling and machine translation.

  • 8 Conclusion: TRIME uses data batching and memory construction during training to leverage long-range contexts and external memory at testing time.The approach includes three model instantiations: TRIMELM, TRIMELMlong, and TRIMELMext.
  • 8 Conclusion: The method adds little computational overhead and does not modify model architectures, enabling compatibility with other neural models and techniques.
  • 8 Conclusion: TRIME is evaluated across language modeling and machine translation settings, with larger language models and other text-generation tasks identified as future work.

Limitations

The paper identifies limits involving inference efficiency, model and architecture coverage, dataset scale, task scope, and experimental model size. These boundaries motivate future evaluation on larger memories, datasets, architectures, tasks, and language models.

  • Limitations: External-memory retrieval improves performance but reduces inference efficiency, making the model 10× slower when incorporating external memory.The issue may become more serious as the external memory grows; coarser memories and compression are suggested as possible remedies.
  • Limitations: The experiments mainly use Transformer-based models, with an additional adaptation to SRU++, while compatibility with other architectures remains future work.
  • Limitations: Machine-translation evaluation is limited to the small IWSLT’14 dataset, whose external memory contains 4M tokens, leaving larger datasets unevaluated.
  • Limitations: The paper studies language modeling and machine translation, while adapting TRIME to natural-language-understanding tasks such as text classification remains open.
  • Limitations: The largest tested model has 247M parameters, far below the hundreds of billions in state-of-the-art autoregressive language models.

Ethical Considerations

For external-memory experiments, the paper constructs memory from the training corpus but notes that retrieval can leak information from the corpus. It therefore recommends using a public corpus for external memory.

  • Ethical Considerations: External memory is constructed from the training corpus in the experiments, although it could instead be built from any corpus.
  • Ethical Considerations: Retrieving from an external datastore can cause information leakage from its corpus, so the paper recommends using a public corpus.
  • Ethical Considerations: The inference procedure can combine local, long-term, and external memories, retrieving the top-K closest terms when external memory is large.

C Dataset Statistics and Tasks

The experiments cover word- and character-level language modeling, machine translation, domain adaptation, architectural variants, and small-model ablations. Dataset statistics, model configurations, and retrieval settings define these evaluation conditions.

  • C Dataset Statistics and Tasks: The benchmark suite includes WIKITEXT-103, ENWIK8, IWSLT’14, and BOOKSCORPUS, spanning language modeling, machine translation, and domain adaptation.
  • C Dataset Statistics and Tasks: WIKITEXT-103 is a 103M-token word-level language-modeling dataset evaluated with perplexity, while ENWIK8 is a 100M-character dataset evaluated with bits per character.
  • C Dataset Statistics and Tasks: IWSLT’14 DE→EN contains 170K translation pairs and is evaluated with BLEU using a Transformer encoder-decoder model.
  • C Dataset Statistics and Tasks: BOOKSCORPUS uses 100M/250K/250K train/dev/test tokens to test domain adaptation without retraining or fine-tuning WIKITEXT-103-trained models.
  • C Dataset Statistics and Tasks: The SRU++ adaptation uses local and previous-segment long memory, while inference can access long-term memories of 6,144 or 12,228 tokens depending on training segment length.
  • C Dataset Statistics and Tasks: A 7M-parameter, 8-layer, 4-head Transformer is used for ablations comparing baselines, batching, memory back-propagation, local memory, and retrieval.

G Additional Analysis

Additional analyses show that TRIME benefits from its memory-aware batching and training design across local, external, and frequency-specific settings. Retrieval quality, memory mixing, and back-propagation each materially affect performance.

  • Ablation study on TRIMELMext: Random batching raises perplexity to 45.71 (+4.21), while disabling back-propagation to memory representations substantially worsens performance.The ablation tests both BM25-based segment packing and gradient propagation into memory representations.
  • Effectiveness of using local memory: TRIMELM significantly outperforms baselines across all tested local-memory segment lengths.The analysis evaluates TRIMELM with different segment lengths L and reports consistent gains over baselines.
  • Retrieval performance on external memory: TRIME retrieves external memories more accurately than kNN-LM, explaining its lower final perplexity when external memory is incorporated.External-memory experiments retrieve the top K keys by nearest-neighbor search, using K = 1024 in the described setting.
  • Perplexity breakdown for different frequencies: TRIMELM and TRIMELMlong improve rare-word perplexity, whereas TRIMELMext improves perplexity across all frequency buckets.Rare words are defined as tokens with development-set frequency ≤1k; frequent-word results for TRIMELM and TRIMELMlong are similar or slightly worse than the Transformer baseline.
  • Memory mixing probability: Increasing local-memory dropout mitigates external-memory underuse, with p = 0.9 selected for the main experiments.At p = 0, the model learns to rely only on local memory and ignores external memory during training.
Loading 2205.12674v3…