Token Pooling-Aware Training for Multi-Vector Embedding Models

A search model can be trained to expect the compression that will later be applied to its index, and in Tom Aarsen's experiment on the Hugging Face blog that recovered more than half of the search quality the compression had cost. The models in question are multi-vector embedding models such as ColBERT, or those built with the Sentence Transformers library's `MultiVectorEncoder` (explained here earlier). Instead of turning a whole document into one vector, a list of numbers that encodes its meaning, they keep one vector per token, roughly one per word piece. At search time each query token picks its best-matching document vector, a rule called MaxSim, and those best matches are added into the document's score. That lets one part of a query match one passage and another part match a different passage, but it makes indexes large. Aarsen's example: 200 vectors per document take 12.8 GB per million documents even with each coordinate stored in 4 bits, against 3.2 GB if each document is cut to 50 vectors.

Token pooling does that cutting by merging token vectors that are already similar. The method in Sentence Transformers uses Ward clustering, which starts with one group per vector and repeatedly merges the two groups whose combination keeps the groups most compact, then replaces each group with its average. With `pool_factor=4`, each document keeps about a quarter of its vectors, and queries are left unchanged. The weakness is that an ordinarily trained model never learned to work with merged vectors, so a distinction that mattered for retrieval can blur when its tokens are averaged together. Pooling-aware training adds that practice: every training batch computes the retrieval loss twice, once with full document vectors and once after pooling to a quarter, and averages the two, so the same encoder is corrected for mistakes it makes in both forms. The library packages this as `MultiVectorTokenPoolingLoss`, a wrapper around an existing retrieval loss. Because the clustering step makes yes-or-no grouping decisions that cannot be learned through directly, training passes the error from each merged vector straight back to the token vectors it was built from.

To measure the effect, Aarsen trained two models from the same starting model, ModernBERT-base, on the same 120,000 MS MARCO training examples with matching settings, so only the loss differed. He scored them on the NanoBEIR retrieval benchmark with nDCG@10, a measure of how well the top ten results are ordered. With pooling to a quarter, which stores 74.8% fewer document vectors than no pooling, the ordinary model scored 49.03 and the pooling-aware model 51.23. That is 94.6% against 97.8% of each model's own unpooled score. The pooling-aware model also scored slightly higher without pooling, 52.39 against 51.84, so the extra objective did not cost full-size quality in this run. The advantage held at settings the model was not trained on, including fixed caps per document and merge-distance thresholds calibrated to an average vector count. It was largest at the harshest cap of four vectors per document, 43.00 against 37.22.

This is one matched pair of runs with one random seed and no hyperparameter search, scored on one benchmark, so the size of the gain is an observation rather than a general rule. The post builds on earlier pooling-aware work by Stefan Josef and LightOn. Its practical contribution is that the method is now a few lines added to a standard training script. For a team whose index size is what limits a multi-vector deployment, fine-tuning with the pooled loss is a cheap experiment: the pooled version reuses the token vectors already computed, adding clustering and scoring work but no second pass through the model.