9.1 Embeddings & Representation Learning¶

Dense text embeddings are the connective tissue of modern information retrieval. Without them, matching a user’s natural-language question to a relevant document requires either exact keyword overlap or hand-crafted rules — both brittle in practice. With them, a 768-dimensional vector encodes semantic meaning so precisely that “What is the boiling point of water?” and “At what temperature does H₂O vaporize?” land microseconds apart in vector space. This chapter develops every piece of the pipeline: why dense vectors work, how they are trained with contrastive objectives, how the field settled on architectures like bi-encoders with pooled transformer representations, and how to evaluate them rigorously with MTEB. It ends with a direct bridge to the RAG systems that consume these embeddings at query time.

If you are looking for the token embeddings inside a transformer (the learned lookup table that maps token IDs to vectors at the input layer), that is covered in Embeddings & The Input Pipeline. This chapter is about sentence- and document-level representation: a single fixed-size vector that summarizes an entire passage.

Why Dense Representations?¶

The classical information retrieval (IR) approach is sparse retrieval: represent each document as a high-dimensional bag-of-words vector (TF-IDF or BM25), then compute dot products at query time. BM25 remains a strong baseline, but it fails on vocabulary mismatch — if the query says “automobile” and the document says “car,” zero overlap means zero score.

Dense retrieval projects documents and queries into a low-dimensional (128–1024 dimensions) continuous space trained so that semantically related pairs are close regardless of word choice. The trade-off is computation: sparse retrieval can be done with an inverted index that touches only non-zero entries; dense retrieval requires either an exhaustive scan or an approximate nearest neighbor (ANN) index (see Vector Databases & Approximate Nearest Neighbor Search).

The two paradigms are not mutually exclusive. Hybrid search — combining BM25 scores with dense scores via reciprocal rank fusion or learned weighting — consistently outperforms either alone. That is covered in Chunking, Reranking & Hybrid Search.

The Bi-Encoder Architecture¶

Dense embedding models overwhelmingly use the bi-encoder design:

Both query and document are encoded independently by the same (or separate) transformer. The final similarity score is computed as a dot product or cosine similarity at retrieval time. This independence is the key property that enables pre-computation: all document embeddings can be computed offline and stored in an ANN index. Only the query needs to be embedded at inference time.

The cross-encoder alternative runs both query and document through the same forward pass with joint attention, producing higher-quality relevance scores but at \(O(n)\) inference cost per document. Cross-encoders are used as rerankers on a short candidate list; they cannot be the first-stage retriever. See Chunking, Reranking & Hybrid Search for cross-encoder reranking.

From Token Representations to Sentence Embeddings: Pooling¶

A transformer operating on a sequence of \(n\) tokens produces \(n\) contextual vectors, each of shape \(d\). We need a single vector. The pooling layer performs this reduction.

CLS Pooling¶

BERT-style models prepend a special [CLS] token. After the final layer, the vector at position 0 is used as the sequence representation. BERT was pretrained with a next-sentence prediction (NSP) objective that trained [CLS] to carry sentence-level signal. In practice, raw BERT CLS embeddings are mediocre for semantic similarity — the representation space is anisotropic (vectors cluster near a cone, not spread across the hypersphere). Models like SBERT fine-tune the encoder specifically for sentence embedding, dramatically improving CLS quality.

Mean Pooling¶

Average the final-layer token representations, excluding padding tokens:

where \(m_i \in \{0,1\}\) is the attention mask and \(\mathbf{h}_i\) is the hidden state at position \(i\). Mean pooling is empirically stronger than CLS pooling for most bi-encoder tasks, and it is the default in models like sentence-transformers (Reimers & Gurevych, 2019).

Max Pooling and Weighted Pooling¶

Max pooling takes the element-wise maximum across token positions — capturing the strongest activation for each latent feature. It is less common in production. Weighted pooling assigns position or attention-based weights before averaging; some models learn these weights as part of fine-tuning.

import torch
import torch.nn.functional as F


def mean_pool(token_embeddings: torch.Tensor,
              attention_mask: torch.Tensor) -> torch.Tensor:
    """
    Compute attention-mask-weighted mean of token embeddings.

    Args:
        token_embeddings: shape (batch, seq_len, hidden_dim)
        attention_mask:   shape (batch, seq_len), 1 for real tokens

    Returns:
        sentence_embeddings: shape (batch, hidden_dim)
    """
    # Expand mask to hidden_dim so we can multiply element-wise
    mask_expanded = attention_mask.unsqueeze(-1).float()  # (B, L, 1)

    # Zero out padding positions
    sum_embeddings = (token_embeddings * mask_expanded).sum(dim=1)  # (B, H)

    # Count real tokens per sample (clamp to avoid division by zero)
    sum_mask = mask_expanded.sum(dim=1).clamp(min=1e-9)             # (B, 1)

    return sum_embeddings / sum_mask  # (B, H)


def normalize_embeddings(embeddings: torch.Tensor) -> torch.Tensor:
    """L2-normalize so dot product == cosine similarity."""
    return F.normalize(embeddings, p=2, dim=-1)

After pooling, embeddings are typically L2-normalized so that dot products equal cosine similarities, which lie in \([-1, 1]\) and are symmetric and bounded.

Contrastive Learning & the InfoNCE Objective¶

