Extended Mind Transformers let each query token retrieve relevant memory tokens during generation through the Transformer's existing key-value attention structure.
2
The method produces citations tied to the tokens the model actually retrieved, and it can retrieve more information when uncertainty indicates that an answer may be wrong.
3
Extended Mind Transformers work without fine-tuning when used with modern relative position methods, but stride length, top K, and retrieval regularization need careful tuning.
Summary
Phoebe Klett argues that long context and RAG make retrieval decisions too early or at too coarse a level. Long context places all available material in the prompt, which can waste computation and confuse the model. RAG selects documents before generation, usually with an external embedding search. Extended Mind Transformers move retrieval into the Transformer. Memory tokens pass through the model once, their key-value representations are saved, and each query token retrieves and attends to a selected set during generation. Klett describes experiments on a counterfactual retrieval benchmark with contexts from about 2,000 to 16,000 tokens. EMT remains competitive with fine-tuned long-context models and performs better when combined with some RAG context in one experiment. The method also exposes token-level retrieval for causal citations. Uncertainty can trigger regeneration with more memory. Klett covers stride length, top K, similarity masking, and removal of unknown tokens, then points to open-source models, code, and data.
Long context and RAG make retrieval choices at the wrong time or level
Klett says pretrained language models need application-specific information, but long context and RAG handle it differently. Long-context methods extend the model's input window, often through expensive fine-tuning, then place all the material in the prompt. Irrelevant information can confuse the model, and a code question usually needs only some function definitions. RAG reduces the prompt, but an external retriever chooses what matters once, before generation. It also uses a representation that may differ from the way the model reasons about the data. EMT separates memory from the inference query and lets retrieval happen during generation.
EMT uses the Transformer's own key-value representations for retrieval
Extended Mind Attention changes the Transformer's attention mechanism. In each decoder layer, the model represents data as key-value pairs, so Klett says the retrieval mechanism is already present. EMT passes memory tokens through the model and saves those key-value representations. During generation, each query token uses cosine similarity to retrieve a chosen number of memory tokens and attend to them. The retrieved tokens can differ from one generation step to the next. This makes retrieval more granular than selecting a document set before the model starts producing an answer.
Relative position methods let retrieved tokens work without fine-tuning
Retrieved tokens create a position problem because Transformers need information about token order and distance. Earlier work had to fine-tune models with absolute position embeddings so they could use retrieved tokens. Klett says newer relative position methods make EMT possible without further fine-tuning. Her implementations use rotary position embeddings, found in Llama models, and ALiBi linear biases, used by MPT models. Rotary embeddings encode distance through angles between vectors. ALiBi instead reduces the influence of information that is farther away.
The counterfactual benchmark tests whether models use supplied facts
Normal Computing's benchmark uses question-answer pairs with contexts ranging from about 2,000 to 16,000 tokens. To avoid testing facts the model may have memorized, the researchers replace the real answer with a plausible but incorrect alternative. For the question about who wrote 'These Boots Were Made for Walkin',' they replace Lee Hazelwood with Terry Allen. The model must retrieve the answer supplied at inference time rather than rely on its pretrained knowledge. Klett compares EMT with fine-tuned models and a base Llama model using interpolated position embeddings.
EMT stays competitive on long retrieval and can use RAG as extra context
On the counterfactual retrieval task, the base model extrapolates reasonably at shorter extensions but falls off at 16,000 tokens because its position embeddings cannot extend that far. Klett says fine-tuned long-context models perform worse than EMT on shorter inputs, suggesting that long-context fine-tuning can reduce attention quality there. EMT remains competitive with those models through 16,000 tokens, despite not being fine-tuned in the experiment. In another setup, the model receives some RAG context but still relies mainly on internal retrieval, and Klett reports that it outperforms GPT-4 in that comparison.
Retrieved tokens provide citations tied to the model's generation
RAG citations can show that a fact appeared somewhere in the input, but Klett describes them as post-hoc rationalization because they do not prove what information caused the model to generate the answer. EMT can report the memory tokens retrieved during generation. In her example, the model answers when Alexander Grothendieck received French citizenship. The retrieved tokens include the date 1971 and parts of the mathematician's name. Those token-level links show which memory the model attended to when producing the date, giving users more direct evidence about the answer's source.
Uncertainty can trigger another retrieval pass with more memory
EMT exposes token-level output entropy, and Klett says other uncertainty measures can also control retrieval. If the model is uncertain about a generated token, it can regenerate that step while reading more information from the memory cache. In her example, the default retrieval amount produces 1993 instead of the correct answer, 1971. When the uncertain steps are regenerated with more memory, the model gets the answer right. This can reduce hallucinations while avoiding the cost of retrieving a large amount of memory for every token.
Stride, top K, and regularization determine practical EMT behavior
Stride length controls how memory representations are created from a model trained on a fixed context length. A smaller stride gives each token more high-quality representations but requires more computation. Top K controls how many key-value pairs each query token can retrieve. Longer memories generally benefit from a larger value, but Klett says it should depend on memory length. Similarity masking removes retrieved tokens below a similarity threshold, with 0.25 given as an example. For RoPE models, removing memory entries associated with unknown tokens helps prevent messy data from producing misleading matches.
"We can look up exactly which tokens were retrieved from those memories and used during generation."09:38
Who should watch
You are building a RAG system and need citations that show which input tokens caused an answer, rather than only which documents were retrieved.
Your application has long memories but should avoid putting every available token into the prompt at every generation step.
You want to experiment with retrieval inside a decoder-only Transformer without fine-tuning, and need guidance on stride, top K, and filtering.
