lightonai / LateOn-hpool-regularized

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 to use ColBERT retrieval with a smaller index footprint through hierarchical pooling. 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.

In practice, training uses multiple pooling budgets ( [4, 8, 16, 32, 64, 128, 300] ) rather than a single fixed target. This encourages the model to become broadly more poolable instead of overfitting to one compression point.

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

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, where the model is optimized aggressively for pooling, full-token performance remains essentially unchanged.

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

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 difference though, is that AGC learns a fixed compression budget: it's trained to compress to 32 tokens, and that's what it does. 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 dendogram and put the cut point where you want.

Adaptive Pooling

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.

With fixed pooling, every document is compressed to the same number of vectors. 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 becomes more efficient.

This gives the final operating point:

98.7% retention at approximately 5× compression.

Cumulative Improvement

Overall, the model improves retention by 21.6 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 slighlty worse.

---

When Should You Use This Model?

Use LateOn-hpool-regularized if you want a ColBERT with a smaller index and plan to use hierarchical pooling to reduce document vectors while being able to set compression budget at inference time.

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

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