The core training objective for dense embedding models is contrastive learning: pull representations of semantically similar pairs together, push dissimilar ones apart. The dominant loss function is InfoNCE (Information Noise-Contrastive Estimation), also known as the NT-Xent loss in the SimCLR literature.

InfoNCE Derivation¶

Given a mini-batch of \(N\) (query, positive document) pairs \(\{(q_i, d_i^+)\}_{i=1}^N\), the model produces normalized embeddings \(\mathbf{e}_{q_i}\) and \(\mathbf{e}_{d_j}\) for all queries and all documents. The InfoNCE loss treats all \(N-1\) other documents in the batch as in-batch negatives for query \(i\):

where \(\tau > 0\) is the temperature hyperparameter. The loss is a cross-entropy over \(N\) classes: class \(i\) is the correct document for query \(i\). Minimizing it encourages \(\mathbf{e}_{q_i} \cdot \mathbf{e}_{d_i^+}\) to be the maximum among all \(N\) dot products in row \(i\).

Why temperature matters. A small \(\tau\) (e.g., 0.01–0.1) sharpens the softmax, creating a steep gradient signal and forcing the model to distinguish very fine-grained differences. A large \(\tau\) (near 1.0) makes the loss soft and nearly uniform. Typical values: 0.05–0.1 for image encoders (SimCLR used 0.07), 0.02–0.05 for text retrieval.

Scaling with batch size. With batch size \(N\), each query sees \(N-1\) negatives. Doubling \(N\) roughly doubles the number of negatives per query, which empirically improves representation quality — this is why contrastive learning loves large batches. DPR (Karpukhin et al., 2020) used batch size 128; modern models often train with effective batch sizes in the thousands via gradient accumulation.

Full In-Batch Negative Loss Implementation¶

import torch
import torch.nn.functional as F


def infonce_loss(
    query_emb: torch.Tensor,   # (B, D), already L2-normalized
    doc_emb: torch.Tensor,     # (B, D), already L2-normalized
    temperature: float = 0.05,
) -> torch.Tensor:
    """
    Symmetric InfoNCE loss with in-batch negatives.

    Each query is matched to the corresponding document (diagonal).
    All other documents in the batch serve as negatives.

    Args:
        query_emb:   Normalized query embeddings.
        doc_emb:     Normalized positive-document embeddings.
        temperature: Softmax temperature (lower = sharper contrast).

    Returns:
        Scalar loss (mean over batch).
    """
    B = query_emb.size(0)

    # Similarity matrix: (B, B)
    # sim[i][j] = cosine_sim(query_i, doc_j)
    sim = torch.matmul(query_emb, doc_emb.T) / temperature  # (B, B)

    # Targets: each query i matches document i (diagonal)
    labels = torch.arange(B, device=query_emb.device)

    # Query-to-document direction
    loss_q2d = F.cross_entropy(sim, labels)

    # Document-to-query direction (symmetric)
    loss_d2q = F.cross_entropy(sim.T, labels)

    return (loss_q2d + loss_d2q) / 2


# ---- Quick sanity check ----
if __name__ == "__main__":
    torch.manual_seed(42)
    B, D = 4, 768

    # Random normalized embeddings
    q = F.normalize(torch.randn(B, D), dim=-1)
    d = F.normalize(torch.randn(B, D), dim=-1)

    loss = infonce_loss(q, d, temperature=0.05)
    print(f"Loss (random init): {loss.item():.4f}")   # ~ log(B) ≈ 1.386

    # Perfect embeddings: q[i] == d[i]
    d_perfect = q.clone()
    loss_perfect = infonce_loss(q, d_perfect, temperature=0.05)
    print(f"Loss (perfect align): {loss_perfect.item():.6f}")  # ≈ 0.0

Hard Negatives: Going Beyond In-Batch Sampling¶

In-batch negatives are easy to implement but are often too easy — randomly sampled passages are unlikely to be genuinely confusable with the query. Hard negatives are passages that are superficially relevant but do not actually answer the query.

Mining Hard Negatives¶

BM25-mined negatives. Run BM25 retrieval on the query; take high-ranked documents that are not the gold positive. These share vocabulary with the query but differ in meaning.

Dense-mined negatives. Use an earlier checkpoint of the embedding model to retrieve top-\(k\) passages; use non-positive ones as negatives. This “ANN negative mining” strategy, used in DPR and ANCE (Xiong et al., 2021), iteratively updates the negatives as the model improves. Training with the model’s own hard negatives dramatically accelerates convergence.

Cross-encoder filtered negatives. Apply a high-quality cross-encoder reranker to a large candidate set, then use passages that the cross-encoder scores as irrelevant as negatives. These are the hardest and most informative.

Triplet Loss vs. InfoNCE¶

An alternative to InfoNCE is the triplet loss with a margin \(m\):

where \(s\) is cosine similarity. This only considers one negative per query per step. In practice, InfoNCE with large \(N\) outperforms triplet loss because it simultaneously considers many negatives, providing a richer gradient signal.

False Negatives and Denoised Negatives¶

A subtle failure mode: the mining procedure mislabels genuinely relevant documents as negatives (false negatives). The gold label says “this passage is not the answer,” but it is — and the gradient actively pushes a correct pair apart. This is the dominant source of noise in hard-negative training, and it gets worse the harder you mine, because the highest-ranked mined candidates are precisely the ones most likely to be unlabelled positives. RocketQA (Qu et al., 2021) was the paper that made this concrete and gave the standard fix. Denoising strategies:

