Dense Retrieval: Semantic Search & Bi-Encoder Implementation › Training Dense Retrievers

Creating effective dense retrievers requires specialized training to produce embeddings where similar queries and documents cluster together. This section covers the key components of the training process, from the mathematical objective that guides learning to the practical considerations of data collection and negative sampling. Understanding training is important for practitioners who fine-tune their own models and for anyone who wants to understand why off-the-shelf models have certain strengths and weaknesses.

The basic challenge in training a dense retriever is supervision: how do we tell the model what "similar" means? Unlike classification tasks where labels are clear-cut, relevance in retrieval is gradational and context-dependent. A document about photosynthesis is highly relevant to "how do plants make food" but somewhat relevant to "renewable energy sources." The training framework needs to capture this nuance.

The goal of dense retriever training is to learn an embedding space where relevant query-document pairs have high similarity and irrelevant query-document pairs have low similarity. The key insight is that we can formulate this as a contrastive learning problem: teach the model to distinguish between documents that satisfy the information need and documents that do not.

Given a query $q$, a relevant (positive) document $d^+$, and irrelevant (negative) documents $\{d^-_1, d^-_2, ..., d^-_k\}$, we want:

$$\text{sim}(E_q(q), E_d(d^+)) > \text{sim}(E_q(q), E_d(d^-_i)) \quad \forall i$$

where:

• $\text{sim}(\cdot)$ is the similarity function (e.g., dot product or cosine)
• $E_q(q), E_d(\cdot)$ are the encoder embeddings for queries and documents respectively
• $d^+$ is the relevant (positive) document
• $d^-_i$ is the $i$-th irrelevant (negative) document
• $\forall i$ indicates the condition should hold for all negative samples

The intuition behind this formulation is geometric: we want the query embedding to lie closer to the positive document embedding than to any negative document embedding in the vector space. The training loss operationalizes this geometric constraint.

The most common loss function is the contrastive loss (also called InfoNCE loss or NT-Xent loss). This function treats retrieval as a multi-class classification task over the set of candidates in a training batch. It first computes the probability that the positive document is the correct match among all candidates (positive plus negatives), then minimizes the negative log of that probability.

To build the loss step by step: we start with the query embedding $\mathbf{q}$ and compute its similarity to the positive document embedding $\mathbf{d}^+$ and each negative embedding $\mathbf{d}^-_i$. We then convert these similarities to a probability distribution using softmax with temperature $\tau$:

$$\begin{aligned} P(d^+|q) &= \frac{\exp(\text{sim}(\mathbf{q}, \mathbf{d}^+) / \tau)}{\exp(\text{sim}(\mathbf{q}, \mathbf{d}^+) / \tau) + \sum_{i=1}^{k} \exp(\text{sim}(\mathbf{q}, \mathbf{d}^-_i) / \tau)} \\ \mathcal{L} &= -\log P(d^+|q) \end{aligned}$$

where:

• $P(d^+|q)$ is the probability that $d^+$ is the correct match for query $q$
• $\mathcal{L}$ is the contrastive loss we minimize during training
• $\mathbf{q}, \mathbf{d}^+, \mathbf{d}^-_i$ are the embeddings for query, positive, and negative documents
• $\tau$ is the temperature parameter that controls the sharpness of the probability distribution
• $k$ is the number of negative samples per positive

The numerator $\exp(\text{sim}(\mathbf{q}, \mathbf{d}^+) / \tau)$ is the exponentiated similarity score for the positive pair. The denominator normalizes this by summing the exponentiated scores for the positive document and all $k$ negative documents. The resulting $P(d^+|q)$ is the probability that the model assigns to the positive document being the correct match, given all the candidates.

The loss $\mathcal{L} = -\log P(d^+|q)$ is minimized when $P(d^+|q)$ approaches 1, which happens when the positive document's similarity is much higher than all negative documents' similarities. Maximizing this probability is exactly what we want: make the positive pair the most similar, by a large margin.

The temperature parameter $\tau$ deserves special attention. When $\tau$ is small (close to 0), the exponential function amplifies differences between similarity scores. This makes the model more confident in its distinctions and giving a sharper training signal. When $\tau$ is large, the probability distribution becomes more uniform. This provides a weaker gradient signal. Finding the right temperature, typically in the range of 0.05 to 0.1 for dense retrieval, is often done through experimentation, though some training frameworks learn it automatically.

Dense retrievers require training data consisting of query-document pairs with relevance labels. The quality and diversity of training data is one of the most important factors determining retrieval performance. Common sources include:

Natural Questions (NQ): Google's dataset of real questions from Google Search users, paired with Wikipedia passages containing answers. This provides a natural distribution of factual queries and their supporting evidence. The queries reflect how real users phrase questions, which is valuable for training a retriever that handles diverse natural language inputs. The limitation is that coverage is restricted to knowledge available on Wikipedia.

