Multi-vector embedding models assign a contextual vector to many tokens per document, giving each query token an independent chance to match document content. This richness improves retrieval expressiveness but increases index size: corpora with millions of documents can quickly produce hundreds of millions of vectors. Token pooling reduces representation size by merging similar token vectors, yet overly aggressive pooling can remove retrieval-relevant information. Pooling-aware training teaches the encoder which information to preserve under pooling so you can retrieve better at the same storage budget.
This article explains how pooling works, gives a short reproducible Sentence Transformers example for compressed retrieval training, and surveys the resulting quality versus storage tradeoffs. In the author’s experiments, with pool_factor=4 a normally trained model retained 94.6% of its unpooled NanoBEIR nDCG@10 after pooling, while pooling-aware training increased retention to 97.8% for the same 4x reduction in document vectors.
Why multi-vector retrieval needs many vectors
A single-vector model compresses a whole document into one vector. Multi-vector models (ColBERT-style) keep contextual token vectors so different query parts can match different document parts. Using MaxSim scoring, each query vector picks its best-matching document vector and scores sum: MaxSim(Q,D) = sum_{q in Q} max_{d in D} q^T d.
Even with residual quantization like Fast-PLAID (2- or 4-bit), storing hundreds of vectors per document accumulates significant storage. Example: 128 dimensions at 4 bits per residual coordinate and 200 vectors requires ~12.8 KB of quantized residuals per document (≈12.8 GB for one million documents). Pooling each document to 50 vectors reduces that to ~3.2 GB at the same dimensionality and precision.
The practical observation is many document vectors are similar: several tokens may describe closely related concepts, so they can be represented together without losing much utility.
Token pooling: HierarchicalTokenPooling
HierarchicalTokenPooling groups similar token embeddings using Ward hierarchical clustering. Start with one group per token, then repeatedly merge the two groups whose union yields the smallest increase in within-group spread. The spread is the sum of squared distances from vectors to their group mean, so Ward favors compact, similar groups.
This merging builds a hierarchy of groups; choosing a level sets the compression ratio. Each group is represented by its mean, then the resulting vectors are normalized for MaxSim scoring.
You can make compression proportional to document length:
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
pooling = HierarchicalTokenPooling(pool_factor=4)
With pool_factor=4, documents keep roughly a quarter of their original vectors, so longer documents retain proportionally more capacity. By default the first token is protected from merging, and pooling is applied only to documents (query embeddings unchanged).
Pooling budget controls
HierarchicalTokenPooling supports three controls:
- max_tokens=16: at most 16 vectors per document (protected tokens counted).
- pool_factor=4: roughly one quarter of each document's original vectors.
- distance_threshold=...: a shared Ward merge-distance threshold that lets vector counts vary with embedding geometry.
The choice matters with variable document lengths. pool_factor keeps representation proportional to length; distance_threshold increases vector count where more distinct embeddings must be preserved.
Calibrating adaptive pooling
Adaptive pooling uses a shared Ward distance threshold to merge more aggressively while preserving more vectors for distinct documents. It can be applied to any multi-vector model, including those not trained with pooling-aware objectives.
Rather than hand-tuning the threshold, calibrate it on a representative sample of your corpus to target an average token count. Using a model's document embeddings on calibration_documents, call HierarchicalTokenPooling.calibrate_distance_threshold(calibration_embeddings, target_tokens=16) to get a threshold you can reuse when indexing. The target is an average, so individual documents can vary and achieved averages can shift on other corpora.
Pooling-aware training principle
Pooling alone can already be useful, but problems arise when a document distinction important for retrieval gets lost by merging. The original training objective did not require the model to perform well with compressed representations. Pooling-aware training fixes this by applying retrieval objectives both before and after pooling.
This is analogous to Matryoshka embeddings that train at multiple embedding widths. For the example in this article, each training batch contributes two losses:
- Retrieval loss with the original (unpooled) document embeddings.
- Retrieval loss after compressing documents to ≈1/4 of their original vectors.
The losses are averaged: L = 1/2 L_unpooled + 1/2 L_pooled. The first term trains the model for full representations; the second arranges token embeddings so useful matches survive pooling. Both objectives update the same encoder, so it receives feedback about retrieval performance before and after compression.
Ward clustering makes discrete grouping decisions, so during training HierarchicalTokenPooling uses a straight-through estimator to pass gradients from pooled vectors back to their member tokens without differentiating cluster assignments. The encoder can thus learn from retrieval errors made with pooled representations.
A loss wrapper in Sentence Transformers
MultiVectorTokenPoolingLoss wraps an existing multi-vector retrieval loss and applies it using supplied poolers. Example setup:
from sentence_transformers import MultiVectorEncoder from sentence_transformers.multi_vector_encoder.losses import ( CachedMultiVectorMultipleNegativesRankingLoss, MultiVectorTokenPoolingLoss, ) from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
model = MultiVectorEncoder("answerdotai/ModernBERT-base", similarity_fn_name="maxsim") pooling = HierarchicalTokenPooling(pool_factor=4)
base_loss = CachedMultiVectorMultipleNegativesRankingLoss(model, mini_batch_size=64) loss = MultiVectorTokenPoolingLoss( model, base_loss, token_poolings=[None, pooling], weights=[0.5, 0.5], )
CachedMultiVectorMultipleNegativesRankingLoss provides the retrieval objective (queries should score their positive documents above in-batch negatives). Gradient caching lets the encoder process small mini-batches while keeping a larger set of negatives. In the author's training script mini-batches of 64 and training batch size 256 were used.
MultiVectorTokenPoolingLoss adds the compression objective: None means one variant uses original embeddings, pooling supplies the factor-4 variant. Pooling variants reuse computed token embeddings so each extra pooler only adds clustering and scoring work without another encoder forward pass. Poolers can be custom BaseTokenPooling implementations that preserve gradients through input embeddings.
The complete training script in the repository includes dataset setup, a trainer, validation loss, and small retrieval evaluations during training. It can be run locally or on cloud GPUs (for example with Hugging Face Jobs) using the provided commands.
Optional kernel acceleration
MaxSim scoring, Ward clustering, and supported backbone ops can be accelerated with kernels (CUDA/Triton). Install training and kernel deps:
pip install -U "sentence-transformers[train,kernels]"
Enable kernels by setting use_kernels=True when constructing the MultiVectorEncoder. Sentence Transformers selects compatible kernels and falls back to existing implementations where replacements are unavailable. The default MaxSim kernel computes scores without building the full token similarity matrix in memory but requires CUDA/Triton; the default Ward kernel runs clustering on CUDA and can be ~30–150x faster than SciPy depending on settings.
Using the trained model
Load the pooling-aware checkpoint and compress document embeddings at inference. Kernel settings aren’t saved with the model, so pass use_kernels=True at load time if desired.
Example (short):
model = MultiVectorEncoder("tomaarsen/mve-ModernBERT-base-token-pooling") pooling = HierarchicalTokenPooling(pool_factor=4)
queries = ["Which planet is known as the Red Planet?"]
... documents ...
query_embeddings = model.encode_query(queries) document_embeddings = model.encode_document(documents, token_pooling=pooling)
document vector counts fall (e.g. from [18,21,19] to [5,6,5]) while ranking may be preserved.
You can change pool_factor at inference to explore alternate storage budgets or set the pooler as the runtime default on MultiVectorEncoder initialization. The pooler’s tasks decide which inputs are compressed (e.g., documents only).
Evaluation results
The author trained two models starting from answerdotai/ModernBERT-base with fresh 128-dimensional token projections:
- tomaarsen/mve-ModernBERT-base-token-pooling-baseline — trained only with the cached contrastive loss.
- tomaarsen/mve-ModernBERT-base-token-pooling — trained with the cached contrastive loss wrapped to include unpooled and factor-4 pooled objectives.
Both runs used one epoch over the same 120,000 MS MARCO triplets, seed 42, identical initialization and optimization settings, token cap 300, and short queries expanded to at least 32 token positions. Only the loss differed; no hyperparameter search was performed.
Evaluation on NanoBEIR@13 reported mean nDCG@10 (%). Results compared standard and pooling-aware training across pool_factors 2–6, unpooled retrieval, fixed max_tokens caps, and adaptive thresholds. Highlights:
- At the trained pool_factor=4, nDCG@10 improves from 49.03 to 51.23 for the pooling-aware model.
- The pooling-aware model stores 74.8% fewer document vectors than unpooled retrieval at pool_factor=4.
- Score retention (pooled mean nDCG@10 divided by the model’s own unpooled mean) rises from 94.6% (ordinary) to 97.8% (pooling-aware) at pool_factor=4.
- Unpooled performance also increased from 51.84 to 52.39.
Gains persist for fixed caps and adaptive pooling despite training using only full and proportional pooling. For one adaptive setting averaging ~29.58 vectors/document, the pooling-aware model retained 97.3% of its unpooled score while storing 84.7% fewer vectors; the ordinary model retained 93.3% at that budget.
The improvement is especially notable at low vector budgets (for example max_tokens=4 showed a larger point gain in the experiment).
Combining compression techniques
If index size is the priority, the suggested order of techniques is:
- Embedding quantization (Fast-PLAID to 4/2/1 bits) to reduce storage per coordinate.
- Token pooling to reduce vectors per document, optionally with pooling-aware fine-tuning.
- Matryoshka token embeddings to reduce token vector dimensionality, optionally with fine-tuning for smaller widths.
If the model’s default dimensionality exceeds 128, consider Matryoshka truncation and fine-tuning earlier. These methods can be combined; evaluate resulting index size and retrieval quality on your corpus and search engine. MultiVectorMatryoshkaLoss can be used to train robustness to truncation and can be wrapped with MultiVectorTokenPoolingLoss to jointly train for fewer dimensions and fewer vectors.
Conclusion
Adding a pooling-aware term to the training objective lets the model actively learn to use smaller document representations. This significantly improves retrieval quality retention at fixed storage budgets, making multi-vector retrieval more practical for large-scale deployments.
References and code
The article builds on works such as Stefan Josef’s Learn to Pool and LightOn’s pooling regularization studies. Code examples, the complete training script, evaluation scripts, and example checkpoints (e.g. tomaarsen/mve-ModernBERT-base-token-pooling and the baseline) are available in the Sentence Transformers repository.
