Dense Retrieval and Bi-Encoders

Encoding and Similarity

A bi-encoder consists of two encoder functions — one for queries, one for documents — that are often weight-shared (same BERT backbone) or sometimes separate networks. Given a query $q$ and a document $d$, the model produces:

Explanation: Each encoder independently maps its input to a fixed-length dense vector of hidden dimension $h$ (typically 768 for BERT-base), with no cross-attention between query and document tokens. The relevance score is then a simple inner product (or cosine similarity):

Explanation: This single dot product is all that happens at retrieval time — the entire semantic matching problem is compressed into one multiply-accumulate operation per candidate document.

Training with In-Batch Negatives

The dominant training objective is the InfoNCE / in-batch negative contrastive loss. Given a batch of $B$ (query, positive document) pairs $\{(q_i, d_i^+)\}_{i=1}^B$, all other documents in the batch act as negatives:

Explanation: For each query $q_i$, we want the score with its positive $d_i^+$ to be the maximum across all $B$ documents in the batch; $\tau$ is a temperature hyperparameter (often 1.0 or learned) that sharpens or softens the distribution. where $\tau$ is temperature, $s(\cdot, \cdot)$ is the dot-product score from Eq. [2], and the denominator sums over all $B$ documents including the positive. With large batch sizes (e.g., 128–2048), this gives $B - 1$ free negatives per positive, making training extremely efficient.

Hard Negative Mining

Standard in-batch negatives are often too easy (random documents are obviously irrelevant). Hard negative mining selects negatives that are semantically close but not relevant:

Explanation: The denominator now includes $K$ hard negatives $d_k^-$ retrieved by BM25 or a previous dense model — documents that look relevant lexically or semantically but are not the true answer, forcing the encoder to learn finer semantic distinctions.

Approximate Nearest Neighbor Retrieval

At inference, given a query embedding $\mathbf{q}$, we find:

Explanation: This is an exhaustive maximum inner product search (MIPS) over the full document corpus $\mathcal{D}$, which in practice is approximated using FAISS HNSW or IVF indexes, reducing search from $O(|\mathcal{D}|)$ to $O(\log |\mathcal{D}|)$ with small recall trade-offs.

PyTorch Implementation Sketch

```python import torch import torch.nn as nn import torch.nn.functional as F from transformers import AutoModel, AutoTokenizer class BiEncoder(nn.Module): def __init__(self, model_name: str = "bert-base-uncased"): super().__init__()

# Weight-shared encoder for query and document self.encoder = AutoModel.from_pretrained(model_name) def encode(self, input_ids, attention_mask):

# Pool [CLS] token representation out = self.encoder(input_ids=input_ids, attention_mask=attention_mask) return out.last_hidden_state[:, 0, :] # shape: (B, H) def forward(self, q_ids, q_mask, d_ids, d_mask): q_emb = self.encode(q_ids, q_mask) # (B, H) d_emb = self.encode(d_ids, d_mask) # (B, H) return q_emb, d_emb

def contrastive_loss(q_emb, d_emb, temperature=1.0): """In-batch negative InfoNCE loss."""

# Normalize for cosine similarity variant q_emb = F.normalize(q_emb, dim=-1) d_emb = F.normalize(d_emb, dim=-1)

# Score matrix: (B, B) — diagonal is positive pairs scores = torch.matmul(q_emb, d_emb.T) / temperature # (B, B)

# Labels: each query matches doc at same index labels = torch.arange(scores.size(0), device=scores.device)

# Cross-entropy over rows and columns (symmetric loss) loss_q = F.cross_entropy(scores, labels) loss_d = F.cross_entropy(scores.T, labels) return (loss_q + loss_d) / 2

# --- Indexing and retrieval with FAISS --- import faiss import numpy as np def build_faiss_index(doc_embeddings: np.ndarray, use_gpu=False): """Build an IVF-Flat FAISS index for MIPS.""" dim = doc_embeddings.shape[1] nlist = min(4096, len(doc_embeddings) // 39) quantizer = faiss.IndexFlatIP(dim) # Inner product index = faiss.IndexIVFFlat(quantizer, dim, nlist, faiss.METRIC_INNER_PRODUCT) if use_gpu: res = faiss.StandardGpuResources() index = faiss.index_cpu_to_gpu(res, 0, index)

# Normalize if using cosine similarity faiss.normalize_L2(doc_embeddings) index.train(doc_embeddings) index.add(doc_embeddings) return index def retrieve(query_emb: np.ndarray, index, top_k=100): faiss.normalize_L2(query_emb) scores, doc_ids = index.search(query_emb, top_k) # (1, k), (1, k) return doc_ids[0], scores[0] ```