• Skip the very top of the ranking. Do not take ranks 1–10 as negatives; sample from ranks ~10–50 instead. Cheap and surprisingly effective.
• Cross-encoder filtering. Score each mined candidate with a strong reranker and drop anything above a confidence threshold (RocketQA’s “denoised hard negatives”).
• Absolute/relative score caps. Reject a candidate whose similarity to the query exceeds some max_score, or that is within a margin of the positive’s own score.
• Soft labels. Instead of a hard 0/1 target, distil the teacher’s score distribution over candidates (the approach BGE-M3 uses for self-knowledge distillation).

These are not just theory — sentence_transformers.util.mine_hard_negatives, shown below, exposes exactly these knobs (range_min, max_score, margin, cross_encoder).

Architecture: Sentence-Transformers and the BERT-Based Bi-Encoder¶

The canonical sentence embedding architecture is SBERT (Sentence-BERT, Reimers & Gurevych, 2019), which fine-tuned a BERT-base backbone on sentence pairs using contrastive/siamese networks. Modern models follow the same structure but with stronger backbones and larger training sets.

Modern choices for the backbone include:

By 2026 the frontier of the MTEB/MMTEB leaderboard is dominated by these LLM-based encoders — Qwen3-Embedding (built on the Qwen3 foundation models, Apache-2.0) is the leading fully-open family, while Google’s Gemini Embedding leads among proprietary API models. The BERT-scale models above remain the pragmatic default when latency and memory dominate.

Training Data Pipelines¶

Modern embedding models are trained in two or three stages:

1. Weakly supervised pre-finetuning on hundreds of millions of (query, passage) pairs from heterogeneous web sources: QA pairs, title-body pairs from web pages, Reddit question-reply threads, etc.
2. Supervised fine-tuning on curated high-quality datasets (Natural Questions, MS MARCO, SNLI, STSb, etc.) with hard negatives.
3. Task-specific fine-tuning for specialized domains (legal, biomedical, code).

Where do the pairs come from if you have none? This is the question that blocks most people trying to adapt an embedder to their own corpus, and it has three standard answers, in increasing order of cost:

• Structural pairs, free from the corpus itself. Title ↔ body, section heading ↔ paragraph, docstring ↔ function, question ↔ accepted answer, citation context ↔ cited abstract. This is exactly what stage 1 above harvests from the open web, and any structured corpus has some version of it.
• Self-supervised pairs from a single sentence. SimCSE (Gao et al., EMNLP 2021) encodes the same sentence twice with dropout active; the two different dropout masks give two different vectors that InfoNCE treats as a positive pair, with the rest of the batch as negatives. No labels at all, and it recovers most of the anisotropy fix. TSDAE (denoising autoencoding of corrupted sentences) is the other common unsupervised option; in sentence-transformers these are losses.MultipleNegativesRankingLoss over duplicated sentences and losses.DenoisingAutoEncoderLoss.
• LLM-generated synthetic queries. Take each passage, prompt an instruct model with “write the question this passage answers,” and keep the pair. This is the Doc2Query / InPars (Bonifacio et al., 2022) / Promptagator (Dai et al., 2023) line of work, and it is how E5-mistral-7b-instruct built training data for tasks with no public dataset. A few tens of thousands of synthetic pairs over your own documents is usually enough to beat a general-purpose model on your own retrieval slice — and it is cheap enough to do in an afternoon.

Whichever source you use, mine hard negatives afterwards against your real corpus; the negatives matter more than the positives.

The sentence-transformers library (Reimers & Gurevych, now maintained at Hugging Face) is the practical starting point for both training and inference.

from sentence_transformers import SentenceTransformer

# Load a pre-trained backbone (already fine-tuned for sentence embedding)
model = SentenceTransformer("BAAI/bge-small-en-v1.5")

# Encode a batch of texts -- returns numpy array by default
texts = [
    "What is the capital of France?",
    "Paris is the capital and most populous city of France.",
    "The Eiffel Tower is located in Paris.",
]
embeddings = model.encode(texts, normalize_embeddings=True)

print(f"Embedding shape: {embeddings.shape}")  # (3, 384)

# Cosine similarity between query and each document
import numpy as np
query_emb = embeddings[0]
doc_embs  = embeddings[1:]
scores = doc_embs @ query_emb  # dot product of L2-normalized = cosine sim
print(f"Scores: {scores}")     # e.g. [0.71, 0.58] — first doc more relevant

Fine-Tuning for Real: the sentence-transformers Trainer¶

Everything this chapter derives from scratch — InfoNCE with in-batch negatives, mined hard negatives, gradient caching, Matryoshka losses — exists as a first-class API in sentence-transformers v3+, which wraps the Hugging Face Trainer. This is what you should actually run; the from-scratch loop later in the chapter is for understanding what it does.

# pip install "sentence-transformers>=3.1" datasets accelerate faiss-cpu
from datasets import Dataset
from sentence_transformers import (
    SentenceTransformer,
    SentenceTransformerTrainer,
    SentenceTransformerTrainingArguments,
    losses,
)
from sentence_transformers.util import mine_hard_negatives

