💬 NLP

Language Model Memory and Memory Models for Language

This paper argues that standard next-token prediction training yields information-poor embeddings unsuitable for memory, and proposes a parallelizable encoder-decoder architecture trained with combined causal and information-retention objectives to achieve high-fidelity memory formation and computational efficiency.

Original authors: Benjamin L. Badger

Published 2026-05-20

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Benjamin L. Badger

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Idea: The "Bad Memory" of AI

Imagine you are reading a very long book. A standard AI (called a Large Language Model) is like a student who reads the book but only remembers the very last sentence they read to guess what comes next. They don't really "store" the whole story in their head; they just keep a running tally of the immediate past.

This paper argues that these AI models have very poor memories of the actual text they have processed. Even if you feed them a huge chunk of text, the "summary" they create in their internal brain (called an embedding) is like a blurry, low-resolution photo. It's good enough to guess the next word, but if you tried to use that photo to recreate the original book word-for-word, you would fail miserably.

The Problem: Why "Next Word" Training Fails Memory

The paper explains that these AI models are trained to do one specific job: predict the next word.

The Analogy: Imagine a game where you have to guess the next word in a sentence. To win, you only need to remember the last few words. You don't need to remember the entire story.
The Result: Because the AI only needs to remember a tiny bit to win the game, it stops trying to remember the rest. It becomes an expert at guessing the next word, but a terrible archivist. It discards most of the information it sees because it doesn't need it for the specific task it was trained on.

The authors found that if you try to force these models to act as a "memory bank" (storing a whole paragraph to be retrieved later), they fail. They simply cannot compress a whole story into a single "memory token" without losing almost all the details.

The Solution: The "Perfect Photographer" (Autoencoders)

The paper introduces a different type of model called an Autoencoder.

The Analogy: Think of an Autoencoder as a high-end photographer. Its only job is to take a picture of an object and then try to reconstruct that exact object from the photo.
The Result: Because its goal is to recreate the exact original image, it learns to take a perfect, high-fidelity photo. It captures every detail.
The Finding: When the authors trained these "photographers" (Autoencoders) to store text, they created perfect memories. They could store a whole paragraph in a tiny space and recreate it perfectly later.

The New Architecture: The "Specialized Team"

Since the standard AI (the "Word Guesser") is bad at memory, but the Autoencoder (the "Photographer") is great at it, the authors propose a new team structure:

The Encoder (The Photographer): This part of the system is trained only to take perfect "photos" (memories) of the text. It freezes its brain once it learns to do this perfectly.
The Decoder (The Storyteller): This is the part that usually predicts the next word. In this new system, it is trained to look at the "photos" (memories) and use them to tell the story or predict the next word.

The "Curriculum" Trick:
The authors found that if you just throw the "Storyteller" into the deep end with these perfect photos, it gets confused. It's used to looking at the text directly.

The Fix: They used a teaching method called Curriculum Learning.
- Step 1: Teach the Storyteller to understand the "photos" (memories) without any text help. It has to learn to read the memory first.
- Step 2: Once it masters reading the memories, give it the text back as well.
- Result: The model learns to use the perfect memories effectively.

Why Do This? (The Benefits)

The paper highlights two main reasons to switch to this "Memory Model":

Speed and Efficiency:
- Current Way: To read a long book, the AI has to re-read every single page every time it wants to guess a new word. It's like re-reading a library to find one fact.
- Memory Way: The AI reads the book once, takes a "snapshot" (memory), and then just looks at the snapshot. This is much faster and uses less computer power.
Arbitrary Access:
- Standard models are like a tape recorder; they can only hear what's playing right now.
- Memory models are like a library card catalog; they can jump to any specific fact stored in the memory, no matter how far back it was.

The Catch: You Can't Just "Guess" Your Way to Memory

The paper concludes with a crucial warning: You cannot simply train a model to predict the next word and expect it to develop a good memory.

The Lesson: Predicting the next word is a "one-way street." It's easy to go from text to a guess, but impossible to go from the guess back to the original text.
The Fix: To get a good memory, you must train the model with a combined goal: "Predict the next word AND be able to recreate the whole story from your memory." If you don't force the model to do both, it will keep its memory "blurry."

Summary

Standard AI: Great at guessing the next word, terrible at remembering the whole story.
The Problem: Training AI to just guess words makes it forget everything else.
The Fix: Train a separate "Photographer" to take perfect snapshots of the text, then teach the "Storyteller" how to use those snapshots.
The Result: A faster, more efficient AI that can actually remember and retrieve specific details from long texts, provided it is trained with the right "memory" goals.

Technical Summary: Language Model Memory and Memory Models for Language

1. Problem Statement

The paper addresses a fundamental gap in the understanding of how Large Language Models (LLMs) store input information. While it is widely assumed that hidden layer embeddings act as a "memory" of the input sequence, the authors argue that this memory is poorly characterized and often insufficient for tasks requiring arbitrary information access.

