Token Pooling-Aware Training for Multi-Vector Embedding Models

Back to Articles Token Pooling-Aware Training for Multi-Vector Embedding Models Published November 11, 2026 Update on GitHub Upvote 2 Tom Aarsen tomaarsen Follow Multi-vector embedding models give each query token its own opportunity to find a match in a document.

This makes them wonderfully expressive for search, but those document vectors add up quickly.A corpus with millions of documents can contain hundreds of millions of vectors!Token pooling makes these representations smaller by merging similar token vectors.

The more aggressively I compress them, however, the easier it becomes to lose something useful.Pooling-aware training teaches the encoder to preserve the information that matters when its vectors are merged, helping you retrieve better from the same storage budget.

In this post, I'll explain how pooling works, train a model for compressed retrieval with Sentence Transformers, and explore the resulting quality and storage tradeoff.With poolfactor=4, my normally trained model retains 94.

6% of its unpooled NanoBEIR nDCG@10, with a 4x reduction in document vectors.Token pooling-aware training raises that retention to 97.8% at the same storage budget!Let's see how.

This post builds on pooling-aware training studied in Stefan Josef's Learn to Pool and LightOn's pooling regularization study, which also explores targeted training and adaptive budgets.

Here, I'll focus on understanding the method and putting it to work with a short, reproducible Sentence Transformers example.

Table of Contents Why Multi-Vector Retrieval Needs So Many Vectors Token Pooling Choosing a Pooling Budget Calibrating Adaptive Pooling Token Pooling-Aware Training A Loss Wrapper in Sentence Transformers Optional Kernel Acceleration Using the Trained Model Evaluation Results Combining Compression Techniques References Code and Models Why Multi-Vector Retrieval Needs So Many Vectors A single-vector embedding model compresses an entire document into one vector.

Multi-vector models such as ColBERT instead keep a collection of contextual token vectors, allowing different parts of a query to match different parts of a document.

For example, a query about a laptop with a long battery life can find evidence about the laptop in one part of a document and battery endurance in another.

With MaxSim scoring, each query vector selects its best matching document vector, and those scores are added together: MaxSim⁡(Q,D)=∑q∈Qmax⁡d∈Dq⊤d \operatorname{MaxSim}(Q, D) = \sum{q \in Q} \max{d \in D} q^\top d MaxSim(Q,D)=q∈Q∑​d∈Dmax​q⊤d Even with the 2-bit or 4-bit residual quantization supported by Fast-PLAID, storing hundreds of vectors per document adds up.

At 4 bits per residual coordinate, 200 vectors of 128 dimensions require 12.8 KB of quantized residuals per document, or 12.8 GB across a million documents.Pooling each document to 50 vectors reduces that to 3.2 GB, at the same dimensionality and precision.

The useful observation is that many document vectors can be similar.Several tokens might describe closely related concepts, creating an opportunity to represent them together.Token Pooling HierarchicalTokenPooling groups similar token embeddings using Ward hierarchical clustering.

Starting with one group per token vector, it repeatedly merges the two groups that cause the smallest increase in how spread out the vectors are within their groups.

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.These merges build a hierarchy of progressively larger groups, letting us choose a level of compression that fits our budget.

We then represent each group by its mean and normalize the resulting vectors again, ready for MaxSim scoring.You can make compression proportional to document length: from sentencetransformers.multivectorencoder.

modules import HierarchicalTokenPooling pooling = HierarchicalTokenPooling(poolfactor=4) With poolfactor=4, documents retain approximately one quarter of their original vectors, so longer documents can keep proportionally more capacity.

By default, the first token is protected from merging and pooling applies only to document embeddings, leaving query embeddings unchanged.

Choosing a Pooling Budget HierarchicalTokenPooling offers three alternative controls: Setting What it expresses maxtokens=16 A maximum of 16 vectors per document, including protected tokens.poolfactor=4 Approximately one quarter of each document's original vectors.distancethreshold=...

A shared Ward merge-distance threshold, allowing vector counts to vary with embedding geometry.This distinction matters when document lengths vary.With poolfactor=4, a longer document can retain approximately one quarter of its vectors, allowing its representation to grow with its content.

With distancethreshold, the vector count can grow when more distinct embeddings need to be preserved.Calibrating Adaptive Pooling Adaptive pooling uses a shared distance threshold to merge similar embeddings more aggressively while preserving more vectors for documents with distinct content.

It can be used with any multi-vector model, including one trained without a pooling-aware objective.Rather than choosing a threshold manually, you can calibrate it for a target average vector count.

Given a representative sample of your corpus in calibrationdocuments: from sentencetransformers import MultiVectorEncoder from sentencetransformers.multivectorencoder.modules import HierarchicalTokenPooling model = MultiVectorEncoder("lightonai/LateOn") calibrationembeddings = model.

encodedocument(calibrationdocuments) threshold = HierarchicalTokenPooling.calibratedistancethreshold(calibrationembeddings, targettokens=16) pooling = HierarchicalTokenPooling(distancethreshold=threshold) # Apply pooling, targeting an average of 16 vectors per document documentembeddings = model.

