Tokens and Embeddings, lesson 12
How embedding models are trained
Train a model to rank each query's passage above the rest of its batch.
The embedding models in the last few lessons' exercises learned from pairs of texts that belong together. This lesson shows how that training works, and when to do your own. You need cosine similarity ("Cosine similarity and the dot product") and softmax with cross-entropy (Neural Networks' "Backprop through softmax and cross-entropy").
By the end, you can:
Part 1. Score every query against every passage
Work out the in-batch contrastive loss for a small batch of pairs.
An embedding model is a transformer whose token vectors get pooled into one, by averaging them (mean pooling, from "Blind spots of cosine similarity") or by keeping a special first token's. Straight out of a plain language model, those vectors match people's sense of similarity poorly.
Training on pairs fixes that: a query and a passage that answers it, like a question and its accepted answer. E5, a family of embedding models from Microsoft, trained on about 270 million pairs from Reddit, Stack Exchange, papers and web pages.
Example 1.1
Three pairs, one grid
A batch holds three help-centre pairs: "reset my password", "cancel my plan" and "download an invoice", each with its article. The model scores every query against every article by cosine. What's the loss for the first query, and for the batch?
- Read the first row0.8, 0.1, 0.2
Its own article, on the diagonal, scores highest.
- Softmax the rowe^0.8 = 2.226, e^0.1 = 1.105, e^0.2 = 1.221 → 0.489, 0.243, 0.268
Each divided by the total, 4.552: a three-way guess at which article goes with the query.
- Take the loss−ln 0.489 = 0.72
Cross-entropy on the query's own article.
- Average the rows(0.72 + 0.85 + 0.92) ÷ 3 = 0.83
Rows 2 and 3, worked the same way, give 0.85 and 0.92.
Answer0.72 for the first query, 0.83 for the batch.
Each row is a small classification task: pick this query's article out of the batch. Lowering the loss pushes the diagonal up and the rest of each row down. The wrong answers come free: they're already in the batch, and the grid is one matrix multiply.
In-batch negatives and contrastive loss
In-batch negatives are the other passages in a batch, used as each query's wrong answers. The contrastive loss is the cross-entropy of picking each query's own passage (score s⁺) from its row, where Σ adds up e^s over the row. van den Oord and colleagues named it InfoNCE in 2018.
loss = −ln(e^s⁺ ÷ Σ e^s)
Look again at row 1. It's already right, 0.8 against 0.1 and 0.2, yet its article gets only 0.489. Cosines can't pass 1, so the scores sit too close for a confident guess, and the loss keeps pushing rows that are fine.
Your turn
In another batch of three, "change my email" scores 0.3, 0.9 and 0.3 against the articles, and its own is the middle one. What's its loss?
Part 2. Temperature sharpens the guess
Divide the scores by a temperature, and say what it does to the loss.
Keep going with Pro
You've done the first part. Pro opens the rest of this lesson, and every other one.
- Temperature sharpens the guess
- Bigger batches and harder negatives
- Fine-tune on your own pairs, then measure
Tokens and Embeddings overviewSpotted a mistake? Tell us.