Token Pooling-Aware Training for Multi-Vector Embedding Models

Token Pooling-Aware Training addresses the high storage costs of multi-vector models like ColBERT by merging similar token embeddings, reducing document vectors by 74.8% while maintaining 97.8% of the original nDCG@10 score on NanoBEIR at pool_factor=4.
The technique employs a HierarchicalTokenPooling module that uses Ward hierarchical clustering to group similar token embeddings, with a pool_factor=4 setting reducing vector counts from an average of 193.22 to 48.68 per document.
Training utilizes a loss wrapper that averages the retrieval objective for original and pooled representations, using straight-through estimators to backpropagate gradients through the discrete clustering decisions.
Kernel acceleration for MaxSim scoring and Ward clustering provides a 30-150x speedup over SciPy, and the approach is implemented in Sentence Transformers with released checkpoints based on answerdotai/ModernBERT-base.

Implementation and Evaluation

Models are trained on 120,000 MS MARCO triplets with a 128-dimensional token projection, evaluated on NanoBEIR@13 using mean nDCG@10.
Pooling-aware training outperforms baseline models across all compression levels, with the largest gains observed at max_tokens=4 where nDCG@10 increases from 37.22 to 43.00.
Adaptive pooling using distance_threshold calibration allows for corpus-specific compression, maintaining 97.3% of unpooled performance while reducing storage by 84.7%.

For a corpus of 1 million documents with 200 vectors each, applying pool_factor=4 reduces total storage from 12.8 GB to 3.2 GB using 4-bit residual quantization.
Recommended compression pipeline order is embedding quantization, token pooling, and Matryoshka token embedding dimensionality reduction.
Inference allows dynamic adjustment of pool_factor to explore different storage budgets without retraining, preserving ranking order for relevant documents.