lightonai/LateOn-regularized

Party is over: regularizing ColBERT models to fix efficient ANN methods

LateOn-regularized

📚 Collection | 📝 Blog

DenseOn | LateOn-regularized | LateOn | PyLate | FastPLAID

---

Overview

LateOn-regularized is a ColBERT (multi-vector) retrieval model built on the LateOn methology, trained by LightOn using PyLate. It applies a novel Straight-Through Estimator (STE) regularization during supervised fine-tuning that makes the model compatible with efficient ANN candidate-generation methods like MUVERA and SMVE, while fully preserving PLAID retrieval performance.

This model is the recommended choice when you plan to use MUVERA or SMVE for candidate generation in your ColBERT retrieval pipeline. If you only use PLAID indexes, LateOn is equivalent and equally recommended.

For full details on the problem, methodology, and analysis, please read our blog post.

Why This Model Exists

Recent ColBERT models (including LateOn) produce highly anisotropic token embeddings: nearly all vectors cluster into a narrow cone with mean pairwise cosine similarity as high as 0.95. While this does not affect PLAID-based retrieval, it breaks ANN methods like MUVERA and SMVE that rely on random projections to separate embeddings. When all vectors point in roughly the same direction, random hyperplanes and rotations cannot distinguish them, making candidate generation essentially random.

Mean-centering the embeddings removes the dominant shared direction and helps significantly, but isn't enough on its own as the residual geometry still doesn't project cleanly.

STE-based regularization solves this by directly optimizing the model to produce discriminative representations after MUVERA/SMVE projection. During training, we build the true projected codes in the forward pass and compute a contrastive loss on them, then use the Straight-Through Estimator to pass gradients through the non-differentiable discretization steps (SimHash bucketing in MUVERA, top-k sparsification in SMVE). The total loss is a weighted combination of the standard MaxSim contrastive loss and the projected-representation contrastive loss.

Surprisingly, the regularization does not make the space more isotropic. Instead, it teaches the model to concentrate discriminative information into fewer effective dimensions that random projections can capture efficiently, the opposite of what isotropy-based intuition would predict. The regularization also transfers across methods (a MUVERA-trained model helps SMVE and vice-versa), across random seeds, and across hyperparameter configurations.

See the blog post for the complete analysis.

Results

Important: All results below use mean-centered representations. We recommend always centering embeddings when using MUVERA or SMVE.

MUVERA / SMVE Retrieval (8 BEIR datasets average NDCG@10)

rk=k means the top-k candidates from MUVERA/SMVE are reranked with full MaxSim.

Model	Center	PLAID	Method	rk=0	rk=50	rk=100	rk=200	rk=400	rk=1000
LateOn	no	55.28	MUVERA	2.89 (-52.4)	9.46 (-45.8)	12.53 (-42.8)	15.87 (-39.4)	20.29 (-35.0)	27.52 (-27.8)
			SMVE	1.51 (-53.8)	7.59 (-47.7)	10.84 (-44.4)	14.24 (-41.0)	18.29 (-37.0)	25.10 (-30.2)
LateOn	yes	55.28	MUVERA	32.66 (-22.6)	46.22 (-9.1)	48.68 (-6.6)	50.01 (-5.3)	50.85 (-4.4)	52.09 (-3.2)
			SMVE	21.18 (-34.1)	37.39 (-17.9)	40.90 (-14.4)	43.52 (-11.8)	45.83 (-9.5)	48.61 (-6.7)
LateOn-regularized (MUVERA)	yes	55.72	MUVERA	40.80 (-14.9)	51.60 (-4.1)	52.37 (-3.4)	53.02 (-2.7)	53.60 (-2.1)	53.87 (-1.9)
			SMVE	37.18 (-18.5)	48.97 (-6.8)	50.45 (-5.3)	51.51 (-4.2)	52.40 (-3.3)	53.50 (-2.2)
LateOn-regularized (SMVE)	yes	54.92	MUVERA	40.22 (-14.7)	50.94 (-4.0)	51.90 (-3.0)	52.67 (-2.3)	53.02 (-1.9)	53.35 (-1.6)
			SMVE	40.39 (-14.5)	50.70 (-4.2)	52.09 (-2.8)	52.92 (-2.0)	53.19 (-1.7)	53.47 (-1.5)

Key takeaways:

• PLAID performance is preserved (55.72 vs 55.28 for MUVERA-reg), confirming that regularization does not conflict with MaxSim requirements.
• Massive improvement for MUVERA/SMVE: centered MUVERA rk=0 goes from 32.66 → 40.80 (+8.1 points), and with rk=200 reranking the gap to full PLAID drops to just 2.7 points (note: the goal of the study was making representations better, not sweeping MUVERA/SMVE parameters to get perfect scores)
• Cross-method transfer: MUVERA-regularized training also boosts SMVE, and vice-versa, showing the regularization improves global compressibility rather than optimizing a single projection subspace.

BEIR (14 datasets, NDCG@10) — PLAID Retrieval

LateOn-regularized maintains the strong PLAID performance of the LateOn family. For full BEIR results, see the LateOn model card.

Related Checkpoints

Model	Description	Link
LateOn-regularized (this card)	STE-regularized for MUVERA/SMVE compatibility	lightonai/LateOn-regularized
LateOn	Base supervised ColBERT model (best for PLAID-only)	lightonai/LateOn
LateOn-unsupervised	Pre-training only	lightonai/LateOn-unsupervised
DenseOn	Single-vector counterpart	lightonai/DenseOn

