After the party comes the free lunch: regularizing ColBERT models for hierarchical pooling

LateOn-hpool-regularized

📚 Collection | 📝 Blog

LateOn-hpool-regularized | LateOn-regularized | LateOn | PyLate | FastPLAID

---

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:

Training at the deployment target gains 1.6 percentage points of retention with virtually no cost to full-token quality. The trade-off: the model is less versatile at other (higher) compression levels.

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.

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 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.

Hierarchical pooling 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 Pooling

Not all documents carry the same density of information. With adaptive pooling, easy-to-compress documents receive fewer vectors while harder documents retain more. The average budget remains the same, but the token allocation is more efficient.

This model with adaptive pooling reaches 98.7% retention at ~5× compression. Combining targeted training with adaptive pooling pushes this to 99.4%.

Cumulative Improvement

Overall, the full stack improves retention by 22.3 percentage points at the 32-token compression point. 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.

---

When Should You Use This Model?

Use LateOn-hpool-regularized if you want a versatile ColBERT model that works well across all compression levels via hierarchical pooling, with the ability to set the budget at inference time.

For maximum retention at a known deployment target, consider training a targeted single-budget variant using the training code in PyLate. See the blog post for details.

Use LateOn-regularized instead if your main goal is compatibility with projection-based candidate generation methods such as MUVERA or SMVE.

Use LateOn if you only care about full-token PLAID retrieval and do not need pooling or projection-oriented compression.

---

Related Checkpoints

---

Model Details

Model Description

• Model type: PyLate ColBERT model
• Base family: LateOn
• Backbone: ModernBERT-base
• Regularization: STE-based contrastive loss on hierarchically pooled document representations
• Training: Multi-budget (λ=1.5, r∈{4, 8, 16, 32, 64, 128, 300})
• Document length: 300 tokens
• Query length: 32 tokens
• Output dimensionality: 128
• Similarity function: MaxSim
• Pooling method: Hierarchical pooling
• Recommended compression mode: Adaptive hierarchical pooling
• Language: English
• License: Apache 2.0

Model Sources

• Documentation: PyLate Documentation
• Repository: PyLate on GitHub
• FastPLAID: FastPLAID on GitHub
• Collection: DenseOn and LateOn

Full Model Architecture

ColBERT(
  (0): Transformer({'max_seq_length': 299, 'do_lower_case': False, 'architecture': 'ModernBertModel'})
  (1): Dense({'in_features': 768, 'out_features': 1536, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity', 'use_residual': True})
  (2): Dense({'in_features': 1536, 'out_features': 768, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity', 'use_residual': True})
  (3): Dense({'in_features': 768, 'out_features': 128, 'bias': False, 'activation_function': 'torch.nn.modules.linear.Identity', 'use_residual': False})
)

---

Usage

Sentence Transformers

This model can be used with Sentence Transformers as a multi-vector (ColBERT-style late interaction) retriever via the MultiVectorEncoder:

pip install "sentence-transformers>=6.0.0"

from sentence_transformers import MultiVectorEncoder

model = MultiVectorEncoder("lightonai/LateOn-hpool-regularized")

query = "Which planet is known as the Red Planet?"
documents = [
    "Venus is often called Earth's twin because of its similar size and proximity.",
    "Mars, known for its reddish appearance, is often referred to as the Red Planet.",
    "Jupiter, the largest planet in our solar system, has a prominent red spot.",
    "Saturn, famous for its rings, is sometimes mistaken for the Red Planet.",
]

query_embeddings = model.encode_query(query)
document_embeddings = model.encode_document(documents)
print(query_embeddings.shape, document_embeddings[0].shape)
# (12, 128) (18, 128)

# MaxSim late-interaction scoring (higher is more relevant)
scores = model.similarity(query_embeddings, document_embeddings)
print(scores)
# tensor([[11.3723, 11.6385, 11.5022, 11.6186]])

PyLate

First install PyLate:

pip install -U pylate

Retrieval with PLAID

Load the ColBERT model and initialize a PLAID index.

from pylate import indexes, models, retrieve

model = models.ColBERT(
    model_name_or_path="lightonai/LateOn-hpool-regularized",
)

index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
    override=True,
)

documents_ids = ["1", "2", "3"]
documents = [
    "document 1 text",
    "document 2 text",
    "document 3 text",
]

documents_embeddings = model.encode(
    documents,
    batch_size=32,
    is_query=False,
    show_progress_bar=True,
    pool_factor=5, # ⚠️ Adapt to your target budget
)

index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

Once the index has been created, it can be reloaded without re-encoding the documents:

index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
)

Retrieve documents for a query:

retriever = retrieve.ColBERT(index=index)

queries_embeddings = model.encode(
    ["query for document 3", "query for document 1"],
    batch_size=32,
    is_query=True,
    show_progress_bar=True,
)

scores = retriever.retrieve(
    queries_embeddings=queries_embeddings,
    k=10,
)

Reranking

You can also use the model for reranking without building a full index.

from pylate import rank, models

model = models.ColBERT(
    model_name_or_path="lightonai/LateOn-hpool-regularized",
)

queries = [
    "query A",
    "query B",
]

documents = [
    ["document A", "document B"],
    ["document 1", "document C", "document B"],
]

documents_ids = [
    [1, 2],
    [1, 3, 2],
]

queries_embeddings = model.encode(
    queries,
    is_query=True,
)

documents_embeddings = model.encode(
    documents,
    is_query=False,
    pool_factor=5, # ⚠️ Adapt to your target budget
)

reranked_documents = rank.rerank(
    documents_ids=documents_ids,
    queries_embeddings=queries_embeddings,
    documents_embeddings=documents_embeddings,
)

---

Citation

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}},
}

Learn to Pool

@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},
}

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{clavie2024reducingfootprintmultivectorretrieval,
  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}},
}

PyLate

@inproceedings{DBLP:conf/cikm/ChaffinS25,
  author       = {Antoine Chaffin and Raphaël Sourty},
  title        = {PyLate: Flexible Training and Retrieval for Late Interaction Models},
  booktitle    = {Proceedings of the 34th ACM International Conference on Information and Knowledge Management},
  year         = {2025},
  url          = {https://github.com/lightonai/pylate},
  doi          = {10.1145/3746252.3761608},
}