Medium Research

Contrastive Loss (InfoNCE)

InfoNCE turns matching embeddings into a classification task. An embedding is a vector representing an example. Each anchor has one designated match among a batch of candidates, and the loss rewards assigning that match a high probability. The objective comes from Contrastive Predictive Coding.

Here row i of anchors matches row i of positives. Compare each anchor with every row of positives using a dot product. Divide these scores by a positive temperature, then apply softmax across each row. The loss for anchor i is the negative natural log of probability (i, i). Other rows of positives serve as negatives for that anchor, even if they happen to contain identical vectors.

For anchors = positives = [[1, 0], [0, 1]] and temperature 1, the score matrix is [[1, 0], [0, 1]]. Each correct match receives probability exp(1) / (exp(1) + exp(0)), about 0.731059. The mean loss is about 0.313262. With only one candidate, its probability is 1 and the loss is zero.

Your task

Implement infonce_loss(anchors, positives, temperature):

Both floating-point inputs have shape (N, d), with N > 0 and d > 0.
temperature is a positive Python float. Assume scaled dot products are finite. A lower temperature emphasizes differences between scores.
Use the input embeddings as given. Do not normalize them.
Average the N anchor-to-candidate losses. Do not add a reverse-direction loss or compare anchors with other anchors.
Return the scalar loss as a Python float. Decimal rounding is optional; numerical comparisons allow normal floating-point error.

Compute log probabilities stably. Subtract each row's maximum score before exponentiating and keep the selected score in those same shifted coordinates. Adding a large maximum back only to subtract it again can lose the small normalization term. For a row of equal scores, the loss should be log(N) even when the common score is very large.

The Python-float return is this exercise's output format. In a training loop, keep the loss as a tensor so automatic differentiation can use it.

Understand the ideas behind this exercise. The Problem tab contains the exact implementation requirements.

Learn vectors that express similarity

An embedding maps an example to a vector. Metric learning trains these vectors so that related examples are close or similar under a chosen comparison. Cosine similarity compares direction:

$$\operatorname{sim}(a,b)=\frac{a\cdot b}{\|a\|\|b\|}.$$

Vectors [1, 0] and [2, 0] have similarity 1 despite different lengths. Vectors [1, 0] and [0, 1] have similarity 0. A zero vector has no direction, so follow the task's rule for a zero denominator or epsilon.

Compare an anchor, a positive, and a negative

Triplet loss uses an anchor a, a related positive p, and an unrelated negative n. A common formula is

$$\max(0,d(a,p)-d(a,n)+m).$$

The margin m specifies how much farther away the negative should be. For positive distance 2, negative distance 4, and margin 1, the loss is zero. If the negative distance is 2.5, the loss is 0.5.

Euclidean distance and squared Euclidean distance give different losses and gradients. Use the definition in the exercise.

Select the positive from several candidates

InfoNCE treats the matching task as classification. Divide similarity scores by a positive temperature to get logits, then apply cross-entropy with the positive candidate as the target. A lower temperature concentrates probability more strongly on the largest scores without changing their order.

Check which axis contains candidates, which index is the positive, and whether to normalize embeddings first. Paired batches often place positives on the diagonal, but only when their ordering supports that convention. Computing a loss in one matching direction differs from averaging losses in both directions.

Detecting runtime…

Contrastive Loss (InfoNCE)

Your task

Learn vectors that express similarity

Compare an anchor, a positive, and a negative

Select the positive from several candidates

Your code, saved here

Stuck?

JAX reference solution