Matryoshka Hash Representations for Model-Aware Compact Semantic Retrieval

원문 초록

Retrieval-augmented generation (RAG) depends on dense retrieval: each document is stored as a learned vector, and queries are answered by nearest-neighbor search. Keeping one full-precision vector per document dominates index cost at corpus scale, so systems replace each vector with a short code of a few bytes (quantization). Standard quantizers such as product quantization (PQ) pick the code that best reconstructs the original vector, but reconstruction is only a proxy for retrieval: at a few bytes per document, a model-aware code trained on the retriever's ranking signal uses the budget better. A single code is even more useful if its short prefixes are each directly searchable, letting a deployment set its efficiency-quality operating point without re-encoding the corpus. But training all prefixes under one objective makes the early bits a compromise across budgets: short codes improve while the full-width code degrades. Binary codes sharpen the conflict, since a gradient step either flips a stored bit or leaves it unchanged. We introduce Matryoshka Hash Representations (MHR), a two-stage procedure that separates full-width training from prefix organization. MHR first learns a longer binary code, then freezes the model and trains zero-initialized residual code adaptors for directly searchable prefixes. Documents are stored at one bit per coordinate, queries keep continuous logits as in PQ, and search runs on FAISS FastScan. Trained on MS MARCO and zero-shot transferred to seven BEIR datasets, one MHR code reaches 0.502 NDCG@10 at 16 bytes and 0.556 at 32 bytes, against 0.404 and 0.524 for the best compact baselines; the gains hold across training seeds, on all 13 public BEIR datasets, and against off-the-shelf binary Matryoshka embeddings with stronger backbones. The same code also helps shortlist candidates for reranking and prune low-storage graph indexes such as LEANN.