Memory Papers
Memory Architecture

Training Language Models with Memory Augmentation

AuthorsZexuan Zhong, Tao Lei, Danqi Chen

arXiv 20222022
arXivPDF

TL;DR

TRIME trains language models with in-batch memories via a contrastive objective, cutting WIKITEXT-103 perplexity from 18.70 to 15.37 with memory augmentation.

SharePost on XLinkedIn

Read our summary here, or open the publisher PDF on the next tab.

THE PROBLEM

Memory-augmented LMs are not trained with their memories (perplexity 18.70 on WIKITEXT-103)

Existing memory-augmented language models often add cache or kNN memories only at test time or from a separately trained encoder, leaving them unoptimized during training. This leads to WIKITEXT-103 perplexity as high as 18.70 for a 247M Transformer baseline.

When continuous cache or kNN-LM are bolted on post hoc, the underlying language model never learns to rely on these memories, limiting gains and wasting non-parametric memory capacity. This hurts long-context modeling and domain adaptation despite large external corpora.

HOW IT WORKS

TRIME — Training with In-batch Memories

TRIME introduces a contrastive training objective that uses in-batch memories, aligning hidden states with both token embeddings and contextual memory representations built from local, long-term, and external memory. TRIME instantiates three concrete models: TRIMELM, TRIMELMlong, and TRIMELMext.

Conceptually, TRIME treats the batch like fast RAM where recent and similar contexts are cached, while a large external datastore behaves like disk accessed via approximate nearest neighbors. The training objective teaches the encoder to treat these two tiers as a unified memory hierarchy.

By directly optimizing similarity between query contexts and memory contexts, TRIME enables language models to exploit memories far beyond the attention window, something a plain context window cannot achieve even with longer sequences.

DIAGRAM

TRIME Inference Flow with Local, Long-term, and External Memory

This diagram shows how TRIME combines local, long-term, and external memories to compute next-token probabilities at inference.

DIAGRAM

TRIME Training and Batching Pipeline

This diagram shows how TRIME constructs in-batch memories via default, consecutive, and BM25-based batching strategies during training.

PROCESS

How TRIME Handles a Language Modeling Step

  1. 01

    Training with In-batch Memories

    TRIME constructs Mtrain from in-batch context target pairs and applies the contrastive objective aligning encoder outputs with token embeddings and memory representations.

  2. 02

    Local Memory Construction

    TRIME builds Mlocal from preceding tokens within each segment, enabling TRIMELM to use in-segment history as neural cache during both training and testing.

  3. 03

    Long-term Memory Construction

    TRIME packs consecutive segments from the same document so TRIMELMlong can access long-range contexts beyond the attention window as Mlong.

  4. 04

    External Memory Construction

    TRIME uses BM25 batching to group lexically similar segments, training TRIMELMext to query a large external datastore Mext via approximate nearest neighbor search.

KEY CONTRIBUTIONS

Key Contributions

  • 01

    Training with In-batch Memories

    TRIME introduces a contrastive objective that aligns encoder representations with token embeddings and in-batch contextual memories, improving WIKITEXT-103 perplexity from 18.70 to 17.76 using only local memory.

  • 02

    Long-term Memory via Batching

    TRIME's consecutive segment batching lets TRIMELMlong exploit 15k–25k token contexts, reducing WIKITEXT-103 perplexity from 25.87 to 22.66 without changing the Transformer architecture.

  • 03

    External Memory Adaptation

    TRIME combines BM25-based batching and in-batch contrastive training so TRIMELMext improves kNN-LM on WIKITEXT-103 from 16.23 to 15.41 test perplexity using a 247M Transformer.

RESULTS

By the Numbers

Test perplexity

18.70

baseline Transformer on WIKITEXT-103 (247M, L = 3,072)

Test perplexity

17.76

-0.94 vs Transformer (TRIMELM with local memory)

Test perplexity

15.41

-0.82 vs kNN-LM our run with continuous cache (TRIMELMext)

Tokens per second

3600.0

no speed loss vs vanilla LM for TRIMELM and TRIMELMlong

These metrics come from WIKITEXT-103 language modeling, which tests next-token prediction on long Wikipedia articles. The MAIN_RESULT shows that TRIME can leverage memory augmentation to reduce perplexity from 18.70 to as low as 15.41 without architectural changes.

BENCHMARK

By the Numbers

These metrics come from WIKITEXT-103 language modeling, which tests next-token prediction on long Wikipedia articles. The MAIN_RESULT shows that TRIME can leverage memory augmentation to reduce perplexity from 18.70 to as low as 15.41 without architectural changes.

BENCHMARK

Performance of TRIME Models on WIKITEXT-103 (247M models, L = 3,072)

Test perplexity on WIKITEXT-103 for 247M-parameter models with different memory augmentations.

KEY INSIGHT

The Counterintuitive Finding

TRIME shows that even for Transformers, explicitly training with local memory reduces WIKITEXT-103 perplexity from 18.70 to 17.76 without any architectural changes. This contradicts prior claims that continuous cache offers limited benefit for Transformer-based models.

It is surprising because self-attention already has full access to recent tokens, yet TRIME's neural cache style objective still yields a 0.94 perplexity gain. This challenges the assumption that explicit cache mechanisms are redundant once attention sees the same context.

WHY IT MATTERS

What this unlocks for the field

TRIME enables language models to be trained end-to-end with local, long-term, and external memories, turning non-parametric memory into a first-class, optimized component. This gives practitioners a way to scale context effectively without modifying architectures.

Builders can now retrofit existing Transformers with TRIME to get kNN-LM style gains, like reducing WIKITEXT-103 perplexity to 15.41, while keeping training and inference pipelines simple and architecture-agnostic.

~14 min read← Back to papers

Related papers

Memory Architecture

A Control Architecture for Training-Free Memory Use

Yanzhen Lu, Muchen Jiang et al.

· 2026

TAG routes low-confidence steps to uncertainty-based routing, filters them with guarded acceptance with rollback, chooses between bank selection across rule and exemplar memory, and prunes via evidence-based retirement inside a unified control loop. On SVAMP and ASDiv, TAG reaches 81.0% and 85.2% accuracy, improving over the 74.0% and 77.5% no-memory baselines while a compute-matched Retry baseline stays flat.

arXiv:2604.18206Read explainer

Questions about this paper?

Paper: Training Language Models with Memory Augmentation

Answers use this explainer on Memory Papers.

Checking…