Token Pooling-Aware Training for Multi-Vector Embedding Models

Hugging Face 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 pool_factor=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. 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.