If you run search over your own documents with a ColBERT-style model, the index is probably bigger than you would like. These multi-vector models keep a vector for every token rather than one per document, which makes them precise and expensive to store. A new post on the Hugging Face blog shows how Sentence Transformers can now train a model to stay accurate after those vectors are squeezed, recovering most of the quality usually lost to compression.
The storage problem is concrete. At 4 bits per coordinate, 200 vectors of 128 dimensions take 12.8 KB per document, or 12.8 GB across a million documents. Pooling each document down to 50 vectors cuts that to 3.2 GB. The library's HierarchicalTokenPooling merges similar token vectors using Ward clustering, either by a fixed factor or by a distance threshold that can be calibrated to a target average, and it leaves query vectors untouched. Pooling works on any multi-vector model, but a model never trained with it can lose distinctions that matter.
The new piece is a loss wrapper, MultiVectorTokenPoolingLoss, which trains on two versions of each batch: the full document vectors and a version pooled to about a quarter. The idea mirrors Matryoshka embeddings, which train smaller prefixes to stay useful. In the author's experiment, two ModernBERT-base models trained for one epoch on the same 120,000 MS MARCO triplets, differing only in the loss. At a pool factor of 4, the ordinary model kept 94.6% of its unpooled NanoBEIR score; the pooling-aware one kept 97.8%, with nDCG@10 rising from 49.03 to 51.23 while storing 74.8% fewer document vectors. Unpooled quality improved slightly too. A GPU kernel for the clustering is said to be 30-150x faster than SciPy.
For teams using PyLate, LightOn's library for ColBERT-style retrieval built on Sentence Transformers, this is directly relevant: the work builds on LightOn's own study of pooling regularisation, and the training recipe runs on a single A10G through Hugging Face Jobs within four hours. The practical advice is to combine quantisation, pooling and Matryoshka truncation, and to measure index size and quality on your own corpus.
The caveats are the author's own. This is one experiment on one base model with no hyperparameter search, evaluated on NanoBEIR, a small benchmark suite. The gain is clearest at the pool factor the model was trained for, even if it carries over to other settings. Results on long or domain-specific documents may differ.