model = SentenceTransformer("BAAI/bge-base-en-v1.5")

# Step 1: (anchor, positive) pairs. Column ORDER is what matters to the loss,
# not the column names. This is the output of the data pipeline above.
pairs = Dataset.from_dict({
    "anchor":   ["What is the capital of France?", "Who wrote Hamlet?"],
    "positive": ["Paris is the capital of France.", "Hamlet was written by Shakespeare."],
})

# Step 2: mine hard negatives against the corpus with the current model
# (this is exactly the ANCE loop, with the RocketQA denoising knobs).
train_ds = mine_hard_negatives(
    pairs,
    model,
    num_negatives=4,
    range_min=10,          # skip ranks 1-9: most likely to be FALSE negatives
    range_max=50,          # sample from ranks 10-50: hard but probably wrong
    max_score=0.8,         # reject anything suspiciously similar to the query
    margin=0.05,           # ...and anything within 0.05 of the positive's score
    sampling_strategy="top",
    output_format="n-tuple",   # -> anchor, positive, negative_1 ... negative_4
    batch_size=256,
    use_faiss=True,
)

# Step 3: loss. MultipleNegativesRankingLoss IS the InfoNCE we derived:
# it uses the positive plus every other row's positive AND negatives as the
# softmax denominator. `scale` is 1/temperature, so 20.0 == tau = 0.05.
base_loss = losses.MultipleNegativesRankingLoss(model, scale=20.0)

# Wrap it so the same batch is scored at every Matryoshka prefix.
loss = losses.MatryoshkaLoss(model, base_loss,
                             matryoshka_dims=[768, 256, 128, 64])

args = SentenceTransformerTrainingArguments(
    output_dir="bge-base-mydomain",
    num_train_epochs=1,
    per_device_train_batch_size=128,   # bigger batch == more negatives
    learning_rate=2e-5,
    warmup_ratio=0.1,
    bf16=True,
    # Guarantees no two rows in a batch share a text, so in-batch negatives
    # are never accidental duplicates of the positive.
    batch_sampler="no_duplicates",
)

trainer = SentenceTransformerTrainer(
    model=model, args=args, train_dataset=train_ds, loss=loss,
)
trainer.train()
model.save_pretrained("bge-base-mydomain/final")

Two swaps cover most scaling problems. When the batch no longer fits in memory, replace MultipleNegativesRankingLoss with losses.CachedMultipleNegativesRankingLoss(model, scale=20.0, mini_batch_size=32) — that is the GradCache algorithm (discussed below) packaged as a drop-in, and it lets a single GPU train at batch sizes in the thousands. When you want a stronger negative filter, pass cross_encoder=CrossEncoder("BAAI/bge-reranker-v2-m3") to mine_hard_negatives so candidates are rescored by a reranker before being accepted. BAAI’s FlagEmbedding toolkit offers the same pipeline with BGE-specific recipes, and swift/LLaMA-Factory can fine-tune the multi-billion-parameter LLM-based encoders with LoRA when a full fine-tune is out of budget.

Matryoshka Representation Learning (MRL)¶

A practical limitation of fixed-dimension embeddings: the 1024-dim model is always 1024-dim, even when you want to store billions of vectors at lower cost or serve low-latency queries that do not need full precision.

Matryoshka Representation Learning (Kusupati et al., 2022, NeurIPS) trains a single model to produce embeddings that are useful at multiple prefix dimensions simultaneously. Like Russian nesting dolls, the first 64 dimensions alone form a useful embedding, as do the first 128, 256, 512, and 1024.

How MRL Training Works¶

Add a loss term for each dimensionality granularity in a set \(\mathcal{M} = \{32, 64, 128, 256, 512, d\}\):

where \(\mathbf{e}_{:m}\) denotes the first \(m\) dimensions of the full embedding and \(\lambda_m\) are weighting coefficients (often uniform). The total loss is the weighted sum of contrastive losses at each granularity; gradients flow through all prefix slices simultaneously.

The effect: the first few dimensions encode the most discriminative signal; later dimensions refine it. At test time you can truncate to any supported size and re-normalize.

def matryoshka_infonce_loss(
    query_emb: torch.Tensor,      # (B, D) full dimension
    doc_emb:   torch.Tensor,      # (B, D) full dimension
    dims:      list[int] = None,  # prefix dimensions to train
    temperature: float = 0.05,
    weights:   list[float] = None,
) -> torch.Tensor:
    """
    Matryoshka contrastive loss: InfoNCE at each prefix dimension.
    Gradients flow through all prefix slices simultaneously.
    """
    D = query_emb.size(-1)
    if dims is None:
        dims = [32, 64, 128, 256, D]
    if weights is None:
        weights = [1.0] * len(dims)

    total_loss = torch.tensor(0.0, device=query_emb.device)

    for dim, w in zip(dims, weights):
        # Slice to prefix dimension and re-normalize
        q_slice = F.normalize(query_emb[:, :dim], dim=-1)
        d_slice = F.normalize(doc_emb[:, :dim], dim=-1)

        # Standard InfoNCE at this granularity
        loss_at_dim = infonce_loss(q_slice, d_slice, temperature)
        total_loss = total_loss + w * loss_at_dim

    return total_loss / sum(weights)

