Fine-tuning embeddings

Objectives

• Understand when embedding fine-tuning is necessary
• Prepare a pair dataset for training
• Implement fine-tuning with lightweight adapters

When to fine-tune?

Let’s start with the bad news, which is actually good: in most cases, pre-trained embedding models work perfectly fine, and fine-tuning is a waste of effort. Only three situations truly justify it. The first is a domain with highly specialized vocabulary—medicine, law, industry—where the generic model doesn’t know that “MI” and “myocardial infarction” refer to the same thing. The second is hitting a plateau: you’ve optimized chunking, added hybrid search and reranking, and metrics won’t budge. The third is simply having annotated text pairs on hand, similar or dissimilar, sitting idle in your search logs.

Conversely, skip fine-tuning if you have fewer than a thousand annotated pairs: training will memorize your examples without generalizing. Skip it too if performance is already satisfactory—a two-point gain in recall isn’t worth maintaining an extra model. And above all, verify that better chunking or a better prompt won’t solve the problem: it’s almost always cheaper and faster.

Approach 1: Linear adapter (Matryoshka)

The most economical approach leaves the embedding model untouched. You let it produce its vectors, then train a small projection layer on top that warps the space to bring closer together what needs to be close in your domain. OpenAI’s model remains a black box called via API; only a few hundred thousand parameters are trained, on CPU, in minutes.

Training relies on triplets: a query, a document that answers it, and a document that covers the same topic but doesn’t answer it. Look at the medical examples below—the negative for “paracetamol side effects” isn’t a text about football, but a text about paracetamol sold in pharmacies. This is called a hard negative, and it’s what teaches the model something: an obvious negative corrects nothing since the generic model already classified it correctly.

import numpy as np
from openai import OpenAI
from sklearn.model_selection import train_test_split

client = OpenAI()

# Dataset of pairs (query, positive_document, negative_document)
triplets = [
    {
        "query": "paracetamol side effects",
        "positive": "Paracetamol can cause liver disorders...",
        "negative": "Paracetamol is available in pharmacies..."
    },
    {
        "query": "children's aspirin dosage",
        "positive": "For children aged 6 to 12, the recommended dose...",
        "negative": "Aspirin was discovered in the 19th century..."
    },
    # ... more triplets
]

def prepare_embeddings(triplets: list[dict]) -> dict:
    """Generate embeddings for the triplets."""
    all_texts = []
    for t in triplets:
        all_texts.extend([t["query"], t["positive"], t["negative"]])

    response = client.embeddings.create(
        input=all_texts,
        model="text-embedding-3-large"
    )
    embs = np.array(
        [d.embedding for d in sorted(response.data, key=lambda x: x.index)]
    )

    n = len(triplets)
    return {
        "queries": embs[0::3],
        "positives": embs[1::3],
        "negatives": embs[2::3],
    }

The adapter itself is a two-layer network that projects the model’s 3,072 dimensions down to 256, with final normalization so output vectors remain comparable by cosine similarity. The loss function is a triplet margin loss: it equals zero as soon as the query is closer to the positive than to the negative by at least a margin of 0.2, and penalizes proportionally otherwise. In other words, you only teach the model what it hasn’t yet understood, and training naturally stabilizes once the triplets are well separated.

import torch
import torch.nn as nn
import torch.optim as optim

class EmbeddingAdapter(nn.Module):
    """Projection layer to adapt embeddings."""
    def __init__(self, input_dim: int = 3072, output_dim: int = 256):
        super().__init__()
        self.projection = nn.Sequential(
            nn.Linear(input_dim, output_dim),
            nn.ReLU(),
            nn.Linear(output_dim, output_dim),
        )

    def forward(self, x):
        projected = self.projection(x)
        return projected / projected.norm(dim=-1, keepdim=True)

def triplet_loss(anchor, positive, negative, margin=0.2):
    """Triplet margin loss."""
    dist_pos = 1 - torch.cosine_similarity(anchor, positive)
    dist_neg = 1 - torch.cosine_similarity(anchor, negative)
    return torch.relu(dist_pos - dist_neg + margin).mean()

def train_adapter(data: dict, epochs: int = 50):
    """Train the adapter on triplets."""
    model = EmbeddingAdapter()
    optimizer = optim.Adam(model.parameters(), lr=1e-3)

    queries = torch.tensor(data["queries"], dtype=torch.float32)
    positives = torch.tensor(data["positives"], dtype=torch.float32)
    negatives = torch.tensor(data["negatives"], dtype=torch.float32)

    for epoch in range(epochs):
        optimizer.zero_grad()

        q_proj = model(queries)
        p_proj = model(positives)
        n_proj = model(negatives)

        loss = triplet_loss(q_proj, p_proj, n_proj)
        loss.backward()
        optimizer.step()

        if (epoch + 1) % 10 == 0:
            print(f"Epoch {epoch + 1}/{epochs}, Loss: {loss.item():.4f}")

    return model

# Train
data = prepare_embeddings(triplets)
adapter = train_adapter(data)

