Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks

In plain words

BERT is very good at judging how alike two sentences are, but it has to read both sentences together every time, so comparing a new sentence against a big collection means running the model once per pair. That is far too slow for search or clustering.

Sentence-BERT changes the setup so that each sentence is turned into a single list of numbers once. The same network reads each sentence separately, and the outputs are averaged into one vector. Two sentences are then compared with a quick calculation on their vectors, and the vectors can be stored and searched. The network is trained on labelled sentence pairs or triplets so that similar sentences end up with similar vectors. The authors report that its accuracy on similarity benchmarks is competitive with BERT's, with some gaps, while large comparisons drop from many hours to seconds.

TL;DR

• Sentence-BERT (SBERT) adds a pooling step to a pretrained BERT or RoBERTa and fine-tunes it with a siamese structure, so each sentence gets one fixed-size embedding §3.
• Similarity is cosine similarity between two embeddings, which allows nearest-neighbour search and clustering over a collection §1 §3.
• Training uses an NLI softmax head, a cosine regression loss, or a triplet loss. The softmax head is only used during training §3 §6.
• Author-reported: SBERT-NLI-large reaches an average Spearman of 76.55 over seven STS tasks, close to RoBERTa-based SRoBERTa-NLI-large at 76.68, and above InferSent and Universal Sentence Encoder tab. 1. Cross-topic AFS results are weaker than BERT tab. 3.

Problem

A cross-encoder (BERT reads a sentence pair together and outputs a score) must run once for every pair. Finding the most similar pair among 10,000 sentences takes n(n−1)/2 = 49,995,000 inference passes, about 65 hours on a V100 GPU in the authors' setup §1 §2.

BERT outputs, used directly as sentence vectors, are poor for similarity. Averaging BERT's output layer or using the [CLS] token gives average STS correlations of 54.81 and 29.19, below averaged GloVe vectors at 61.32 tab. 1 §4.1. The authors treat this as the motivation: a model that gives independent sentence embeddings is needed for semantic search and clustering §1.

Method

Architecture §3 fig. 1. Two copies of the same pretrained encoder (BERT or RoBERTa, 12 or 24 layers §2) share weights. A pooling layer turns the token outputs of each sentence into one vector §3. The paper compares three pooling strategies: the [CLS] output, the mean of all token outputs (MEAN), and the max over time (MAX) §3. MEAN is the default §3 §3.1.

Objective for labelled pairs (NLI classification) §3. For embeddings u and v with dimension n and k labels, the model computes softmax(W_t (u, v, |u − v|)), where W_t maps the concatenation of size 3n to k logits, and trains with cross-entropy §3. The equation is unnumbered in the paper.

Objective for scored pairs (regression) §3 fig. 2. The cosine similarity of u and v is compared to the gold score with mean squared error §3.

Objective for triplets §3. For an anchor a, a positive p and a negative n, the loss is max(‖s_a − s_p‖ − ‖s_a − s_n‖ + ε, 0) with Euclidean distance and margin ε = 1 §3.

Training on NLI §3.1. SBERT is trained on SNLI (570,000 pairs) and MultiNLI (430,000 pairs) with the three-way softmax objective for one epoch, batch size 16, Adam with learning rate 2e-5, and a linear warm-up over 10% of the training data §3.1.

Training on STS benchmark and on triplets §4.2 §4.4. The regression objective is trained on the STS benchmark training split (5,749 pairs). Two setups are compared: STSb only, and NLI first then STSb §4.2. The triplet objective is trained for one epoch on about 1.8 million Wikipedia-section triplets §4.4.

Inference §3 §6. The softmax head is discarded. Similarity is the cosine of the pooled embeddings of the two sentences §6. The authors also tried negative Manhattan and Euclidean distances and report roughly the same results, without reporting numbers §4.

The pseudocode below follows the official code at code_commit. The mean pooling is pooling.py. The NLI head is softmax.py. The triplet hinge is triplet.py.

# Sentence embedding with mean pooling  [§3]
# Input:  token ids (B, T), attention mask (B, T)
# Output: sentence embeddings (B, n)

def embed(token_ids, mask):
    # Run the pretrained encoder (BERT or RoBERTa)  [§3]
    token_vectors = encoder(token_ids, mask)            # (B, T, n)
    # Average over the real tokens only, ignoring padding  [§3]
    total = (token_vectors * mask[..., None]).sum(dim=1)  # (B, n)
    count = mask.sum(dim=1, keepdims=True)              # (B, 1)
    return total / count                                # (B, n)