MS MARCO: Microsoft's large-scale reading comprehension dataset with approximately 500,000 queries drawn from Bing search logs, each paired with relevant passages from web documents. Its scale and web domain make it the most popular dataset for training general-purpose retrievers. Models trained on MS MARCO generalize reasonably well to many retrieval tasks, which makes it a common pre-training choice before domain-specific fine-tuning.

Synthetic data from LLMs: A recent and increasingly popular approach uses large language models to generate queries for existing documents. Given a passage, an LLM generates realistic questions that the passage would answer. This allows creating training data for any domain without requiring expensive human annotation. The quality of synthetic data has improved substantially as LLMs have become more capable, and recent work shows that models trained on synthetic data can approach the performance of models trained on human-annotated data.

Click logs from production systems: In deployed search or recommendation systems, user clicks provide implicit relevance signals. Documents that users click after submitting a query are treated as positive examples. This provides enormous amounts of training signal at low cost. This approach has been used by large-scale systems like Google, Bing, and Amazon to continuously improve their retrieval models. The challenge is handling noise: users sometimes click irrelevant results out of curiosity or mistake.

Domain-specific curated datasets: For specialized domains like legal, medical, or scientific retrieval, researchers and organizations have created curated datasets. Examples include BEIR (a benchmark with 18 diverse retrieval datasets), BioASQ for biomedical question answering, and MLDR for multilingual dense retrieval. Fine-tuning on these domain-specific datasets is often necessary for high-quality retrieval in specialized applications.

The choice of negative documents significantly impacts training quality, and this is one of the most studied aspects of dense retrieval training. The intuition is simple: the negatives define what "not relevant" means, and the model learns to distinguish positives from negatives. If negatives are too easy (obviously irrelevant), the model learns weak distinctions that do not help at inference time. If negatives are too hard (indistinguishable from positives), the training signal becomes noisy and inconsistent.

Random negatives: Sample random documents from the corpus. This is simple but often too easy. Random documents from a large corpus are typically obviously irrelevant to any given query. A retriever trained only on random negatives learns to distinguish relevant documents from random text, which is easy and does not prepare it for the harder cases where it needs to distinguish between two semantically similar documents, only one of which answers the query.

BM25 negatives: Use BM25 to find documents that contain query terms but are not relevant. These are called "hard negatives" because they share surface features with the query but differ semantically, forcing the model to learn representations that go beyond lexical overlap. For example, for the query "causes of inflation," a BM25 negative might be an article that mentions "inflation" and "causes" frequently but discusses historical inflation rather than economic causes. Training on such negatives teaches the model to distinguish semantic nuances.

In-batch negatives: Use positive documents from other queries in the same training batch as negatives for the current query. This approach is computationally efficient because it requires no additional encoding: the positive document embeddings already computed for other queries can be reused as negatives for the current query. With a batch size of 64, each query has 63 in-batch negatives at no extra encoding cost. The main limitation is that in-batch negatives may be easy (the other queries in the batch are typically about different topics), but with large batch sizes and diverse training data, in-batch negatives can be quite effective.

Self-mined hard negatives: Use the current model to find high-scoring documents that are not labeled as relevant. These are the hardest possible negatives because the current model finds them confusable with true positives. After each training epoch (or after a fixed number of steps), re-rank the corpus using the current model and mine new hard negatives. This iterative mining process progressively makes training harder, keeping the model in a regime where it is always learning something challenging.

Research has consistently shown that combining strategies works best. A typical approach: start training with in-batch negatives plus some BM25 negatives for initial training, then periodically re-mine hard negatives using the current model to refine the embedding space. The original DPR paper used BM25 negatives and in-batch negatives. More recent work like ANCE (Approximate Nearest Neighbor Negative Contrastive Estimation) demonstrates that self-mined hard negatives can substantially improve over BM25 negatives, especially later in training when the model has become good enough that BM25 negatives are no longer challenging.

A typical dense retriever training pipeline proceeds as follows:

1. Initialize query and document encoders from a pre-trained language model (BERT, RoBERTa, or a domain-specific model)
2. Collect training data: query-positive document pairs plus negative documents using one of the strategies above
3. For each training batch, encode all queries and documents using the two encoders
4. Compute the contrastive loss using the similarity scores between query embeddings and all document embeddings in the batch
5. Backpropagate gradients through both encoders and update their weights
6. Periodically evaluate on a held-out validation set to monitor retrieval quality
7. Optionally, re-mine hard negatives using the current model state after every few thousand steps

The training process is computationally intensive because we need to encode many negatives per query to provide a strong training signal, and BERT-based encoders are large models. Techniques like in-batch negatives help by reusing computation across the batch, but training a high-quality dense retriever still typically requires substantial GPU resources and hours to days of compute time. Smaller, more efficient encoder architectures like MiniLM or DistilBERT can reduce training time substantially while maintaining competitive retrieval quality.

An important practical consideration is gradient flow through both encoders. Both the query encoder and the document encoder receive gradients during training, so both are updated simultaneously. This means the document embeddings change at every training step, which invalidates any cached embeddings. In practice, this is handled by recomputing document embeddings periodically (not after every step) or by freezing one encoder and only training the other, depending on the application.