The Pulse
Hugging Face details pooling-aware training for embedding models
Hugging Face has published a training approach for multi-vector embedding models that accounts for compression during retrieval. An experiment compared the approach with ordinary training at the same document-vector budget.
A training method that accounts for token pooling raised a multi-vector embedding model’s mean NanoBEIR nDCG@10 score from 49.03 to 51.23 at a pooling factor of four, according to a Hugging Face post by Tom Aarsen. The experiment compared a model trained with the compression-aware objective against one trained without it.
Token pooling combines similar document vectors into a smaller representation. The training approach applies a retrieval objective both before and after pooling, giving the encoder feedback on its compressed representations as well as its full ones.
What the comparison measured
The comparison used two models trained for one epoch on the same 120,000 MS MARCO triplets. They shared initialization and optimization settings, a token cap and query expansion; the loss function was the difference between the experiments. Aarsen did not run a hyperparameter search.
For each model, the evaluation measured the mean retrieval score with unpooled documents and under several pooling settings. At the factor used in training, the pooling-aware model retained more of its own unpooled score than the ordinary model, according to the post. The evaluation also covered fixed vector caps and adaptive thresholds.
How pooling changes retrieval
A multi-vector embedding model represents a document with a collection of contextual token vectors rather than a single vector for the whole document. During retrieval, each query vector finds its best match among the document vectors, and the matching scores are added together.
The example uses Ward hierarchical clustering to group similar token embeddings. It starts with separate groups, repeatedly merges them according to their internal spread, and represents each resulting group by its mean. The pooled vectors are normalized again for scoring.
Pooling can follow a document’s length, enforce a fixed maximum or use a distance threshold that allows vector counts to vary. The adaptive threshold can be calibrated against document embeddings from a sample of the corpus, with a target average rather than an identical count for every document. Query embeddings remain unchanged under the default pooling configuration.
Training and deployment
The MultiVectorTokenPoolingLoss wrapper applies an existing retrieval loss to the selected pooling variants. In the example, the unpooled and pooled losses receive equal weight. The variants reuse the encoder’s token embeddings rather than requiring a separate encoder pass for each one.
The clustering assignments are discrete, so the training method uses an estimator to pass gradients from pooled vectors back to their member tokens. That gives the encoder feedback from retrieval errors on the compressed representation without differentiating the cluster assignments themselves.
The post also describes using the trained model with a different pooling factor at inference and combining token pooling with other compression methods. Those include reducing the precision of stored coordinates and reducing the dimensionality of each token vector.