Dimension truncation composes with precision reduction, and the two multiply. sentence_transformers.quantization.quantize_embeddings(emb, precision="int8" | "binary") casts a float32 embedding to one byte or one bit per dimension; binary embeddings are compared with Hamming distance and, on MRL-trained models, typically retain a large majority of retrieval quality while cutting memory 32×. Stack that on a 768→128 MRL truncation and a 100M-vector index drops from ~307 GB to under 2 GB — cheap enough to hold in RAM, with a full-precision rescoring pass over the top few hundred candidates to recover the lost recall. The index-side machinery for this (product quantization, rescoring, the recall/memory triangle) is Vector Databases & Approximate Nearest Neighbor Search.

MRL is now standard in leading open embedding models. OpenAI’s text-embedding-3 family supports variable dimensions via MRL-style training. BGE-M3 (Chen et al., 2024) combines MRL with multi-lingual and multi-granularity retrieval, and the 2025-era standard-bearers (Qwen3-Embedding, Gemini Embedding) expose a full MRL range — Qwen3-Embedding, for instance, lets you request any dimension from 32 to 4096 from the same model.

Instruction-Following Embeddings¶

Generic embedding models treat every query identically. But “software developer” is relevant to both a job-listing query and a biography question — the embedding should reflect the task context. Instruction embeddings (E5-mistral-7b-instruct, Instructor, GTE-Qwen) prepend a natural-language task instruction to the query at inference time:

Represent the following sentence for searching relevant passages: <query>
Represent this sentence for retrieval:
  Task: Given a question, find the most relevant academic paper.
  Input: <query>

The instruction is tokenized and encoded together with the query, so the pooled vector is a function of both. In a bidirectional (BERT-style) encoder every query token attends to the instruction tokens directly; in a decoder-only encoder the causal mask means the instruction sits in the prefix, and the last-token representation has attended over all of it. Either way the same backbone lands the query in a different region of the space depending on the stated task. One operational consequence: the instruction must be applied at index time exactly as the model was trained. E5 and BGE prefix documents with passage: and queries with query:; if you index with the wrong prefix (or none), scores degrade silently — no error, just worse retrieval.

This is especially powerful for LLM-based embedders. Models like E5-mistral-7b-instruct (Wang et al., 2023) — and the 2025 frontier models Qwen3-Embedding and Gemini Embedding that superseded it on the leaderboards — use a decoder-only LLM backbone and take the last-token representation instead of mean pooling (since left-to-right autoregressive models produce context-rich representations at the final token, not the first).

from sentence_transformers import SentenceTransformer

# Instruction-aware embedding model
model = SentenceTransformer("intfloat/e5-large-v2")

# Prefix 'query:' vs 'passage:' tells the model the role
queries = ["query: How do transformers handle long sequences?"]
docs    = [
    "passage: Transformers scale quadratically with sequence length in attention.",
    "passage: The history of the Transformer architecture dates to 2017.",
]

q_emb = model.encode(queries, normalize_embeddings=True)
d_emb = model.encode(docs,    normalize_embeddings=True)

scores = q_emb @ d_emb.T
print(scores)  # [[0.73, 0.41]] — first document correctly ranked higher

The instruction mechanism gives the same frozen backbone very different behaviors for different retrieval tasks without any fine-tuning, a form of lightweight task specification.

Evaluation: MTEB¶

The Massive Text Embedding Benchmark (MTEB, Muennighoff et al., 2023) is the standard leaderboard for embedding model evaluation. It covers 58 datasets across 8 task types:

MTEB’s retrieval sub-tasks (the BEIR benchmark, Thakur et al., 2021) are especially important for RAG use cases. Models are evaluated zero-shot on held-out domains, measuring generalization rather than in-distribution performance.

The original MTEB has since grown into MMTEB (Massive Multilingual Text Embedding Benchmark, Enevoldsen et al., 2025), a community-driven expansion covering over 500 quality-controlled tasks across 250+ languages, plus instruction-following, long-document, and code-retrieval tasks. By 2026 the public multilingual leaderboard on MMTEB is the reference most practitioners consult when picking a model.

Running MTEB Evaluation¶

pip install mteb

import mteb
from sentence_transformers import SentenceTransformer

# Load the model as an MTEB-compatible encoder
model_name = "BAAI/bge-small-en-v1.5"
model = mteb.get_model(model_name)

# Run a single retrieval task
tasks = mteb.get_tasks(tasks=["NFCorpus"])
evaluation = mteb.MTEB(tasks=tasks)
results = evaluation.run(model, output_folder=f"results/{model_name}")

# results contains nDCG@10 and other metrics per task

Rather than enumerating tasks by hand, ask mteb for a named benchmark — mteb.get_benchmark("MTEB(eng, v2)") or the multilingual/code suites — and pass the resulting task list to mteb.MTEB(...); the library also caches results so a re-run only evaluates what is missing. For RAG system design, prioritize the BEIR retrieval subset because it directly predicts retrieval quality in your pipeline. Semantic textual similarity (STS) tasks measure something slightly different — fine-grained similarity rather than ranking relevance.

The leaderboard is a shortlist generator, not a decision. Models with a two-point MTEB gap routinely swap places on a specific corpus, and the top of the leaderboard is partly a story about which public benchmarks a model’s training mix overlapped. The decision comes from 100–300 labelled (query, gold-passage) pairs from your own domain, scored with Recall@\(k\) and nDCG@\(k\) — a day of annotation that is worth more than any leaderboard column. You can register that set as a custom AbsTaskRetrieval in mteb and get the same metrics and result format as the public tasks; see Chunking, Reranking & Hybrid Search for stage-by-stage retrieval measurement.

