Training Sentence Embeddings on 1 Billion Data Pairs

To advance Sentence Embedding technology that maps sentences into meaning-bearing vectors, this discloses the results and methodology of a large-scale training project using 1 billion training pairs.

Key Technical Features:

• Infrastructure and Framework: Efficient large-scale training was performed using the JAX/Flax framework and 7 TPU v3-8 units.
• Model Architecture: Based on a Transformer model, applying a Pooling operation to contextualized word vectors to derive the final sentence representation.
• Training Strategy: Adopted a Contrastive Learning approach using Multiple Negative Ranking Loss (InfoNCE). This method increases similarity for positive pairs within a batch while pushing away the remaining samples (in-batch negatives).
• Similarity Function: Compares the characteristics of cosine similarity and dot-product, and applied a scaling factor during actual training to adjust the score differences.