Create README.md

README.md

@@ -0,0 +1,349 @@
+---
+tags:
+- ColBERT
+- PyLate
+- sentence-transformers
+- sentence-similarity
+- feature-extraction
+pipeline_tag: sentence-similarity
+library_name: PyLate
+license: apache-2.0
+language:
+- en
+---
+
+<div align="center">
+
+<img src="https://cdn-uploads.huggingface.co/production/uploads/609bbe2f4932693ca2009d6a/wAztcBQeYSvCcj52mVK_w.png" alt="LightOn" width="512">
+
+<h1 align="center">Party is over: regularizing ColBERT models to fix efficient ANN methods</h1>
+
+<h3 align="center">LateOn-regularized</h3>
+
+[![Website](https://img.shields.io/badge/LightOn-Website-blue?logo=google-chrome)](https://lighton.ai)
+[![LinkedIn](https://img.shields.io/badge/LightOn-LinkedIn-0A66C2?logo=linkedin)](https://www.linkedin.com/company/lighton/)
+[![X](https://img.shields.io/badge/@LightOnIO-X-black?logo=x)](https://x.com/LightOnIO)
+
+  📚 [Collection](https://huggingface.co/collections/lightonai/denseon-and-lateon) | 📝 [Blog](https://huggingface.co/blog/lightonai/lateon-regularization)
+
+</div>
+
+
+
+<p align="center">
+<a href="https://huggingface.co/lightonai/DenseOn">DenseOn</a> |
+<a href="https://huggingface.co/lightonai/LateOn-regularized">LateOn-regularized</a> |
+<a href="https://huggingface.co/lightonai/LateOn">LateOn</a> |
+<a href="https://github.com/lightonai/pylate">PyLate</a> |
+<a href="https://github.com/lightonai/fast-plaid">FastPLAID</a>
+</p>
+
+---
+
+## Overview
+
+**LateOn-regularized** is a ColBERT (multi-vector) retrieval model built on the [LateOn](https://huggingface.co/lightonai/LateOn) methology, trained by [LightOn](https://lighton.ai) using [PyLate](https://github.com/lightonai/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](https://arxiv.org/abs/2405.19504) and [SMVE](https://www.topk.io/blog/20260311-smve-multi-vector-retrieval), **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](https://huggingface.co/lightonai/LateOn) is equivalent and equally recommended.
+
+For full details on the problem, methodology, and analysis, please read our [blog post](https://huggingface.co/blog/lightonai/lateon-regularization).
+
+## 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](https://huggingface.co/blog/lightonai/lateon-regularization) 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](https://huggingface.co/lightonai/LateOn).
+
+## Related Checkpoints
+
+| Model | Description | Link |
+|:---|:---|:---|
+| **LateOn-regularized** *(this card)* | STE-regularized for MUVERA/SMVE compatibility | [`lightonai/LateOn-regularized`](https://huggingface.co/lightonai/LateOn-regularized) |
+| LateOn | Base supervised ColBERT model *(best for PLAID-only)* | [`lightonai/LateOn`](https://huggingface.co/lightonai/LateOn) |
+| LateOn-unsupervised | Pre-training only | [`lightonai/LateOn-unsupervised`](https://huggingface.co/lightonai/LateOn-unsupervised) |
+| DenseOn | Single-vector counterpart | [`lightonai/DenseOn`](https://huggingface.co/lightonai/DenseOn) |
+
+## Model Details
+
+### Model Description
+- **Model Type:** PyLate ColBERT model
+- **Base model:** [LateOn-unsupervised](https://huggingface.co/lightonai/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](https://lightonai.github.io/pylate/)
+- **Repository:** [PyLate on GitHub](https://github.com/lightonai/pylate)
+- **Hugging Face:** [PyLate models on Hugging Face](https://huggingface.co/models?library=PyLate)
+
+### 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:
+
+```bash
+pip install -U pylate
+```
+
+### Retrieval
+
+Use this model with PyLate to index and retrieve documents. The index uses [FastPLAID](https://github.com/lightonai/fast-plaid) for efficient similarity search.
+
+#### Indexing documents
+
+Load the ColBERT model and initialize the PLAID index, then encode and index your documents:
+
+```python
+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:
+
+```python
+# 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:
+
+```python
+# 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:
+
+```python
+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
+
+```bibtex
+@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}},
+}
+```
+
+#### DenseOn and LateOn
+
+```bibtex
+@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
+
+```bibtex
+@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
+
+```bibtex
+@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"
+}
+```
+