# NLI training step with the softmax head  [§3] [§3.1]
# Input:  ids_a, ids_b (B, T); labels (B,) in {0, 1, 2}
# Output: scalar cross-entropy loss

def nli_loss(ids_a, ids_b, labels):
    # Embed both sentences with the same (tied) encoder  [§3] [fig. 1]
    u = embed(ids_a)                                    # (B, n)
    v = embed(ids_b)                                    # (B, n)
    # Concatenate u, v and |u - v|; no product u * v  [§3] [§6]
    features = concat([u, v, abs(u - v)], axis=1)       # (B, 3n)
    # Linear layer W_t of shape (3n, k) gives logits  [§3]
    logits = features @ W_t                             # (B, k)
    # Cross-entropy over the k labels  [§3]
    return cross_entropy(logits, labels)

# Triplet training step  [§3] [§4.4]
# Input:  anchor, positive, negative token ids; margin epsilon (paper: 1)
# Output: scalar hinge loss

def triplet_loss(anchor, positive, negative, margin):
    # Embed the three sentences with the shared encoder  [§3]
    s_a = embed(anchor)                                 # (B, n)
    s_p = embed(positive)                               # (B, n)
    s_n = embed(negative)                               # (B, n)
    # Euclidean distances from the anchor  [§3]
    dist_pos = norm(s_a - s_p)                          # (B,)
    dist_neg = norm(s_a - s_n)                          # (B,)
    # Hinge: the positive must be at least margin closer  [§3]
    return mean(max(dist_pos - dist_neg + margin, 0))

# Similarity at inference, softmax head not used  [§3] [§6]
# Input:  two sentences
# Output: cosine similarity in [-1, 1]

def similarity(sentence_a, sentence_b):
    # Embed each sentence separately  [§3]
    u = embed(sentence_a)                               # (n,)
    v = embed(sentence_b)                               # (n,)
    # Cosine similarity of the two vectors  [§3]
    return dot(u, v) / (norm(u) * norm(v))              # scalar

Hyperparameters

Results

Author-reported. Setups: STS evaluation with cosine similarity and Spearman ρ × 100, no STS training data for unsupervised rows §4.1 tab. 1; STSb test set §4.2 tab. 2; AFS with 10-fold CV and cross-topic evaluation §4.3 tab. 3; Wikipedia section triplets §4.4 tab. 4; SentEval 10-fold logistic regression §5 tab. 5; ablation on STSb dev §6 tab. 6; speed on STSb sentences §7 tab. 7.

Table 1 tab. 1: STS tasks, Spearman ρ × 100, no STS training data for these rows §4.1.

Table 2 tab. 2: STS benchmark test set, Spearman ρ × 100, mean ± std over 10 seeds where given §4.2.

Table 3 tab. 3: Argument Facet Similarity, average Pearson r and Spearman ρ §4.3.

Table 4 tab. 4: Wikipedia section triplets, accuracy §4.4.

Table 5 tab. 5: SentEval, accuracy, 10-fold cross-validation §5.

Table 6 tab. 6: Ablation on STSb dev (cosine, Spearman) and NLI classification (accuracy) §6.

Table 7 tab. 7: Sentences per second on STS benchmark sentences, higher is better §7.

Headline claims in the text: the closest pair among 10,000 sentences drops from about 65 hours with BERT to about 5 seconds with SBERT (about 10,000 embeddings plus about 0.01 s of cosine computation) §1 §8. On GPU throughput, the smart-batched SBERT-base comes out roughly 9% ahead of InferSent and roughly 55% ahead of the Universal Sentence Encoder §7 tab. 7.

Paper vs code

Checked against the official repository UKPLab/sentence-transformers at code_commit (shallow clone of the default branch at HEAD, a later version of the library than the 2019 release). Training example scripts for the paper's tasks are present under examples/ at this commit.