Practical Training Recipe¶

Here is a concrete end-to-end training loop for a fine-tuned bi-encoder, illustrating all the moving pieces:

import torch
import torch.nn.functional as F
from torch.utils.data import DataLoader, Dataset
from transformers import AutoTokenizer, AutoModel
from dataclasses import dataclass
from typing import Optional


@dataclass
class BiEncoderConfig:
    model_name: str = "BAAI/bge-base-en-v1.5"
    max_length: int = 512
    embedding_dim: int = 768
    temperature: float = 0.05
    batch_size: int = 64
    learning_rate: float = 2e-5
    warmup_steps: int = 1000
    matryoshka_dims: Optional[list] = None  # e.g. [128, 256, 768]


class EmbeddingModel(torch.nn.Module):
    def __init__(self, config: BiEncoderConfig):
        super().__init__()
        self.config = config
        self.encoder = AutoModel.from_pretrained(config.model_name)
        self.tokenizer = AutoTokenizer.from_pretrained(config.model_name)

    def encode(self, texts: list[str]) -> torch.Tensor:
        """Tokenize, encode, mean-pool, and L2-normalize a list of texts."""
        encoded = self.tokenizer(
            texts,
            max_length=self.config.max_length,
            padding=True,
            truncation=True,
            return_tensors="pt",
        ).to(next(self.parameters()).device)

        # Forward pass through the transformer backbone
        outputs = self.encoder(**encoded)

        # Mean pool over token dimension (excluding padding)
        embeddings = mean_pool(outputs.last_hidden_state, encoded["attention_mask"])

        # L2 normalize
        return F.normalize(embeddings, p=2, dim=-1)

    def forward(self, queries: list[str], docs: list[str]) -> torch.Tensor:
        """Compute InfoNCE loss for a batch of (query, doc) pairs."""
        q_emb = self.encode(queries)
        d_emb = self.encode(docs)

        if self.config.matryoshka_dims is not None:
            return matryoshka_infonce_loss(
                q_emb, d_emb,
                dims=self.config.matryoshka_dims,
                temperature=self.config.temperature,
            )
        else:
            return infonce_loss(q_emb, d_emb, self.config.temperature)


class PairDataset(Dataset):
    """Simple dataset of (query, positive_document) string pairs."""

    def __init__(self, pairs: list[tuple[str, str]]):
        self.pairs = pairs

    def __len__(self):
        return len(self.pairs)

    def __getitem__(self, idx):
        return self.pairs[idx]


def collate_fn(batch):
    queries, docs = zip(*batch)
    return list(queries), list(docs)


def train_one_epoch(model, loader, optimizer, scheduler, device):
    model.train()
    total_loss = 0.0

    for step, (queries, docs) in enumerate(loader):
        optimizer.zero_grad()

        loss = model(queries, docs)
        loss.backward()

        # Gradient clipping prevents spikes, especially important with
        # large batches and small temperatures
        torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=1.0)

        optimizer.step()
        scheduler.step()

        total_loss += loss.item()
        if step % 100 == 0:
            print(f"  Step {step:5d} | loss {loss.item():.4f} | "
                  f"lr {scheduler.get_last_lr()[0]:.2e}")

    return total_loss / len(loader)


# Example usage (pseudo-code, requires actual data and GPU):
# config = BiEncoderConfig(matryoshka_dims=[128, 256, 768])
# model = EmbeddingModel(config).to("cuda")
# dataset = PairDataset(my_training_pairs)
# loader = DataLoader(dataset, batch_size=config.batch_size,
#                     shuffle=True, collate_fn=collate_fn)
# optimizer = torch.optim.AdamW(model.parameters(), lr=config.learning_rate)
# scheduler = get_linear_schedule_with_warmup(optimizer, config.warmup_steps, total_steps)

Key Training Decisions¶

Gradient caching. With large batch sizes, storing all intermediate activations for both query and document towers can overflow GPU memory. Gradient caching (GradCache, Gao et al., 2021) runs three passes instead of one: (1) encode every sub-batch under torch.no_grad() and keep only the embeddings; (2) compute the contrastive loss on the full set of embeddings and cache \(\partial \mathcal{L}/\partial \mathbf{e}\) for each one — a tiny \((B, D)\) tensor; (3) re-encode each sub-batch with gradients and backprop the cached \(\partial \mathcal{L}/\partial \mathbf{e}\) into the encoder, accumulating. Activation memory is now set by the sub-batch, not the batch, at the cost of roughly one extra forward pass. It is mathematically exact — not an approximation. Reference implementation: luyug/GradCache; in practice you get it for free via losses.CachedMultipleNegativesRankingLoss.

Multi-GPU synchronization. With distributed training, in-batch negatives from other GPUs can be gathered via all-gather before computing the loss, effectively multiplying the batch size by the number of GPUs with no per-GPU memory increase. This “all-gather trick” is used in SimCLR, CLIP, and DPR follow-ups.

Temperature as a learned parameter. CLIP (Radford et al., 2021) learned the log-temperature as a scalar parameter, allowing the model to adapt the sharpness of its similarity distribution during training. This avoids manual tuning and is increasingly common.

