Token Pooling-Aware Training for Multi-Vector Embedding Models

Understanding Token Pooling in Multi-Vector Retrieval

Multi-vector embedding models, such as ColBERT, represent documents as collections of contextual token vectors rather than a single vector. This approach allows for high-expressivity search via MaxSim scoring, where query tokens are matched against the most relevant document tokens. However, this granularity leads to significant storage requirements, as a corpus with millions of documents can generate hundreds of millions of vectors.

Token pooling mitigates this by merging similar token vectors into fewer, representative vectors. While standard pooling can be applied to any model, it often leads to information loss because the encoder was not originally trained to perform retrieval using compressed representations. Pooling-aware training addresses this by explicitly teaching the model to maintain retrieval accuracy when its vectors are merged.

The Mechanism of Pooling-Aware Training

Pooling-aware training functions similarly to Matryoshka embeddings by applying a retrieval objective to both full and compressed representations simultaneously. During training, each batch computes two losses:

  1. Unpooled Loss: Retrieval using the original, full set of document embeddings.
  2. Pooled Loss: Retrieval after compressing documents to a target vector count (e.g., using a pool_factor of 4).

The total loss is an average of these two components. This forces the encoder to structure its token embeddings such that critical information survives the clustering process. The HierarchicalTokenPooling module uses Ward hierarchical clustering to group vectors, and the training process employs a straight-through estimator to pass gradients back to the original token embeddings despite the discrete nature of cluster assignments.

Implementation and Performance

Developers can implement this using the MultiVectorTokenPoolingLoss wrapper within the Sentence Transformers library. By wrapping an existing retrieval loss, this approach allows for efficient training without requiring separate encoder passes for different pooling variants.

Key Performance Gains

Evaluations using ModernBERT-base on the NanoBEIR benchmark demonstrate that pooling-aware training significantly outperforms standard models at the same storage budget:

  • Retention: At a pool_factor=4 (a 4x reduction in vectors), pooling-aware training achieves 97.8% retention of unpooled nDCG@10 scores, compared to 94.6% for normally trained models.
  • Flexibility: The benefits of this training extend beyond the specific pool_factor used during training, showing consistent improvements across various fixed caps and adaptive pooling thresholds.
  • Adaptive Calibration: Users can calibrate a distance threshold to achieve a target average vector count, allowing the model to dynamically allocate more vectors to documents with higher information density.

Recommended Compression Strategy

For applications where index size is a primary constraint, the following tiered approach is recommended:

  1. Quantization: Apply techniques like Fast-PLAID (2-bit or 4-bit residual quantization) to reduce storage per coordinate.
  2. Token Pooling: Use pooling-aware fine-tuning to reduce the total number of document vectors.
  3. Matryoshka Embeddings: Utilize Matryoshka token embeddings to reduce the dimensionality of individual vectors.

These techniques are complementary and can be combined by wrapping multiple loss functions, such as MultiVectorMatryoshkaLoss and MultiVectorTokenPoolingLoss, to optimize for both vector count and dimensionality simultaneously.

Sources