Sentence Embeddings & Contrastive Loss

A developer asks "How do I rotate an API key?" Two passages look plausible. One says "Create a replacement key, deploy it, then revoke the old key." Another discusses a billing export token. They're topic-adjacent, but only one answers the question. A keyword-heavy search can still send the developer to the wrong guide.

The rotation guide should rank above the lookalike. A sentence embedding maps each string to one fixed-width vector. Contrastive learning trains that geometry so an intended match lands nearer the query than a confusing candidate. We'll build from pooling token vectors to calculating a batch loss, then test whether better training scores translate into better retrieval.

Keep this ordering in view: the rotation guide should score above the billing-token passage. A lower contrastive loss without better held-out rankings can mean the batch got easier rather than the retriever got better.

From word embeddings to sentence embeddings

Word embeddings gave individual tokens numerical coordinates. Retrieval has to compare variable-length queries and passages with one reusable vector. A sentence encoder performs that compression while trying to preserve the distinctions its task needs from nearby texts.

Averaging context-free word vectors is a useful baseline because it's fast and transparent. It also isolates two jobs that are easy to conflate: mean pooling collapses token states into one vector, while the training objective teaches that vector what relevance means.

Mean pooling inside a trained sentence encoder can work well. Averaging alone is not an explicit retrieval objective. Sentence-level training gives the model a target for those rankings; contrastive loss is one way to provide it.

Why pooling alone isn't enough

Mean pooling of word embeddings

Before reaching for a larger encoder, predict what a commutative operation will do to "compiler calls linker" and "linker calls compiler." This tiny baseline uses a local vector table, so you can run it and see exactly what averaging throws away:

word_vectors: dict[str, tuple[float, float, float]] = {
    "compiler": (0.9, 0.1, 0.0),
    "calls": (0.8, 0.2, 0.1),
    "linker": (0.7, 0.1, 0.2),
    "build": (0.1, 0.9, 0.2),
    "failed": (0.2, 0.8, 0.1),
}

def mean_pool(
    sentence: str,
    vectors: dict[str, tuple[float, float, float]],
) -> tuple[float, float, float]:
    tokens = [token.lower() for token in sentence.split()]
    token_vectors = [vectors[token] for token in tokens if token in vectors]

    if not token_vectors:
        raise ValueError("at least one known token is required")
    if any(len(vector) != 3 for vector in token_vectors):
        raise ValueError("this table requires three-dimensional vectors")

    return tuple(
        sum(vector[dim] for vector in token_vectors) / len(token_vectors)
        for dim in range(len(token_vectors[0]))
    )

pooled_a = mean_pool("compiler calls linker", word_vectors)
pooled_b = mean_pool("linker calls compiler", word_vectors)

print("A:", tuple(round(value, 3) for value in pooled_a))
print("B:", tuple(round(value, 3) for value in pooled_b))
print("Same vector:", pooled_a == pooled_b)

A: (0.8, 0.133, 0.1)
B: (0.8, 0.133, 0.1)
Same vector: True

The final line exposes the failure: both sentences produce the same vector because averaging static word vectors ignores order. This baseline skips unknown tokens and rejects an all-unknown input. Mean pooling contextual Transformer states is different: those states already depend on surrounding words and positions. Pooling them doesn't automatically erase the encoder's order information.

What averaging throws away

Reversing two tokens is the clearest collision. The same operation also loses other signals a retrieval boundary may need:

• Word order: "compiler calls linker" and "linker calls compiler" get the same vector because addition is commutative ($A+B = B+A$).
• Common-word dilution: Frequent words and boilerplate can wash out rare, informative tokens.
• No context: A static vector for "key" can't tell an API token from KMS key material.

[CLS] token from BERT

Taking the [CLS] (Classification) token from BERT (Bidirectional Encoder Representations from Transformers) and treating it as a sentence vector is a tempting shortcut. BERT's original pair architecture feeds both texts through one Transformer for a joint task. Its [CLS] state isn't trained by default so that independently encoded sentences can be ranked by cosine similarity.