Bridging to RAG¶

Dense embeddings are the front door of every Retrieval-Augmented Generation (RAG) system. The full pipeline:

Several design decisions at the embedding layer propagate through the entire system:

Asymmetric query/document length. Queries are short (5–20 tokens); documents are long (100–500 tokens). Some models use separate query and document encoders (DPR-style) or a single encoder with different prefix instructions (E5/BGE). Matching the encoder’s training distribution to your actual query and document lengths matters.

Chunking policy. A document is split into chunks before embedding. The chunk size (128–512 tokens) and overlap determine how much context each embedded unit contains. This is the first point where RAG quality is determined — see Chunking, Reranking & Hybrid Search.

Index refresh latency. When new documents arrive, their embeddings must be computed and inserted into the ANN index. With asynchronous pipelines, there is a window where new content is not yet retrievable. High-throughput embedding inference on GPU batches can reduce this lag to seconds. In production you do not call model.encode in your application process: you run a dedicated embedding server. Hugging Face’s Text Embeddings Inference (TEI) is the standard open-source choice — a Rust/Candle server with continuous batching, Flash-Attention kernels, and an OpenAI-compatible /embed endpoint, deployable as ghcr.io/huggingface/text-embeddings-inference and serving bi-encoders, rerankers, and sequence classifiers alike. infinity is a comparable Python alternative, and both vLLM and SGLang can serve LLM-based encoders through their pooling/embedding APIs, which matters when your embedder is a 7B model (see vLLM: Architecture, PagedAttention & Internals).

Embedding drift. If the embedding model is updated (e.g., fine-tuned on domain data), all stored document embeddings become stale and must be recomputed. This “re-indexing” cost — on the order of millions of API calls or hours of GPU time — is a real operational concern in production RAG.

Token budget vs. retrieved quality. Retrieving \(k=20\) chunks provides more coverage but costs more LLM context tokens. The embedding model’s precision (ratio of relevant chunks among top-\(k\)) directly determines whether context budget is spent on signal or noise. LLMs with long-context windows (see Long-Context Pretraining & Context Extension) relax this tension but do not eliminate it.

For the next stage — indexing these embeddings and querying them efficiently at scale — see Vector Databases & Approximate Nearest Neighbor Search. For the full RAG architecture, see Retrieval-Augmented Generation Architectures.

Where this lands in the capstone. Stack-100M’s narrow research agent uses a dense retriever as one of its two tools, and the accuracy decomposition there shows the retrieval term dominates end-to-end accuracy — which is why the agent pairs a ~100M generator with a 300–600M embedder (EmbeddingGemma-300M or Qwen3-Embedding-0.6B via sentence-transformers, FAISS IndexFlatIP for a corpus this small). Yes, the retriever ends up larger than the model it serves; that is the correct allocation when grounding is load-bearing. See A Narrow Auto-Research Agent: ReAct, Tool-Use & Retrieval by Distillation for the wired-up version, including the hybrid BM25 + dense fusion.

Further Reading¶

• Reimers & Gurevych, Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks, EMNLP 2019 — the foundational bi-encoder fine-tuning paper.
• Karpukhin et al., Dense Passage Retrieval for Open-Domain Question Answering (DPR), EMNLP 2020 — established in-batch negative training for retrieval.
• Xiong et al., Approximate Nearest Neighbor Negative Contrastive Estimation for Dense Text Retrieval (ANCE), ICLR 2021 — ANN-mined hard negatives.
• Gao et al., SimCSE: Simple Contrastive Learning of Sentence Embeddings, EMNLP 2021 — dropout as a data augmentation for contrastive pre-training; the standard unlabelled bootstrap.
• Wang & Isola, Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hypersphere, ICML 2020 — the theory behind why InfoNCE de-anisotropizes an embedding space.
• Qu et al., RocketQA: An Optimized Training Approach to Dense Passage Retrieval, NAACL 2021 — cross-encoder-denoised hard negatives and cross-batch negatives.
• Gao, Yao & Callan, Scaling Deep Contrastive Learning Batch Size under Memory Limited Setups (GradCache), RepL4NLP 2021 — exact large-batch contrastive training on one GPU.
• Thakur et al., BEIR: A Heterogeneous Benchmark for Zero-shot Evaluation of Information Retrieval Models, NeurIPS 2021 — the standard retrieval evaluation suite.
• Muennighoff et al., MTEB: Massive Text Embedding Benchmark, EACL 2023 — comprehensive multi-task embedding leaderboard.
• Kusupati et al., Matryoshka Representation Learning, NeurIPS 2022 — multi-granularity embedding training.
• Wang et al., Text Embeddings by Weakly-Supervised Contrastive Pre-training (E5), 2022 — large-scale weakly supervised bi-encoder pre-training.
• sentence-transformers library (Reimers, Hugging Face) — the practical go-to for embedding model training and inference.
• MTEB Python package (Hugging Face) — easy reproducible evaluation of any embedding model.

Exercises¶

1. (Conceptual.) A colleague proposes replacing your dense bi-encoder retriever with a cross-encoder, arguing that cross-encoders score relevance more accurately. Explain (a) why the bi-encoder, not the cross-encoder, is the design that makes first-stage retrieval over millions of documents tractable, and (b) what specifically breaks if you try to use a cross-encoder as the first-stage retriever. Then describe the standard way both architectures are used together in a production RAG pipeline.

