After the party comes the free lunch: regularizing ColBERT models to enhance pooling capabilities and reduce index footprint

Hierarchical pooling already lets you cut ColBERT storage in half without hurting retrieval. We now show that STE-based regularization, the same trick we used to fix MUVERA/SMVE, pushes this much further: 99.4% retention at 5× compression, without hurting full-token performance. The key insight is that the same as for MUVERA/SMVE: the best way to achieve compressability is to directly train for it.

🤗 Models: LateOn-hpool-regularized | LateOn-regularized | LateOn

Table of Contents

• Hierarchical Pooling: The Quiet Workhorse
• A Regularization That Keeps on Giving
• Same Trick, New Target
• Results: From 77% to 98% Retention
• Targeted Training: Matching Budget to Objective
• How Does This Compare to Learned Pooling?
• Adaptive Budgets: Not All Documents Are Created Equal
• The Full Picture
• Takeaways and What's Next

Hierarchical pooling: the quiet workhorse

Late interaction models like ColBERT store one embedding per token in every document. That's what gives them their expressive power, but it also makes them expensive. A 300-token document doesn't get one vector; it gets 300. Multiply that by millions of documents and the storage bill adds up fast.

In our earlier work, we introduced a simple idea: what if many of those 300 tokens are saying roughly the same thing? If two token embeddings are nearly identical, you can merge them without losing much information (especially as the ColBERT embeddings are very anisotropic). Hierarchical pooling does exactly this. You build a dendrogram over the token embeddings in each document via hierarchical clustering, then iteratively merge the closest clusters and pool them until you hit your target budget. It turned out to be remarkably effective: we could cut storage in half with virtually no performance loss, and push further with graceful degradation. The idea has since proven effective across modalities as well. For more information about hierarchical pooling, please refer to the original blogpost.

We can push the compression further than cutting it in half. However, at aggressive compression ratios (e.g pooling a document down to 32 tokens), the performance degrades more and more.

We recently showed in our regularization work that the geometry of the embedding space itself was the bottleneck for efficient retrieval methods like MUVERA and SMVE, and that Straight-Through Estimator (STE)-based regularization could reshape it. A natural question emerged: if the regularization makes embeddings more compressible for random projections, does it also make them more compressible for pooling?

A regularization that keeps on giving

Originally, we just wanted to check if the model that was regularized for MUVERA/SMVE benefited from it for hierarchical pooling retention, without any pooling-specific training at all.

The answer was a pleasant surprise. Retention at 32 tokens jumped from 77.1% to 88.3%, a free 11-point improvement. The MUVERA regularization, which was never about clustering or pooling, had reshaped the space in a way that made tokens more mergeable.

This is not entirely surprising in hindsight. In the previous post, we discovered that STE regularization concentrates each document's tokens into fewer effective dimensions rather than spreading them out. Fewer effective dimensions means more redundancy between tokens, which is precisely what hierarchical pooling exploits. The regularization was teaching the model to say the same thing with fewer independent directions, and pooling is just another way of capitalizing on that redundancy.

But 88.3% is not 100%. Could we do better by regularizing directly for the target?

Same trick, new target

If you checked our MUVERA/SMVE work, you know the drill. Hierarchical pooling, like SimHash partitioning and top-k selection, has a gradient-breaking operation: the cluster assignment. Which tokens get merged depends on a discrete clustering decision, and small changes in the embeddings won't shift that decision until a threshold is crossed, making the gradient zero almost everywhere.

We apply the same STE trick. In the forward pass, we perform the real hierarchical clustering and pooling, compute MaxSim on the pooled representations, and get a proper loss. In the backward pass, we pretend the cluster assignment was an identity and let the gradient flow through. The model receives a signal that says: "if your tokens were pooled this way, here's how the retrieval score would change," and it learns to arrange its embeddings so that the pooling preserves what matters. The total loss is a weighted combination between the full MaxSim loss and the loss on the pooled representation, defined by λ.

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: from 77% to 98% retention

We trained several configurations, varying λ from 0.5 to 1.5, starting from the SMVE-regularized checkpoint and adding the hierarchical pooling regularization on top. During training, we used compression targets of [4, 8, 16, 32, 64, 128, 300] tokens at each step rather than fixing a single target. 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.

