What it actually does
Nothing here specifies what an embedding should be. It specifies only that paired things should end up closer than unpaired things, and the geometry falls out of enforcing that over millions of pairs. This is why the technique is defined almost entirely by its data: the encoder is ordinary, and the consequential choices are what counts as a positive pair, where the negatives come from, and the temperature that decides how hard the loss pushes on the closest wrong answer. Get the pairs right and a small encoder beats a much larger one. Get them wrong and the embeddings look fine on held-out pairs from the same distribution while failing on everything real, because the model learned to solve the sampling artefact rather than the task.
- A meaningful notion of positive pair exists in the data, since that pairing is the entire supervision signal
- The negatives approximate the candidate distribution the system will face at serving time
- A single geometry can serve every downstream use, which fails when different consumers need different notions of similarity
- Similarity is symmetric and transitive enough for a metric space to represent it — asymmetric relations such as entailment fit badly
- The corpus can be re-embedded whenever the encoder changes, since the two are inseparable operations
Loss function
A cross-entropy over similarities: identify the true partner among a set of candidates. Two things are worth reading carefully. The denominator contains the negatives, so the loss is defined relative to a candidate set rather than in absolute terms — the same embedding is good or bad depending on what it is compared against, which is why the negative-sampling policy belongs in the model definition and not in the data loader. And the temperature is not a nuisance parameter: expanding the gradient shows that its reciprocal scales the weight on each negative, so a small temperature makes the loss almost entirely about the single closest wrong answer. That is the hard-negative behaviour people attribute to sampling tricks, and a substantial part of it is just the temperature.
- embeddings of a positive pair, almost always L2-normalized so similarity is a cosine
- the negatives for anchor i — the design decision that dominates everything else here
- temperature; small values concentrate the gradient on the single hardest negative
- cosine similarity after normalization, which bounds the logits and makes the temperature meaningful
- anchors in the batch; with in-batch negatives the batch size IS the negative count
How it is minimized
AdamW with the largest batch that fits, since in-batch negatives make batch size a capacity parameter rather than a throughput knob
- softmax weight on negative j — concentrated on the nearest negative when tau is small
- scales the whole gradient, which is why temperature and learning rate interact strongly
- a weighted centroid of the negatives; the anchor is pushed away from it
- the positive, pulled toward unconditionally and with weight 1 regardless of temperature
The gradient makes the mechanism legible in a way the loss does not. Each anchor is pulled toward its positive with unit weight and pushed away from a softmax-weighted centroid of its negatives, and the temperature controls only how peaked that weighting is. At high temperature all negatives contribute roughly equally and training is stable but slow to sharpen; at low temperature the gradient is dominated by the nearest negative, which is the informative one and also the one most likely to be a mislabelled positive. That single sentence is the whole tension of the technique. It also explains the batch-size effect: with in-batch negatives the candidate set is the batch, so a larger batch is a harder and better-estimated problem rather than merely a faster one — which is why published results at batch sizes in the thousands do not reproduce at a hundred, and why the memory bank and momentum-encoder tricks exist at all.
temperature
The most consequential knob. Small values concentrate the gradient on the hardest negative, which is both the signal and the risk (0.01 to 0.2)
batch size
With in-batch negatives this is the negative count, so it is a capacity parameter rather than a throughput knob (256 to 32768)
hard-negative ratio
Share of mined negatives against random ones. Too high and false negatives dominate the gradient (0 to 0.5)
embedding dimension
Trades retrieval quality against index memory and query latency; gains flatten well before most defaults (64 to 1024)
augmentation strength
Defines what invariance is learned. This is the actual task specification in the self-supervised setting (task-specific)
projection head depth
Discarded after training. The representation before it usually transfers better than the one the loss was applied to (0 to 2 layers)
Three characteristic failures, and the first two look nothing alike. Representation collapse is the classic one: all embeddings converge to a narrow cone or a single point, the loss plateaus at log of the batch size, and the diagnostic is the mean pairwise similarity rather than the loss — which is why alignment and uniformity are tracked as separate metrics. It arises when the negatives are too easy or too few, and normalization plus enough negatives usually prevents it. The second failure is subtler and more common in practice: false negatives. With in-batch sampling, some negatives are genuine positives that happen to be in the batch, and at a small temperature the gradient concentrates precisely on those, so the model is explicitly trained to separate things that belong together. The symptom is good held-out pair accuracy with poor retrieval, which is easy to misread as an indexing problem. The third is a shortcut failure: the encoder solves the augmentation rather than the semantics — matching on colour histogram or on a tokenizer artefact — and the tell is that accuracy collapses under a mild distribution shift that should not matter.
Encoding is O(B·C) for encoder cost C, and the similarity matrix adds O(B²·d), which is negligible until batch sizes reach the thousands and then becomes a real memory term at B² floats. Retrieval at serving time is O(log n) with an approximate index against O(n·d) for exact search, and that gap is the entire reason this technique is preferred to a cross-encoder despite being less accurate pair-for-pair.
How, where, and why it is used
Time-Series Forecasting
AdaptedLearn a window embedding by treating temporally adjacent or augmentation-related windows as positives, then use the embedding for retrieval-based forecasting: find the k nearest historical windows and combine their continuations. The embedding replaces a hand-built similarity measure, which is the actual contribution — the forecast itself is a weighted average of what happened next.
- Retrieval-based forecasting over a large historical panel, where analogues exist and are worth finding
- Regime identification, where the embedding clusters windows into states a downstream model conditions on
- Cold-start series with no history of their own but many similar series to borrow from
- Similarity search over a series library for exploration and quality control
The analogue story is genuinely useful where the panel is large enough that near-duplicates of the current window exist. It also has a property most forecasting methods lack: the retrieved neighbours are an explanation, so a forecast comes with the historical episodes that produced it, which is frequently worth more to a reviewer than an accuracy point. Two honest limits. Defining positives is harder than in language or vision — temporal adjacency conflates "similar" with "nearby", so the embedding can learn a clock rather than a shape, and this is the default failure. And a properly tuned direct forecaster usually beats retrieval on plain accuracy, so the case for this rests on interpretability, cold-start, or a genuinely enormous analogue library rather than on error metrics.
- Normalize each window before encoding, or the embedding sorts by level rather than by shape
- Use augmentations that preserve shape — jitter, scaling, masking — rather than adjacency alone, which teaches a clock
- Exclude windows overlapping the anchor from the negatives; they are false negatives by construction and the temperature will find them
- Keep the split strictly temporal, since retrieval from the future is the easiest leak in this whole design
Rolling-origin backtesting of the retrieval forecast against a direct forecaster, which is the comparison that decides whether the interpretability is free or paid for. Separately, inspect retrieved neighbours by eye on a sample — an embedding that retrieves windows from the same month rather than the same shape is a clock, and no aggregate metric says so clearly.
- Adjacency-only positives, which teach recency rather than similarity
- Overlapping windows as negatives, which are false negatives the small temperature then punishes hardest
- Retrieval across a non-temporal split, which leaks the future into the neighbour set
- No direct-forecaster baseline, so the accuracy cost of the interpretability is never measured
Anomaly Detection
AdaptedFit the embedding on normal data only, then score by distance to the k nearest normal neighbours in embedding space — a point far from everything the encoder has seen packed together is anomalous. The embedding does the work of turning a high-dimensional distance that means nothing into one that means something.
- High-dimensional inputs where raw-space distance is meaningless — images, logs, sensor windows
- Detection with abundant normal data and effectively no labelled anomalies
- Settings where the nearest normal neighbours must be shown as evidence for the alert
- Novelty detection over a growing reference set, where new normal modes can be added without retraining
A reasonable fit with one important caveat about what the score is. The technique earns its place because a learned metric makes nearest-neighbour detection viable in dimensions where Euclidean distance is uninformative, and because the retrieved neighbours are a natural explanation for an alert — the reviewer sees what the system compared against. The caveat is that this is a distance, not a probability: it has no calibration, the threshold is chosen from a quantile of a normal-data score distribution and nothing more, and it does not compose with other detectors that produce genuine likelihoods. The characteristic failure is subtle and worth naming: whatever invariances the augmentations built in become blind spots, so if the augmentation included brightness jitter the embedding cannot see a brightness anomaly. That means the augmentation policy silently determines which anomalies are detectable, and it is rarely reviewed with that in mind.
- Fit on verified-normal data only; contaminated training data teaches the encoder that anomalies belong in the normal cone
- Review augmentations as a list of things the detector will be blind to, because that is exactly what they are
- L2-normalize before indexing, since the index and the loss must use the same metric
- Set the threshold from a quantile of normal-data scores, and state plainly that it is not calibrated
Precision at the alert volume the reviewers can actually handle, rather than AUC — the operating point is a staffing constraint and AUC hides it. Score stability under benign shift matters as much as detection rate here, because an embedding that drifts moves every threshold with it.
- Treating the distance as a probability; it is uncalibrated and does not combine with likelihood-based detectors
- Augmentations that make the anomaly of interest invariant, which is a blind spot introduced on purpose and forgotten
- Contaminated normal data, which is the failure that most flatters offline metrics
- A stale reference set, so legitimate new normal modes register as anomalies
Optimization
Primary fitThree separate optimization problems live here. The loss itself is a constrained geometry problem — pull positives together, spread everything else out, on a sphere. The negative-sampling policy is a sampling-design problem with a sharp trade at its centre. And serving is an index problem, since the whole point of a dual encoder over a cross-encoder is precomputable embeddings.
- Temperature as the knob that decides how much of the gradient the hardest negative receives
- Batch size as a capacity parameter, because in-batch negatives make the batch the candidate set
- Hard-negative mining against false-negative risk, which is the central and unavoidable trade
- Alignment against uniformity as two measurable quantities the single loss value conflates
Worth studying for two lessons that transfer well beyond embeddings. The first is that a loss defined relative to a sampled candidate set makes the sampler part of the model, and no amount of architecture work substitutes for getting it right — this is the clearest example in this reference of the data policy dominating the model. The second is that decomposing a single loss into two measurable properties, alignment and uniformity, turns an opaque number into a diagnosis: collapse shows as uniformity going to zero while the loss sits at log of the batch size, and nothing about the loss alone distinguishes that from healthy training. Both lessons generalize to any objective with a sampled denominator.
- Track alignment and uniformity separately; the loss value alone cannot distinguish collapse from convergence
- Tune temperature and learning rate together, since the gradient scales with the reciprocal of temperature
- Cap the hard-negative share and audit mined negatives for false positives, which are what a small temperature punishes hardest
- Measure retrieval with the production index rather than exact search, because the index is part of the system
Recall at k through the actual approximate index, since exact-search numbers overstate what production will do. Alongside it, the mean pairwise similarity as a collapse guard, and a scan of mined hard negatives for false positives — a few minutes of reading the mined pairs routinely finds a systematic error the metrics cannot.
- Tuning temperature by loss value, which improves as the geometry degenerates
- Reporting exact-search recall for a system that will serve through an approximate index
- Mining ever-harder negatives without auditing them, which trains the model to separate true positives
- Comparing across batch sizes without saying so, since batch size changes the problem rather than the speed
Natural Language
Primary fitTrain a dual encoder on query-document pairs — from click logs, question-answer data or synthetic generation — then index document embeddings and retrieve by nearest neighbour at query time. The two towers may share weights or not; sharing is usually better with limited data and worse when queries and documents have genuinely different surface form.
- First-stage retrieval in search and in retrieval-augmented generation, where the candidate pool is too large to score pairwise
- Semantic deduplication and near-duplicate detection at corpus scale
- Clustering and topic discovery over documents with no labels
- Zero-shot classification by embedding label descriptions and retrieving the nearest
The dominant approach for first-stage retrieval, and the reason is structural: document embeddings are precomputed, so query time is one encoder pass plus an index lookup regardless of corpus size. A cross-encoder is more accurate on any single pair and cannot be precomputed, which is why the standard architecture is contrastive retrieval followed by cross-encoder reranking of the top candidates — the two are complements rather than alternatives. The honest caveats are about what embeddings are bad at. Exact matching on rare terms, identifiers and numbers is a known weakness, which is why hybrid retrieval with a lexical index remains standard practice rather than a legacy holdover. And these models are sensitive to the query-document asymmetry in their training data, so a model trained on short keyword queries degrades on long natural-language ones in a way that is invisible until it is measured.
- Match the query and document prefixes or instructions the checkpoint was trained with; they are part of the model
- Use hybrid retrieval with a lexical index, since embeddings reliably miss exact rare-term matches
- Mine hard negatives from the current index rather than sampling randomly, and audit them for false negatives
- Rerank the top candidates with a cross-encoder; the two-stage design exists because each stage is bad at the other job
Recall at the cutoff the reranker actually consumes — recall at 100 when 100 are reranked, not at 10 — since the first stage exists to feed the second. Measure through the production index, and break results out by query length and by rare-term presence, which is where the failures hide.
- Omitting the instruction prefix the checkpoint expects, which degrades everything with no error
- Pure dense retrieval where exact identifier matching matters
- Query-distribution mismatch between training and serving, invisible until measured by query type
- Evaluating at a cutoff unrelated to how many candidates the downstream stage consumes
Computer Vision
Primary fitEither self-supervised, with two augmented views of the same image as the positive pair, or paired across modalities using image-caption data. The augmentation policy in the first case and the caption distribution in the second are the actual task specification — they define what the encoder is required to treat as the same thing.
- Pretraining vision encoders without labels, then fine-tuning or linear-probing for a downstream task
- Image retrieval and visual near-duplicate detection at scale
- Zero-shot classification by embedding class descriptions and retrieving the nearest
- Cross-modal search where text queries retrieve images and the reverse
Established practice, and the instructive part is how visible the augmentation dependence is: strong colour jitter produces embeddings invariant to colour, which is exactly right for object recognition and exactly wrong for anything where colour is the signal. Nowhere else in this reference is a hyperparameter so directly a task specification. Contrastive pretraining also outperforms masked reconstruction under linear probing while being weaker for dense prediction, so the two objectives rank differently depending on the downstream task and neither is generally better. The cross-modal variant is the more interesting case because the supervision is nearly free at web scale, and its weakness follows from the same fact: captions describe salient objects, so these embeddings are strong on object identity and weak on counting, spatial relations and fine-grained attributes.
- Choose augmentations as a statement of what must be treated as identical, because that is what they are
- Use large batches or a memory bank; the negative count is a capacity parameter here more than anywhere
- Discard the projection head at inference — the layer beneath it transfers better than the one the loss touched
- Do not expect counting or spatial reasoning from caption-supervised embeddings; the supervision never contained it
Linear probe and full fine-tuning reported separately, because they rank pretraining objectives differently and a single number hides that. For retrieval, recall at k through the production index, plus a deliberate check on the attributes the augmentations made invariant.
- Augmentations that destroy the signal the downstream task needs
- Small batches without a memory bank, which starves the loss of negatives and invites collapse
- Keeping the projection head, which usually costs transfer quality
- Assuming fine-grained or relational understanding from caption supervision that never taught it
Recommendation & Ranking
Primary fitEmbed users and items into one space using interactions as positive pairs, then retrieve candidates by nearest neighbour against the user embedding. This is candidate generation, not ranking — a separate and usually much heavier model scores the retrieved few hundred.
- Candidate generation over catalogues too large to score item by item
- Cold-start items, embedded from content features rather than from interaction history
- Related-item and complementary-item retrieval from co-occurrence
- Session-based retrieval where the recent sequence is encoded as the query
Standard architecture for candidate generation, for the same precomputation reason as text retrieval. The domain-specific complication is popularity bias, and it is severe: random negatives are dominated by unpopular items, so the model learns that popular items are universally good and the embedding space collapses toward a popularity ordering. The standard correction is sampling negatives proportional to popularity with a logit correction, and it is not optional — without it the retrieved set is a bestseller list with extra steps. The second complication is that interaction positives conflate interest with exposure: a user clicked what was shown, so the embedding partly learns the previous ranker’s behaviour, which is a feedback loop that offline metrics computed on logged data cannot see at all.
- Sample negatives by popularity with a logit correction; uncorrected random negatives produce a popularity ranking
- Include content features so cold-start items can be embedded without interaction history
- Refresh item embeddings on a schedule, since catalogue drift moves the geometry under a static index
- Treat the exposure-versus-interest confound explicitly, because logged-data metrics cannot surface it
Recall of relevant items at the candidate-set size the ranker consumes, split by item popularity so the tail is visible, followed by an online test — offline metrics on logged data are measuring the previous ranker as much as this one. Track catalogue coverage of the retrieved set, since a collapsed embedding retrieves the same few thousand items for everyone.
- Uncorrected random negatives, which bake in popularity bias
- Confusing candidate generation with ranking and expecting calibrated scores from a distance
- A stale item index after catalogue turnover
- Believing offline metrics on logged data, which reflect the incumbent ranker’s exposure decisions
What it costs to run
Dominated by batch size rather than by model size, because in-batch negatives make the batch the candidate set — which pushes training toward large-memory hardware or a memory bank even for modest encoders. Fine-tuning an existing embedding model on domain pairs is hours on one accelerator and is usually the right starting point. Illustrative, not a measured benchmark.
One encoder pass per query plus an index lookup, with document embeddings precomputed. That is the entire architectural argument for a dual encoder over a cross-encoder, and it holds: query latency is independent of corpus size, while index memory grows linearly with it.
Re-embedding the corpus is the expensive operation, so the cadence is set by how often that can be afforded rather than by how often the model improves. Note the coupling that surprises people: a new encoder invalidates every stored embedding, so model updates and full re-indexing are one operation and must be planned together.
- Mean pairwise similarity of a fixed probe set as a collapse guard — the loss will not tell you
- Recall at k on a frozen labelled probe set, measured through the production index rather than exact search
- Query-embedding norm and nearest-neighbour distance distribution, which shift when queries move out of domain
- Index staleness: the fraction of corpus items embedded by an older encoder version, which should normally be zero
- Score-distribution drift on normal data where the embedding feeds anomaly detection, since every threshold moves with it
- Changing the encoder invalidates every stored embedding. Model updates and full corpus re-indexing are a single operation, and treating them separately produces a silently mixed index
- Embeddings from different models are not comparable even at identical dimension. Mixing two versions in one index yields nonsense with no error anywhere
- The similarity metric must match training. An index configured for inner product over embeddings trained with cosine degrades quietly rather than failing
- The instruction or prefix a checkpoint was trained with is part of the model; omitting it at query time degrades retrieval with no error
- A distance is not a probability. Thresholds are quantiles of an observed score distribution, they do not transfer across encoder versions, and they do not compose with likelihood-based scores
- Exact-search evaluation overstates production recall, sometimes substantially. Measure through the index that will actually serve
- Augmentation-induced invariances are permanent blind spots, and they are rarely written down anywhere a future maintainer will find them
Pros and cons, in context
Precomputable embeddings, so query cost is independent of corpus size
The whole architectural argument for a dual encoder. One encoder pass plus an index lookup at any scale, which is why this is the first stage of essentially every large retrieval system.
Supervision from pairing alone, with no labels required
Co-occurrence, augmentation and click logs all provide pairs, which makes the technique available where labels are not. The trade is that the pairing definition becomes the most consequential decision in the system.
One geometry serves retrieval, clustering, deduplication and detection
A single trained encoder supports several unrelated consumers, which is unusual and a genuine operational advantage — provided they all want the same notion of similarity, which is worth checking.
Retrieved neighbours are a built-in explanation
Every result comes with the items it was closest to, which for anomaly review and forecast justification is frequently worth more than an accuracy point.
The negative-sampling policy dominates the architecture
The loss is defined relative to a sampled candidate set, so the sampler is part of the model. Two teams with the same encoder and different negatives get materially different systems, and no architecture work compensates.
False negatives are punished hardest by design
In-batch negatives include genuine positives, and a small temperature concentrates the gradient on exactly those. The symptom is good pair accuracy with poor retrieval, which is easy to misdiagnose as an indexing problem.
Representation collapse, invisible in the loss
Embeddings converge to a narrow cone while the loss plateaus at log of the batch size. Diagnosable only from mean pairwise similarity or a uniformity metric, which is why both are tracked separately.
Batch size is a capacity parameter, not a throughput knob
A larger batch is a harder and better-estimated problem, so published results at batch sizes in the thousands do not reproduce at a hundred. Memory banks and momentum encoders exist to work around this rather than to speed anything up.
Weak at exact matching
Rare terms, identifiers and numbers are reliably missed, which is why hybrid retrieval with a lexical index is standard practice rather than legacy baggage.
The score is a distance with no calibration
Thresholds are quantiles of an observed distribution, do not transfer across encoder versions, and do not compose with likelihood-based scores. Anything needing a probability needs a different model.
Augmentation choices become permanent blind spots
Whatever the augmentations made invariant, the embedding cannot see. That silently determines which anomalies are detectable and which distinctions are representable, and it is almost never documented.
Intuitive, idiomatic, optimized
The same algorithm in three languages, each progressing through Kent Beck's three stages. Every optimization names the lever it pulls and the trade-off it accepts — cost profiles are illustrative, not measured benchmarks.
The math, transcribed literally. No dependencies.
"""InfoNCE contrastive embeddings, transcribed from the objective.
No ML library. The loss written out so the two halves can be read against
each other:
L = -mean_i log[ exp(sim(z_i, z_i+)/tau)
/ (exp(sim(z_i, z_i+)/tau) + sum_j exp(sim(z_i,z_j)/tau)) ]
The thing to read for: the DENOMINATOR contains the negatives. The loss is
defined relative to a candidate set, not in absolute terms, so the same
embedding is good or bad depending on what it is compared against. That is
why the negative-sampling policy belongs in the model definition and not
in the data loader.
"""
import math
def dot(left, right):
return sum(left[idx] * right[idx] for idx in range(len(left)))
def l2_normalize(vector):
"""Project onto the unit sphere.
Not cosmetic. Normalization bounds every similarity to [-1, 1], which
is what makes the temperature a meaningful scale rather than an
arbitrary one, and it is the main structural defence against collapse.
"""
norm = math.sqrt(dot(vector, vector))
if norm < 1e-12:
return [0.0] * len(vector)
return [value / norm for value in vector]
def encode(features, w_hidden, w_out):
"""The encoder is deliberately boring: one hidden layer and a project.
Everything consequential about this technique is in the pairs and the
negatives, not here. Swapping this for a transformer changes the
quality and changes none of the reasoning below it.
"""
hidden = [
max(0.0, dot(row, features)) for row in w_hidden
]
projected = [dot(row, hidden) for row in w_out]
return l2_normalize(projected)
def info_nce_loss(anchors, positives, negatives_for, temperature):
"""The loss, one anchor at a time.
`negatives_for(i)` returns the candidate set for anchor i. Passing it
as a function rather than a matrix is the honest signature: the
sampling policy is an argument to the loss, because it IS part of the
loss.
"""
if temperature <= 0.0:
raise ValueError("temperature must be positive")
total = 0.0
for idx, anchor in enumerate(anchors):
positive_logit = dot(anchor, positives[idx]) / temperature
logits = [positive_logit]
for negative in negatives_for(idx):
logits.append(dot(anchor, negative) / temperature)
# Log-sum-exp over positive AND negatives. The positive appears in
# the denominator too -- it is one candidate among many, and the
# loss asks the model to pick it out.
top = max(logits)
log_z = top + math.log(sum(math.exp(value - top) for value in logits))
total += log_z - positive_logit
return total / len(anchors)
def in_batch_negatives(embeddings, anchor_index):
"""Every other item in the batch is a negative.
Free and extremely effective, and it carries the technique's most
common quiet failure: some of these ARE positives that happen to share
the batch. At a small temperature the gradient concentrates on exactly
those, so the model is explicitly trained to separate things that
belong together. The symptom is good held-out pair accuracy with poor
retrieval, which usually gets misdiagnosed as an indexing problem.
"""
return [
embedding
for other, embedding in enumerate(embeddings)
if other != anchor_index
]
def info_nce_gradient(anchor, positive, negatives, temperature):
"""The gradient, because it explains the mechanism better than the loss.
dL/dz_i = (1/tau) * [ sum_j p_ij * z_j - z_i+ ]
Each anchor is pulled toward its positive with UNIT weight and pushed
away from a softmax-weighted centroid of its negatives. The
temperature controls only how peaked that weighting is:
- high tau: all negatives contribute about equally. Stable, slow.
- low tau: the nearest negative dominates. That negative is the most
informative one AND the most likely to be a mislabelled positive.
That single trade is the whole tension of the technique.
"""
width = len(anchor)
logits = [dot(anchor, positive) / temperature]
for negative in negatives:
logits.append(dot(anchor, negative) / temperature)
top = max(logits)
exps = [math.exp(value - top) for value in logits]
total = sum(exps)
probabilities = [value / total for value in exps]
# Weighted centroid of the candidates, positive included.
centroid = [0.0] * width
candidates = [positive] + list(negatives)
for weight, candidate in zip(probabilities, candidates):
for k in range(width):
centroid[k] += weight * candidate[k]
return [
(centroid[k] - positive[k]) / temperature for k in range(width)
]
def alignment_and_uniformity(anchors, positives, alpha=2.0, t=2.0):
"""Decompose the loss into two things you can actually diagnose.
The loss value alone cannot distinguish collapse from convergence.
These two can:
alignment -- mean squared distance between positive pairs. Lower is
better; it measures whether pairs are pulled together.
uniformity -- log mean exp(-t * squared distance) over ALL pairs.
Lower (more negative) means better spread.
Collapse shows as uniformity going to ~0 while the loss sits at
log(batch_size) and looks completely healthy.
"""
n_items = len(anchors)
align = 0.0
for idx in range(n_items):
squared = sum(
(anchors[idx][k] - positives[idx][k]) ** 2
for k in range(len(anchors[idx]))
)
align += squared ** (alpha / 2.0)
align /= n_items
pair_total = 0.0
pair_count = 0
for left in range(n_items):
for right in range(left + 1, n_items):
squared = sum(
(anchors[left][k] - anchors[right][k]) ** 2
for k in range(len(anchors[left]))
)
pair_total += math.exp(-t * squared)
pair_count += 1
uniformity = math.log(pair_total / pair_count) if pair_count else 0.0
return {"alignment": align, "uniformity": uniformity}
def retrieve(query, corpus, top_k):
"""Nearest-neighbour retrieval by exact search.
Exact search here so the loss and the retrieval use provably the same
metric. In production this becomes an approximate index, and the
difference is not negligible: exact-search recall overstates what the
served system does, sometimes substantially.
"""
scored = [(dot(query, embedding), idx) for idx, embedding in enumerate(corpus)]
scored.sort(reverse=True)
return [idx for _, idx in scored[:top_k]]
Encoding is O(B·C) for encoder cost C and the similarity matrix adds O(B²·d), which is negligible until batch sizes reach the thousands and then becomes a real memory term at B² floats. Illustrative, not a measured benchmark: this version recomputes the log-sum-exp per anchor in Python and does exact retrieval, so it demonstrates the objective rather than any realistic cost.
How this language actually wants it written.
"""Contrastive embeddings with the sampler as a declared component.
The design point: the loss is defined relative to a sampled candidate set,
so the sampler is part of the model. Hiding it in a data loader hides the
most consequential decision in the system. Here it is a protocol with a
name, and the two things it can get wrong -- too few negatives and false
negatives -- are both checkable.
"""
from __future__ import annotations
from contextlib import contextmanager
from dataclasses import dataclass
from typing import Final, Iterator, Protocol, Sequence
import numpy as np
from numpy.typing import NDArray
# Above this mean pairwise cosine, the embedding occupies a narrow cone and
# the geometry has effectively degenerated.
COLLAPSE_SIMILARITY: Final[float] = 0.9
MIN_TEMPERATURE: Final[float] = 1e-3
class RepresentationCollapse(RuntimeError):
"""Embeddings have converged to a narrow cone.
The loss will NOT show this: it plateaus at log(batch_size) and looks
like healthy convergence. Only mean pairwise similarity or a
uniformity metric distinguishes the two, which is why both are
computed here rather than left to the caller.
"""
class DegenerateTemperature(ValueError):
"""Temperature at or below zero, or small enough to overflow.
The gradient scales with 1/tau, so a very small temperature both
concentrates all the gradient on the nearest negative and makes the
logits numerically unstable.
"""
class MetricMismatch(ValueError):
"""The index metric disagrees with the training metric.
A silent production failure: an index configured for raw inner product
over embeddings trained with cosine similarity degrades quietly rather
than erroring, and the symptom looks like a quality regression.
"""
class TooFewNegatives(ValueError):
"""A candidate set small enough that the loss barely constrains anything.
With in-batch negatives the batch size IS the negative count, which is
why results at batch sizes in the thousands do not reproduce at a
hundred: the larger batch is a harder problem, not a faster one.
"""
class NegativeSampler(Protocol):
"""Where the negatives come from.
Declared as a protocol because this is the component that decides what
the model learns. Two teams with the same encoder and different
samplers have materially different systems.
"""
@property
def name(self) -> str:
...
def candidates(
self, embeddings: NDArray[np.float32], anchor: int
) -> NDArray[np.int32]:
...
@dataclass(frozen=True, slots=True)
class InBatchSampler:
"""Every other item in the batch. Free, effective, and quietly risky.
Some of these candidates are genuine positives that happen to share
the batch. At a small temperature the gradient concentrates on exactly
those, so the model is trained to separate things that belong
together. `known_duplicates` exists so a caller who can identify them
is able to exclude them.
"""
known_duplicates: tuple[tuple[int, int], ...] = ()
@property
def name(self) -> str:
return "in-batch"
def candidates(
self, embeddings: NDArray[np.float32], anchor: int
) -> NDArray[np.int32]:
excluded = {anchor}
for left, right in self.known_duplicates:
if left == anchor:
excluded.add(right)
elif right == anchor:
excluded.add(left)
return np.array(
[idx for idx in range(embeddings.shape[0]) if idx not in excluded],
dtype=np.int32,
)
@dataclass(frozen=True, slots=True)
class MinedHardSampler:
"""Hard negatives mined from the current index, capped deliberately.
The cap is the point. Ever-harder negatives improve the gradient right
up until the hardest candidates are predominantly false negatives, and
then the model is being trained to separate true positives. There is
no way to detect that from the loss, so the share is bounded and the
mined pairs are meant to be audited by eye.
"""
hard_indices: tuple[tuple[int, ...], ...]
max_hard_fraction: float = 0.5
def __post_init__(self) -> None:
if not 0.0 <= self.max_hard_fraction <= 1.0:
raise ValueError(
f"max_hard_fraction must lie in [0, 1], got {self.max_hard_fraction}"
)
@property
def name(self) -> str:
return f"mined-hard(cap={self.max_hard_fraction})"
def candidates(
self, embeddings: NDArray[np.float32], anchor: int
) -> NDArray[np.int32]:
n_items = embeddings.shape[0]
random_pool = [idx for idx in range(n_items) if idx != anchor]
budget = int(len(random_pool) * self.max_hard_fraction)
hard = [idx for idx in self.hard_indices[anchor] if idx != anchor][:budget]
remaining = [idx for idx in random_pool if idx not in set(hard)]
return np.array(hard + remaining[: len(random_pool) - len(hard)], dtype=np.int32)
@dataclass(frozen=True, slots=True)
class EmbeddingGeometry:
"""Alignment and uniformity, which the loss value conflates.
Returned rather than computed on request, because a diagnostic that
has to be assembled by hand does not get watched and this failure is
silent in every other signal.
"""
alignment: float
uniformity: float
mean_pairwise_similarity: float
max_pairwise_similarity: float
@property
def has_collapsed(self) -> bool:
return self.mean_pairwise_similarity > COLLAPSE_SIMILARITY
def raise_if_collapsed(self) -> None:
"""Guard clause for a training loop that should stop now."""
if not self.has_collapsed:
return
raise RepresentationCollapse(
f"mean pairwise cosine {self.mean_pairwise_similarity:.3f} exceeds "
f"{COLLAPSE_SIMILARITY}; uniformity {self.uniformity:.3f} -- the loss "
"value will look healthy regardless"
)
@dataclass(frozen=True, slots=True)
class ContrastiveResult:
loss: float
geometry: EmbeddingGeometry
sampler_name: str
negatives_per_anchor: int
class ContrastiveObjective:
"""InfoNCE with the sampler injected and the geometry reported."""
def __init__(
self,
sampler: NegativeSampler,
temperature: float = 0.07,
min_negatives: int = 16,
) -> None:
if temperature <= MIN_TEMPERATURE:
raise DegenerateTemperature(
f"temperature {temperature} at or below {MIN_TEMPERATURE}; the "
"gradient scales with 1/tau and the logits will overflow"
)
self._sampler = sampler
self._temperature = temperature
self._min_negatives = min_negatives
def __call__(
self,
anchors: NDArray[np.float32],
positives: NDArray[np.float32],
) -> ContrastiveResult:
# Guard clauses first, cheapest to most expensive.
if anchors.shape != positives.shape:
raise ValueError(
f"anchor shape {anchors.shape} != positive shape {positives.shape}"
)
if anchors.shape[0] < 2:
raise TooFewNegatives(
f"batch of {anchors.shape[0]} leaves no negatives at all"
)
anchors = l2_normalize(anchors)
positives = l2_normalize(positives)
n_anchors = anchors.shape[0]
total = 0.0
negative_count = 0
for anchor_index in range(n_anchors):
candidates = self._sampler.candidates(anchors, anchor_index)
if candidates.size < self._min_negatives:
raise TooFewNegatives(
f"sampler {self._sampler.name} produced {candidates.size} "
f"negatives for anchor {anchor_index}, below "
f"{self._min_negatives}"
)
negative_count = candidates.size
anchor = anchors[anchor_index]
positive_logit = float(anchor @ positives[anchor_index])
negative_logits = anchors[candidates] @ anchor
logits = np.concatenate(
([positive_logit], negative_logits)
) / self._temperature
shifted = logits - logits.max()
log_z = float(np.log(np.exp(shifted).sum()) + logits.max())
total += log_z - positive_logit / self._temperature
return ContrastiveResult(
loss=total / n_anchors,
geometry=measure_geometry(anchors, positives),
sampler_name=self._sampler.name,
negatives_per_anchor=negative_count,
)
def l2_normalize(embeddings: NDArray[np.float32]) -> NDArray[np.float32]:
"""Project onto the unit sphere. Shared by the loss AND the index.
One function, used by both, so they cannot disagree about the metric --
which is the silent production failure this guards against.
"""
norms = np.linalg.norm(embeddings, axis=-1, keepdims=True)
return embeddings / np.maximum(norms, 1e-12)
def measure_geometry(
anchors: NDArray[np.float32],
positives: NDArray[np.float32],
alpha: float = 2.0,
t: float = 2.0,
) -> EmbeddingGeometry:
"""Alignment, uniformity, and the two similarity statistics.
Cheap enough to compute every step, and the only signal that separates
a collapsed embedding from a converged one.
"""
squared_positive = np.sum((anchors - positives) ** 2, axis=-1)
alignment = float(np.mean(squared_positive ** (alpha / 2.0)))
similarity = anchors @ anchors.T
upper = np.triu_indices_from(similarity, k=1)
pairwise_similarity = similarity[upper]
squared_all = np.maximum(2.0 - 2.0 * pairwise_similarity, 0.0)
uniformity = float(np.log(np.mean(np.exp(-t * squared_all))))
return EmbeddingGeometry(
alignment=alignment,
uniformity=uniformity,
mean_pairwise_similarity=float(np.mean(pairwise_similarity)),
max_pairwise_similarity=float(np.max(pairwise_similarity)),
)
@contextmanager
def embedding_index(
corpus: NDArray[np.float32], metric: str = "cosine"
) -> Iterator[NDArray[np.float32]]:
"""Own the index for the duration of a query session, then release it.
A context manager because the index is a real resource: it is
frequently a memory-mapped file or an on-device buffer, and the version
coupling matters -- a new encoder invalidates every stored embedding,
so an index outliving its encoder is a correctness bug rather than a
stale cache.
"""
if metric != "cosine":
raise MetricMismatch(
f"index metric {metric!r} but embeddings are trained with cosine; "
"an inner-product index over normalized vectors is equivalent, "
"anything else degrades silently"
)
normalized = l2_normalize(corpus)
try:
yield normalized
finally:
del normalized
def retrieve(
query: NDArray[np.float32],
index: NDArray[np.float32],
top_k: int,
) -> Sequence[int]:
"""Exact nearest neighbour, which overstates production recall.
Kept exact here so the metric provably matches the loss. Anything
serving at scale uses an approximate index, and the recall gap between
the two is not negligible -- measure through the index that will
actually serve.
"""
if top_k <= 0:
raise ValueError(f"top_k must be positive, got {top_k}")
scores = index @ l2_normalize(query[None, :])[0]
return np.argsort(-scores)[:top_k].tolist()
What changed: The negative-sampling policy becomes an explicit protocol rather than an implicit loop, because the loss is defined relative to a candidate set and hiding the sampler in the data loader misrepresents where the modelling happens. Alignment and uniformity are returned alongside the loss rather than computed by a caller who remembers to, since collapse is invisible in the loss value and a diagnostic nobody assembles is a diagnostic nobody watches. Specific exceptions separate the failures that actually occur: collapse detected from mean pairwise similarity, a temperature outside the range where the gradient is numerically sane, and an embedding-metric mismatch between training and the index. The encoder and the index are given one shared normalization step so they cannot disagree about the metric, which is a silent production failure otherwise.
Same asymptotics as the naive version, with the per-anchor similarity computed as one matvec rather than a Python loop. Illustrative, not a measured benchmark: the substantive change is that a collapsed geometry now raises and the sampler is a named component, so the decision that dominates retrieval quality is visible in the call site instead of buried in a data loader.
Faster, and the justification for each change.
"""Contrastive embeddings as one Gram matrix and one masked softmax.
With in-batch negatives, every anchor's candidate set is the same batch.
So the natural form of the loss is not a loop over anchors at all -- it is
a single B x B similarity matrix, masked, then one row-wise cross-entropy
where the correct class for row i is column i.
That matrix is also the memory cost of the technique. It grows as B^2, and
since batch size is a CAPACITY parameter here rather than a throughput
knob -- a larger batch is a harder and better-estimated problem, not just
a faster one -- the B^2 term is what actually bounds how good the model
can get on given hardware. That is why memory banks and momentum encoders
exist: they add negatives without adding to this matrix.
"""
from __future__ import annotations
from dataclasses import dataclass
from typing import Final
import numpy as np
from numpy.typing import NDArray
COLLAPSE_SIMILARITY: Final[float] = 0.9
MASK_VALUE: Final[float] = -1e4 # additive; float32-safe, unlike -inf
@dataclass(slots=True)
class ContrastiveBuffers:
"""Allocated once per batch shape and overwritten in place.
Held at the largest shape seen, so memory does not shrink when a
smaller batch follows a larger one -- an acceptable trade for a
training loop that does no per-step allocation.
"""
similarity: NDArray[np.float32]
mask: NDArray[np.float32]
row_max: NDArray[np.float32]
log_z: NDArray[np.float32]
@classmethod
def for_batch(cls, max_batch: int) -> ContrastiveBuffers:
return cls(
similarity=np.empty((max_batch, 2 * max_batch), dtype=np.float32),
mask=np.zeros((max_batch, 2 * max_batch), dtype=np.float32),
row_max=np.empty(max_batch, dtype=np.float32),
log_z=np.empty(max_batch, dtype=np.float32),
)
def normalize_inplace(embeddings: NDArray[np.float32]) -> NDArray[np.float32]:
"""L2-normalize in place on a contiguous float32 block.
Contiguity and a single dtype are not cosmetic here: the Gram matmul
below is the dominant BLAS call, and a strided or mixed-precision input
forces a copy on every step.
"""
if embeddings.dtype != np.float32 or not embeddings.flags.c_contiguous:
embeddings = np.ascontiguousarray(embeddings, dtype=np.float32)
norms = np.linalg.norm(embeddings, axis=-1, keepdims=True)
np.maximum(norms, 1e-12, out=norms)
embeddings /= norms
return embeddings
def info_nce(
anchors: NDArray[np.float32],
positives: NDArray[np.float32],
buffers: ContrastiveBuffers,
temperature: float = 0.07,
false_negative_pairs: NDArray[np.int32] | None = None,
) -> tuple[float, NDArray[np.float32]]:
"""One Gram matrix, one masked softmax, one cross-entropy.
Layout: column j of the similarity block is positive j for j < B, and
anchor j-B for j >= B. So row i has its true positive at column i, and
every other column -- both blocks -- is a negative. That single layout
gives 2B-1 negatives per anchor from one B x 2B matmul.
The returned matrix is reused by the geometry metrics and by
hard-negative mining, which is the fusion that matters: recomputing it
three times would triple the dominant cost.
"""
if temperature <= 0.0:
raise ValueError(f"temperature must be positive, got {temperature}")
anchors = normalize_inplace(anchors)
positives = normalize_inplace(positives)
n_anchors = anchors.shape[0]
similarity = buffers.similarity[:n_anchors, : 2 * n_anchors]
np.matmul(anchors, positives.T, out=similarity[:, :n_anchors])
np.matmul(anchors, anchors.T, out=similarity[:, n_anchors:])
similarity /= np.float32(temperature)
mask = buffers.mask[:n_anchors, : 2 * n_anchors]
mask.fill(0.0)
# An anchor against itself in the second block is not a negative, it is
# the same vector: similarity 1 regardless of training.
np.fill_diagonal(mask[:, n_anchors:], MASK_VALUE)
# Known false negatives, masked additively rather than by rebuilding
# index arrays per anchor. These are the candidates a small temperature
# punishes hardest, so excluding the ones that can be identified is
# the single highest-value correction available.
if false_negative_pairs is not None and false_negative_pairs.size > 0:
left = false_negative_pairs[:, 0]
right = false_negative_pairs[:, 1]
mask[left, right] = MASK_VALUE
mask[right, left] = MASK_VALUE
mask[left, n_anchors + right] = MASK_VALUE
mask[right, n_anchors + left] = MASK_VALUE
similarity += mask
# Fused log-sum-exp: subtract the row max in place, then one log of one
# sum rather than materializing exp() as a separate array.
row_max = buffers.row_max[:n_anchors]
np.max(similarity, axis=1, out=row_max)
similarity -= row_max[:, None]
log_z = buffers.log_z[:n_anchors]
np.log(np.exp(similarity).sum(axis=1), out=log_z)
# The correct class for row i is column i. No gather needed.
positive_logits = np.diagonal(similarity[:, :n_anchors])
loss = float(np.mean(log_z - positive_logits))
return loss, similarity
def geometry_from_similarity(
similarity: NDArray[np.float32],
n_anchors: int,
temperature: float,
t: float = 2.0,
) -> dict[str, float]:
"""Alignment and uniformity from the matrix already computed.
Reuses the shifted matrix rather than recomputing cosines, which is the
whole point of passing it back: the geometry diagnostics are then
nearly free and can run every step.
Collapse shows as mean similarity near 1 and uniformity near 0 while
the loss sits at log(2B-1) and looks entirely healthy. No other signal
separates those two states.
"""
anchor_block = similarity[:, n_anchors:] * np.float32(temperature)
upper = np.triu_indices(n_anchors, k=1)
pairwise = anchor_block[upper]
squared = np.maximum(2.0 - 2.0 * pairwise, 0.0)
uniformity = float(np.log(np.mean(np.exp(-t * squared))))
positive_block = similarity[:, :n_anchors] * np.float32(temperature)
positive_similarity = np.diagonal(positive_block)
alignment = float(np.mean(np.maximum(2.0 - 2.0 * positive_similarity, 0.0)))
mean_similarity = float(np.mean(pairwise))
return {
"alignment": alignment,
"uniformity": uniformity,
"mean_pairwise_similarity": mean_similarity,
"collapsed": float(mean_similarity > COLLAPSE_SIMILARITY),
}
def mine_hard_negatives(
similarity: NDArray[np.float32],
n_anchors: int,
top_k: int,
exclude_self: bool = True,
) -> NDArray[np.int32]:
"""The hardest in-batch negatives, from the same matrix again.
Worth being explicit about what these are. They are the candidates the
gradient already concentrates on at a small temperature, so mining
them is mostly a way to find them for AUDIT rather than a way to
increase their weight. Reading a sample of mined pairs by eye
routinely finds a systematic false-negative pattern that no aggregate
metric surfaces.
"""
block = similarity[:, n_anchors:].copy()
if exclude_self:
np.fill_diagonal(block, MASK_VALUE)
# argpartition rather than argsort: only the top-k boundary matters.
partitioned = np.argpartition(-block, top_k, axis=1)[:, :top_k]
return partitioned.astype(np.int32)
def batch_size_is_capacity(batch_size: int, embedding_dim: int) -> dict[str, float]:
"""What the B^2 term costs, and why memory banks exist.
Illustrative arithmetic, not a benchmark. The ratio is the point: the
similarity matrix overtakes the embedding block as the memory
consumer once the batch exceeds the embedding dimension, and since a
larger batch is a strictly harder problem for this loss, that crossover
is where the technique stops scaling on-device and starts needing a
memory bank or a momentum encoder to add negatives off the matrix.
"""
embedding_bytes = float(batch_size) * embedding_dim * 4.0
similarity_bytes = float(batch_size) * 2.0 * batch_size * 4.0
return {
"negatives_per_anchor": float(2 * batch_size - 1),
"embedding_mib": embedding_bytes / (1024.0 ** 2),
"similarity_mib": similarity_bytes / (1024.0 ** 2),
"similarity_share": similarity_bytes / (embedding_bytes + similarity_bytes),
}
What changed: The per-anchor loop becomes one B-by-B similarity matrix and a single vectorized cross-entropy, which is the natural form of the objective — with in-batch negatives, every anchor’s candidate set is the same batch, so the whole loss is one masked softmax over a Gram matrix. That matrix is the memory cost of the technique at scale and is computed once and reused for the loss, the geometry diagnostics and the hard-negative mining rather than three times. Normalization is done in place on a contiguous float32 block so the matmul hits BLAS at full efficiency, the log-sum-exp is fused, and false negatives are excluded by additive masking rather than by rebuilding index arrays per anchor.
Vectorize to NumPy/BLAS (@, np.dot, einsum)
The whole loss becomes one Gram matrix and one masked softmax instead of a per-anchor loop of matvecs
Trade-off: Peak memory grows as the square of the batch size, which becomes the binding constraint at the batch sizes this loss actually wants
Fuse adjacent operations so intermediates are never materialized
One similarity matrix serves the loss, the geometry metrics and hard-negative mining rather than being recomputed for each
Trade-off: The three consumers now share a mutable buffer, so the masking order matters and an in-place edit for one silently affects the others
Ensure contiguous memory layout and a single dtype
Normalized embeddings form one contiguous float32 block, so the Gram matmul runs at full BLAS efficiency
Trade-off: A copy is required when the caller supplies a non-contiguous or mixed-precision view, which costs bandwidth on every step
Pre-allocate output arrays and use in-place operations
The similarity matrix and mask are allocated once per batch shape and overwritten, so a training loop does no per-step allocation
Trade-off: The buffers are held at the largest batch shape seen, so memory does not shrink when a smaller batch follows a larger one
The loss becomes one B-by-2B matmul plus a row-wise softmax, so encoder cost dominates again and the similarity term is O(B²·d) arithmetic against O(B²) memory. Illustrative, not a measured benchmark: that quadratic memory is the real ceiling, because batch size here is a capacity parameter rather than a throughput knob, and it is why memory banks exist to add negatives without adding to this matrix.