Token level matching boosts memory retrieval accuracy in language models
Beyond the Final Embedding: Retrieving Memories from Token-Level Evidence
Information RetrievalComputation and Language
Summary
When virtual assistants try to find helpful information stored in past conversations or documents, they usually compare summaries of the whole memory to the query. This can miss important details that only appear in small parts of the text. The authors show that looking closely at individual words or tokens improves finding the right memories. Their method matches query words to specific words in the memory, combining this with overall similarity, and works well without extra training. This approach helps assistants answer questions more accurately, especially when long-term memory is involved.
What this means in practice
- •For virtual assistant developers: Improve the accuracy of memory retrieval in conversational AI by using token-level matching alongside traditional embeddings.
- •For enterprise search engineers: Enhance document retrieval systems to find relevant passages by matching query tokens directly to stored content without retraining models.
Authors
Han Luo, Zhan Zhuang, Xu He, Jujia Zhao, Yu Li, Suqin Yuan, Dong Li, Lei Feng
Abstract
Memory-augmented assistants retrieve evidence from stored conversations and documents to answer queries. Most embedding-based retrievers compare single-vector representations of queries and memories, even when relevance depends on a small part of the content. Such local evidence can be matched to queries using the contextual token representations already computed by the same embedding models. Our method uses sparse token matching to improve retrieval without additional training. Each query token retrieves a small set of matching tokens across candidate memories. The matching scores are aggregated for each memory and combined with whole-memory similarity to rank memories. An optional low-rank query adapter can be trained with supervision while keeping the encoder and stored memory representations fixed. Our method improves retrieval across eight benchmarks and multiple model families without additional training. Query adaptation further increases the gains to 13.3 and 14.3 percentage points in average Recall@10 and NDCG@10 on Qwen3-Embedding-4B. These retrieval gains translate into higher question-answering accuracy on long-term memory benchmarks.