Memory as a Modality for LLMs

Abstract

Large language models are increasingly deployed as reasoning engines and autonomous agents that interact with users, tools, documents, and environments over extended horizons. Yet these systems lack native, persistent memory, limiting their ability to accumulate experience, adapt across tasks, and reuse knowledge from prior interactions. Existing mechanisms—retrieval-augmented generation (RAG) and KV-cache compression—either retrieve isolated text fragments or compress a single active context, but do not learn reusable representations from accumulated experience.

We introduce MEMORILLA, a memory-as-modality architecture that compresses conversations, documents, and agent trajectories into compact latent memory embeddings directly usable by a frozen LLM. Through gradient-trained attention, MEMORILLA compresses entire document collections into a fixed set of query-conditioned memory embeddings and preserves global structure from prior contexts rather than retrieving isolated snippets. A single trained memory module matches or exceeds the strongest RAG baseline on seven of eight benchmarks at 29–875× fewer context tokens, reaching 92.30% accuracy on PersonaMem-v2 where RAG achieves 31.72%, and exceeding RAG by +7.60 EM on TriviaQA with 372× fewer tokens.¹

Background

Many applications require a model to draw on information beyond its pretraining data—a user’s conversation history, an agent’s past trajectories, or a large reference corpus. Retrieval has become the dominant solution: relevant text is pulled into context, underpinning everything from RAG to KV-cache compression to text-based agent memory. In all of these, memory is treated as information to be fetched rather than learned, and the LLM never builds a representation of the experience it accumulates.

In humans, memory is supported not by perfect recall of past events but by an evolving internal model of what has been experienced. Analogously, a learned corpus representation native to the LLM could encode an entire document collection once, serve diverse queries from a compact set of embeddings, and preserve cross-document patterns that single-passage retrieval misses.

Building such a module for a frozen LLM is constrained on three axes:

• Embedding-space compatibility.Since the decoder is frozen, the encoded representation must lie within the LLM’s existing embedding space to remain interpretable.
• Query adaptivity.The same encoded corpus must serve diverse downstream queries—question-conditional encodings overfit to one query, while question-agnostic ones (mean pooling, fixed soft prompts) lack the needed granularity.
• Scale. Inputs span hundreds of documents per entity, where mean pooling collapses content into a single centroid and learned-query compression degrades as input length grows.

Existing methods address subsets of these constraints; none satisfies all three at once. What remains missing is a reusable, query-adaptive, multi-document memory module that operates entirely within a frozen LLM’s embedding space.

The MEMORILLA Architecture

MEMORILLA treats accumulated context as a learned latent representation rather than text to be retrieved. For each inference step it takes a collection of N source document chunks D = {d₁, …, d_N} (conversation histories, agent trajectories, or reference material) and a textual query q. A learned memory module compresses D into a fixed set of K embedding vectors, conditioned on q, which are prepended to the embedded query and consumed by a frozen decoder—much as a vision encoder feeds a vision-language model. Memory is computed once and reused across all subsequent queries, so inference cost is fixed in the number of memory tokens rather than the original context length.

Encoders

Both the query encoder and the document encoder use a frozen Qwen3-Embedding-4B model (hidden dimension h = 2560), encoding queries and documents separately. Each chunk is truncated to 128 tokens on a sentence boundary and mapped to a single embedding vector. All chunk embeddings are precomputed and cached, amortizing encoding cost across every query for a given source.

The memory module

The memory module is the sole trainable component. It refines K learnable prototype vectors into the output memory through four stages:

1. Query-prototype initialization. K shared, learnable prototypes (Xavier-uniform initialized) seed the memory slots.
2. Question conditioning. Prototypes attend to the query embedding via multi-head cross-attention, steering each slot toward question-relevant information while preserving slot diversity through the residual connection.
3. Iterative refinement.Two cross-attention blocks let the memory slots extract information from the normalized document embeddings, each with a 4×-expansion feedforward sub-layer, 8 attention heads, and dropout 0.1.
4. Projection. A final LayerNorm and a learned linear layer produce the output memory vectors.

Integrating memory into a frozen LLM

We add a single special <|memory|> token to the vocabulary, repeat it Ktimes early in the input sequence, and replace each position’s embedding with the corresponding memory vector. Rotary position embeddings extend naturally over the concatenated sequence, and generation proceeds through standard autoregressive decoding. The decoder is one of the frozen Qwen3-4B or Qwen3-8B models.

Training

Only the memory module is trained; the encoders and decoder stay frozen and gradients flow exclusively through f_mem. The model maximizes the conditional log-likelihood of the answer tokens, with loss computed only over answer positions. Training follows a two-stage pipeline: pretraining on 2M English Wikipedia paraphrase examples for general-purpose compression, then finetuning on a mixture of ten datasets (reading comprehension, open-domain QA, dialogue summarization) spanning over 1.5M QA pairs.²

Results

We evaluate a singletrained memory module across eight benchmarks spanning three axes—personalization, long-document QA, and knowledge-intensive QA—all using a frozen Qwen3-Embedding-4B encoder and a Qwen3-4B/8B decoder, with K = 16 memory tokens.

• Personalization. On PersonaMem-v2, MEMORILLA reaches 92.30%accuracy versus 31.72% for RAG-50 and 87.66% for xRAG. It surpasses ComMer on LaMP-4 and LaMP-7 while using 8× fewer compression tokens—all at 29–875× fewer context tokens than RAG-50.
• Long-document QA. On NarrativeQA, MEMORILLA compresses stories averaging 70k tokens into 16 memory tokens and outperforms RAG-1 and xRAG. On TriviaQA it reaches 63.20 EM, +7.60 over RAG-50at 372× fewer tokens.
• Knowledge-intensive QA. On FactKG, MEMORILLA hits 84.33% accuracy, +5.94over RAG-50 at 104× compression. On PubMedQA—where a handful of self-contained abstracts fit the context window—RAG retains an edge.

Performance scales favorably with both memory capacity and decoder size: moving from a 4B to an 8B decoder amplifies MEMORILLA’s advantage over every baseline.

What Does MEMORILLA Learn?

Per-question analysis reveals a consistent picture of what compression preserves and where it falls short:

• Concepts over surface tokens.On NarrativeQA, MEMORILLA beats RAG-50 by +9.89pp on “why” and +7.07pp on “how” questions, but trails on “who” and “which,” where the answer is a specific entity name that retrieval can locate verbatim.
• Relational reasoning.On FactKG, MEMORILLA’s advantage grows monotonically with entity count and reasoning depth—RAG-50 drops from 78.7 to 51.8 as relational hops increase from 1 to 4, while MEMORILLA stays in the 80–89 range.
• Complementary to retrieval. The two methods solve largely non-overlapping questions. Their union answers 94.3% of FactKG questions and boosts PubMedQA accuracy by +21.0% over either method alone.

This mirrors the neuroscience distinction between semantic and episodic memory: human memory likewise preserves abstract structure while losing surface detail over time.

Why It Matters

MEMORILLA reframes memory from a retrieval layer into a learned modality for frozen LLMs. A single trainable module compresses hundreds of documents into 16 latent tokens, matches or exceeds RAG-50 across seven benchmarks at a fraction of the context cost, and transfers across personalization, long-document QA, and fact verification without task-specific changes.

More broadly, memory as a modality opens the door to persistent, adaptive systems—agents that carry forward compact representations of past trajectories, personalize to users without storing raw histories, and improve through continual learning. The task-transfer results suggest memory modules could become reusable components shared across applications, much as vision encoders became standard in vision-language models.