Efficient Multivector Retrieval with Token-Aware Clustering and Hierarchical Indexing

1. Problem Statement

Multivector retrieval models (e.g., ColBERT) have achieved state-of-the-art effectiveness in Information Retrieval (IR) by encoding documents as sets of token-level vectors and using "late interaction" (MaxSim) to compute relevance. However, their deployment is hindered by:

• High Computational and Memory Costs: Storing and searching billions of token vectors requires massive resources.
• Inefficient Clustering: Current solutions rely on standard $k$-means clustering to generate centroid vectors for compression and candidate gathering.
  • Scalability Bottleneck: $k$-means scales poorly with dataset size and the number of clusters ($k$), often requiring tens of hours or expensive GPUs to train even moderately large sets (e.g., 262k centroids).
  • Token Imbalance: Standard $k$-means treats all vectors equally. Consequently, frequent tokens (stopwords) dominate the loss function and consume the majority of centroids, while rare, highly discriminative tokens are underrepresented, degrading retrieval quality.

2. Methodology

The authors propose Tachiom, a retrieval system comprising two core innovations: Token-Aware Clustering (Tac) and a Hierarchical Index for optimized retrieval.

A. Token-Aware Clustering (Tac)

Tac addresses the limitations of standard $k$-means by exploiting the structural property of token identity in multivector embeddings. Instead of clustering all vectors globally, Tac decomposes the problem into independent per-token subproblems.

• Centroid Allocation Strategy: Tac allocates the centroid budget based on token frequency and semantic variance, rather than uniform vector distribution.
  • Tail Handling: Extremely rare tokens ($n < \mu$) are guaranteed at least 1 centroid; small tokens ($\mu \le n < \tau$) get 2.
  • Damped Scoring: For active tokens, a weight $w_j$ is calculated as $\sqrt{n_j} \cdot s_j$, where $n_j$ is frequency and $s_j$ is semantic variance (component-wise variance). The square root dampens the influence of high-frequency tokens, preventing them from monopolizing the budget.
  • Bounding & Reconciliation: Hard floors and ceilings ensure minimum representation for rare tokens and prevent over-allocation. The final budget is reconciled to match the target $k$.
• Parallelization: By solving independent $k_j$-means problems for each token type, Tac enables massive parallelization, drastically reducing computational complexity compared to global $k$-means.
• Residual Compression: After clustering, residuals (vector - centroid) are normalized and compressed using Product Quantization (PQ).

B. Hierarchical Index for Retrieval (Tachiom)

Tachiom leverages the high-resolution centroids generated by Tac to optimize the retrieval pipeline in two phases:

1. Gather Phase (Centroid-Only Scoring):
   • An Hnsw graph is built over the centroids.
   • Each centroid maintains an inverted list of document IDs.
   • For a query, the system traverses the Hnsw graph to find nearest centroids, retrieves associated documents, and computes an approximate MaxSim score using only centroid interactions.
   • This bypasses expensive token-level computations entirely during the initial candidate selection.
2. Refine Phase (PQ-Optimized MaxSim):
   • The top candidates are refined using both centroids and PQ-compressed residuals.
   • Cache-Optimized Layout: The authors introduce a specialized memory layout for PQ distance tables. Instead of scattered memory accesses for each query token, distances are organized into Macro-blocks (subspaces), Centroid blocks, and Query token micro-blocks. This allows retrieving distances for all query tokens against a specific centroid as a single contiguous memory block, significantly improving cache locality and SIMD utilization.

3. Key Contributions

1. Token-Aware Clustering (Tac): A novel algorithm that decomposes global clustering into per-token subproblems. It theoretically guarantees a speedup lower-bounded by the ratio of total vocabulary weight to the maximum single-token weight.
2. Tachiom Architecture: A retrieval system that combines Hnsw-based centroid gathering with a cache-optimized PQ layout tailored for late interaction, enabling efficient full scoring.
3. Theoretical & Empirical Validation: The paper provides a theoretical lower bound for Tac's speedup and demonstrates that the method scales to millions of centroids while maintaining or improving retrieval effectiveness.

4. Experimental Results

Experiments were conducted on Ms Marco-v1 (8.8M passages) and LoTTE (2.4M passages) using ColBERTv2 embeddings.

• Clustering Efficiency:
  • Tac achieved up to 247$\times$ faster training than standard $k$-means (Faiss) and 84$\times$ faster than Faiss with MKL acceleration.
  • Tac successfully clustered 600M vectors into 262k clusters in 8 minutes, whereas competitors hit a 24-hour timeout before reaching 131k–524k clusters.
  • Tac scales seamlessly to over 4 million centroids in ~100 minutes.
• Retrieval Effectiveness:
  • Tac achieved equal or superior MRR@10 compared to standard $k$-means at fixed centroid budgets, proving that prioritizing rare tokens improves approximation quality.
• Retrieval Speed:
  • Tachiom achieved up to 9.8$\times$ faster end-to-end retrieval compared to state-of-the-art systems (Emvb, Warp, Igp).
  • It matched the peak effectiveness of Emvb (which uses expensive supervised PQ variants) while using standard PQ and achieving a 5.5$\times$ speedup over Emvb.

5. Significance

This work fundamentally shifts the paradigm for multivector retrieval efficiency:

• Scalability: It removes the computational bottleneck of clustering, allowing systems to utilize millions of centroids (finer granularity) which were previously infeasible to train.
• Quality-Efficiency Trade-off: By aligning centroid allocation with retrieval objectives (favoring discriminative rare tokens) rather than pure reconstruction error, it achieves better effectiveness with fewer resources.
• Practical Deployment: The combination of fast training and ultra-fast inference (via Hnsw gathering and optimized PQ) makes high-performance multivector retrieval viable for large-scale production environments.

The authors have released their Rust implementation of Tac and Tachiom on GitHub to ensure reproducibility.