Sentence Transformers Maintainer Tom Aarsen: Pooling-Aware Training Preserves 97.8% Quality Under 4x Compression for Multi-Vector Embeddings
Multi-vector embedding models, such as ColBERT, provide each query token with its own independent opportunity to find a match within a document. This makes them highly expressive for search, but the cost is that each document generates hundreds of vectors. A corpus containing millions of documents may hold hundreds of millions of vectors; even with 4-bit residual quantization via Fast-PLAID, 200 vectors of 128 dimensions require 12.8 KB per document, totaling 12.8 GB across a million documents.
Token pooling compresses representations by merging similar token vectors through Ward hierarchical clustering. For example, with pool_factor=4, a document retains only about one-quarter of its original vectors. The more aggressive the compression, the easier it becomes to lose discriminative information critical for retrieval—a risk because standard training objectives never required the encoder to adapt to such compression.
In a blog post published on November 11, 2026, Sentence Transformers maintainer Tom Aarsen introduced and implemented a pooling-aware training method. Its core concept resembles Matryoshka embeddings: applying retrieval objectives simultaneously across multiple representation granularities. Specifically, each training batch contributes two losses—the retrieval loss for original document vectors and the retrieval loss for pooled document vectors—averaged with equal weights of 0.5.
In implementation, MultiVectorTokenPoolingLoss wraps existing multi-vector retrieval losses (such as CachedMultiVectorMultipleNegativesRankingLoss) and accepts a set of token_poolings parameters. None represents the original vectors, while HierarchicalTokenPooling(pool_factor=4) represents the compressed vectors. During training, Ward clustering uses a straight-through estimator to pass gradients from pooled vectors back to their member tokens without differentiating the cluster assignments themselves, allowing the encoder to learn from retrieval errors made with the compressed representations.
Aarsen trained two models on answerdotai/ModernBERT-base: one using only standard contrastive loss, and another using pooling-aware loss. Both were trained for one epoch on 120,000 MS MARCO triplets with seed 42, matching initialization and optimization settings exactly, differing only in the loss function. In the NanoBEIR@13 evaluation, the standard model retained 94.6% of its unpooled nDCG@10 at pool_factor=4 (49.03 vs. 51.84), while the pooling-aware model retained 97.8% (51.23 vs. 52.39).
Notably, the benefits of pooling-aware training extend beyond the pool_factor=4 used during training. Gains persist under fixed caps (reaching +3.48 percentage points at max_tokens=8 and +5.78 at max_tokens=4) and adaptive distance thresholds, despite training utilizing only full representations and proportional pooling conditions.
The applicability boundaries of this method warrant attention: all figures come from a single self-test by the author, the unpooled baseline itself was relatively low (51.84/52.39), the training scale was limited (one epoch, 120K triplets), and the evaluation set was NanoBEIR@13 rather than the complete BEIR benchmark. Aarsen himself noted that no hyperparameter search was conducted. The method has been integrated into the Sentence Transformers library, with code and model weights (tomaarsen/mve-ModernBERT-base-token-pooling) available for reproduction.
Sources:huggingface.co