That joint path also has the wrong serving shape for a large corpus. Reimers and Gurevych estimated about 50 million pair inferences, or roughly 65 hours on a V100, to find the nearest pair in 10,000 sentences. SBERT (Sentence-BERT) instead encodes each sentence once, then compares the cached vectors with cosine similarity.^{[1]Reference 1Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks.https://arxiv.org/abs/1908.10084}

The same paper measured how the shortcuts behave. Across seven semantic textual similarity (STS) tasks, raw [CLS] averaged 29.19 and mean-pooled BERT states averaged 54.81, reporting Spearman rank correlation multiplied by 100. Averaged GloVe (Global Vectors) scored 61.32. These are historical baselines; both pooling choices can work after sentence-level training.^{[1]Reference 1Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks.https://arxiv.org/abs/1908.10084}

Anisotropy, collapse, and rankings

One possible obstacle is anisotropy: embedding directions occupy some regions more densely than others. Ethayarajh measured this in contextual word representations from BERT, ELMo, and GPT-2, with substantial differences between layers.^{[2]Reference 2How Contextual are Contextualized Word Representations? Comparing the Geometry of BERT, ELMo, and GPT-2 Embeddings.https://arxiv.org/abs/1909.00512} That result doesn't establish the cause of every sentence-retrieval failure.

Uniformly distributed unit vectors have expected pairwise cosine zero. A shared directional component can instead make many cosines positive, as the constructed cone below illustrates. Representation collapse is a more extreme failure: different inputs receive identical or nearly identical representations, losing useful distinctions. Anisotropy alone doesn't mean those distinctions have vanished. Also, selecting [CLS] is not mean pooling; it selects one contextual state.

A high score baseline is not sufficient evidence that retrieval is broken. The intfloat/e5-base-v2 model card describes cosines concentrated around 0.7–1.0 and emphasizes relative ordering, rather than absolute scores.^{[3]Reference 3E5-base-v2 model cardhttps://huggingface.co/intfloat/e5-base-v2} Ask whether the right passages outrank the distractors on held-out queries. Inspect score distributions alongside Recall@K; a cosine of 0.9 is neither a relevance probability nor a universal acceptance threshold.

Alignment vs uniformity on the unit hypersphere

Wang and Isola analyzed two properties of normalized representations on the unit hypersphere $\mathcal{S}^{d-1}$:^{[4]Reference 4Understanding Contrastive Representation Learning through Alignment and Uniformity on the Hyperspherehttps://arxiv.org/abs/2005.10242}

1. Alignment: Positive pairs $(x, x^+)$ should map nearby. One alignment metric is the expected squared Euclidean distance: $\mathcal{L}_{\text{align}}(f) \triangleq \mathbb{E}_{(x, x^+) \sim p_{\text{pos}}} \left[ \|f(x) - f(x^+)\|_2^2 \right]$ Because $\|u - v\|_2^2 = \|u\|_2^2 + \|v\|_2^2 - 2(u^\top v) = 2 - 2(u^\top v)$ for unit vectors, minimizing alignment loss directly maximizes cosine similarity $u^\top v \to 1$.

2. Uniformity: A second metric rewards spreading sampled representations across the sphere: $\mathcal{L}_{\text{uniform}}(f; t) \triangleq \log \mathbb{E}_{x, y \overset{i.i.d.}{\sim} p_{\text{data}}} \left[ \exp\left(-t \|f(x) - f(y)\|_2^2\right) \right], \quad t > 0$ Nearby pairs contribute more to this Gaussian kernel average than distant pairs. Minimizing it discourages crowding. This is a distributional property, not a requirement to separate every pair, including semantic positives.

For a fixed $\tau>0$, a shared unit-normalized encoder, and independently sampled negatives from $p_{\text{data}}$, their population contrastive loss has this infinite-negative limit: $\lim_{K \to \infty} \mathcal{L}_{\text{InfoNCE}}(f; \tau, K) - \log K = -\frac{1}{\tau} \mathbb{E}_{(x, x^+)} [f(x)^\top f(x^+)] + \mathbb{E}_x \left[ \log \mathbb{E}_{y} \left[ \exp\left( \frac{f(x)^\top f(y)}{\tau} \right) \right] \right]$

The first term is alignment up to scaling and an additive constant. The second is related to, but not equal to, the uniformity metric above: here the outer expectation contains a logarithm; in $\mathcal{L}_{\text{uniform}}$ the logarithm sits outside the pairwise expectation. The paper relates their minimizers under its assumptions. It does not prove that every finite minibatch produces uniformly distributed, useful sentence embeddings.

Contrastive learning for sentence embeddings

The core idea

Now turn the retrieval failure into a training target. For each query, name a passage that should be nearby and candidates that should not. Contrastive training reshapes the space so that ordering becomes part of the representation instead of a hope at serving time.

Cosine similarity measures the angle between two vectors and ignores their length. For two unit vectors, it's their dot product: +1 is the same direction, 0 is orthogonal, and a negative value is opposing.

Those values aren't a meaning score by themselves. Only an evaluated embedding model makes cosine ranking useful. The next lesson studies that scoring contract in detail.

SimCSE (Simple Contrastive Learning of Sentence Embeddings) showed how small the construction can be: pass the same sentence through the encoder twice, let dropout supply two noisy views, and treat those views as a positive pair.^{[5]Reference 5SimCSE: Simple Contrastive Learning of Sentence Embeddings.https://arxiv.org/abs/2104.08821} The model gets a positive without a second labeled sentence.

E5 later trained single-vector text embeddings contrastively from a large weakly supervised pair corpus.^{[6]Reference 6Text Embeddings by Weakly-Supervised Contrastive Pre-training.https://arxiv.org/abs/2212.03533} That same pattern, rather than a pooling rule alone, is what gives a vector space a retrieval contract.

The batch-ranking objective used below is InfoNCE: for one anchor, the true match should outrank every other candidate in the batch.

InfoNCE loss

You want the matching API-doc passage close to the query and every non-match farther away. InfoNCE writes that ranking problem as classification: for this query, which passage in the batch is the right one? Before calculating, predict the answer for $q_1$: $p_1$ should win because 0.90 is above 0.20.

Walk through a batch of two query-passage pairs. After you normalize the embeddings, cosine similarities are just dot products:

For query $q_1$, the true match is $p_1$ (similarity 0.90). The other passage in the batch, $p_2$, acts as an in-batch negative (similarity 0.20). InfoNCE turns that ordering into probability: it wants the model to make $p_1$ more likely than $p_2$.

For this worked row, choose a sharp temperature $\tau = 0.05$ and compute the loss contribution for $q_1$ step by step:

1. Scale the similarities into logits: positive logit = $0.90 / 0.05 = 18.0$, negative logit = $0.20 / 0.05 = 4.0$
2. Exponentiate (this turns logits into unnormalized probabilities): $\exp(18.0) \approx 65{,}659{,}969$, $\exp(4.0) \approx 54.6$
3. Normalize with softmax into a probability for the positive: $65{,}659{,}969 / (65{,}659{,}969 + 54.6) \approx 0.99999917$
4. Take negative log: $-\log(0.99999917) \approx 0.00000083$ (tiny loss; the model is already very confident)

If the model were wrong ($q_1$ similarity to $p_1$ only 0.20, to $p_2$ 0.90), the positive probability would drop to about $8.3 \times 10^{-7}$ and the loss would jump to roughly 14 nats. The optimizer would receive a strong gradient pushing the correct pair closer.

The one-positive, query-to-passage InfoNCE form for a batch of $N$ positive pairs is:^{[7]Reference 7Representation Learning with Contrastive Predictive Coding.https://arxiv.org/abs/1807.03748}

$\mathcal{L} = -\frac{1}{N} \sum_{i=1}^{N} \log \frac{\exp(\text{sim}(z_i, z_i^+) / \tau)}{\sum_{j=1}^{N} \exp(\text{sim}(z_i, z_j^+) / \tau)}$

Reading the formula

For each example $i$, the numerator gives the true match $z_i^+$ its score. The denominator adds that positive to every candidate in the batch, turning the row into a multiple-choice question.

Temperature $\tau$ controls how sharply those logits separate. The loss therefore asks for one concrete change: make the true pair more likely than every alternative in the batch.

Where:

• $z_i, z_i^+$ are embeddings of a positive pair (e.g., query and relevant document)
• $\tau$ is the temperature parameter
• All non-matching examples in that denominator act as in-batch negatives

This copy-runnable implementation keeps the same calculation visible instead of hiding the matrix math behind a framework. Production training code would vectorize this in PyTorch or another tensor library. The loop below makes the denominator explicit:

from math import exp, isfinite, log, sqrt

def normalize(vector: list[float]) -> list[float]:
    if not vector or not all(isfinite(value) for value in vector):
        raise ValueError("expected a nonempty finite vector")
    scale = max(abs(value) for value in vector)
    if scale == 0:
        raise ValueError("vector must be nonzero")
    scaled = [value / scale for value in vector]
    norm = sqrt(sum(value * value for value in scaled))
    return [value / norm for value in scaled]

def dot(left: list[float], right: list[float]) -> float:
    if len(left) != len(right):
        raise ValueError("vector widths must match")
    return sum(a * b for a, b in zip(left, right))

def row_cross_entropy(logits: list[float], correct: int) -> float:
    if not logits or not all(isfinite(value) for value in logits):
        raise ValueError("expected nonempty finite logits")
    if type(correct) is not int or not 0 <= correct < len(logits):
        raise ValueError("target index outside row")
    peak = max(logits)
    return log(sum(exp(value - peak) for value in logits)) + (peak - logits[correct])

def info_nce_loss(
    query_vectors: list[list[float]],
    positive_vectors: list[list[float]],
    temperature: float = 0.2,
) -> float:
    if not isfinite(temperature) or temperature <= 0:
        raise ValueError("temperature must be finite and positive")
    if not query_vectors or len(query_vectors) != len(positive_vectors):
        raise ValueError("expected equal nonempty paired batches")
    queries = [normalize(vector) for vector in query_vectors]
    positives = [normalize(vector) for vector in positive_vectors]
    losses: list[float] = []

    for row, query in enumerate(queries):
        logits = [dot(query, candidate) / temperature for candidate in positives]
        losses.append(row_cross_entropy(logits, row))

    return sum(losses) / len(losses)

query_vectors = [[1.0, 0.0, 0.0], [0.0, 1.0, 0.0]]
positive_vectors = [
    [0.90, 0.15, sqrt(1 - 0.90**2 - 0.15**2)],
    [0.20, 0.85, sqrt(1 - 0.20**2 - 0.85**2)],
]

loss = info_nce_loss(query_vectors, positive_vectors, temperature=0.05)
score_pos = dot(normalize(query_vectors[0]), normalize(positive_vectors[0]))
score_neg = dot(normalize(query_vectors[0]), normalize(positive_vectors[1]))
extreme_loss = row_cross_entropy([1000.0, 986.0], correct=0)

print("loss:", f"{loss:.8f}")
print("q1 positive score:", round(score_pos, 4))
print("q1 negative score:", round(score_neg, 4))
print("stable large-logit loss:", f"{extreme_loss:.8f}")
print("tied extreme logits:", f"{row_cross_entropy([1e308, 1e308], 0):.6f}")
print("tiny vector:", [round(value, 6) for value in normalize([1e-300, 1e-300])])

loss: 0.00000083
q1 positive score: 0.9
q1 negative score: 0.2
stable large-logit loss: 0.00000083
tied extreme logits: 0.693147
tiny vector: [0.707107, 0.707107]

The equation is often expanded into raw exponentials on paper. Code should subtract the maximum before exponentiating. Keep the target gap separate from the log-sum term: forming peak + log_sum first can lose the small term when peak is enormous. Two tied logits must give $\log 2$, even when both are $10^{308}$. Scaling before taking a vector norm also avoids squaring tiny coordinates into zero or huge coordinates into infinity. This small implementation expects numeric coordinates; it is not a general input parser.

Tensor shapes and one backward pass

In a real encoder, token states have shape [batch, tokens, hidden]. Mask padding before mean pooling, then normalize the pooled vectors if the loss uses cosine. Use the pooling and special-token rules of the chosen model; this toy simply includes both non-padding positions. Its fixed states are invented, not Transformer outputs.

import torch
import torch.nn.functional as F

torch.set_default_dtype(torch.float64)
states = torch.tensor([
    [[1., 0.], [0., 1.], [99., 99.]],
    [[0., 2.], [0., 0.], [99., 99.]],
], requires_grad=True)
mask = torch.tensor([[1., 1., 0.], [1., 0., 0.]])
counts = mask.sum(dim=1, keepdim=True)
if torch.any(counts == 0):
    raise ValueError("each sequence needs a non-padding token")
pooled = (states * mask.unsqueeze(-1)).sum(dim=1) / counts
print("masked means:", pooled.detach().tolist())

# This is a separate two-pair cosine matrix, with intentionally modest margins.
queries = F.normalize(pooled, dim=-1)
passages = F.normalize(torch.tensor([[1., 0.], [0., 1.]]), dim=-1)
logits = queries @ passages.T / 0.2  # [2, 2], not elementwise multiplication
targets = torch.arange(len(queries))  # column i is the positive for row i
loss = F.cross_entropy(logits, targets)
loss.backward()

print("mean row loss:", f"{loss.item():.6f}")
print("padding gradient:", states.grad[:, 2, :].tolist())
print("first token gradient:", [round(v, 6) for v in states.grad[0, 0].tolist()])

masked means: [[0.5, 0.5], [0.0, 2.0]]
mean row loss: 0.349931
padding gradient: [[-0.0, 0.0], [0.0, 0.0]]
first token gradient: [-0.883883, 0.883883]

The first row initially ties both passages, so its loss is $\log 2$, while the second strongly favors its positive. Gradients reach the unmasked token states through pooling and normalization; padding receives zero gradient. An actual training loop would pass this gradient into encoder parameters and call the optimizer. This example performs no model training or network download.

Dual (symmetric) InfoNCE

The formula above is query → candidate only: each query row treats other positives in the batch as negatives. Many dual-encoder and multimodal recipes also run the reverse direction and average both:

On the same $N \times N$ similarity matrix $S_{ij} = \text{sim}(z_i^{q}, z_j^{d}) / \tau$:

• Row CE ($\mathcal{L}_{q \to d}$): for each query $i$, the correct column is $i$ (query ranks its matched document).
• Column CE ($\mathcal{L}_{d \to q}$): for each document $j$, the correct row is $j$ (document ranks its matched query).

CLIP uses this symmetric objective for image↔text.^{[8]Reference 8Learning Transferable Visual Models From Natural Language Supervision.https://arxiv.org/abs/2103.00020} Asymmetric serving doesn't forbid symmetric training. The reverse loss adds the assumption that each passage has one correct query in this batch, which may be wrong when several queries share a relevant passage. Choose the objective from the labels and compare held-out query-to-document retrieval. Bi-encoder towers may share weights, as in a Siamese encoder, or use separate parameters.

Triplet loss

InfoNCE compares one anchor with a candidate pool. A second contrastive objective, triplet loss, focuses on one anchor-positive-negative triplet instead, making the required distance gap explicit:

$\mathcal{L} = \max(0, d(a, p) - d(a, n) + m)$

Where:

• $d(\cdot, \cdot)$ is the Euclidean distance between embeddings
• $m$ is a margin hyperparameter
• $a$ is the anchor, $p$ is the positive, $n$ is the negative

The loss enforces that the anchor must be closer to the positive than to the negative by at least margin $m$: $d(a, p) + m \leq d(a, n)$. If this constraint is already satisfied, the loss is zero. The margin prevents the model from wasting capacity pushing already-distant negatives even farther away.

Worked example: computing triplet loss by hand

Consider three sentences about API-key rotation. Before doing the arithmetic, predict the outcome: with the positive at distance 0.2 and the negative at 0.5, a margin of 0.1 should leave no work, while a larger margin could still leave a gap to close:

After encoding, suppose the distances are:

• $d(a, p) = 0.2$ (the paraphrase is close)
• $d(a, n) = 0.5$ (the hard negative is farther, but not by much)

With margin $m = 0.1$, plug into the formula:

$0.2 - 0.5 + 0.1 = -0.2$

$\max(0, -0.2) = 0$

The loss is zero because the model already satisfies the margin constraint: the positive is closer than the negative by more than 0.1. Now imagine a bad model where $d(a, p) = 0.5$ and $d(a, n) = 0.3$ (the negative is closer than the positive):

$0.5 - 0.3 + 0.1 = 0.3$

$\max(0, 0.3) = 0.3$

A non-zero loss tells the optimizer to push the anchor and positive together while pushing the negative away until the gap exceeds the margin.

Comparing the learning signal

Triplet loss makes the required distance gap explicit. InfoNCE makes competition across a candidate pool explicit. Neither choice removes the need to inspect your labels and training batches:

1. Comparison unit: One triplet term contains one negative, but a training step can batch many triplets, mine within the batch, or aggregate multiple negatives per anchor.
2. Inactive examples: A satisfied triplet has zero hinge loss. Many satisfied triplets can provide little signal, so the margin and sampling policy matter.
3. Candidate competition: InfoNCE scores several candidates in one softmax. A negative's score gradient depends on its probability in that pool; an easy row can still saturate.

InfoNCE has no explicit distance margin $m$. It has other design choices, including temperature, pool size, and negative sampling. Sentence Transformers still documents both triplet objectives and batch-ranking objectives; select a loss that fits your data, then compare held-out retrieval.^{[9]Reference 9Sentence Transformers loss referencehttps://www.sbert.net/docs/package_reference/sentence_transformer/losses.html}

Temperature parameter τ and gradient dynamics

Temperature $\tau$ controls the sharpness of the softmax probability distribution over candidate similarities. It's not just a scaling constant: it directly governs the magnitude of gradients and controls how aggressively the loss concentrates learning pressure on hard negatives.

Let $s_j = \text{sim}(q, c_j)$ denote the cosine similarity between query $q$ and candidate $c_j$. The softmax probability assigned to candidate $c_j$ is: $P_j = \frac{\exp(s_j / \tau)}{\sum_k \exp(s_k / \tau)}$

For one row, differentiating its InfoNCE loss with respect to the similarity scores gives:

• For the true positive candidate $p^+$: $\frac{\partial \mathcal{L}}{\partial s^+} = -\frac{1}{\tau} (1 - P^+)$
• For any negative candidate $n_j^-$: $\frac{\partial \mathcal{L}}{\partial s_j^-} = \frac{1}{\tau} P_j$

These two derivatives expose the twin mathematical roles of temperature $\tau$:

1. Scale and saturation: The factor $\frac{1}{\tau}$ increases as $\tau$ decreases, but $P_j$ changes too. Gradients can shrink on an easy row as its positive probability approaches 1. A mean over $N$ rows adds a factor $1/N$ to each row's contribution.
2. Hard-negative weighting (hardness-aware gradients): Compare the gradient penalties applied to two distinct negative candidates $n_1^-$ and $n_2^-$ with similarities $s_1^-$ and $s_2^-$: $\frac{\partial \mathcal{L} / \partial s_1^-}{\partial \mathcal{L} / \partial s_2^-} = \frac{P_1}{P_2} = \exp\left( \frac{s_1^- - s_2^-}{\tau} \right)$

This relative penalty grows exponentially with the similarity difference divided by $\tau$:

• Low temperature (such as $\tau = 0.05$): Suppose a hard negative scores $s_1^- = 0.72$ and an easy negative scores $s_2^- = 0.20$ (a gap of $\Delta s = 0.52$). The gradient ratio is: $\exp\left(\frac{0.52}{0.05}\right) = \exp(10.4) \approx 32{,}860$ This is a ratio of partial derivatives with respect to two scores, not a ratio of encoder parameter updates. Those updates also depend on the encoder's Jacobians, shared parameters, other rows, and optimizer.
• High temperature (such as $\tau = 1.00$): For that exact same gap of $0.52$, the gradient ratio is: $\exp\left(\frac{0.52}{1.00}\right) \approx 1.68$ The score-gradient contributions are more evenly distributed. Whether that improves retrieval depends on which negatives and labels are trustworthy.

Temperature tradeoffs and failure modes

Choosing temperature is therefore a balance between sharpness and stability:

• Large $\tau$: Probabilities flatten for fixed scores. In the limit $\tau\to\infty$, score gradients vanish because of $1/\tau$. This doesn't imply every finite larger temperature gives a smaller gradient.
• Small $\tau$: A clearly correct row can have almost no gradient, while a tie or misranked negative can have a large gradient. A false negative that wins substantial softmax mass then receives unwanted separation pressure. Per-row score-gradient magnitudes are bounded by $1/\tau$; their ratio to another negative can be much larger.

Tune $\tau$ against held-out retrieval metrics and gradient behavior, rather than universal cutoffs. E5 documents $\tau = 0.01$.^{[3]Reference 3E5-base-v2 model cardhttps://huggingface.co/intfloat/e5-base-v2} Sentence Transformers' MultipleNegativesRankingLoss defaults to scale=20, equivalent to $\tau=0.05$ when using cosine.^{[9]Reference 9Sentence Transformers loss referencehttps://www.sbert.net/docs/package_reference/sentence_transformer/losses.html} This is a particular loss default, not a temperature shared by all SBERT training recipes.

For the worked similarity gap ($0.90 - 0.20 = 0.70$), calculate the gradient as well as the probability:

Here the larger temperature gives a larger absolute gradient. Both the gap and the candidate pool matter. The rounded probability at 0.01 hides a nonzero mathematical gradient; floating-point subtraction 1 - P can round it to zero.

from math import exp, isfinite

def probability_of_positive(
    positive_similarity: float,
    negative_similarity: float,
    temperature: float,
) -> float:
    if not isfinite(temperature) or temperature <= 0:
        raise ValueError("temperature must be finite and positive")
    if not all(isfinite(s) for s in (positive_similarity, negative_similarity)):
        raise ValueError("similarities must be finite")
    scaled_gap = (positive_similarity - negative_similarity) / temperature
    if scaled_gap >= 0:
        return 1.0 / (1.0 + exp(-scaled_gap))
    small = exp(scaled_gap)
    return small / (1.0 + small)

for temperature in (0.01, 0.05, 0.10, 1.00):
    probability = probability_of_positive(0.90, 0.20, temperature)
    # Compute the tiny complementary probability directly, not as 1 - P.
    negative_probability = probability_of_positive(0.20, 0.90, temperature)
    gradient = negative_probability / temperature
    print(f"tau={temperature:.2f}: P(positive)={probability:.6f}, |grad|={gradient:.9g}")

tau=0.01: P(positive)=1.000000, |grad|=3.97544974e-29
tau=0.05: P(positive)=0.999999, |grad|=1.66305606e-05
tau=0.10: P(positive)=0.999089, |grad|=0.00911051194
tau=1.00: P(positive)=0.668188, |grad|=0.331812228

Where the pairs come from

A contrastive objective has no way to decide what a pair means. You need texts that should match, texts that shouldn't, and enough variety that the model can't succeed by memorizing topic words. Supervised NLI pairs and SimCSE's dropout views make two useful constructions concrete.

Supervised: fine-tuning on NLI data

Natural Language Inference (NLI) labels whether a hypothesis follows from a premise (entailment), conflicts with it (contradiction), or does neither (neutral). Entailment is directional. It isn't a promise that two texts are interchangeable.

SBERT trained pooled Siamese representations with NLI supervision and evaluated sentence similarity. Its NLI classification setup feeds the two embeddings and their absolute difference to a classifier. Sentence-level training need not use InfoNCE.^{[1]Reference 1Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks.https://arxiv.org/abs/1908.10084}

Supervised SimCSE instead uses entailment as a positive and the corresponding contradiction as a hard negative.^{[5]Reference 5SimCSE: Simple Contrastive Learning of Sentence Embeddings.https://arxiv.org/abs/2104.08821} Whether either recipe transfers to your retrieval task needs separate evaluation.

The labels say which sentences should meet; SBERT's architecture makes that comparison cheap. It runs both sentences through the same encoder with shared weights, pools each into a single vector, then trains on those pooled embeddings.^{[1]Reference 1Sentence-BERT: Sentence Embeddings using Siamese BERT-Networks.https://arxiv.org/abs/1908.10084}

At inference you keep only the single-sentence path. That shared-weight Siamese network makes precomputing document embeddings and nearest-neighbor search practical.

Self-supervised: SimCSE

No labeled NLI data? SimCSE changes where the positive comes from.^{[5]Reference 5SimCSE: Simple Contrastive Learning of Sentence Embeddings.https://arxiv.org/abs/2104.08821} Pass the same sentence through the encoder twice with different dropout masks. Dropout zeros different units each time, so you get two slightly different vectors for the same sentence.

Those views are positives; every other sentence in the batch is a negative. The construction turns an unlabeled batch into a ranking problem, with the same false-negative risk we will audit below.

On the paper's STS suite, unsupervised SimCSE with BERT-base reached 76.3 average Spearman correlation multiplied by 100. The supervised NLI variant reached 81.6. Dropout becomes the data augmentation: removing it or sharing the same mask between views substantially worsened the reported results. These are findings for that recipe, not a theorem that every dropout-free encoder collapses.^{[5]Reference 5SimCSE: Simple Contrastive Learning of Sentence Embeddings.https://arxiv.org/abs/2104.08821} Use training mode to construct these stochastic views; ordinary embedding inference disables dropout for repeatable outputs.

Data augmentation

Once dropout makes a safe positive view, other augmentations become hypotheses about meaning preservation:

• Dropout masks (SimCSE): different mask patterns per forward pass
• Verified paraphrases or back-translations: use only after checking that technical scope and exceptions survive
• Domain pairs: mine resolved duplicate questions that cite the same approved docs passage
• Reranking with cross-encoders: score candidate paraphrases before accepting them as positives

Avoid casual word deletion or insertion for technical docs: dropping "not," a version condition, or an authorization requirement changes the rule while incorrectly labeling the pair positive.

In-batch negatives in practice

The batch can supply a negative pool without another labeling pass. For a batch of $N$ positive pairs, each anchor has one matching positive, and the other $N-1$ candidates act as negatives. Dense Passage Retrieval trained dual encoders this way for open-domain QA.^{[10]Reference 10Dense Passage Retrieval for Open-Domain Question Answering.https://arxiv.org/abs/2004.04906}

Larger batches increase the chance of informative competitors, but they also increase the chance of false negatives: another row may cite the same relevant passage while the loss treats it as wrong. Different passage IDs can also be relevant to the same query. Deduplicate, mask known false negatives, or use a multi-positive loss; an identity check alone doesn't establish non-relevance.

An $N=1$ batch without extra negatives gives probability 1 and loss 0 regardless of embedding quality. Larger batches create an $N\times N$ score matrix, with quadratic score-storage cost. Distributed training can gather embeddings across GPUs, but must handle gradient propagation deliberately. Plain gradient accumulation doesn't create more in-batch negatives unless the implementation explicitly reuses embeddings across microbatches.

import torch
import torch.nn.functional as F

batch = [
    {"query": "How do I rotate an API key?", "doc_id": "api-key-rotation-v3"},
    {"query": "How do I replace an old API token?", "doc_id": "api-key-rotation-v3"},
    {"query": "How do service account tokens expire?", "doc_id": "service-token-lifecycle-v2"},
]

false_negatives = []
for anchor_index, anchor in enumerate(batch):
    for candidate_index, candidate in enumerate(batch):
        if anchor_index == candidate_index:
            continue
        if candidate["doc_id"] == anchor["doc_id"]:
            false_negatives.append(
                f"row {anchor_index} treats row {candidate_index} as negative"
            )

print("false negatives:", false_negatives)
print("action: deduplicate shared-doc positives before InfoNCE")

# Isolate row 0: columns 0 and 1 are copies of its approved passage.
for exclude_duplicate in (False, True):
    scores = torch.tensor([0.8, 0.8, 0.1], dtype=torch.float64, requires_grad=True)
    logits = scores / 0.2
    if exclude_duplicate:
        logits = logits.masked_fill(torch.tensor([False, True, False]), -torch.inf)
    loss = F.cross_entropy(logits.unsqueeze(0), torch.tensor([0]))
    loss.backward()
    print(f"exclude duplicate={exclude_duplicate}: loss={loss.item():.6f}; "
          f"score gradients={[round(value, 6) for value in scores.grad.tolist()]}")

false negatives: ['row 0 treats row 1 as negative', 'row 1 treats row 0 as negative']
action: deduplicate shared-doc positives before InfoNCE
exclude duplicate=False: loss=0.708133; score gradients=[-2.537185, 2.462815, 0.074371]
exclude duplicate=True: loss=0.029750; score gradients=[-0.146561, 0.0, 0.146561]

The duplicate passage receives a positive score derivative in the unmasked row. Gradient descent therefore tries to lower that score even though it is a relevant copy. Masking only that verified duplicate removes its gradient and lets the row learn against the lifecycle passage. The two duplicate encodings can share parameters, so their opposing score derivatives need not produce independent vector movements. The unmasked loss also has a $\log 2$ floor when both duplicate scores stay equal and the unrelated candidate's probability vanishes. No optimizer can make one identical copy win with probability 1. This is a constructed score calculation, not measured encoder training. Masking an off-diagonal positive differs from the multi-positive objectives below, which retain it as a target.

Hard negative mining

Why hard negatives matter

Suppose the model already rejects "GPU driver compatibility" for an API-key query. That example contributes little new signal. Hard negatives such as a KMS-key guide share the broad vocabulary but answer a different intent, forcing the encoder to learn the boundary that matters.

Mining strategies and false negative contamination

Mining can supply informative competitors when your existing pool becomes easy. It is not universally required: sufficiently varied in-batch negatives can remain useful. Inspect actual scores and held-out errors. Mining also risks false negative contamination: an unannotated passage may actually answer the question, so treating it as a negative trains against the intended relevance relation.

Four core negative mining strategies balance informativeness against contamination risk:

1. In-batch negatives

In-batch negatives use the other passages in the current training mini-batch. For $B$ positive pairs, $2B$ text encodings supply $B(B-1)$ off-diagonal comparisons without encoding additional texts.^{[10]Reference 10Dense Passage Retrieval for Open-Domain Question Answering.https://arxiv.org/abs/2004.04906} Encodings can be batched; this does not mean $2B$ separate model calls. Computing and storing the similarity matrix still costs work and memory.

• Advantage: Reuses passage representations already needed for the positive pairs.
• Limitation: Batch composition controls the competitors. They may be easy, informative, or false negatives; they are not necessarily independent random corpus samples.

2. BM25 hard negatives

BM25 lexical retrieval searches the corpus using the query keywords to find passages that share vocabulary without answering the underlying intent:

Query: "How do I rotate an API key?"
BM25 Hard Negative: "Rotate KMS encryption key material"  # shares "rotate" and "key" but answers a different intent
Easy Negative: "GPU driver compatibility matrix"          # shares zero relevant terms

BM25 negatives force the encoder to look beyond superficial word matching and learn syntax, modifiers, and technical intent. However, lexical search frequently retrieves valid paraphrases or unannotated synonyms, requiring careful filtering.

3. Cross-encoder mined negatives and MarginMSE

Cross-encoders use joint self-attention over $[q; d]$ and can be useful teachers when validated on the target task. A teacher score remains an estimate, not verified relevance. One possible mining and distillation pipeline is:

1. Candidate retrieval: Retrieve the top 100 candidates for query $q$ using BM25 and a baseline bi-encoder.
2. Cross-encoder scoring: Score all 100 candidates with a trained cross-encoder.
3. Candidate auditing: Inspect candidates ranked highly by the bi-encoder but poorly by the teacher. Use relevance judgments where available; audit uncertain cases before assigning hard-negative labels.
4. Soft margin distillation (MarginMSE): Instead of forcing a binary 0/1 contrastive target, train the bi-encoder to match the score difference (margin) predicted by the cross-encoder: $\mathcal{L}_{\text{MarginMSE}} = \left( (\text{score}_{student}(q, p^+) - \text{score}_{student}(q, n^-)) - (\text{score}_{teacher}(q, p^+) - \text{score}_{teacher}(q, n^-)) \right)^2$ Average this squared error over the training triples. Hofstätter and colleagues introduced this margin-based distillation objective.^{[11]Reference 11Improving Efficient Neural Ranking Models with Cross-Architecture Knowledge Distillationhttps://arxiv.org/abs/2010.02666} The student inherits the teacher's target margins, including mistakes. Check scoring ranges: a cosine margin lies in $[-2,2]$, so it cannot match an arbitrary unbounded teacher-logit margin. An unnormalized dot-product student or an explicitly validated target scaling avoids that particular mismatch.

A zero distillation loss can still teach the wrong order. Suppose the approved passage is relevant, but the teacher assigns it 0.2 and a non-match 0.8. A student copying the margin $-0.6$ obtains zero MarginMSE while ranking the non-match first. Teacher agreement is a training target; held-out relevance remains the outcome to measure.

4. Iterative (self-adversarial) mining

As a bi-encoder changes, different competitors can become informative. Iterative mining re-encodes and re-indexes the corpus with a newer checkpoint, retrieves high-ranked candidates, then audits them for another training round. Keep evaluation queries separate from the mining loop; teacher filtering does not eliminate annotation errors.

Avoiding false negative contamination

Relevance annotations can be incomplete: one labeled passage does not establish that all other passages are wrong. A valid but unlabeled competitor with substantial softmax probability receives an unwanted negative score gradient $P_j/\tau$. The magnitude depends on the entire row, not on temperature alone.

To protect against false negative contamination:

• Known-positive handling: Deduplicate repeated passages, mask known positive off-diagonal entries, or group positives for a compatible objective. Different document IDs can still answer the same query; a shared topic alone does not prove relevance.
• Teacher filtering: A validated score ceiling can remove suspicious high-scoring candidates from a negative pool. It can also miss valid answers or remove useful negatives. Estimate those errors from labeled samples instead of treating the threshold as a proof.
• Multi-positive objectives: Mark verified positives explicitly. The chosen SupCon formulation averages their negative log-softmax probabilities; it does not simply sum positive probability mass inside one logarithm.^{[12]Reference 12Supervised Contrastive Learninghttps://arxiv.org/abs/2004.11362}

Two positives, two different objectives

For one anchor, write $P_p=\exp(s_p/\tau)/\sum_a\exp(s_a/\tau)$ over the allowed candidates. A SupCon-style row uses

$\mathcal{L}_{\text{mean-log}}=-\frac{1}{|\mathcal{P}|}\sum_{p\in\mathcal{P}}\log P_p.$

A different, positive-mass objective uses $-\log\sum_{p\in\mathcal{P}}P_p$. Predict what happens when the model puts almost all probability on one positive and neglects the other:

from math import isclose, log

# Constructed candidate probabilities: two positives and one negative.
positive_probabilities = (0.999, 0.0005)
negative_probability = 0.0005
assert isclose(sum(positive_probabilities) + negative_probability, 1.0)
mass_loss = -log(sum(positive_probabilities))
mean_log_loss = -sum(log(p) for p in positive_probabilities) / 2
print(f"positive-mass loss: {mass_loss:.6f}")
print(f"mean positive log-loss: {mean_log_loss:.6f}")

positive-mass loss: 0.000500
mean positive log-loss: 3.800951

The mass loss is almost satisfied even though the second positive has probability 0.0005. Averaging the log probabilities still penalizes neglecting it. These objectives encode different goals; their raw loss values are not directly comparable quality scores. The original SupCon setting excludes the anchor itself from its candidate denominator and uses other same-class examples as positives. Retrieval needs query-specific relevance labels, rather than assuming every same-topic passage is an answer.

The mining loop below illustrates one filter: lexical overlap stands in for BM25, invented scores stand in for a teacher, and supplied relevance labels protect known positives. It does not discover whether those labels are correct.

def tokens(text: str) -> set[str]:
    return {part.strip("?.!,").lower() for part in text.split()}

def lexical_overlap(query: str, candidate: str) -> int:
    return len(tokens(query) & tokens(candidate))

def fixture_score(query: str, candidate: str) -> float:
    query_terms = tokens(query)
    candidate_terms = tokens(candidate)

    if "api" in query_terms and "api" in candidate_terms:
        return 0.92
    if "service" in candidate_terms and "key" in candidate_terms:
        return 0.76
    if "kms" in candidate_terms and "key" in candidate_terms:
        return 0.68
    return 0.05

def mine_hard_negatives(
    query: str,
    corpus: list[dict[str, str | bool]],
    top_k: int = 2,
) -> list[str]:
    if type(top_k) is not int or top_k < 0:
        raise ValueError("top_k must be a nonnegative integer")
    if any(type(row["relevant"]) is not bool for row in corpus):
        raise ValueError("relevance labels must be booleans")
    known_non_matches = [
        row
        for row in corpus
        if not row["relevant"]
        and lexical_overlap(query, str(row["text"])) > 0
    ]
    ranked = sorted(
        known_non_matches,
        key=lambda row: fixture_score(query, str(row["text"])),
        reverse=True,
    )
    return [str(row["text"]) for row in ranked[:top_k]]

query = "How do I rotate an API key?"
corpus = [
    {"text": "Where can I replace an old API token?", "relevant": True},
    {"text": "Rotate KMS encryption key material", "relevant": False},
    {"text": "How do I rotate a service account signing key?", "relevant": False},
    {"text": "GPU driver compatibility matrix", "relevant": False},
]

print(mine_hard_negatives(query, corpus))

['How do I rotate a service account signing key?', 'Rotate KMS encryption key material']

Hard negatives in the InfoNCE denominator

Mining changes the loss only after those candidates appear in the softmax. Production trainers usually do one of two things:

1. Expand the candidate pool: for each query, score positive + in-batch negatives + K mined hard negatives and run InfoNCE over that longer list.
2. Share a mined candidate bank: append passages as extra candidate columns, with an explicit positive-column index for each query. Don't add unlabeled query rows or assume every new column is a positive diagonal entry.

Either way, a mined hard negative is just another term in $\sum_j \exp(\text{sim}(z_i, z_j) / \tau)$. The toy below keeps the positive and one in-batch negative, then adds a mined hard negative so you can see the extra denominator mass raise the loss:

from math import exp, isfinite, log

def row_cross_entropy(logits: list[float], correct: int) -> float:
    if not logits or not all(isfinite(value) for value in logits):
        raise ValueError("expected nonempty finite logits")
    if type(correct) is not int or not 0 <= correct < len(logits):
        raise ValueError("target index outside row")
    peak = max(logits)
    return log(sum(exp(v - peak) for v in logits)) + (peak - logits[correct])

# Cosine similarities for one query at temperature 0.05
temperature = 0.05
positive = 0.90
in_batch_negative = 0.20
mined_hard_negative = 0.72

easy_logits = [positive / temperature, in_batch_negative / temperature]
hard_logits = easy_logits + [mined_hard_negative / temperature]

print("loss without mined hard neg:", round(row_cross_entropy(easy_logits, 0), 6))
print("loss with mined hard neg:", round(row_cross_entropy(hard_logits, 0), 6))

loss without mined hard neg: 1e-06
loss with mined hard neg: 0.026958

The positive score didn't change. The mined candidate added competing probability mass, so the loss rose. This isn't evidence that retrieval improved: no weights changed in this calculation. It shows how a verified hard negative supplies additional training pressure.

Bi-encoder vs cross-encoder

The loss trains a vector space. Serving still has to decide when the query and the document are allowed to interact. That placement decision, not the brand name of the encoder, sets latency and index size.

Bi-encoder (dual encoder)

A bi-encoder encodes the query and document independently, then compares them with dot product or cosine similarity. Documents can be pre-encoded and indexed. At query time you encode the query once, then use an approximate nearest neighbor (ANN) search index.

That speed comes from keeping the texts apart. Without cross-attention, the single-vector score can miss phrase order, negation, and fine token alignment.

Cross-encoder

A cross-encoder concatenates query and document and runs them through one Transformer. Full attention can model phrase-level interactions, and on a suitable shortlist that often improves precision.

The same interaction makes it expensive: each scored pair needs joint encoding. For fixed sequence lengths, scoring $N$ documents requires $O(N)$ pair-encoding work per query. Batching and parallelism affect model-call count and latency, but they don't turn these into reusable query-independent document vectors. Measure throughput before choosing a shortlist.

Late interaction: ColBERT and the MaxSim operator

A single-vector score compresses a passage into a fixed-size representation before seeing the query. It can preserve order and negation, but a given model may lose distinctions that matter for retrieval. A cross-encoder instead lets both texts interact through joint self-attention, at the cost of query-dependent pair encoding.

ColBERT (Contextualized Late Interaction over BERT)^{[13]Reference 13ColBERT: Efficient and Effective Passage Search via Contextualized Late Interaction over BERT.https://arxiv.org/abs/2004.12832} offers another tradeoff using late interaction:

Instead of pooling into a single vector, ColBERT retains multiple contextual token vectors. The original implementation filters document punctuation and augments queries with mask positions, so these are not simply one vector for every literal word:

1. The query encoder outputs normalized token vectors $E_q \in \mathbb{R}^{|q| \times d}$ (typically projected to $d = 128$).
2. The document encoder outputs normalized token vectors $E_d \in \mathbb{R}^{|d| \times d}$.
3. Relevance is computed using the MaxSim operator: $\text{Score}(q, d) = \sum_{i=1}^{|q|} \max_{j=1}^{|d|} E_{q,i} \cdot E_{d,j}^\top$

For each query token $i$, MaxSim finds the single document token $j$ that produces the maximum dot product, and sums those peak alignments.

What MaxSim preserves

• Soft token alignment: Each retained query position gets its own best match. Contextual states can distinguish API keys from door keys, but identical token strings do not guarantee near-1 cosine, and contextualization does not guarantee correct disambiguation.
• Offline document indexing: Because document tokens don't interact with query tokens inside the transformer, documents are encoded and indexed once offline.
• Query-time vector scoring: Candidate generation and token comparisons replace joint query-document Transformer encoding. Index traversal, decompression, and memory transfers still contribute to latency.

MaxSim does not enforce a one-to-one alignment: several query positions may select the same document position. Its final sum also does not directly enforce phrase order; order information must come through the contextual vectors. Test modifier and negation queries rather than assuming this scoring rule handles them correctly.

Index storage and ColBERTv2 compression

The raw vector payload scales with token count × token-vector width × bytes per coordinate. At equal coordinate precision, 100 or 200 vectors of width 128 use about 16.7 or 33.3 times the payload of one width-768 vector, before index overhead. A flat “100×” rule ignores vector width and compression.

ColBERTv2 stores a nearest-centroid ID and a quantized residual for each token vector, reconstructing an approximation as centroid plus residual. Its paper describes 1- or 2-bit residual coordinates and reports substantial compression relative to earlier late-interaction representations.^{[14]Reference 14ColBERTv2: Effective and Efficient Retrieval via Lightweight Late Interactionhttps://arxiv.org/abs/2112.01488} That is not a guarantee of single-vector index size or unchanged quality on your corpus.

This is a budget tradeoff:

• Bi-encoder: A useful starting point when one vector per passage meets recall and RAM requirements.
• Late interaction (ColBERT): A candidate when token-level scoring improves held-out retrieval enough to justify a larger, more complex index.
• Cross-encoder: A candidate for reranking a bounded shortlist when measured ordering gains justify joint encoding.

similarities = {
    "api": {"rotate": 0.12, "api": 0.93, "key": 0.08},
    "key": {"rotate": 0.19, "api": 0.07, "key": 0.91},
}

best_by_query_token = {
    query_token: max(document_scores.values())
    for query_token, document_scores in similarities.items()
}
score = sum(best_by_query_token.values())

print("best token scores:", best_by_query_token)
print("MaxSim score:", round(score, 2))

best token scores: {'api': 0.93, 'key': 0.91}
MaxSim score: 1.84

A production reranking pattern

Serving stacks often combine the two. A reranker spends the expensive cross-encoder pass only on candidates the bi-encoder admitted. The first stage sets the available recall; the second aims to improve ordering within that set. If the correct passage never enters the shortlist, reranking can't recover it.

The small function below sorts invented scores for one fixed query. It doesn't run either model or demonstrate a measured cross-encoder improvement. Reducing candidate_k to 1 leaves only the KMS passage, so even the favorable reranking scores can't recover a relevant document.

from math import isfinite

documents = [
    {
        "id": "kms-key-rotation",
        "text": "How do I rotate a KMS encryption key?",
        "bi_score": 0.84,
        "cross_score": 0.21,
    },
    {
        "id": "api-key-rotation",
        "text": "API key rotation runbook",
        "bi_score": 0.81,
        "cross_score": 0.96,
    },
    {
        "id": "replace-old-api-token",
        "text": "Create a replacement key, deploy it, then revoke the old key.",
        "bi_score": 0.78,
        "cross_score": 0.91,
    },
    {
        "id": "gpu-driver-notes",
        "text": "GPU driver compatibility notes",
        "bi_score": 0.10,
        "cross_score": 0.05,
    },
]

def search_with_rerank(
    corpus: list[dict[str, str | float]],
    candidate_k: int = 3,
    top_k: int = 2,
) -> list[str]:
    if any(type(k) is not int or k < 0 for k in (candidate_k, top_k)):
        raise ValueError("cutoffs must be nonnegative integers")
    seen = set()
    for doc in corpus:
        doc_id = doc["id"]
        if type(doc_id) is not str or not doc_id or doc_id in seen:
            raise ValueError("document IDs must be nonempty unique strings")
        seen.add(doc_id)
        for name in ("bi_score", "cross_score"):
            score = doc[name]
            if type(score) not in (int, float) or not isfinite(score):
                raise ValueError("scores must be finite numbers")
    candidates = sorted(corpus, key=lambda doc: doc["bi_score"], reverse=True)[
        :candidate_k
    ]
    reranked = sorted(
        candidates,
        key=lambda doc: doc["cross_score"],
        reverse=True,
    )
    return [doc["id"] for doc in reranked[:top_k]]

results = search_with_rerank(documents)
print(results)
print("shortlist of one:", search_with_rerank(documents, candidate_k=1))
bad_scores = [dict(doc) for doc in documents]
bad_scores[0]["bi_score"] = float("nan")
try:
    search_with_rerank(bad_scores)
except ValueError as error:
    print("invalid ranking:", error)

['api-key-rotation', 'replace-old-api-token']
shortlist of one: ['kms-key-rotation']
invalid ranking: scores must be finite numbers

Instruction-tuned embeddings and asymmetric prefixes

The problem with asymmetric search and task ambiguity

Information retrieval is structurally asymmetric. Search queries are typically short, telegraphic, and interrogative ("how do I rotate an API key?", "403 expired token"). Document passages are long, declarative, and descriptive ("Create a replacement key, deploy it, then revoke the old key.").

Sharing encoder weights does not force questions and answers to identical vectors. However, an encoder used for several tasks needs to know which similarity relation you want. Should “Python logging error” cluster with programming topics, or retrieve a troubleshooting passage? Role prefixes and task instructions can make that distinction explicit when the model was trained to use them.

Task-specific prefixes and asymmetric instruction tuning

These named checkpoints and APIs illustrate different input contracts, checked October 10, 2026:

• intfloat/e5-base-v2: Its card specifies "query: " for queries and "passage: " for documents in asymmetric retrieval. For symmetric similarity, both inputs use "query: ".^{[3]Reference 3E5-base-v2 model cardhttps://huggingface.co/intfloat/e5-base-v2}
• BAAI/bge-base-en-v1.5: Its card recommends a query instruction for short-query-to-passage retrieval, with plain passages. It also documents use without instructions and recommends comparing both settings on your task.^{[15]Reference 15BGE-base-en-v1.5 model cardhttps://huggingface.co/BAAI/bge-base-en-v1.5}
• Original INSTRUCTOR API: Accepts [instruction, text] pairs for domain and task conditioning.^{[16]Reference 16INSTRUCTOR embedding usagehttps://github.com/xlang-ai/instructor-embedding} Instructions are inputs the model was trained to interpret, not a promise that arbitrary prompts improve retrieval.

The operational danger of prefix mismatch

Follow the selected checkpoint's model card, rather than applying one family-wide rule. E5-base-v2 warns that omitting its prefixes degrades performance; BGE-v1.5 describes a more flexible instruction setting. Neither establishes a universal size of recall loss or mathematically separate query and passage subspaces. Record prompt formatting alongside the checkpoint, pooling, width, and scoring function. If document formatting changes, regenerate the affected corpus vectors and evaluate the resulting index.

def format_e5_pair(query: str, passage: str) -> tuple[str, str]:
    """Asymmetric retrieval inputs for intfloat/e5-base-v2."""
    return f"query: {query}", f"passage: {passage}"

def format_bge_v15_pair(query: str, passage: str) -> tuple[str, str]:
    """Recommended instructed retrieval setting for BAAI/bge-base-en-v1.5."""
    instruction = "Represent this sentence for searching relevant passages: "
    return f"{instruction}{query}", passage

def format_instructor_input(instruction: str, text: str) -> list[str]:
    """One instruction/text pair for the original INSTRUCTOR encode API."""
    return [instruction, text]

e5_query, e5_passage = format_e5_pair(
    "how do I rotate an API key?",
    "Create a replacement key, deploy it, then revoke the old key.",
)

bge_query, bge_passage = format_bge_v15_pair(
    "how do I rotate an API key?",
    "Create a replacement key, deploy it, then revoke the old key.",
)

instructor_example = format_instructor_input(
    "Represent the developer question for retrieving API docs:",
    "how do I rotate an API key?",
)

assert e5_query.startswith("query: ")
assert e5_passage.startswith("passage: ")
assert bge_query.startswith("Represent this sentence")
assert bge_passage == "Create a replacement key, deploy it, then revoke the old key."
assert instructor_example[0].startswith("Represent")
assert instructor_example[1] == "how do I rotate an API key?"
# Input-format check only; no model is loaded or called.

Matryoshka representation learning (MRL)

What arbitrary truncation can lose

Suppose your vector budget allows only 128 dimensions per document. Slicing a width-768 embedding to vector[:128] saves coordinates, but it can discard distinctions needed to rank passages. The amount of quality loss is an empirical question; collapse is not inevitable.

Predict the ranking after cutting these constructed unit vectors to one coordinate. Does re-normalization restore the lost sign in the second coordinate?

# Continue info_nce.py for normalize and dot. No trained model is used.
query = [0.6, 0.8]
match = [0.6, 0.8]
distractor = [0.6, -0.8]
print("full cosines:", [round(dot(query, row), 2) for row in (match, distractor)])
short_query = normalize(query[:1])
print("prefix cosines:", [dot(short_query, normalize(row[:1])) for row in (match, distractor)])
print("a prefix can stay unit:", dot([1.0], [1.0]) == 1.0)
try:
    normalize([0.0])  # Valid full vector [0, 1] has a zero first-coordinate prefix.
except ValueError:
    print("zero prefix rejected")

full cosines: [1.0, -0.28]
prefix cosines: [1.0, 1.0]
a prefix can stay unit: True
zero prefix rejected

Re-normalizing restores length, not information. These two prefixes tie because their distinguishing coordinate was discarded. A truncation trained for your target width aims to avoid losing the distinctions the task needs.

The nested representation mechanism

Matryoshka representation learning (MRL) supervises selected prefix widths so one representation can serve several capacity budgets.^{[17]Reference 17Matryoshka Representation Learning.https://arxiv.org/abs/2205.13147} Named after Russian nesting dolls, it encourages useful shorter representations nested inside a larger vector. The original paper applies the idea to several objectives, including classification; it is not restricted to sentence-level InfoNCE.

For a cosine-based contrastive application, choose prefix widths $\mathcal{D} = \{d_1, d_2, \dots, d_M\}$, such as $\{32,64,128,256,512,768\}$. At each selected width, slice both embeddings, normalize nonzero prefixes, and apply the contrastive objective:

$\mathcal{L}_{\text{MRL}} = \sum_{d \in \mathcal{D}} w_d \, \mathcal{L}_{\text{InfoNCE}}\left( \frac{z_{1:d}}{\|z_{1:d}\|_2}, \frac{z^+_{1:d}}{\|z^+_{1:d}\|_2}; \tau \right)$

Length and scoring after a cut

For a unit full vector, a prefix has norm at most 1. It stays 1 if every discarded coordinate is zero, and it can become zero if all retained coordinates are zero. For dot-product search intended to equal cosine, normalize each nonzero prefix. A cosine function may perform that normalization internally. If a model instead uses unnormalized dot products, follow and evaluate that scoring contract rather than silently replacing it with cosine.

Hierarchical two-stage vector search

MRL enables a high-throughput deployment pattern called hierarchical (adaptive) retrieval:

1. Stage 1 (Short-vector search): Index a supported 64- or 128-dimensional prefix and retrieve a shortlist, perhaps 500 candidates. Ten million width-64 float32 vectors have a raw payload of 2.56 GB, versus 30.72 GB at width 768 (decimal GB). Graph links, IDs, metadata, and other index overhead add storage.
2. Stage 2 (Full-width rescoring): Fetch the full vectors for those candidates and rescore with the selected model's scoring function. Full width and numerical precision are separate choices; quantized vectors remain approximate.
3. Measure the outcome: Compare first-stage Recall@500, final ranking quality, latency, and total storage against full-width retrieval. The shortlist cannot recover omitted passages. Neither 99% accuracy retention nor a particular speedup follows from the dimensions alone.

# Continue the prefix example and reuse info_nce.py's normalized loss.
def matryoshka_loss(
    embeddings_a: list[list[float]],
    embeddings_b: list[list[float]],
    dims: tuple[int, ...] = (2, 4, 6),
) -> float:
    if not embeddings_a or len(embeddings_a) != len(embeddings_b):
        raise ValueError("expected equal nonempty paired batches")
    width = len(embeddings_a[0])
    if any(len(row) != width for row in embeddings_a + embeddings_b):
        raise ValueError("all embedding widths must match")
    if not dims or len(set(dims)) != len(dims):
        raise ValueError("choose distinct prefix widths")
    if any(type(dim) is not int or not 1 <= dim <= width for dim in dims):
        raise ValueError("prefix width outside embedding")
    losses = []

    for dim in dims:
        truncated_a = [row[:dim] for row in embeddings_a]
        truncated_b = [row[:dim] for row in embeddings_b]
        losses.append(info_nce_loss(truncated_a, truncated_b))

    return sum(losses) / len(losses)

queries = [[1.0, 0.0, 0.9, 0.1, 0.5, 0.2], [0.0, 1.0, 0.1, 0.9, 0.2, 0.5]]
positives = [[0.95, 0.05, 0.85, 0.15, 0.45, 0.25], [0.05, 0.95, 0.15, 0.85, 0.25, 0.45]]

loss = matryoshka_loss(queries, positives)
for dim in (2, 4, 6):
    print(f"dim={dim} loss={matryoshka_loss(queries, positives, (dim,)):.4f}")
print(round(loss, 4))

dim=2 loss=0.0088
dim=4 loss=0.0147
dim=6 loss=0.0225
0.0153

Why it matters

Evaluation: STS and MTEB

Semantic textual similarity (STS)

STS asks whether embedding similarities preserve human judgments of relatedness. Datasets like STS-B (STS Benchmark) provide sentence pairs rated by annotators. These developer-themed examples and ratings are invented to illustrate the format:

"A CI build failed" / "A test run failed" => 4.5
"A CI build failed" / "A password was reset" => 1.2

Compute cosine similarity for every pair, then calculate Spearman rank correlation against the human ratings. High correlation means similar ordering on those pairs. It doesn't calibrate cosine as a probability or guarantee relevant passages will rank above thousands of corpus distractors.

MTEB (Massive text embedding benchmark)

As models improved, optimizing only for STS stopped being enough. A model that's excellent at pairwise relatedness can still fail at retrieval or clustering. The Massive Text Embedding Benchmark (MTEB) was introduced to make those task differences visible.^{[18]Reference 18MTEB: Massive Text Embedding Benchmark.https://arxiv.org/abs/2210.07316}

The original paper evaluated 33 models on 58 datasets grouped into 8 task categories and found that no one model dominated every category. Much of the retrieval slice comes from BEIR.^{[19]Reference 19BEIR: A Heterogeneous Benchmark for Zero-shot Evaluation of Information Retrieval Models.https://arxiv.org/abs/2104.08663}

The table below is the original snapshot, not today's catalog. The 2025 MMTEB paper expanded evaluation to over 500 quality-controlled tasks across more than 250 languages.^{[20]Reference 20MMTEB: Massive Multilingual Text Embedding Benchmarkhttps://arxiv.org/abs/2502.13595} That collection is not the same as every selected leaderboard benchmark. Current MTEB offers named benchmark scopes and filters for language, task, modality, and model type.^{[21]Reference 21MTEB Leaderboardhttps://leaderboard.mteb.org/} Record the exact scope, dataset revisions, splits, model revision, and input formatting before comparing scores.

The original reranking category orders supplied passage lists using embedding similarities. It does not automatically evaluate the cross-encoder pipeline discussed above.^{[18]Reference 18MTEB: Massive Text Embedding Benchmark.https://arxiv.org/abs/2210.07316}

Choosing a model in practice

Use a leaderboard average to screen candidates, not to make the deployment decision. Start with the operational questions that define your own retrieval contract:

• Does the model expect plain text, query/passage prefixes, or explicit instructions?
• Which shorter widths are supported, and do their measured retrieval results meet your requirements?
• How well does it handle your language mix, domain jargon, and query length distribution?
• What are the latency, throughput, and memory costs once you batch and index it at production scale?
• Do you still need BM25 or a cross-encoder reranker to hit Recall@K and NDCG (normalized discounted cumulative gain) targets?

A leaderboard delta doesn't establish that a model change solves poor chunking, shortlist misses, or mislabeled negatives. Evaluate the whole retrieval path. The original MTEB comparison found no model dominating every task category.^{[18]Reference 18MTEB: Massive Text Embedding Benchmark.https://arxiv.org/abs/2210.07316} Compare candidates within the same benchmark scope, then measure API-doc retrieval, languages, latency, and storage on your own workload.^{[21]Reference 21MTEB Leaderboardhttps://leaderboard.mteb.org/} A reranker is worth adding when its measured gains fit your budget; it is not a mandatory stage for every application.

Measure coverage and ordering separately

For a query with two relevant passages, retrieving one of them gives Recall@2 of 1/2, not 1. Reciprocal rank asks a different question: how far down is the first relevant result? Average it across queries to get mean reciprocal rank (MRR). These fixed rankings expose both cases:

retrieval_cases = [
    ({"rotation", "replacement"}, ["kms", "rotation", "replacement"]),
    ({"expiry"}, ["expiry", "rotation", "kms"]),
    ({"revocation"}, ["kms", "expiry", "rotation"]),
]

def metrics_at_k(relevant: set[str], ranking: list[str], k: int) -> tuple[float, float]:
    if not relevant or type(k) is not int or k < 1:
        raise ValueError("need relevant judgments and a positive integer cutoff")
    if len(set(ranking)) != len(ranking):
        raise ValueError("ranking must contain unique document IDs")
    selected = ranking[:k]
    recall = len(relevant & set(selected)) / len(relevant)
    reciprocal_rank = next(
        (1 / rank for rank, doc in enumerate(selected, 1) if doc in relevant), 0.0
    )
    return recall, reciprocal_rank

results = [metrics_at_k(relevant, ranking, 2) for relevant, ranking in retrieval_cases]
for i, (recall, rr) in enumerate(results, 1):
    print(f"query {i}: recall@2={recall:.2f}, RR@2={rr:.2f}")
print("mean recall@2:", f"{sum(r for r, _ in results) / len(results):.2f}")
print("MRR@2:", f"{sum(rr for _, rr in results) / len(results):.2f}")

query 1: recall@2=0.50, RR@2=0.50
query 2: recall@2=1.00, RR@2=1.00
query 3: recall@2=0.00, RR@2=0.00
mean recall@2: 0.50
MRR@2: 0.50

Moving rotation from rank 2 to rank 1 improves that query's reciprocal rank from 0.5 to 1 without changing Recall@2. Adding replacement to the top two improves recall. Neither action recovers the third query's absent relevant passage.

For a real experiment, hold out queries and near-duplicate groups from training and negative mining, record the corpus version, and keep relevance judgments fixed when comparing models. Incomplete judgments can label useful results as misses. Evaluate token truncation and passage boundaries too: no pooling rule can preserve a sentence that was cut off before encoding.

Key libraries and tools

Building embedding-based systems requires the right tooling:

Mastery check

Key concepts

• alignment and uniformity in embedding space
• InfoNCE numerator, denominator, and temperature
• hard negatives vs easy negatives
• bi-encoder vs cross-encoder vs late interaction
• reranking as recall first, then precision
• Matryoshka prefix training and evaluation of shorter widths

Evaluation rubric

• Foundational: Derives the InfoNCE objective and explains what the numerator, denominator, and temperature do.
• Intermediate: Distinguishes static averaging, contextual pooling, and sentence-level training, including padding masks and normalization.
• Intermediate: Explains why hard negatives matter more than random negatives once the model already separates broad topics.
• Advanced: Compares bi-encoders, cross-encoders, and late-interaction models by latency, indexability, accuracy, and storage.
• Advanced: Explains ColBERT's MaxSim scoring and why it keeps more token-level signal than a single document vector.
• Advanced: Explains Matryoshka embeddings and when shorter prefixes are worth the storage-accuracy tradeoff.
• Advanced: Designs a two-stage production retrieval pipeline with recall and latency budgets defended quantitatively.

Follow-up questions

Common pitfalls

Raw [CLS] is treated like a search-ready sentence embedding

• Symptom: Nearly every query-document pair gets suspiciously similar cosine scores.
• Possible cause: Raw BERT [CLS] lacks sentence-retrieval training. Repeated vectors or a scoring bug can also flatten results; high absolute cosine alone does not diagnose failure.
• Fix: Inspect vectors, masks, and matrix shapes, then compare a retrieval-trained encoder on held-out rankings.

Negatives stay too easy

• Symptom: Training loss falls, but recall on realistic queries barely moves.
• Possible cause: The supplied negatives are already separated, or the training labels fail to represent held-out queries.
• Fix: Inspect loss-active rows and retrieval errors; add audited confusing non-matches when those errors justify mining.

The reranker is asked to save missing recall

• Symptom: The reranker looks good in pairwise inspection, yet the right document is often absent in production results.
• Cause: The correct passage never entered the shortlist.
• Fix: Measure first-stage Recall@K separately. Improve retrieval or increase the shortlist only within a measured latency budget.

Dimensions are shortened blindly

• Symptom: Index storage drops as expected, but retrieval quality falls off a cliff.
• Possible cause: The discarded coordinates contained useful distinctions, or prefix normalization and indexing used an inconsistent scoring contract.
• Fix: Use Matryoshka-trained or provider-documented shortening controls and benchmark each target width.

Task conditioning is ignored

• Symptom: One embedding model works for clustering but underperforms on retrieval.
• Possible cause: The selected checkpoint expected query/passage prefixes or instructions, but every input was embedded as plain text.
• Fix: Follow the model card format for retrieval, clustering, and classification separately.

Try it yourself

These exercises let you verify your understanding without needing a GPU cluster.

Exercise 1: compute triplet loss by hand

Given an anchor $a$, positive $p$, and negative $n$ with distances $d(a,p) = 0.3$ and $d(a,n) = 0.7$, compute the triplet loss for margins $m = 0.1$ and $m = 0.5$. In which case does the model still have work to do?

Solution sketch: For $m = 0.1$: $0.3 - 0.7 + 0.1 = -0.3$, so $\max(0, -0.3) = 0$. The margin is already satisfied. For $m = 0.5$: $0.3 - 0.7 + 0.5 = 0.1$, so $\max(0, 0.1) = 0.1$. The larger margin forces the model to pull the positive even closer or push the negative farther away.

Exercise 2: spot the embedding mistake

A teammate gets nearly identical scores for every query-document pair using raw BERT [CLS]. What should you inspect before blaming the embedding geometry, and which model comparison would be useful?

Solution sketch: Check that distinct texts produce distinct vectors, inspect attention masks and pooling, and verify that scoring computes the intended query-by-document matrix. Then compare held-out rankings using a retrieval-trained encoder with its documented input format. Raw [CLS] lacks that training contract, but close scores alone do not prove anisotropy is the cause or guarantee a model swap fixes the pipeline.

Exercise 3: design a two-stage retrieval pipeline

You have 2 million documentation chunks and a latency budget of 200 ms per query. Query encoding takes 10 ms; sequential cross-encoder scoring takes 15 ms per pair. Assume no batching or parallelism. Why does scoring the full corpus miss the budget, and what shortlist might fit?

Solution sketch: $2{,}000{,}000 \times 15\,\text{ms} = 30{,}000{,}000\,\text{ms} \approx 8.3$ hours per query. The cross-encoder is far too slow for the full corpus. Reserve 10 ms for query encoding and choose a top-10 shortlist only if ANN lookup and overhead fit inside the remaining 40 ms: reranking then costs $10 \times 15\,\text{ms} = 150\,\text{ms}$, for at most 200 ms total. If Recall@10 is inadequate, the budget or reranker throughput must change; silently widening to top 100 violates the requirement.