MS MARCO

MS MARCO Passage Ranking is a large dataset to train models for information retrieval. It consists of about 500k real search queries from Bing search engine with the relevant text passage that answers the query. This page shows how to train Multi-Vector Encoder (ColBERT-style) models on this dataset so that they can be used for searching text passages given queries (key words, phrases or questions).

There are pre-trained models available, which you can directly use without the need of training your own models. For more information, see: Pretrained Models.

MultiVectorMultipleNegativesRankingLoss

Training code: training_contrastive.py

When we use MultiVectorMultipleNegativesRankingLoss, we provide triplets: (query, positive_passage, negative_passage) where positive_passage is the relevant passage to the query and negative_passage is a non-relevant passage, mined with BM25 in the sentence-transformers/msmarco-bm25 dataset. Every query token is compared against every passage token, and the resulting MaxSim score of the (query, positive_passage) pair is optimized to be higher than the scores against the negative passage and against all other in-batch passages.

Losses that use in-batch negatives benefit heavily from larger batch sizes. If GPU memory is the bottleneck, training code: training_cached_contrastive.py demonstrates the same recipe with CachedMultiVectorMultipleNegativesRankingLoss, which reaches much larger effective batch sizes at a small speed cost via GradCache.

MultiVectorDistillKLDivLoss

Training code: training_kd.py

The strongest late-interaction models (e.g. ColBERTv2, GTE-ModernColBERT) are trained with knowledge distillation: instead of binary relevance, the model learns to reproduce the score distribution of a stronger teacher model over N candidate documents per query. The lightonai/ms-marco-en-bge dataset provides per-query candidate document IDs together with teacher scores from a BGE model. resolve_ids() resolves the IDs against the query and document texts on the fly, and MultiVectorDistillKLDivLoss minimizes the KL divergence between the (softmaxed) teacher scores and the student’s MaxSim scores. The script also builds its model explicitly via modules=... in the style of modern PyLate models like lightonai/LateOn: a residual bottleneck head before the final projection, a query length cap, and a punctuation skiplist. It trains fp32 weights under bf16 autocast with FlashAttention-2, and runs evaluation under the same autocast.

The train split can also be streamed (streaming=True) to skip the initial download: apply resolve_ids via IterableDataset.map with an explicit features= (streaming map cannot infer the output schema, and the trainer requires one). The resolve_ids() documentation shows the exact recipe. The query and document lookup datasets are random-access joins, so they must stay regular (materialized) datasets.