LateOn-hpool-regularized

📚 Collection | 📝 Blog

Overview

LateOn-hpool-regularized is a ColBERT-style multi-vector retrieval model trained by LightOn using PyLate. It is built on the LateOn family and regularized specifically to make hierarchical pooling much more effective.

Late interaction models are powerful because they store one embedding per document token and score with MaxSim. The downside is index size: long documents require many vectors. We previously introduced hierarchical pooling to reduce this footprint by clustering and merging similar token embeddings inside each document. This model is trained so that such pooling preserves retrieval quality even at aggressive compression ratios.

The headline result is simple:

LateOn-hpool-regularized reaches 98.7% retention at a 32-token document budget, corresponding to roughly 5× compression, with negligible impact on full-token retrieval performance.

This model is the recommended LateOn checkpoint when you want a versatile ColBERT that works well across all compression levels. It is trained with multiple pooling budgets ([4, 8, 16, 32, 64, 128, 300]) so you can set the compression ratio at inference time. If you know your deployment target, training specifically for that budget pushes retention even further (see Targeted Training below).

If you only use full-token PLAID retrieval without pooling, LateOn or LateOn-regularized remain strong choices.

For the complete methodology and analysis, refer to our blog post.

Method

Hierarchical pooling contains a discrete operation: cluster assignment. Tokens are merged according to a hierarchical clustering tree, and this assignment is not differentiable. To train through it, we use a Straight-Through Estimator.

During training:

1. The model encodes query and document tokens normally.
2. Document token embeddings are hierarchically clustered and pooled in the forward pass.
3. Retrieval scores are computed on the pooled document representation.
4. The loss combines the standard full-token MaxSim loss with a pooled-representation loss.
5. During the backward pass, the Straight-Through Estimator lets gradients flow through the non-differentiable pooling decision.

The total loss is a weighted sum of the full-token MaxSim contrastive loss and the pooled-representation contrastive loss.

This released checkpoint uses multi-budget training with targets [4, 8, 16, 32, 64, 128, 300] at each step, which encourages the model to become broadly more poolable instead of overfitting to one compression point.

[!IMPORTANT] Important note: while we converged towards this approach very naturally due to our previous projection work and having introduced the hierarchical pooling approach, it is important to note that an independent researcher (Stefan Josef) did the same kind of study a few months ago. Although we did not have this work in mind while converging back to the idea, we were totally aware of the work and actually chatted with the author during the LIR workshop. Thus, we would like to highlight that the paternity of the idea should be given to Stefan and not to us. You can see these notes as an independent rediscovery, highlighting how much the idea makes sense from the beginning, and most importantly, a scaling up of the study to more data. Stefan is an independent researcher whose compute capacity was limited (as stated in the paper). When scaled to LateOn data, the results become even stronger, so we are happy to share them (alongside the model), but again, please have a look at the original work

Results

All of the results can be found in this space. The results are only on a subset of the 7 smaller BEIR datasets to make experiments easier and match AGC setup (SciFact, NFCorpus, ArguAna, TREC-Covid, FiQa, Touché, SciDocs). Given the original results of the hpool paper, we do not expect much different from running the whole BEIR suit.

Hierarchical Pooling Retention at 32 Tokens

The table below reports retention at a 32-token document budget, measured against the full-token LateOn reference.

The regularization strongly improves pooling retention while preserving full-token retrieval quality. Even at λ=1.5 (this model), where the model is optimized aggressively for pooling, full-token performance remains essentially unchanged.

Targeted Training

If you know your deployment compression budget, you can train specifically for that target instead of using multi-budget training. With targeted r=32 training at λ=0.7: