Multi-vector embedding models, like ColBERT, offer high precision in search, but the number of vectors stored per document quickly drives up the cost of the index. A grouping technique, token pooling, reduces this size, and adapted training allows, according to the source, for better preservation of result quality at an equal storage budget. With pool_factor=4, a normally trained model retains 94.6% of its uncompressed nDCG@10 on NanoBEIR, compared to 97.8% for the model trained while accounting for pooling.
Why Multi-Vector Models Store So Many Vectors
A classic embedding model summarizes an entire document into a single vector. Multi-vector models, on the contrary, retain a set of contextual vectors, one per token (a unit of text, often a word or part of a word), which allows different parts of a query to match different parts of a document. The source gives the example of a query about a long-lasting laptop, which might find a passage about the device in one part of the document and a passage about battery life in another.
This functionality comes at a cost. According to the presented calculation, with 4-bit quantization per residual coordinate, 200 vectors of 128 dimensions represent 12.8 KB per document, or 12.8 GB for one million documents. Reducing each document to 50 vectors brings this total down to 3.2 GB, with identical dimension and precision.
The method is based on an observation: many vectors from the same document can be close, as several tokens often describe similar concepts. HierarchicalTokenPooling, the Sentence Transformers library, groups these vectors via Ward clustering, which merges at each step the two groups causing the smallest increase in internal dispersion. Each group is then replaced by its average and normalized. With pool_factor=4, a document retains about a quarter of its initial vectors, so a longer document keeps proportionally more vectors.
Training the Model for Compression
The stronger the compression, the easier it is to lose useful information. Training that accounts for pooling aims precisely to preserve important information when vectors are merged. Each training batch combines two losses: search with the original document embeddings, and search after compression to about a quarter of the vectors. These two losses are averaged, and both objectives update the same encoder.
As clustering makes discrete decisions, the source indicates that the training uses a straight-through estimator, which propagates gradients from the pooled vectors to the tokens that compose them, without differentiating the cluster assignments. The MultiVectorTokenPoolingLoss class encapsulates an existing search loss, here CachedMultiVectorMultipleNegativesRankingLoss, with weights of 0.5 for each variant. The starting model is answerdotai/ModernBERT-base, trained on 120,000 MS MARCO triplets for one epoch.
Results on NanoBEIR
At pool_factor=4, nDCG@10 increases from 49.03 to 51.23 for the model trained with pooling, while reducing the number of document vectors by 74.8% compared to an uncompressed search. The uncompressed model also improves, from 51.84 to 52.39. For an average of 29.58 vectors per document, adaptive pooling retains 97.3% of the uncompressed score of the trained model, compared to 93.3% for the normally trained model.
To reduce the size of an index, the source first suggests quantizing the embeddings, then token pooling, possibly with fine-tuning that accounts for pooling, and finally Matryoshka embeddings to reduce the dimension of each vector. These approaches can be combined. The work on which the article is based is described in the study Learn to Pool by Stefan Josef.