# Save
torch.save(adapter.state_dict(), "medical_adapter.pt")

One operational point you can’t miss: the day you deploy this adapter, every vector already stored in your database must pass through the projection again, otherwise you’d be comparing vectors from two different spaces. Plan for complete corpus reindexing in your production rollout.

Approach 2: Fine-tuning with Sentence Transformers

When the adapter isn’t enough, you need to modify the embedding model itself, which requires an open model that you host—here BAAI/bge-large-en-v1.5. Sentence Transformers makes this accessible: you describe examples as text pairs with a continuous label, where 1.0 means “these two texts should be close” and 0.2 means “they look superficially similar, push them apart.” The CosineSimilarityLoss then pushes the calculated cosine similarity toward the provided label.

Three epochs and one hundred warmup steps are conservative values, and for good reason: on a modest dataset, training too long destroys the model’s general knowledge, making it excellent on your examples and mediocre everywhere else.

from sentence_transformers import (
    SentenceTransformer, InputExample, losses
)
from torch.utils.data import DataLoader

# Load a base model
model = SentenceTransformer("BAAI/bge-large-en-v1.5")

# Prepare training examples
train_examples = [
    InputExample(
        texts=["paracetamol side effects",
               "Paracetamol can cause liver disorders"],
        label=1.0
    ),
    InputExample(
        texts=["paracetamol side effects",
               "Paracetamol is sold in pharmacies"],
        label=0.2
    ),
]

train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=16)
train_loss = losses.CosineSimilarityLoss(model)

# Train
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=3,
    warmup_steps=100,
    output_path="./finetuned_medical_model"
)

# Use the fine-tuned model
embeddings = model.encode([
    "abdominal pain after taking medication",
    "adverse drug effects"
])

Preparing a quality dataset

Here lies the real obstacle: nobody has a thousand annotated pairs lying around. The workaround is to build them from what you already have, namely your documents. You submit each document to the LLM asking for the questions it answers precisely, and you get question-document pairs without any manual annotation. A temperature of 0.7 ensures the three questions generated for the same document aren’t three variants of the same sentence, and lstrip cleans up the numbering the model adds spontaneously.

def generate_training_pairs(
    documents: list[str],
    n_per_doc: int = 3
) -> list[dict]:
    """Generate query-document pairs for training."""
    pairs = []

    for doc in documents:
        response = client.chat.completions.create(
            model="gpt-5.6-terra",
            messages=[{
                "role": "user",
                "content": (
                    f"Here is a document:\n\n{doc[:500]}\n\n"
                    f"Generate {n_per_doc} questions that this document "
                    f"answers precisely. One question per line."
                )
            }],
            temperature=0.7
        )
        questions = response.choices[0].message.content.strip().split("\n")

        for q in questions[:n_per_doc]:
            q = q.strip().lstrip("0123456789.-) ")
            if q:
                pairs.append({"query": q, "document": doc})

    return pairs

Evaluating fine-tuning

Training with decreasing loss proves nothing: the only question that matters is whether search has improved. The function below measures recall@5 on the same set of queries, with and without the adapter, simply by passing adapter=None to get the baseline. Both the corpus and queries pass through the projection when provided, which is the condition for both vectors to live in the same space.

Always run this comparison before deploying. It happens regularly that an adapter trained on too few triplets degrades recall, and without measurement you’d ship a regression thinking you shipped an improvement.

def evaluate_before_after(
    queries: list[str],
    relevant_docs: list[list[str]],
    corpus: list[str],
    adapter=None
):
    """Compare performance before and after fine-tuning."""
    resp = client.embeddings.create(
        input=corpus, model="text-embedding-3-large"
    )
    corpus_embs = np.array(
        [d.embedding for d in sorted(resp.data, key=lambda x: x.index)]
    )

    if adapter:
        corpus_embs_adapted = adapter(
            torch.tensor(corpus_embs, dtype=torch.float32)
        ).detach().numpy()
    else:
        corpus_embs_adapted = corpus_embs

    recall_scores = []
    for query, relevants in zip(queries, relevant_docs):
        q_resp = client.embeddings.create(
            input=query, model="text-embedding-3-large"
        )
        q_emb = np.array(q_resp.data[0].embedding)

        if adapter:
            q_emb = adapter(
                torch.tensor(q_emb, dtype=torch.float32).unsqueeze(0)
            ).detach().numpy()[0]

        scores = corpus_embs_adapted @ q_emb
        top_5 = set(np.argsort(scores)[-5:])
        relevant_ids = set(
            i for i, d in enumerate(corpus) if d in relevants
        )
        recall = len(top_5 & relevant_ids) / len(relevant_ids)
        recall_scores.append(recall)

    return np.mean(recall_scores)

Summary

• Embedding fine-tuning improves performance on specialized domains
• A linear adapter is the simplest method and often sufficient
• Sentence Transformers enables deeper fine-tuning of the full model
• A minimum of 1,000 annotated pairs is recommended
• Always evaluate before and after to validate the gain