Hugging Face Sentence Embedding Models with 1B Training Pairs

TL;DR

Hugging Face developed state-of-the-art general-purpose sentence embedding models by training on a massive corpus of up to 1 billion sentence pairs. This was achieved using JAX/Flax frameworks and Google TPU v3-8 infrastructure to optimize contrastive learning at scale.

Training Methodology

Model Architecture

Sentence embedding models map sentences to vectors of real numbers to capture semantic meaning for applications like clustering, text mining, and question answering. Because the set of all possible sentences is infinite, these models use a composition module—typically a Transformer followed by a pooling operation over contextualized word vectors—to compute the final representation.

Multiple Negative Ranking Loss (MNRL)

The models were trained using a contrastive training method known as Multiple Negative Ranking Loss (also referred to as InfoNCE or NTXentLoss). This approach uses in-batch negatives to optimize the embedding space:

  1. Dataset Composition: The training set consists of pairs $(a_i, p_i)$ with close meanings (e.g., query-answer pairs, duplicate questions, or cited paper titles).
  2. Objective: The model is trained to map positive pairs $(a_i, p_i)$ to close vectors while pushing unmatched pairs $(a_i, p_j)$ where $i \neq j$ to distant vectors.
  3. Similarity Functions: The model computes a similarity matrix between all pairs in a batch. The similarity function used is either Cosine-Similarity or Dot-Product.
    • Cosine-Similarity: Normalized vectors ensure the highest similarity is with the vector itself (1) and are proportional to Euclidean distance, making it compatible with k-means clustering.
    • Dot-Product: Can be faster with some approximate nearest neighbor methods, but does not work with k-means clustering and allows other vectors to have higher dot-products than the vector's similarity to itself.

To prevent score differences from becoming too small, a scaling factor $C$ (typically $C=20$) is applied to the similarity score: $sim_{scaled}(a, b) = C * sim(a, b)$.

Optimizing Embedding Quality

Batch Size and Hard Negatives

Batch composition is critical for contrastive learning performance. Hugging Face identified three primary levers for improving quality:

  • Batch Size: Larger batch sizes generally correlate with better model performance.
  • Hard Negatives: Including samples $p_j$ that are semantically close to the positive sample $p_i$ but are not the correct match (e.g., "What is the capital of France?" vs. "What is the capital of the US?") forces the model to learn more precise semantic distinctions.
  • Cross-Dataset Batches: By mixing at least two different datasets within a single batch, the model learns a global structure between different topics rather than just local structures within a single topic.

Infrastructure and Data Scale

The project utilized 7 TPUs v3-8 to handle the matrix multiplications required for large-scale contrastive learning. The training data consisted of concatenated datasets totaling up to 1 billion sentence pairs.

Results and Applications

Hugging Face trained 20 general-purpose Sentence Transformer models based on architectures including Mini-LM, RoBERTa, DistilBERT, and MPNet, achieving state-of-the-art (SOTA) results on multiple general-purpose sentence similarity evaluation tasks.

Along with the models, eight specialized datasets for Sentence-Similarity, Question Answering, and Gender Evaluation were released. These embeddings enable several practical applications:

  • Sentence Similarity: Comparing the semantic closeness of two texts using cosine similarity.
  • Asymmetric QA: Determining the likelihood that a candidate passage answers a specific query.
  • Search and Clustering: Retrieving nearby answers to a query using dot-product distance.
  • Gender Bias Evaluation: Identifying inherent gender bias in training sets by comparing model similarity scores for gendered pronouns in occupation-based anchor texts.

Sources