The core problem identified is that standard causal language models, trained solely on next-token prediction, produce embeddings that contain very little input information. This is because the objective function (predicting the next token) is non-invertible and many-to-one; the model only needs to retain enough information to predict the immediate next token, not the entire history. Consequently, these models fail to form high-fidelity memories capable of arbitrary lookup, limiting their utility as memory modules for other models or for extending context windows efficiently.

2. Methodology

2.1 Information Measurement and Inversion

To quantify the information retained in embeddings, the authors propose an inversion framework. Instead of using gradient descent (which they found ineffective for single embeddings), they train a decoder model to reconstruct the input sequence from the encoder's final hidden layer embedding.

Metrics: Two primary metrics are used:
1. Entropy Ratio ( $H_r$ ): Measures the fraction of input information retained by comparing the cross-entropy loss of the reconstruction against theoretical bounds.
2. Token Accuracy: The proportion of input tokens correctly identified by the decoder (Hamming metric).
Architectures Tested: The study evaluates Causal LLMs, Retrieval models, and Autoencoders (trained for input regeneration) across various scales (from small Mixers/Transformers to large LLMs like BERT, Llama, and Qwen).

2.2 Memory Model Architectures

The paper introduces a parallelizable Encoder-Decoder Memory Model architecture designed to separate information retention from information extraction:

Encoder: Maps sequences of tokens to a fixed-size embedding (memory).
Decoder: Takes these memory embeddings (and potentially additional tokens) to predict next tokens.
Parallelization: Unlike recurrent memory models, this architecture allows the encoder to process input chunks in parallel, reducing inference complexity from $O(n^2)$ to approximately $O(n^{4/3})$ or better depending on chunking strategies.

2.3 Training Strategies

The authors explore several training paradigms to improve memory formation:

Combined Objectives: Training models with a linear combination of next-token prediction loss and a "copy" loss (where the model must reproduce a sequence it has seen).
Curriculum Learning: A staged approach where decoders first learn to process information-rich embeddings (via a "blank copy" task where no tokens are provided) before being introduced to token embeddings.
Frozen Encoders: Training a high-fidelity autoencoder encoder first, freezing it, and then training the decoder to utilize these embeddings. This aims to decouple the difficult task of forming memories from the task of decoding them.

3. Key Results

3.1 Causal Models Have Poor Memory

Experiments confirm that causal language models and retrieval models exhibit minimal information retention. Even with increased model size and compute, causal embeddings retain very little input information (e.g., token accuracy around 5% for 512-token contexts). In contrast, autoencoders trained specifically for input regeneration achieve near-perfect memory formation (token accuracy >90% for 512 tokens).

3.2 Invertibility and Data Dependence

The study establishes that while strict mathematical invertibility is impossible for standard architectures due to non-invertible operations (e.g., ReLU, LayerNorm), models are "practically invertible" on finite datasets if the output dimension is high enough. However, causal training does not encourage the formation of these invertible, information-rich embeddings. Furthermore, autoencoders learn non-trivial, data-dependent encodings rather than trivial token-to-vector mappings.

3.3 Success of Combined Objectives and Curriculum Learning

Combined Training: Models trained with both next-token prediction and copy objectives successfully learn to form and decode information-rich memories while maintaining causal prediction capabilities.
Curriculum Learning: Decoders trained via curriculum learning (starting with blank copy tasks) can effectively utilize frozen, high-fidelity encoder embeddings. This approach outperforms training with trainable encoders on difficult tasks like "blank copy" (reconstructing sequences without token inputs).
Benchmark Performance: Models trained with combined objectives show notable improvements on information-intensive benchmarks (e.g., LongBench, Multi-Document QA, Code Completion) compared to purely causal models, though general language understanding benchmarks (e.g., ARC, HellaSwag) see little change.

3.4 Architectural Compatibility

The paper notes that "Mixer" architectures are generally more efficient encoders for information retention, while Transformers are more efficient decoders. However, the authors found that embeddings from frozen Mixer encoders are often incompatible with Transformer decoders, leading to decreased training efficiency, suggesting that the learned manifolds of these architectures differ significantly.

4. Significance and Claims

The paper claims to provide a basis for understanding why memory models have historically underperformed full-context Transformers: causal training objectives are inherently unsuited for high-fidelity memory formation.

Key contributions include:

Characterization of Memory: Demonstrating that standard LLM embeddings are information-poor and that this is a result of the non-invertible nature of next-token prediction.
New Training Paradigm: Proposing that accurate memory formation requires combined objective functions (specifically including reconstruction or copy tasks) and potentially curriculum learning.
Architectural Proposal: Introducing a parallelizable encoder-decoder memory model that decouples memory formation from memory usage, offering computational efficiencies at inference (reduced KV cache, lower time-to-first-token).
Implications for Retrieval: The findings suggest that current embedding-based retrieval systems are fundamentally limited because standard embedding models do not retain sufficient information to allow for arbitrary lookup of content within a text chunk.

The authors conclude that while next-token prediction alone is insufficient for creating usable memory, models can be trained to form and decode accurate memories, enabling more efficient and capable language systems that can access arbitrary input information.