encodedocument(documentstoindex, tokenpooling=pooling) Calibration only needs document embeddings and returns a threshold you can reuse when indexing.The target is an average, so individual documents can receive more or fewer than 16 vectors.The achieved average can also change on another corpus.

Token Pooling-Aware Training Pooling can already be useful on an ordinary multi-vector model.The difficulty appears when a document contains a distinction that matters for retrieval, but merging its vectors makes that distinction harder to recover.

The original training objective never required the model to work with these compressed representations.But, you can teach it to do exactly that.If you've used Matryoshka embeddings, the training 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.

For my example, each training batch contributes two losses: Retrieval using the original document embeddings.Retrieval after compressing documents to approximately one quarter of their original vectors.

I average them: L=12Lunpooled+12Lpooled \mathcal{L} = \tfrac{1}{2}\mathcal{L}{\mathrm{unpooled}} + \tfrac{1}{2}\mathcal{L}{\mathrm{pooled}} L=21​Lunpooled​+21​Lpooled​ The first term teaches the model to retrieve with its full representations.

The second teaches it to arrange its token embeddings so that useful matches survive pooling.Both objectives update the same encoder, giving it feedback about retrieval before and after compression.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.

A Loss Wrapper in Sentence Transformers MultiVectorTokenPoolingLoss wraps an existing multi-vector retrieval loss and applies it using your chosen poolers: from sentencetransformers import MultiVectorEncoder from sentencetransformers.multivectorencoder.

losses import ( CachedMultiVectorMultipleNegativesRankingLoss, MultiVectorTokenPoolingLoss, ) from sentencetransformers.multivectorencoder.

modules import HierarchicalTokenPooling model = MultiVectorEncoder("answerdotai/ModernBERT-base", similarityfnname="maxsim") pooling = HierarchicalTokenPooling(poolfactor=4) baseloss = CachedMultiVectorMultipleNegativesRankingLoss(model, minibatchsize=64) loss = MultiVectorTokenPoolingLoss( model, baseloss, tokenpoolings=[None, pooling], weights=[0.

5, 0.5], ) CachedMultiVectorMultipleNegativesRankingLoss supplies the retrieval objective.It teaches each query to score its relevant document above the negative documents in the batch.By default, it uses model.

similarity, using the model's similarity setting (MaxSim in this example) and kernel settings.A model configured with similarityfnname="meanmaxsim" would also use MeanMaxSim during training.

Gradient caching allows the encoder to process small mini-batches while retaining a larger set of in-batch negatives.My training script uses mini-batches of 64 with a training batch size of 256.MultiVectorTokenPoolingLoss adds the compression objective.

Here, None means that one variant uses the original embeddings, while pooling supplies the factor-4 variant.The two weights average their losses.

Pooling variants reuse the computed token embeddings, so each additional pooler adds clustering and scoring work without requiring a separate encoder pass for that variant.The poolers can also be custom BaseTokenPooling implementations that preserve gradients through their input embeddings.

You can adapt the training objective to the representation you want to deploy.Replace the factor-4 pooler with your chosen deployment configuration, or include several poolers and choose weights for their losses.

The complete training script includes the dataset, trainer, validation loss, and small retrieval evaluations during training.From a checkout of the repository, run both variants locally with uv: uv run examples/multivectorencoder/training/tokenpooling/trainingcachedcontrastive.

py --baseline uv run examples/multivectorencoder/training/tokenpooling/trainingcachedcontrastive.py The script declares its training and kernel dependencies, so uv installs them for you.Compatible kernels are enabled automatically.You can also train on a cloud GPU with Hugging Face Jobs.

Install the CLI and log in with a token that can run Jobs and upload models: pip install -U huggingfacehub hf auth login hf jobs uv run \ --flavor a10g-small \ --timeout 4h \ --secrets HFTOKEN \ -- https://raw.githubusercontent.

com/huggingface/sentence-transformers/main/examples/multivectorencoder/training/tokenpooling/trainingcachedcontrastive.py This runs the pooling-aware variant on an A10G GPU with a four-hour time limit.

Passing HFTOKEN as a secret lets the script upload the trained model to your Hub account when training finishes.To run the ordinary-loss baseline as well, repeat the command with --baseline after the script URL.

Optional Kernel Acceleration You can accelerate MaxSim scoring, Ward clustering, and supported backbone operations with kernels.

Install the training and kernel dependencies: pip install -U "sentence-transformers[train,kernels]" In the training example above, replace the model and pooler initialization before constructing the loss: model = MultiVectorEncoder( "answerdotai/ModernBERT-base", device="cuda", similarityfnname="maxsim", usekernels=True, ) pooling = HierarchicalTokenPooling(poolfactor=4) Hierarchical poolers inherit the model's kernel settings by default, so you only need to enable kernels on the model.

Sentence Transformers selects compatible kernels for your setup and uses the existing implementation wherever a replacement isn't available.For exampl