← Back to DevBytes

Embedding Model Fine-Tuning: Adapting to Domain-Specific Vocabulary

Embedding Model Fine-Tuning: Adapting to Domain-Specific Vocabulary

Embedding models are the backbone of modern retrieval-augmented generation (RAG) pipelines, semantic search systems, and recommendation engines. Off-the-shelf models like sentence-transformers/all-MiniLM-L6-v2 or BAAI/bge-large-en are trained on general-purpose corpora and perform well on everyday language. However, when your application operates in a specialized domain — medical records, legal contracts, financial filings, or internal engineering wikis — these models often struggle to capture the semantic nuance of domain-specific vocabulary. Fine-tuning embeddings on in-domain data bridges this gap, producing vectors that better reflect the relationships between concepts your users actually care about.

What Embedding Fine-Tuning Actually Means

Fine-tuning an embedding model means taking a pre-trained encoder (typically a BERT-family transformer) and continuing its training on a dataset that reflects your target domain. The objective is to adjust the model's representation space so that semantically similar in-domain texts land close together, while dissimilar texts are pushed apart. Unlike generative fine-tuning, which optimizes next-token prediction, embedding fine-tuning optimizes a contrastive or triplet loss that directly shapes the geometry of the vector space.

The most common training signals are:

Why Domain Adaptation Matters

General embedding models are trained on web text, Wikipedia, and similar broad sources. They have never seen your company's product codes, your clinical abbreviations, or the specific way your legal team phrases indemnification clauses. This causes three concrete problems:

Fine-tuning addresses all three by reshaping the embedding space around your domain's actual semantic structure. Empirically, teams report 15–40% improvements in retrieval metrics like nDCG@10 and recall@k after domain adaptation.

Approaches to Fine-Tuning

1. Supervised Contrastive Learning

The most reliable approach when you have labeled query-document pairs. You train with MultipleNegativesRankingLoss, which treats each pair as positive and all other pairs in the batch as negatives. This is efficient and works well with as few as a few thousand pairs.

2. Triplet Loss with Hard Negatives

When you can annotate hard negatives — documents that look relevant but aren't — triplet loss (anchor, positive, negative) produces sharper representations. This is the approach used by high-performing models like E5 and BGE.

3. Unsupervised Domain Adaptation

If you lack labeled pairs but have a corpus of domain text, you can use techniques like TSDAE (Transformer-based Denoising Auto-Encoder) or contrastive tension. These are weaker than supervised methods but still meaningfully outperform off-the-shelf models.

Practical Implementation

Below is a complete, runnable example using the sentence-transformers library. We'll fine-tune all-MiniLM-L6-v2 on synthetic medical QA pairs to demonstrate the workflow.

Installing Dependencies

pip install sentence-transformers datasets torch

Preparing the Training Data

Your dataset should contain at minimum a query and a positive column. Hard negatives are optional but recommended.

import pandas as pd
from datasets import Dataset

# Synthetic medical domain pairs
data = [
    {"query": "What is myocardial infarction?", "positive": "Myocardial infarction, commonly known as a heart attack, occurs when blood flow to the heart muscle is blocked."},
    {"query": "Symptoms of pulmonary embolism", "positive": "Pulmonary embolism presents with sudden shortness of breath, chest pain, and rapid heart rate."},
    {"query": "How is type 2 diabetes diagnosed?", "positive": "Type 2 diabetes is diagnosed through fasting blood glucose tests, HbA1c levels, and oral glucose tolerance tests."},
    {"query": "Treatment for community-acquired pneumonia", "positive": "Community-acquired pneumonia is typically treated with empiric antibiotics such as macrolides or fluoroquinolones."},
    {"query": "What causes atrial fibrillation?", "positive": "Atrial fibrillation is caused by abnormal electrical signals in the atria, often linked to hypertension and coronary disease."},
    {"query": "Define chronic kidney disease staging", "positive": "Chronic kidney disease is staged from G1 to G5 based on glomerular filtration rate and kidney damage markers."},
    {"query": "Mechanism of action of metformin", "positive": "Metformin reduces hepatic glucose production and improves insulin sensitivity through AMPK pathway activation."},
    {"query": "Risk factors for stroke", "positive": "Major stroke risk factors include hypertension, atrial fibrillation, diabetes, smoking, and hyperlipidemia."},
]

df = pd.DataFrame(data)
dataset = Dataset.from_pandas(df)

Training with MultipleNegativesRankingLoss

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

# Load the pre-trained base model
model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")

# Convert to InputExample objects
train_examples = [
    InputExample(texts=[row["query"], row["positive"]])
    for _, row in df.iterrows()
]

# DataLoader with a reasonable batch size
# Larger batches give more in-batch negatives, improving contrastive signal
train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=8)

# MultipleNegativesRankingLoss uses every other pair in the batch as a negative
train_loss = losses.MultipleNegativesRankingLoss(model)

