Multi-vector embedding models give every query token a chance to find its own match in a document. This makes them wonderfully expressive for search—different parts of a query can latch onto different parts of a document. But there's a problem: those document vectors add up fast.
A corpus with millions of documents can contain hundreds of millions of vectors. At 4 bits per coordinate with 128 dimensions, 200 vectors per document means 12.8 KB per document, or 12.8 GB across a million documents. Token pooling merges similar token vectors to reduce this footprint—but aggressive compression risks losing information that matters for retrieval.
Pooling-aware training is the clever workaround: teach the encoder to preserve what matters when its vectors will be merged. The results are striking. With a pool_factor=4, a normally trained model retains 94.6% of its unpooled NanoBEIR nDCG@10 score. With pooling-aware training, that retention jumps to 97.8% at the same 4x storage reduction.
Why Multi-Vector Models Need Compression
Single-vector embeddings compress an entire document into one vector. Multi-vector models like ColBERT keep a collection of contextual token vectors instead, enabling fine-grained matching.
MaxSim scoring is the secret sauce: each query vector selects its best matching document vector, and those scores are summed. A query about "laptop with long battery life" can find evidence about the laptop in one part of a document and battery endurance in another.
Even with 2-bit or 4-bit residual quantization from Fast-PLAID, storing hundreds of vectors per document is expensive. Pooling each document from 200 vectors down to 50 reduces storage from 12.8 GB to 3.2 GB across a million documents—same dimensionality, same precision, four times fewer vectors.
The insight is simple: many document vectors are similar. Several tokens might describe closely related concepts, creating an opportunity to represent them together without losing much.
How Hierarchical Token Pooling Works
The HierarchicalTokenPooling module groups similar token embeddings using Ward hierarchical clustering. It starts with one group per token vector, then repeatedly merges the two groups that cause the smallest increase in within-group spread.
This spread is measured by the sum of squared distances from each vector to its group's mean, so Ward's method favors compact groups of similar vectors. The merges build a hierarchy of progressively larger groups, letting you choose a compression level that fits your storage budget.
Each group is then represented by its mean, and the resulting vectors are normalized again for MaxSim scoring.
Controlling the Budget
You can make compression proportional to document length with pool_factor=4, which retains approximately one quarter of each document's original vectors. Longer documents keep proportionally more capacity.
Alternatively, you can use:
max_tokens=16: cap each document at 16 vectorsdistance_threshold=...: use a shared Ward merge-distance threshold, allowing vector counts to vary with embedding geometry
The distinction matters when document lengths vary. With pool_factor, a longer document can retain more vectors, allowing its representation to grow with its content. With distance_threshold, the vector count can grow when more distinct embeddings need to be preserved—this is adaptive pooling.
Calibrating Adaptive Thresholds
Rather than choosing a threshold manually, you can calibrate it for a target average vector count using a representative sample of your corpus. The calibration returns a threshold you can reuse when indexing.
The target is an average, so individual documents can receive more or fewer vectors. The achieved average can also shift on a different corpus. Importantly, adaptive pooling works with any multi-vector model, including ones trained without a pooling-aware objective.
Training Models to Survive Compression
Pooling can be useful on an ordinary multi-vector model, but problems appear when a document contains a distinction that matters for retrieval and merging its vectors makes that distinction harder to recover. The original training objective never required the model to work with compressed representations.
Pooling-aware training fixes this by teaching the model to arrange its token embeddings so that useful matches survive pooling.
If you've used Matryoshka embeddings, the principle will feel familiar. Matryoshka applies a retrieval objective at several embedding widths, teaching smaller prefixes to remain useful. Pooling-aware training applies the objective to representations with fewer token vectors.
Each training batch contributes two losses:
- Retrieval using the original document embeddings
- Retrieval after compressing documents to approximately one quarter of their vectors
These are averaged: the first term teaches the model to retrieve with its full representations, the second teaches it to arrange token embeddings so useful matches survive pooling. Both objectives update the same encoder, giving it feedback about retrieval before and after compression.
The Gradient Challenge
Ward clustering makes discrete decisions about which tokens belong together. During training, HierarchicalTokenPooling uses a straight-through estimator to pass gradients from pooled vectors back to their member tokens, without differentiating the cluster assignments. The encoder can then learn from retrieval errors made with its pooled representations.
Implementation in Sentence Transformers
The MultiVectorTokenPoolingLoss wrapper makes this straightforward. It wraps an existing multi-vector retrieval loss and applies it using your chosen poolers:
from sentence_transformers.multi_vector_encoder.losses import (
CachedMultiVectorMultipleNegativesRankingLoss,
MultiVectorTokenPoolingLoss,
)
from sentence_transformers.multi_vector_encoder.modules import HierarchicalTokenPooling
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 supplies the retrieval objective, teaching each query to score its relevant document above negative documents in the batch. Gradient caching allows small mini-batches while retaining a larger set of in-batch negatives.
The loss wrapper can handle multiple pooling configurations simultaneously—you could train with pool_factor=2, pool_factor=4, and unpooled representations all contributing to a single weighted loss.
Results: Storage vs Quality Tradeoff
The evaluation compares normal training against pooling-aware training across multiple pool factors on NanoBEIR nDCG@10.
With pool_factor=4 (4x compression), the normally trained model retains 94.6% of its unpooled performance. Pooling-aware training pushes that to 97.8% at the same storage budget.
The gap widens at higher compression. At pool_factor=6, normal training drops to around 91% retention, while pooling-aware training maintains stronger performance. The model learned to organize its token embeddings so that aggressive merging preserves what matters for retrieval.
This isn't just about raw scores—it's about deployment economics. If you're building a search system over tens of millions of documents, the difference between 12.8 GB and 3.2 GB of vectors per million documents is the difference between expensive and tractable.
Combining Compression Techniques
Token pooling stacks with other compression methods. You can combine it with:
- Residual quantization (2-bit or 4-bit) to reduce coordinate precision
- Dimensionality reduction to shrink vector width
- Matryoshka-style training to enable lower-dimensional prefixes
Each technique attacks a different dimension of the storage problem. Token pooling reduces the number of vectors per document. Quantization reduces bits per coordinate. Dimensionality reduction reduces coordinates per vector.
The compounding effect is powerful: 4x fewer vectors, 4-bit quantization instead of 32-bit, and 64 dimensions instead of 128 could yield a 64x storage reduction with careful training.
The Broader Picture
This work builds on Stefan Josef's Learn to Pool paper and LightOn's pooling regularization study, which also explores targeted training and adaptive budgets. The Sentence Transformers implementation makes the technique immediately accessible for practitioners.
What I find elegant here is the symmetry with other compression-aware training methods. Matryoshka taught us that models can learn to preserve information in dimension prefixes. Quantization-aware training showed that models can learn to tolerate precision loss. Now pooling-aware training demonstrates that models can learn to arrange their representations for graceful merging.
The pattern is clear: if you tell the model during training what kind of compression it will face at inference, it can adapt. The encoder isn't just learning to embed—it's learning to embed in a way that survives your deployment constraints.
For teams building retrieval systems, this is practical magic. You don't have to choose between expressiveness and efficiency anymore. Train with pooling awareness, compress aggressively, and keep most of your quality. The storage budget becomes a training parameter, not a post-hoc constraint you impose on an unwilling model.