Creating Text Embedding Models
Adapted from: Jay Alammar & Maarten Grootendorst, Hands-On Large Language Models (O'Reilly, 2024), Chapter 10 — rewritten and condensed for teaching, with code examples adapted from the original.
Every earlier lesson in this course used embedding models — for classification (Lesson 4), for representing tokens and sentences (Lesson 2), for dense retrieval (Lesson 8) — without ever asking how those models actually get built. This lesson closes that gap: how do you train an embedding model from scratch, and — far more commonly in practice — how do you fine-tune an existing one for your own data?
What Makes a Good Embedding, Again?
Recall the core idea from Lesson 2: an embedding model converts unstructured text into a numeric vector, and a good embedding model places semantically similar texts near each other in that vector space while pushing dissimilar texts apart.
The important nuance is that "similar" isn't one fixed thing — it's defined entirely by what the model was trained to notice. A model trained to group text by topic will place two movie reviews near each other regardless of whether they're positive or negative, as long as they're both about movies. A model trained instead for sentiment will do the opposite: it'll pull all the negative reviews together and all the positive reviews together, even across completely different topics. Neither notion of "similar" is more correct than the other — they're just optimized for different downstream tasks. Whatever notion of similarity you want, you get there by choosing the right training examples.
Check: what "similar" means
Defining similarity for an embedding model
1 / 1Why might two embedding models trained on the same reviews dataset place the exact same pair of documents in very different relative positions?
Pick an answer to see if you are right.
Contrastive Learning
The dominant technique behind training and fine-tuning embedding models is contrastive learning: teach the model by showing it pairs of documents labeled as similar or dissimilar, rather than trying to define "similarity" as some abstract rule up front.
The intuition connects to a broader idea sometimes called contrastive explanation: understanding why something is the case often depends on understanding what it's being contrasted against, not just the thing in isolation. Asking a model "is this a dog?" in isolation gives it much weaker signal than asking "is this a dog, and not this cat sitting right next to it?" — the contrast is what actually teaches the model which features matter. Contrastive learning applies exactly this principle to embeddings: by repeatedly showing a model pairs that are alike and pairs that aren't, it gradually learns which features distinguish one from the other — and, just as importantly, which features to ignore.
Lesson 2 covered word2vec's skip-gram and negative-sampling setup — training on real neighboring-word pairs (positive) contrasted against randomly sampled words (negative). That's a textbook example of contrastive learning, applied years before the term became closely associated with sentence embeddings specifically.
The most influential framework for applying contrastive learning to sentence-level embeddings is sentence-transformers.
Check: contrastive learning
Contrastive learning
1 / 2What is the core mechanism contrastive learning uses to teach a model what "similar" means?
Pick an answer to see if you are right.
SBERT: From Cross-Encoders to Bi-Encoders
Before sentence-transformers, the standard way to compare two sentences with BERT was a cross-encoder: concatenate both sentences (separated by a special token), feed them through BERT together, and attach a classification head that outputs a single similarity score.
Cross-encoders are accurate, but they don't scale to search or clustering at all. Comparing every pair among 10,000 sentences means n·(n-1)/2 ≈ 50 million forward passes — completely impractical. Worse, a cross-encoder never actually produces a reusable embedding for a single sentence on its own; it only ever scores a pair, together, in one pass.
The alternative — simply averaging BERT's word-level outputs, or using its [CLS] token, to get a "sentence embedding" — turns out to perform worse than a much older, simpler baseline: averaging plain GloVe word vectors.
The bi-encoder fix
Sentence-transformers (also called SBERT) solves this with a different architecture entirely:
- Drop the classification head.
- Add a mean pooling layer on top of BERT's token outputs, averaging them into one fixed-size vector — a real, reusable sentence embedding.
- Train with a Siamese (weight-tied) setup: two identical copies of the same BERT model — really just one model, called twice — each embed one sentence from a pair. The two resulting embeddings (plus their difference) feed into a softmax classifier during training, whose job is to predict the relationship between the two sentences.
This architecture is called a bi-encoder: each sentence is encoded independently, so embeddings can be precomputed once and reused across any number of comparisons — the same nearest-neighbor search machinery from Lesson 8 works directly on top of it. Cross-encoders generally edge out bi-encoders on raw accuracy for a single (query, document) comparison, but only a bi-encoder scales to comparing one query against a large archive.
Both architectures fundamentally rely on contrastive learning — the difference is architectural (joint vs. independent encoding), not in the underlying training philosophy.
Check: cross-encoders vs. bi-encoders
SBERT and the bi-encoder architecture
1 / 2Why does a cross-encoder scale poorly to search over a large archive?
Pick an answer to see if you are right.
Creating an Embedding Model
Building an embedding model from scratch needs three ingredients: labeled data that expresses similarity/dissimilarity, a base Transformer to embed individual pieces of text, and a loss function to optimize against.
Generating contrastive examples from NLI data
Natural language inference (NLI) datasets are a convenient, ready-made source of similarity signal. Each example is a (premise, hypothesis) pair labeled as entailment (the hypothesis follows from the premise), contradiction (they conflict), or neutral (neither).
For example, given the premise "The chef seasoned the soup with fresh basil and set it to simmer,":
- "Basil was added to the soup while it cooked on the stove" — entailment (same event, different words)
- "The chef refused to put any herbs in the soup" — contradiction
- "The chef has worked in restaurants for ten years" — neutral (plausible, but not implied)
Entailment pairs make natural positive examples for contrastive learning; contradiction (and often neutral) pairs make natural negatives.
This lesson uses the MNLI corpus (part of the GLUE benchmark) — 392,702 labeled sentence pairs — working with a 50,000-pair subset to keep training fast. Smaller subsets do make training less stable, so a larger slice (or the full dataset) is preferable whenever training time allows:
from datasets import load_dataset
# 0 = entailment, 1 = neutral, 2 = contradiction
train_dataset = load_dataset("glue", "mnli", split="train").select(range(50_000))
train_dataset = train_dataset.remove_columns("idx")
Training with softmax loss
The base Transformer here is plain bert-base-uncased — every layer stays trainable, since freezing layers tends to hurt final performance more than it saves compute:
from sentence_transformers import SentenceTransformer
embedding_model = SentenceTransformer("bert-base-uncased")
The original sentence-transformers paper trained with softmax loss — predicting the entailment/neutral/contradiction label from the two sentence embeddings:
from sentence_transformers import losses
train_loss = losses.SoftmaxLoss(
model=embedding_model,
sentence_embedding_dimension=embedding_model.get_sentence_embedding_dimension(),
num_labels=3,
)
To measure progress, evaluation uses the STSB dataset — human-labeled sentence pairs with similarity scores from 1–5, rescaled to 0–1:
from sentence_transformers.evaluation import EmbeddingSimilarityEvaluator
val_sts = load_dataset("glue", "stsb", split="validation")
evaluator = EmbeddingSimilarityEvaluator(
sentences1=val_sts["sentence1"],
sentences2=val_sts["sentence2"],
scores=[score / 5 for score in val_sts["label"]],
main_similarity="cosine",
)
Training itself follows the now-familiar Hugging Face Trainer pattern:
from sentence_transformers.training_args import SentenceTransformerTrainingArguments
from sentence_transformers.trainer import SentenceTransformerTrainer
args = SentenceTransformerTrainingArguments(
output_dir="base_embedding_model",
num_train_epochs=1,
per_device_train_batch_size=32,
per_device_eval_batch_size=32,
warmup_steps=100,
fp16=True,
eval_steps=100,
logging_steps=100,
)
trainer = SentenceTransformerTrainer(
model=embedding_model, args=args,
train_dataset=train_dataset, loss=train_loss, evaluator=evaluator,
)
trainer.train()
| Argument | What it controls |
|---|---|
num_train_epochs | Full passes over the training data (kept at 1 here purely for speed) |
per_device_train/eval_batch_size | Samples processed together per step — higher generally means faster training/eval |
warmup_steps | Steps spent linearly ramping the learning rate up from zero |
fp16 | Trains using 16-bit floats instead of 32-bit, cutting memory use and often speeding up training |
Evaluating afterward (evaluator(embedding_model)) reports several similarity-metric variants; the headline number used throughout this lesson is Pearson cosine — cosine similarity between centered vectors, from 0 to 1. This softmax-trained baseline lands around 0.59–0.60 — a starting point the rest of the lesson tries to beat.
Check: NLI data and the softmax baseline
NLI data and training an embedding model from scratch
1 / 2Why do entailment and contradiction labels in an NLI dataset make convenient contrastive-learning signal?
Pick an answer to see if you are right.
Evaluating Beyond a Single Benchmark
STSB alone is a narrow slice of "good at embeddings." The Massive Text Embedding Benchmark (MTEB) exists precisely to broaden that picture — spanning 8 task types across 58 datasets and 112 languages, with a public leaderboard for comparing models:
from mteb import MTEB
evaluation = MTEB(tasks=["Banking77Classification"])
results = evaluation.run(model)
Beyond accuracy, MTEB also reports evaluation time for each task — a deliberate design choice, since embedding models used for things like real-time semantic search need to be fast, not just accurate. Running a model against the entire MTEB suite can take hours on typical hardware, which is exactly why this lesson sticks to the much faster STSB benchmark for every comparison that follows.
Loss Functions
Softmax loss illustrates the historical starting point, but it's not the strongest option available today. Two loss functions consistently perform well across many use cases: cosine similarity loss and multiple negatives ranking (MNR) loss.
Cosine similarity loss
This loss directly compares the model's cosine similarity between two sentence embeddings against a labeled similarity score between 0 and 1 — a very direct, intuitive optimization target for semantic similarity tasks.
To reuse the MNLI data with this loss, entailment gets mapped to a similarity score of 1, while both neutral and contradiction map to 0:
mapping = {2: 0, 1: 0, 0: 1} # (neutral/contradiction)=0, (entailment)=1
train_dataset = Dataset.from_dict({
"sentence1": train_dataset["premise"],
"sentence2": train_dataset["hypothesis"],
"label": [float(mapping[label]) for label in train_dataset["label"]],
})
train_loss = losses.CosineSimilarityLoss(model=embedding_model)
Training otherwise proceeds exactly as before. The result: Pearson cosine ≈ 0.72 — a substantial jump over softmax loss's ~0.59, from changing nothing but the loss function.
Multiple negatives ranking (MNR) loss
MNR loss (also known as InfoNCE or NT-Xent loss) works on triplets: an anchor, a genuine positive, and a negative. Rather than requiring hand-picked negatives, it builds them from in-batch negatives — mismatching a positive pair from one example with a positive pair from a different example in the same batch. Given a batch of anchors and positives, every anchor's cosine similarity gets computed against every positive in the batch (not just its own), and the whole thing is optimized as a classification problem: pick out the true match from among all the batch's candidates, via cross-entropy over the similarity scores.
mnli = load_dataset("glue", "mnli", split="train").select(range(50_000))
mnli = mnli.remove_columns("idx").filter(lambda x: x["label"] == 0) # entailment only
train_dataset = {"anchor": [], "positive": [], "negative": []}
soft_negatives = list(mnli["hypothesis"])
random.shuffle(soft_negatives)
for row, soft_negative in zip(mnli, soft_negatives):
train_dataset["anchor"].append(row["premise"])
train_dataset["positive"].append(row["hypothesis"])
train_dataset["negative"].append(soft_negative)
train_dataset = Dataset.from_dict(train_dataset)
train_loss = losses.MultipleNegativesRankingLoss(model=embedding_model)
Filtering down to entailment-only pairs shrinks the dataset from 50,000 to roughly 16,875 rows — but the result is the best score yet: Pearson cosine ≈ 0.80, clearly ahead of both softmax (0.59) and cosine similarity loss (0.72).
🧪 In-batch classification, illustrated with fixed similarity scores
MNR loss tends to benefit from larger batch sizes — a bigger batch means more competing candidates per anchor, which makes the classification task genuinely harder and pushes the model toward finer-grained representations.
Easy, semi-hard, and hard negatives
In-batch negatives built this way are almost always easy negatives — usually completely unrelated to the anchor, which makes the model's job easy and caps how much it can actually learn. What helps far more is a hard negative: something topically very close to the correct answer, but still wrong.
Take the question "What is the population of Kyoto?":
- Positive: "Kyoto has a population of roughly 1.4 million people."
- Easy negative: "The Eiffel Tower was completed in 1889." — totally unrelated.
- Semi-hard negative: "Kyoto is famous for its historic temples and gardens." — about the right city, wrong topic.
- Hard negative: "Osaka, located near Kyoto, has a population of about 2.7 million." — same topic (population), a nearby related place, and still the wrong answer.
Sourcing negatives roughly falls into three tiers of effort: easy (random sampling, as above), semi-hard (using a pretrained embedding model's cosine similarity to find topically-related-but-not-quite-right candidates), and hard (manual labeling, or using a generative model to construct or judge candidate pairs). Hard negatives make training genuinely harder — and that difficulty is usually what drives the largest quality gains.
Check: loss functions and negative sampling
Loss functions and negative sampling
1 / 3Why did switching from softmax loss to cosine similarity loss improve the STSB score, using the exact same base model and dataset?
Pick an answer to see if you are right.
Fine-Tuning an Embedding Model
Training from a raw bert-base-uncased checkpoint, as above, is expensive and rarely necessary. In practice, it's far more common to start from an already-trained sentence-transformers model and fine-tune it toward your specific data or domain.
Supervised fine-tuning from a pretrained checkpoint
The training loop is identical to the MNR loss example — only the starting checkpoint changes, from bert-base-uncased to a model already trained for sentence embeddings, like all-MiniLM-L6-v2 (small, fast, and a strong general-purpose starting point):
embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
train_loss = losses.MultipleNegativesRankingLoss(model=embedding_model)
# ...same SentenceTransformerTrainingArguments and SentenceTransformerTrainer as before
This pushes the STSB score to ≈ 0.85 — the best result in the lesson. That said, this comparison isn't perfectly apples-to-apples: the pretrained checkpoint had already seen the entire MNLI dataset (and more) during its own original training, while this fine-tuning run only added 50,000 additional examples on top. Still, the workflow itself — take a strong pretrained embedding model, fine-tune on your own labeled pairs — is exactly the technique that matters most in practice.
Check: fine-tuning from a pretrained model
Fine-tuning a pretrained embedding model
1 / 1What is the main practical advantage of fine-tuning a pretrained sentence-transformers model instead of training bert-base-uncased from scratch?
Pick an answer to see if you are right.
Augmented SBERT: Stretching a Small Labeled Dataset
Real fine-tuning datasets are often far smaller than the hundreds of thousands (or billions) of pairs large embedding models are originally trained on. Augmented SBERT is a technique specifically for this situation — squeezing more value out of a small labeled ("gold") dataset by using it to automatically label a much larger set of new pairs.
The process:
- Fine-tune a cross-encoder on the small, fully-labeled gold dataset. Cross-encoders are slower, but more accurate — a good fit for this one-time labeling task, since it's not being used for large-scale search.
- Gather new, unlabeled sentence pairs — real unlabeled data if you have it, or, failing that, synthetically constructed pairs (even something as simple as recombining premises and hypotheses from your existing gold set into new combinations).
- Label those new pairs using the fine-tuned cross-encoder, producing a much larger silver dataset — labeled, but not ground truth, since the labels come from model predictions rather than human annotation.
- Train the bi-encoder (the actual embedding model you want) on gold + silver combined.
# Step 1: fine-tune a cross-encoder on the small gold dataset
from sentence_transformers.cross_encoder import CrossEncoder
cross_encoder = CrossEncoder("bert-base-uncased", num_labels=2)
cross_encoder.fit(train_dataloader=gold_dataloader, epochs=1, warmup_steps=100)
# Steps 2-3: label a much larger set of pairs with the trained cross-encoder
output = cross_encoder.predict(pairs, apply_softmax=True)
silver["label"] = output.argmax(axis=1)
# Step 4: train the bi-encoder on gold + silver combined
data = pd.concat([gold, silver], ignore_index=True).drop_duplicates(
subset=["sentence1", "sentence2"], keep="first"
)
train_dataset = Dataset.from_pandas(data, preserve_index=False)
train_loss = losses.CosineSimilarityLoss(model=embedding_model)
# ...train as usual
Using just 10,000 gold-labeled examples (20% of the 50,000 used in the cosine similarity loss example) plus their silver-augmented expansion, this reaches ≈ 0.71 — nearly matching the 0.72 achieved with 5x as much fully-labeled data. Comparing a model trained on gold-only versus gold+silver is a good way to sanity-check how much value the silver data is actually adding for your specific dataset.
Randomly recombining premises and hypotheses to build new candidate pairs skews heavily toward dissimilar pairs — most random combinations just won't be related. A better approach: use a pretrained embedding model to do a rough semantic-search pass first, retrieving top-k similar candidates for each sentence, then hand those to the cross-encoder for labeling. It's still an approximation (that embedding model wasn't trained on your data), but it's a much better starting point than pure random pairing.
Check: Augmented SBERT
Augmented SBERT
1 / 3What problem does Augmented SBERT specifically address?
Pick an answer to see if you are right.
Unsupervised Learning: TSDAE
Sometimes there's no labeled data at all. TSDAE (Transformer-based Sequential Denoising Auto-Encoder) is one of several unsupervised approaches (alongside SimCSE, Contrastive Tension, and GPL) for training an embedding model without any labels whatsoever — and it's particularly effective for domain adaptation.
The idea: take a sentence, delete a percentage of its words to create a "damaged" version, then run the damaged sentence through an encoder (with a pooling layer) to produce a single embedding. A decoder then tries to reconstruct the original, undamaged sentence purely from that embedding. The core assumption: the better the embedding captures the sentence's actual meaning, the more accurately the decoder can rebuild it.
This resembles masked language modeling, but instead of predicting individual masked tokens, the whole original sentence has to be reconstructed from a single compressed vector — a much harder, and in some ways more informative, target.
🧪 The noising step, concretely
After training, only the encoder is kept — the decoder was purely a training aid, useful only for forcing the encoder to produce embeddings good enough to reconstruct from:
from sentence_transformers.datasets import DenoisingAutoEncoderDataset
flat_sentences = mnli["premise"] + mnli["hypothesis"]
damaged_data = DenoisingAutoEncoderDataset(list(set(flat_sentences)))
from sentence_transformers import models, SentenceTransformer
# TSDAE specifically found [CLS] pooling more effective than mean pooling,
# since mean pooling loses positional information that [CLS] retains.
word_embedding_model = models.Transformer("bert-base-uncased")
pooling_model = models.Pooling(
word_embedding_model.get_word_embedding_dimension(), "cls"
)
embedding_model = SentenceTransformer(modules=[word_embedding_model, pooling_model])
from sentence_transformers import losses
# Encoder and decoder share weights (tied), and reconstruction happens
# without ever being told exactly where the deleted words were.
train_loss = losses.DenoisingAutoEncoderLoss(embedding_model, tie_encoder_decoder=True)
Trained this way — on entirely unlabeled sentences — TSDAE reaches ≈ 0.70 on STSB, remarkably close to some of the labeled techniques covered earlier, despite using zero labels at all.
Domain adaptation
TSDAE's real strength shows up in domain adaptation: adapting an embedding model to a new domain (legal text, medical text, internal company documents) that looks quite different from whatever it originally trained on. The standard recipe, adaptive pretraining, chains two of this lesson's techniques together:
- Run TSDAE (or masked language modeling) on a large pile of unlabeled, in-domain text — this is cheap, since it needs no labels, and it exposes the model to the target domain's vocabulary and style.
- Fine-tune that domain-adapted model using supervised training or Augmented SBERT — this data doesn't strictly need to come from the target domain itself, since the model already picked up the domain's characteristics in step 1.
Check: TSDAE and domain adaptation
TSDAE and domain adaptation
1 / 3What does TSDAE's decoder actually try to reconstruct during training?
Pick an answer to see if you are right.
Chapter Summary
- An embedding model's notion of "similar" is entirely defined by its training data — the same architecture can be steered toward topic, sentiment, or any other axis just by choosing different training pairs.
- Contrastive learning — training on labeled similar/dissimilar pairs — underlies essentially every technique in this lesson, from word2vec through modern sentence embeddings.
- Sentence-transformers (SBERT) replaced the expensive cross-encoder architecture with a bi-encoder: mean-pooled, independently-computed, reusable sentence embeddings from a Siamese, weight-tied BERT setup.
- Loss function choice matters enormously on identical data: softmax loss (~0.59 Pearson cosine) < cosine similarity loss (~0.72) < multiple negatives ranking loss (~0.80) in this lesson's running comparison — and hard negatives (topically close, but wrong) teach a model far more than easy, random ones.
- Fine-tuning a pretrained sentence-transformers checkpoint is almost always more practical than training from a raw BERT checkpoint, and reached the strongest score in this lesson (~0.85).
- Augmented SBERT stretches a small labeled ("gold") dataset by using a fine-tuned cross-encoder to auto-label a much larger ("silver") dataset for bi-encoder training — recovering most of the performance of a much larger fully-labeled dataset.
- TSDAE trains a strong embedding model with zero labels, by corrupting sentences and learning to reconstruct them — and chaining it before supervised fine-tuning (adaptive pretraining) is a standard recipe for domain adaptation.
Embedding models are only one half of the fine-tuning story. The next lesson turns to fine-tuning representation models directly for classification — covering both full BERT fine-tuning and the broader techniques that make adapting large pretrained models to a specific task practical.