# Fine-tune
model.fit(
    train_objectives=[(train_dataloader, train_loss)],
    epochs=3,
    warmup_steps=10,
    show_progress_bar=True,
    output_path="./fine-tuned-medical-embeddings",
)

print("Fine-tuning complete. Model saved to ./fine-tuned-medical-embeddings")

Adding Hard Negatives with Triplet Loss

For stronger results, annotate hard negatives — documents that share vocabulary with the query but are not the correct answer.

from sentence_transformers import InputExample, losses
from sentence_transformers.evaluation import TripletEvaluator

triplet_data = [
    InputExample(
        texts=[
            "What is myocardial infarction?",  # anchor
            "Myocardial infarction is the death of heart muscle due to ischemia.",  # positive
            "Myocarditis is inflammation of the heart muscle, often viral in origin.",  # hard negative
        ]
    ),
    InputExample(
        texts=[
            "Treatment for community-acquired pneumonia",
            "Community-acquired pneumonia is treated with macrolides or fluoroquinolones.",
            "Hospital-acquired pneumonia requires broad-spectrum antibiotics like vancomycin.",
        ]
    ),
    # ... more triplets
]

triplet_dataloader = DataLoader(triplet_data, shuffle=True, batch_size=8)
triplet_loss = losses.TripletLoss(
    model,
    distance_metric=losses.TripletDistanceMetric.COSINE,
    triplet_margin=0.5,
)

model.fit(
    train_objectives=[(triplet_dataloader, triplet_loss)],
    epochs=5,
    warmup_steps=10,
    output_path="./fine-tuned-medical-embeddings-triplet",
)

Evaluating the Fine-Tuned Model

Always evaluate against a held-out set using InformationRetrievalEvaluator, which computes nDCG, MRR, and recall@k.

from sentence_transformers import SentenceTransformer
from sentence_transformers.evaluation import InformationRetrievalEvaluator

# Load fine-tuned model
fine_tuned = SentenceTransformer("./fine-tuned-medical-embeddings")

# Evaluation set: queries, corpus, and relevance mapping
queries = {
    "q1": "What causes hypertension?",
    "q2": "How is sepsis treated?",
    "q3": "Side effects of statins",
}

corpus = {
    "d1": "Hypertension is caused by increased peripheral vascular resistance and fluid retention.",
    "d2": "Sepsis treatment involves early antibiotics, fluid resuscitation, and vasopressors.",
    "d3": "Statins may cause myalgia, elevated liver enzymes, and rarely rhabdomyolysis.",
    "d4": "Hypertension management includes ACE inhibitors and lifestyle modification.",
    "d5": "Septic shock is defined by persistent hypotension despite fluid resuscitation.",
}

relevant_docs = {
    "q1": {"d1", "d4"},
    "q2": {"d2", "d5"},
    "q3": {"d3"},
}

evaluator = InformationRetrievalEvaluator(
    queries=queries,
    corpus=corpus,
    relevant_docs=relevant_docs,
    name="medical-eval",
    show_progress_bar=True,
)

results = evaluator(fine_tuned)
print(f"nDCG@10: {evaluator.ndcg_at_k['medical-eval'][10]:.4f}")
print(f"Recall@5: {evaluator.recall_at_k['medical-eval'][5]:.4f}")
print(f"MRR@10: {evaluator.mrr_at_k['medical-eval'][10]:.4f}")

Using the Fine-Tuned Model in a RAG Pipeline

from sentence_transformers import SentenceTransformer
import numpy as np

model = SentenceTransformer("./fine-tuned-medical-embeddings")

# Index your document corpus
documents = [
    "Myocardial infarction requires immediate reperfusion therapy.",
    "Type 2 diabetes management focuses on glycemic control and cardiovascular risk reduction.",
    "Atrial fibrillation increases stroke risk fivefold and requires anticoagulation.",
]

doc_embeddings = model.encode(documents, normalize_embeddings=True, convert_to_numpy=True)

# Encode a user query
query = "How do you treat a heart attack?"
query_embedding = model.encode([query], normalize_embeddings=True, convert_to_numpy=True)

# Retrieve top-k by cosine similarity
scores = query_embedding @ doc_embeddings.T
top_k_idx = np.argsort(scores[0])[::-1][:3]

for idx in top_k_idx:
    print(f"Score: {scores[0][idx]:.4f} | Document: {documents[idx]}")

Best Practices

Conclusion

Fine-tuning embedding models for domain-specific vocabulary is one of the highest-leverage improvements you can make to a retrieval or RAG system. The workflow is straightforward: collect or generate in-domain pairs, mine hard negatives, train with a contrastive loss, and evaluate rigorously on held-out data. Even modest fine-tuning — a few thousand pairs and three to five epochs — can produce measurable gains in retrieval quality that directly translate into better downstream LLM answers. By treating your embedding model as a first-class, versioned component of your system rather than a static dependency, you build retrieval infrastructure that evolves with your domain and your users' needs.

— Ad —

Google AdSense will appear here after approval

← Back to all articles