2. (Quantitative.) A transformer produces the following final-layer token representations for a 3-token sequence with hidden dimension \(d = 2\): \(\mathbf{h}_1 = [1, 3]\), \(\mathbf{h}_2 = [3, 1]\), \(\mathbf{h}_3 = [10, 0]\). The attention mask is \(\mathbf{m} = [1, 1, 0]\) (the third token is padding). (a) Compute the mean-pooled sentence embedding using the masked formula from the chapter, then L2-normalize it. (b) Now recompute the mean ignoring the mask (averaging all three tokens) and normalize. © Report the cosine similarity between the two resulting unit vectors and state the lesson.

3. (Quantitative.) You are training with the symmetric in-batch InfoNCE loss from the chapter at temperature \(\tau = 0.1\) and batch size \(B = 2\). The similarity matrix of L2-normalized embeddings is

Compute the symmetric loss \((\mathcal{L}_{q\to d} + \mathcal{L}_{d\to q})/2\) that infonce_loss returns. (Useful values: \(e^2 = 7.389\), \(e^3 = 20.086\), \(e^8 = 2980.96\), \(e^9 = 8103.08\).)

4. (Implementation.) The chapter implements mean_pool but only mentions max pooling in prose. Implement a masked max_pool(token_embeddings, attention_mask) with the same signature and shape conventions as mean_pool (input (B, L, H) and (B, L), output (B, H)). It must take the element-wise maximum over real token positions only. Explain why simply calling .max(dim=1) on the raw tensor is wrong, and how you handle padding correctly.

import torch


def max_pool(token_embeddings: torch.Tensor,
             attention_mask: torch.Tensor) -> torch.Tensor:
    """
    Element-wise max over real (non-padding) token positions.

    Args:
        token_embeddings: shape (batch, seq_len, hidden_dim)
        attention_mask:   shape (batch, seq_len), 1 for real tokens

    Returns:
        sentence_embeddings: shape (batch, hidden_dim)
    """
    # Expand mask to hidden_dim so it broadcasts over features
    mask_expanded = attention_mask.unsqueeze(-1).bool()  # (B, L, 1)

    # Send padding positions to -inf so they never win the max
    masked = token_embeddings.masked_fill(~mask_expanded, float("-inf"))

    # Max over the sequence (token) dimension
    return masked.max(dim=1).values  # (B, H)

5. (Quantitative + conceptual.) You store embeddings for a 100-million-chunk corpus in an ANN index using float32 (4 bytes per dimension). Your model was trained with Matryoshka Representation Learning and produces full 768-dimensional embeddings. (a) Compute the raw vector storage for the full 768-dim embeddings, and for MRL-truncated 128-dim embeddings. (b) What is the memory-savings ratio? © After truncating each stored vector to its first 128 dimensions, why must you re-normalize before doing cosine/dot-product search, and why does MRL make the 128-dim prefix usable at all (unlike truncating an ordinary embedding)?

6. (Implementation, harder.) The chapter’s infonce_loss uses only in-batch negatives. Extend it to infonce_with_hard_negatives(query_emb, pos_emb, neg_emb, temperature) where each query has one positive (pos_emb, shape (B, D)) and its own set of \(K\) mined hard negatives (neg_emb, shape (B, K, D)), all L2-normalized. Build the per-query logit row as [positive, K hard negatives], make the positive the target class, and return the mean cross-entropy. Explain how this connects to the “BM25-mined / ANN-mined negatives” discussion in the chapter, and why hard negatives give a stronger gradient than random in-batch negatives.

import torch
import torch.nn.functional as F


def infonce_with_hard_negatives(
    query_emb: torch.Tensor,   # (B, D), L2-normalized
    pos_emb:   torch.Tensor,   # (B, D), L2-normalized
    neg_emb:   torch.Tensor,   # (B, K, D), L2-normalized
    temperature: float = 0.05,
) -> torch.Tensor:
    """
    InfoNCE where each query has one positive and K per-query hard negatives.
    The positive is class 0 in each row of logits.
    """
    B, D = query_emb.shape

    # Positive similarity: elementwise dot -> (B, 1)
    pos_logits = (query_emb * pos_emb).sum(dim=-1, keepdim=True)

    # Hard-negative similarities: batched matmul (B,K,D) x (B,D,1) -> (B,K)
    neg_logits = torch.bmm(neg_emb, query_emb.unsqueeze(-1)).squeeze(-1)

    # Concatenate: (B, 1 + K); column 0 is the positive
    logits = torch.cat([pos_logits, neg_logits], dim=1) / temperature

    # Target class is 0 for every query
    labels = torch.zeros(B, dtype=torch.long, device=query_emb.device)

    return F.cross_entropy(logits, labels)


# ---- Sanity check ----
if __name__ == "__main__":
    torch.manual_seed(0)
    B, K, D = 4, 3, 768
    q = F.normalize(torch.randn(B, D), dim=-1)
    pos = F.normalize(torch.randn(B, D), dim=-1)
    neg = F.normalize(torch.randn(B, K, D), dim=-1)
    print(infonce_with_hard_negatives(q, pos, neg, 0.05).item())  # ~ log(1+K)