Train a Sentence Embedding Model with 1B Training Pairs

Summary

Sentence embeddings map sentences into dense vectors to capture semantic meaning for downstream tasks such as clustering, asymmetric question answering, and text search. To train general-purpose models at scale, the team utilized a dataset compiled from multiple sources containing up to one billion sentence pairs. Training leveraged contrastive learning with Multiple Negative Ranking Loss and in-batch negatives, running across seven TPU v3-8 accelerators with JAX and Flax. Batches combined samples from at least two datasets to balance hard negatives within topics against broader cross-topic semantic structure. The project resulted in twenty trained models—including variants of RoBERTa, DistilBERT, and MPNet—alongside eight specialized evaluation datasets and an interactive demo covering sentence similarity, search, and gender bias.

Context

Sentence embeddings map sentences to real-number vectors to capture generic semantics for downstream tasks like clustering, text mining, and question answering. However, training high-performing general-purpose models requires large-scale datasets and compute infrastructure capable of handling large batch sizes and contrastive learning objectives.

Approach / What changed

The team trained models on 7 TPU v3-8 devices using JAX and Flax on a corpus of 1 billion sentence pairs compiled from multiple datasets. Training utilized Multiple Negative Ranking Loss (InfoNCE) with in-batch negatives and scaled similarity scores (C=20). To maximize quality, batches used large sizes, incorporated hard negatives, and mixed samples from at least two datasets simultaneously to capture both within-topic and cross-topic structures.

Takeaways

• Multiple Negative Ranking Loss uses in-batch negatives, treating unmatched sentence pairs within a batch as negative examples evaluated via cross-entropy loss against a ground-truth similarity matrix.
• Batch quality was optimized by using large batch sizes, incorporating hard negatives from the same dataset, and mixing at least two datasets per batch to learn both intra-topic and inter-topic structures.
• The initiative produced 20 general-purpose Sentence Transformers models (such as Mini-LM, RoBERTa, DistilBERT, and MPNet) trained on 1 billion pairs, achieving state-of-the-art results on sentence similarity tasks.

Related reading