• Triplet margin: discrepancy. Paper uses ε = 1 §3. The TripletLoss default is triplet_margin = 5 triplet.py. The Wikipedia-sections example creates TripletLoss(model) without a margin argument, so it runs with 5 training_wikipedia_sections.py. Pass triplet_margin=1 to match the paper.
• Optimizer: discrepancy. Paper says Adam §3.1. The fit default is torch.optim.AdamW with weight decay 0.01 fit_mixin.py. The NLI example sets neither the optimizer nor the learning rate training_nli.py, so the 2e-5 of the paper is not in the example. Whether the Trainer default matches was not verified.
• Warm-up: matches the examples, not the fit default. Paper uses 10% of training §3.1. The examples pass warmup_steps=0.1 training_nli.py training_wikipedia_sections.py. The old fit default is a fixed 10000 steps fit_mixin.py.
• NLI head: matches. Default concatenation is (u, v, |u − v|) with the product off softmax.py §3 §6. Cross-entropy is the default loss softmax.py. The classifier is a linear layer over 3n inputs softmax.py. Batch size 16, one epoch and num_labels=3 match training_nli.py training_nli.py training_nli.py.
• Regression: matches with a caveat. MSE is the default loss on cosine similarity cosine_similarity.py. The default score transform is the identity, so labels are compared to a cosine in [−1, 1] without rescaling cosine_similarity.py. Paper gold labels are 0–5 §4.1; how they were rescaled is not stated.
• Pooling: matches. MEAN is the default pooling.py. MAX is a masked max over tokens pooling.py.
• STSb and triplet examples. STSb example: batch 16, 4 epochs training_stsbenchmark.py. Wikipedia example: batch 16, one epoch training_wikipedia_sections.py.
• Mixed precision: not in the paper. The NLI example sets fp16=True training_nli.py. The paper does not mention precision.

Pitfalls

• Triplet margin differs. Paper ε = 1 §3; library default 5 triplet.py. Using the default changes the loss scale and the margin's effect on training.
• Optimizer differs. Paper "Adam" §3.1; library default AdamW with weight decay 0.01 fit_mixin.py. Weight decay is not in the paper.
• Regression label scale is not stated. STS labels are 0–5 §4.1, cosine is in [−1, 1], and the default transform is the identity cosine_similarity.py. Rescale labels to the cosine range before training, and say so in your own code.
• STSb epochs for SBERT are not stated. Table 2's caption gives four epochs for the BERT systems only tab. 2. The paper says nothing about SBERT's STSb epochs or batch size; the example uses 4 epochs and 16.
• Seeds and variance. STSb rows use 10 seeds with ± std tab. 2 §4.2. Other tables report single values, and the variance across seeds is not stated for them.
• Pooling matters. MAX pooling is much worse on STSb (69.92 vs 87.44 for MEAN) tab. 6. CLS is slightly worse than MEAN tab. 6. Keep MEAN unless you have a reason.
• The classification head is for training only. At inference only the pooled vectors and cosine are used §6. Do not keep the NLI head in the serving path.
• Concatenation choice matters for NLI. Adding u * v to (u, v, |u − v|) lowers NLI accuracy from 80.78 to 80.44 tab. 6, and the product-only variants are much worse tab. 6. The paper says the |u − v| term is the most important one §6.
• Cross-topic AFS is weak. SBERT-AFS-base drops to ρ = 50.65 in cross-topic evaluation, below BERT-AFS-base at 57.23 tab. 3. The authors suggest that more than two training topics may be needed §4.3. The 10-fold result is close to BERT tab. 3.
• Cosine vs other distances. Manhattan and Euclidean results were "roughly the same" §4; no numbers are reported in the paper.
• Timings are hardware- and library-specific. The 65 h and 5 s figures are estimates on the authors' setup §1 §7. Table 7 mixes frameworks: SBERT and InferSent are PyTorch and USE is TensorFlow Hub §7, so the ratios mix framework and model differences. Smart batching is an inference-time optimisation §7.
• Max sequence length and training hardware are not stated. A model with a shorter or longer token limit will produce different embeddings and different costs.
• Equations are unnumbered in the source. Cite by section (for example §3). Do not assume an equation number.
• Code version drift. The commit is the current library, not the 2019 release, and no release tag was checked. The example scripts may have changed since the paper.

Lineage

• Builds on: BERT (arXiv:1810.04805) §2 and RoBERTa (arXiv:1907.11692) §2, which are fine-tuned rather than trained from scratch §2. Triplet networks are cited to FaceNet (arXiv:1503.03832) §3.
• Compared against: average GloVe embeddings tab. 1; InferSent (Conneau et al., 2017) tab. 1; Universal Sentence Encoder (arXiv:1803.11175) tab. 1; SentEval (arXiv:1803.05449) for transfer tasks §5.
• Superseded by: not stated in this wiki.