Matryoshka Representation Learning (MRL) trains embeddings as hierarchical representations: shorter prefixes preserve a useful approximation of the full embedding, while additional dimensions progressively improve precision. This lets vector search trade representation size for speed, memory, and recall. section 13.7.4 section 13.7
Use MRL in a retrieval pipeline by:
- Generating full document and query embeddings with an MRL-capable model. section 13.7.4
- Slicing embeddings to a selected prefix, such as one-half, one-quarter, or one-eighth of the original dimensions. section 13.7.4
- Indexing the reduced vectors for faster approximate retrieval, then reranking candidates with fuller or higher-precision vectors when recall matters. section 13.7 section 13.7.5
Benchmarking dimension choices
For 18,456 original 1,024-dimensional vectors, a 512-dimensional MRL index improved search time by 49.15% and reduced index size by 50%, with 0.7022 recall and 1.0 reranked recall. At 128 dimensions, search improved by 84.35% and index size by 87.5%, but reranked recall fell to 0.64. section 13.7.4
Listing 13.26 Benchmarking MRL embeddings at different thresholds
def get_mrl_embeddings(embeddings, num_dimensions):
mrl_embeddings = numpy.array( #1
list(map(lambda e: e[:num_dimensions], embeddings))) #1
return mrl_embeddings
def index_mrl_embeddings(doc_embeddings, num_dimensions, mrl_index_name):
mrl_doc_embeddings = get_mrl_embeddings(doc_embeddings, num_dimensions)
print(f"{mrl_index_name} embeddings shape:", mrl_doc_embeddings.shape)
mrl_index = index_full_precision_embeddings( #2
mrl_doc_embeddings, mrl_index_name) #2
return mrl_index
print(f"Original embeddings shape:", doc_embeddings.shape)
original_dimensions = doc_embeddings.shape[1] #3
for num_dimensions in [original_dimensions//2, #4
original_dimensions//4, #5
original_dimensions//8]: #6
mrl_index_name = f"mrl_embeddings_{num_dimensions}"
mrl_index = index_mrl_embeddings(doc_embeddings,
num_dimensions,
mrl_index_name)
mrl_queries = get_mrl_embeddings(query_embeddings,
num_dimensions)
evaluate_search(full_index, mrl_index, mrl_index_name, #7
query_embeddings, mrl_queries) #7
evaluate_rerank_search(full_index, #8
mrl_index, query_embeddings, mrl_queries) #8
Inspect the Python pattern for slicing embeddings to a chosen dimension count and benchmarking multiple MRL index sizes. section 13.7.4
MRL is complementary to ANN search, quantization, and reranking rather than a replacement for them. Combining these techniques can greatly reduce search time and index size, but aggressive compression can lower initial recall and requires evaluation. section 13.7.5 section 13.9
