Tom Aarsen described the method on the Hugging Face blog on 11 November and released training and evaluation scripts along with two example models. Aarsen says the work builds on Stefan Josef’s Learn to Pool and on research by LightOn. Aarsen’s own aim is to explain the method and provide a reproducible example.

A multi-vector model such as ColBERT stores a separate vector for every token in a document. At search time, each query token picks the best-matching document vector, and the scores are summed. The vectors add up: in Aarsen’s calculation, the 4-bit quantised residuals of 128-dimensional vectors take up 12.8 GB for a collection of 1 million documents with 200 vectors per document. With 50 vectors, 3.2 GB is enough.

How pooling and training work

Users can shrink the index by merging similar vectors within a document (token pooling). HierarchicalTokenPooling in Sentence Transformers groups them with Ward’s hierarchical clustering and replaces each group with its mean. The number of vectors can be limited by a ratio (pool_factor), a maximum or a distance threshold. The threshold can be calibrated to produce a target average number of vectors. Because calibration targets the average, the count varies from document to document, and on another dataset even the average may change. Queries are not pooled by default.

A standard model is not trained to work with pooled vectors, so merging can erase a distinction that matters for retrieval. In pooling-aware training, the loss function MultiVectorTokenPoolingLoss computes the retrieval loss at selected pooling levels. Aarsen’s example used two levels, unpooled and pooled to about a quarter, weighted equally. The group assignment is not differentiated; instead, gradients pass from the pooled vector straight to the group’s members (straight-through estimator). The pooled versions are built from vectors already computed, so the model does not need a separate pass for them.

Results from a single comparison run

Aarsen trained two models from ModernBERT-base with the same settings on the same 120,000 MS MARCO examples for one epoch. Only the loss function differed, and no hyperparameters were tuned. The metric was mean nDCG@10 across NanoBEIR’s 13 datasets.

In the table, the pooling-aware model scored higher at every pooling level tested, with the largest gap at the tightest limit.

Source: Aarsen, Hugging Face blog. NanoBEIR, 13 datasets, mean nDCG@10 (%).

The table cannot show whether the advantage holds on other datasets, because the models were tested only on NanoBEIR. Each model had a single training run, and the figures are the author’s own.

Pooled with a factor of 4, the pooling-aware model retained 97.8 per cent of its own unpooled score and the standard model 94.6 per cent. The nDCG@10 scores then differed by 2.2 points (49.03 and 51.23), with 74.8 per cent fewer vectors than without pooling. The advantage also appeared with maximum limits and calibrated thresholds, even though training used only ratio-based pooling.

Outlets including zglg.work and AI Breaking Wire repeat the blog’s figures. Whether the advantage holds in much larger datasets and production indexes has yet to be established, zglg.work wrote.

The method is part of the Sentence Transformers library, and pooling can also be used with conventionally trained multi-vector models. The blog advises running the training script on a local GPU or one from Hugging Face Jobs. For optional GPU acceleration, the library picks compatible GPU kernels and uses the standard implementation elsewhere; the default MaxSim scoring kernel, for example, requires CUDA and Triton.