The first thing to notice is that the regularization strongly helps. At the lowest weight (λ=0.5), the retention jumps to 94.9%, without degrading full tokens performance (actually slightly enhancing them). The second is that pushing λ further increases the retention while still preserving the full token performance. At λ=1.5, we can reach 97.9% retention at 32 tokens! The model can be optimized for extreme pooling without sacrificing its native retrieval quality.

It is worth noting that these results build on the MUVERA-regularized checkpoint. When applying the same hpool regularization directly on top of the unregularized LateOn, the retention was always slightly lower than what is achieved from the MUVERA base. This confirms that the two regularizations compound rather than substitute: the MUVERA step reshapes the geometry broadly, and the hpool-specific loss refines it further for pooling. Exploring the combination of the two (SMVE/MUVERA projections on top of pooled embeddings) and understanding what these regularizations really do in the embedding space is a very promising research direction.

Targeted training: matching budget to objective

The results above used multi-budget training: at each step, the model train for all cutoff in [4, 8, 16, 32, 64, 128, 300]. This gives excellent flexibility (one model for any compression level) but spreads the training signal across many operating points.

A natural question: what if you know your deployment target? If you're going to run at 32 tokens, does training specifically for r=32 do better?

Training specifically for r=32 gains 1.6 percentage points of retention @ 32 tokens over multi-budget training at the same λ. The model concentrates all its regularization effort on the exact compression level it will face at deployment. However, as expected, it does a bit worse at lower cutoff it has not been trained for.

This creates a practical trade-off:

• Multi-budget training is the right default when you want one model that works well across all compression levels: you don't know your deployment target yet, or you want flexibility.
• Targeted training is strictly better when your compression budget is fixed at deployment time.

How does this compare to learned pooling?

A fair question: if we're going to train for pooling anyway, why not learn the pooling function itself? The recent Attention-Guided Clustering (AGC) method does exactly that, replacing hierarchical clustering with a learned attention mechanism that decides which tokens to merge (and how to merge them).

On the 7 BEIR datasets they evaluated, AGC reaches an average retention of 97.4% at 32 tokens. Our hierarchical pooling regularization at λ=1.5 reaches 97.9%, while being trained for all the cutoffs and thus being very versatile.

A few caveats: the training data and base model differ, so the comparison is not apples-to-apples. Our base model is also stronger (0.506 vs 0.462 NDCG@10). Take the numbers with a grain of salt, the point is not to claim superiority but to show that a simpler, more flexible method can reach comparable territory.

One important aspect of hierarchical pooling is that it sets the budget at inference time. You can choose to compress more or less your documents with the same model and methodology, to set the quality-storage trade off. You compute the dendrogram and put the cut point where you want.

Adaptive budgets: not all documents are created equal

This flexibility opens up another opportunity: not all documents carry the same density of information. A document that repeats the same concept in different phrasings has highly redundant tokens that can be aggressively merged. A document covering many distinct topics has tokens that are far apart in embedding space and resist compression.

Fixed-budget pooling treats them the same, which is wasteful: it under-compresses the easy documents and over-compresses the hard ones, and it's the over-compressed outliers that hurt retrieval the most.

Adaptive pooling addresses this by setting a threshold on the inter-cluster distance instead of a fixed token count. We compute the threshold that achieves the target budget on average across the corpus, then apply it per-document. Easy documents get compressed below the budget; hard documents keep more tokens. The total storage is the same, the distribution of the budget is different: it goes where it is needed.

The results confirm the intuition, especially at aggressive compression ratios where the fixed approach starts to degrade the results (or on less easily compressed datasets):

The model trained for all the cutoff reaches 98.7%, near-lossless by any practical standard. Combining fixed cutoff training with adaptive budgeting reaches 99.4% retention at 5× compression, less than 0.3 NDCG@10 points from the full-token model.

Interestingly, targeted training benefits more from adaptive pooling than multi-budget training does. One plausible explanation is that the model trained at a single compression level learns tighter clusters at that level, and adaptive budgeting can then redistribute the savings from easy documents more effectively.

We also tried training the model directly for adaptive pooling (using EMA-based target thresholds), but found it no better than fixed-budget training with adaptive inference. The model doesn't need to know about adaptive budgets during training, it just needs a well-shaped space, and the adaptive allocation handles the rest at inference time.

The full picture

Putting together every ingredient, from baseline to final:

