Improving Long-Context Retrieval with Multi-Prefix Embedding
For practitioners needing efficient long-context retrieval, MPE offers a practical trade-off between detail and storage without requiring token-level labels.
Multi-Prefix Embedding (MPE) partitions documents into chunks with EOS tokens, extracts embeddings at prefix boundaries in a single forward pass, and achieves competitive or superior retrieval performance on MLDR-en, BrowseComp-Plus, and LongEmbed benchmarks compared to single-vector, independent-chunk, and multi-vector baselines, while enabling source attribution.
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.