Fine-tuning an embedding model

The highest return in this course

An embedding model is 20 to 350 million parameters. It trains in minutes on free hardware. It needs hundreds of pairs rather than thousands of examples. And because it sets the recall ceiling for everything downstream, improving it improves every generator you ever put behind it.

This is the cheapest substantial win available in applied fine-tuning, and it is the one most often skipped in favour of training the expensive model that sits after it.

The data is pairs

You need queries paired with the passage that answers them. Nothing else is required:

{"anchor": "how long do refunds take", "positive": "Refunds are processed within 5 working days of the returned item arriving at our warehouse…"}

Three ways to get them:

• Click logs. A search query and the result somebody opened. Free, plentiful, noisy.
• Generated. For each passage in your corpus, ask a model to write three questions it answers. Cheap, and it covers your whole corpus rather than only what people have searched for. Filter by checking that your current retriever ranks the passage highly for the generated question; if it does not, keep the pair, because that is exactly the case worth learning.
• Existing question-answer content. FAQs, support macros, documentation with headings phrased as questions.

Five hundred pairs is enough to see a real change. Five thousand is comfortable.

The loss, and why batch size is the hyperparameter

The standard objective is multiple-negatives ranking loss. For each anchor in a batch, its own positive should score higher than every other positive in the batch, which serve as negatives for free.

The consequence is unusual and important: the batch size is the number of negatives. A batch of 16 asks the model to pick the right passage out of 16. A batch of 128 asks it to pick out of 128, which is a much harder task and produces a much better model.

So unlike almost everywhere else in this course, here you want the batch as large as you can fit. And when you cannot fit it, there is a technique for exactly this:

from sentence_transformers.losses import CachedMultipleNegativesRankingLoss
loss = CachedMultipleNegativesRankingLoss(model, mini_batch_size=16)

Gradient caching computes embeddings in small chunks without gradients, then replays them, giving the gradient of an enormous batch with the memory of a small one. It makes batch sizes of 512 or more reachable on a single modest card.

Hard negatives, and the trap in them

In-batch negatives are random, so most are trivially irrelevant. The model learns to separate refunds from shipping and never learns to separate a refund from an exchange.

Hard negatives fix this: for each query, retrieve the top results with your current model and keep the ones that are not the gold passage.

from sentence_transformers.util import mine_hard_negatives
triplets = mine_hard_negatives(dataset, model, range_min=10, range_max=50,
                               max_score=0.8, num_negatives=3)

Those range_min and max_score arguments are the trap being handled. The very top results are frequently also correct (your corpus has several passages about refunds) and training on a correct passage as a negative teaches the model to push apart two things that belong together. Mining from rank 10 to 50 rather than from rank 1, and capping similarity, avoids most of it. Read twenty mined negatives by hand before training on fifty thousand.

Matryoshka, if you will truncate

Matryoshka loss trains the model so that the first 64, 128 and 256 dimensions of its output are each independently useful. You then store 768-dimensional vectors for reranking and search over 128-dimensional ones, cutting index memory by six times for a small quality cost. It is one extra wrapper around your loss and it is worth it on any index above a few million vectors.

Evaluate on retrieval, not on loss

from sentence_transformers.evaluation import InformationRetrievalEvaluator
ev = InformationRetrievalEvaluator(queries, corpus, relevant_docs)

Recall at 5 and 10, NDCG at 10, on held-out queries against your whole corpus, not against a small sample of it. Retrieving the right passage from 200 candidates is a different and much easier task than retrieving it from 200,000, and a number from the first tells you nothing about the second.

The free path, end to end

Base models: all-MiniLM-L6-v2 at 22M parameters for speed, bge-base-en-v1.5 or e5-base-v2 at around 110M for quality, and multilingual variants where you need them. All free, all on the Hub.

Training 5,000 pairs on a 110M model takes a few minutes on a free T4, and the 22M model trains on a laptop CPU. The resulting model is a few hundred megabytes, serves at thousands of embeddings per second, and runs anywhere.

There is no cheaper intervention in this field with a comparable effect on a retrieval system’s quality.