22.3 percentage points of improved retention, from a model that lost nearly a quarter of its performance at 32 tokens to one that loses barely over half a percent with no impact on full-token retrieval quality.

Note that, although the gain from MUVERA regularization seems very large, directly training the base model for hierarchical pooling yields very similar results, although slightly worse. The two regularizations compound: the MUVERA step reshapes the geometry broadly, and the hpool-specific loss refines it further.

We release the multi-budget λ=1.5 model as it offers versatility with very high retention (97.9% fixed, 98.7% adaptive). For users who know their target, single-budget training pushes the envelope further.

Takeaways and what's next

Four things to remember:

1. STE regularization generalizes beyond its training target. A model regularized for MUVERA/SMVE already pools much better, because of the underlying geometric change. Concentrating information into fewer effective dimensions seems to be what all compression methods exploit, but better understanding all of the effects at play is a promising avenue of research.

2. Targeted regularization compounds. Adding an hpool-specific STE loss on top of the projection regularization pushes retention from 88% to 97–99%, without hurting the full token performance. Most of the hpool gains come from the hpool-specific loss, but prior projection regularization is a bit helpful.

3. Targeted training beats multi-budget at the target point. When you know your deployment compression budget, training specifically for that budget gains retention over multi-budget training. Combined with adaptive pooling, this reaches 99.4% retention, effectively lossless.

4. Adaptive budgeting is free performance. Redistributing tokens from easy documents to hard ones costs nothing at training time and adds up to 2 percentage points at aggressive compression ratios.

These results reinforce a picture that's been building across our recent work: ColBERT embedding spaces have much lower effective dimensionality than their nominal dimension suggests. Hierarchical pooling, MUVERA and SMVE all exploit this in different ways. The STE regularization doesn't create the redundancy, it organizes it in the shape, making it accessible to whatever compression method you throw at the space afterward.

We believe the next step is understanding how far this can be pushed. Can we regularize for pooling and projections simultaneously? Can we reach truly lossless compression at even higher ratios? And can these insights about effective dimensionality inform how we design the next generation of late interaction models from scratch?

Get started now:

• 🤗 Models: LateOn-hpool-regularized | LateOn-regularized | LateOn
• 🛠️ Tools: PyLate | FastPLAID

Acknowledgments

We would like to acknowledge Stefan Josef, who independently explored the idea of training ColBERT models for hierarchical pooling before us. Our work here scales up the study, but the core insight deserves to be attributed to him. See the original paper.

Citation

Learn to Pool: Lightweight Fine-Tuning for Flexible Multi-Vector Compression

@inproceedings{learntopool,
  author       = {Stefan Josef},
  title        = {Learn to Pool: Lightweight Fine-Tuning for Flexible Multi-Vector Compression},
  booktitle    = {Proceedings the 1st Late Interaction Workshop (LIR) @ ECIR 2026},
  year         = {2025},
}

Hierarchical Pooling Regularization

@misc{chaffin2026hpoolregularization,
  title={After the party comes the free lunch: regularizing ColBERT models to enhance pooling capabilities and reduce index footprint},
  author={Chaffin, Antoine},
  year={2026},
  howpublished={\url{https://huggingface.co/blog/lightonai/lateon-hpool-regularization}},
}

LateOn Regularization

@misc{chaffin2026lateonregularization,
  title={Party is over: regularizing ColBERT models to fix efficient ANN methods},
  author={Chaffin, Antoine},
  year={2026},
  howpublished={\url{https://huggingface.co/blog/lightonai/lateon-regularization}},
}

Hierarchical Pooling

@misc{clavié2024reducingfootprintmultivectorretrieval,
      title={Reducing the Footprint of Multi-Vector Retrieval with Minimal Performance Impact via Token Pooling}, 
      author={Benjamin Clavié and Antoine Chaffin and Griffin Adams},
      year={2024},
      eprint={2409.14683},
      archivePrefix={arXiv},
      primaryClass={cs.IR},
      url={https://arxiv.org/abs/2409.14683}, 
}

DenseOn and LateOn

@misc{sourty2026denseonlateon,
  title={DenseOn with the LateOn: Open State-of-the-Art Single and Multi-Vector Models},
  author={Sourty, Raphael and Chaffin, Antoine and Weller, Orion and Demoura, Paulo and Chatelain, Amelie},
  year={2026},
  howpublished={\url{https://huggingface.co/blog/lightonai/denseon-lateon}},
}