Hugging Face ships token pooling-aware training for multi-vector retrieval models

Hugging Face published a technical walkthrough on November 11 describing token pooling-aware training for multi-vector embedding models, and released the supporting pieces alongside it: a loss wrapper in Sentence Transformers, a full training script, an evaluation script, and two example checkpoints on the Hugging Face Hub.

The problem is storage. Multi-vector models such as ColBERT keep a contextual vector for every token in a document, which makes retrieval expressive but expensive at corpus scale. The post puts numbers on it: at 128 dimensions and 4-bit residual quantization, 200 vectors per document cost 12.8 KB, or 12.8 GB across a million documents; pooling each document down to 50 vectors cuts that to 3.2 GB at the same precision.

Pooling itself is handled by HierarchicalTokenPooling, which uses Ward hierarchical clustering to merge similar token vectors, represents each cluster by its mean, renormalizes the result, and leaves MaxSim scoring intact. Users can compress proportionally with pool_factor, cap vectors per document with max_tokens, or set a shared merge-distance threshold; the first token is protected by default and only document embeddings are pooled.

The training change is what the post is really about. Instead of compressing a model after the fact, the encoder is trained with the knowledge that its document vectors will later be merged. MultiVectorTokenPoolingLoss wraps the cached contrastive loss with unpooled and factor-4 objectives, so the model learns to preserve what matters under compression. Optional kernel acceleration is also provided.

The experiments are deliberately controlled. Both models start from answerdotai/ModernBERT-base with a fresh 128-dimensional token projection and train for one epoch on the same 120,000 MS MARCO triplets, seed 42, a 300-token cap, and short queries expanded to at least 32 token positions. Only the loss differs, and no hyperparameter search was run.

On NanoBEIR@13, at pool_factor=4 the ordinary model scores 49.03 mean nDCG@10 and the pooling-aware model 51.23, against unpooled scores of 51.84 and 52.39. That lifts score retention from 94.6% to 97.8% while storing 74.8% fewer document vectors. The gains carry over to fixed caps and calibrated adaptive thresholds: at 29.58 vectors per document on average, the pooling-aware model keeps 97.3% of its unpooled score with 84.7% fewer vectors, versus 93.3% for the ordinary model.

Aarsen positions the method inside a broader compression stack rather than as a replacement for one: quantization first, then token pooling, then Matryoshka truncation, with the losses combinable. For teams running ColBERT-style retrieval, the practical appeal is a reproducible way to shrink indexes without switching architectures. The open question is whether the retention holds on much larger corpora and production indexes.