Model Details

Model Description

• Model Type: PyLate ColBERT model
• Base model: LateOn-unsupervised (ModernBERT-base)
• Regularization: STE-based contrastive loss on MUVERA projected representations
• Document Length: 300 tokens
• Query Length: 32 tokens
• Output Dimensionality: 128 tokens
• Similarity Function: MaxSim
• Recommended: Use mean-centered embeddings for MUVERA/SMVE
• Language: English
• License: Apache 2.0

Model Sources

• Documentation: PyLate Documentation
• Repository: PyLate on GitHub
• Hugging Face: PyLate models on Hugging Face

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 the PyLate library:

pip install -U pylate

Retrieval

Use this model with PyLate to index and retrieve documents. The index uses FastPLAID for efficient similarity search.

Indexing documents

Load the ColBERT model and initialize the PLAID index, then encode and index your documents:

from pylate import indexes, models, retrieve

# Step 1: Load the ColBERT model
model = models.ColBERT(
    model_name_or_path="lightonai/LateOn-regularized",
)

# Step 2: Initialize the PLAID index
index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
    override=True,  # This overwrites the existing index if any
)

# Step 3: Encode the documents
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,  # Ensure that it is set to False to indicate that these are documents, not queries
    show_progress_bar=True,
)

# Step 4: Add document embeddings to the index by providing embeddings and corresponding ids
index.add_documents(
    documents_ids=documents_ids,
    documents_embeddings=documents_embeddings,
)

Note that you do not have to recreate the index and encode the documents every time. Once you have created an index and added the documents, you can re-use the index later by loading it:

# To load an index, simply instantiate it with the correct folder/name and without overriding it
index = indexes.PLAID(
    index_folder="pylate-index",
    index_name="index",
)

Retrieving top-k documents for queries

Once the documents are indexed, you can retrieve the top-k most relevant documents for a given set of queries. To do so, initialize the ColBERT retriever with the index you want to search in, encode the queries and then retrieve the top-k documents to get the top matches ids and relevance scores:

# Step 1: Initialize the ColBERT retriever
retriever = retrieve.ColBERT(index=index)

# Step 2: Encode the queries
queries_embeddings = model.encode(
    ["query for document 3", "query for document 1"],
    batch_size=32,
    is_query=True,  # Ensure that it is set to True to indicate that these are queries
    show_progress_bar=True,
)

# Step 3: Retrieve top-k documents
scores = retriever.retrieve(
    queries_embeddings=queries_embeddings,
    k=10,  # Retrieve the top 10 matches for each query
)

Reranking

If you only want to use the ColBERT model to perform reranking on top of your first-stage retrieval pipeline without building an index, you can simply use rank function and pass the queries and documents to rerank:

from pylate import rank, models

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

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

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

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

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

documents_embeddings = model.encode(
    documents,
    is_query=False,
)

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

Training Details

Regularization Setup

LateOn-regularized starts from the unregularized unsupervised checkpoint and applies the same supervised training step as LateOn, with an additional STE-based regularization signal. The total loss is:

L = (1 - α) × MaxSim contrastive loss + α × projected-representation contrastive loss

The projected representations are built using MUVERA (R=20, k=5, d=16, FDE dim=10,240) in the forward pass, with the Straight-Through Estimator passing gradients through the non-differentiable bucketing step during backpropagation. The random projection matrices are sampled once at the start of training and kept fixed. Despite this, the regularization generalizes across different random seeds and MUVERA/SMVE hyperparameter configurations at evaluation time.

The supervised stage is lightweight — we did not even run the full 3 epochs — demonstrating that effective regularization is achievable with minimal additional compute.

Framework Versions

• Python: 3.11.10
• Sentence Transformers: 5.1.1
• PyLate: 1.3.4
• Transformers: 4.57.5
• PyTorch: 2.9.0+cu128
• Accelerate: 1.12.0
• Datasets: 3.6.0
• Tokenizers: 0.22.1

Citation

BibTeX

LateOn Regularization

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

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{/blog/lightonai/denseon-lateon}},
}

PyLate

@inproceedings{DBLP:conf/cikm/ChaffinS25,
  author       = {Antoine Chaffin and
                  Rapha{\"{e}}l Sourty},
  editor       = {Meeyoung Cha and
                  Chanyoung Park and
                  Noseong Park and
                  Carl Yang and
                  Senjuti Basu Roy and
                  Jessie Li and
                  Jaap Kamps and
                  Kijung Shin and
                  Bryan Hooi and
                  Lifang He},
  title        = {PyLate: Flexible Training and Retrieval for Late Interaction Models},
  booktitle    = {Proceedings of the 34th {ACM} International Conference on Information
                  and Knowledge Management, {CIKM} 2025, Seoul, Republic of Korea, November
                  10-14, 2025},
  pages        = {6334--6339},
  publisher    = {{ACM}},
  year         = {2025},
  url          = {https://github.com/lightonai/pylate},
  doi          = {10.1145/3746252.3761608},
}

Sentence Transformers

@inproceedings{reimers-2019-sentence-bert,
    title = "Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks",
    author = "Reimers, Nils and Gurevych, Iryna",
    booktitle = "Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing",
    month = "11",
    year = "2019",
    publisher = "Association for Computational Linguistics",
    url = "https://arxiv.org/abs/1908.10084"
}