Improving Long-Context Retrieval with Multi-Prefix Embedding

2026-06-22

AI Summary

This work addresses the challenge in long-context retrieval where single-vector embeddings struggle to preserve fine-grained semantics, while multi-vector approaches incur high storage overhead. The authors propose Multi-Prefix Embeddings (MPE), which leverages a causal language model to perform a single forward pass over the full document, segmenting it at end-of-sentence (EOS) boundaries and extracting embeddings at each prefix boundary. Trained solely with document-level labels, MPE enables block-level MaxSim matching that effectively captures cross-block context and inherently supports evidence localization. Experimental results demonstrate that MPE achieves performance on par with or superior to existing single-vector, independently chunked, and multi-vector methods across multiple benchmarks, including MLDR-en, BrowseComp-Plus, and LongEmbed.

Abstract

Long-context retrieval exposes a tension: single-vector embeddings lose fine-grained detail, while token-level multi-vector methods incur prohibitive storage. We propose Multi-Prefix Embedding (MPE), which partitions a document into chunks separated by EOS tokens, encodes the full sequence in a single causal forward pass, and extracts one embedding at each prefix boundary. MPE retains cross-chunk context, enables chunk-level MaxSim matching, and trains with only document-level relevance labels. Experiments on MLDR-en, BrowseComp-Plus, and LongEmbed show that MPE is competitive with or outperforms single-vector, independent-chunk, and multi-vector baselines, while providing a natural source attribution mechanism for locating evidence chunks.

Problem

• long-context retrieval
• single-vector embeddings
• multi-vector methods
• storage overhead
• fine-grained detail

Innovation

• Multi-Prefix Embedding
• long-context retrieval
• chunk-level MaxSim
• cross-chunk context
• source attribution