LateOn-regularized

📚 Collection | 📝 Blog

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.

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 (15 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