MUVERA Python Implementation

The FDE process converts multiple vectors (representing a document or query) into a single fixed-dimension vector through the following steps:

1. Space Partitioning: For each repetition, apply SimHash (multiplying by a random Gaussian matrix) and convert to partition indices using Gray Code. This creates $2^{k_{sim}}$ partitions.
2. Vector Aggregation:
   • For Queries: Sum all vectors belonging to each partition.
   • For Documents: Average all vectors belonging to each partition.
3. Repetition: Repeat the partitioning and aggregation steps with different random seeds.
4. Output: Concatenate the results from all repetitions into a single FDE vector. The final dimension is num_repetitions × num_partitions × projection_dim.

This method uses Locality Sensitive Hashing (LSH) to ensure that vectors close in the original space contribute to the same parts of the FDE vector, resulting in high dot products during comparison.

The FixedDimensionalEncodingConfig dataclass controls the FDE generation process. It is a direct Python mapping of the C++ Protocol Buffer configuration.

Key Parameters:

• dimension (int): The original vector dimension (default: 128).
• num_repetitions (int): Number of independent runs to perform (default: 10).
• num_simhash_projections (int): Controls partition granularity (default: 6).
• seed (int): Random seed for reproducibility (default: 42).
• encoding_type (EncodingType): Determines if vectors are summed or averaged.
• projection_type (ProjectionType): Determines if dimensionality reduction is used.
• projection_dimension (Optional[int]): Target dimension if using reduction.
• fill_empty_partitions (bool): Whether to fill empty partitions.

@dataclass
class FixedDimensionalEncodingConfig:
    dimension: int = 128
    num_repetitions: int = 10
    num_simhash_projections: int = 6
    seed: int = 42
    encoding_type: EncodingType = DEFAULT_SUM
    projection_type: ProjectionType = DEFAULT_IDENTITY
    projection_dimension: Optional[int] = None
    fill_empty_partitions: bool = False
    final_projection_dimension: Optional[int] = None

To generate Fixed Dimensional Encodings (FDE) for queries and documents, use generate_query_fde and generate_document_fde. You must first define a FixedDimensionalEncodingConfig to specify the vector dimension, number of repetitions, and SimHash projections. The resulting FDE vectors can be compared using a dot product to approximate Chamfer similarity.

import numpy as np
from fde_generator import FixedDimensionalEncodingConfig, generate_query_fde, generate_document_fde

# 1. Create configuration
config = FixedDimensionalEncodingConfig(
    dimension=128,              # Vector dimension
    num_repetitions=10,         # Number of independent partitionings
    num_simhash_projections=6,  # Creates 2^6 = 64 partitions
    seed=42
)

# 2. Prepare data
# Query: 32 vectors of 128 dimensions each
query_vectors = np.random.randn(32, 128).astype(np.float32)

# Document: 80 vectors of 128 dimensions each  
doc_vectors = np.random.randn(80, 128).astype(np.float32)

# 3. Generate FDEs
query_fde = generate_query_fde(query_vectors, config)
doc_fde = generate_document_fde(doc_vectors, config)

# 4. Compute similarity (approximates Chamfer similarity)
similarity_score = np.dot(query_fde, doc_fde)
print(f"Similarity: {similarity_score}")

The ColbertFdeRetriever class implements the MUVERA approach. It uses a ColBERT model to generate multi-vector embeddings, which are then compressed into Fixed Dimensional Encodings (FDE). This significantly speeds up search times at a slight cost to recall.

Workflow

1. Initialize: Creates a FixedDimensionalEncodingConfig internally with default parameters.
2. Index: Call .index(corpus). This performs two steps: generating native ColBERT embeddings and then batch-generating FDEs using generate_document_fde_batch.
3. Search: Call .search(query). The query is encoded via ColBERT and then transformed into an FDE using a specific configuration (fill_empty_partitions=False) to perform a fast matrix multiplication against the FDE index.

from main import ColbertFdeRetriever

# Initialize
retriever = ColbertFdeRetriever()

# Index documents
corpus = {"doc1": {"title": "Hello", "text": "World"}}
retriever.index(corpus)

# Search
results = retriever.search("Hello World")
# results: {'doc1': 0.88} (sorted by score descending)