Medium Primitives

Triplet Loss

Triplet loss compares three embeddings (vectors representing examples): an anchor, a related positive, and an unrelated negative. The positive and negative names describe their relationship to the anchor, not the signs of their vector entries.

This exercise uses squared Euclidean distance: subtract two vectors, square each difference, then sum over their coordinates. For each triplet, compute max(0, d_pos - d_neg + margin).

For an anchor [0, 0], positive [1, 0], and negative [2, 0], the squared distances are 1 and 4. With margin 5, the loss is max(0, 1 - 4 + 5) = 2. With margin 2, it is zero.

A triplet has zero loss once the negative's squared distance is at least margin greater than the positive's. Otherwise the loss penalizes the shortfall. Minimizing this objective encourages that separation; it does not guarantee that every triplet will satisfy it after training.

Your task

Implement triplet_loss(anchor, positive, negative, margin):

Each embedding tensor has shape (N, d), with N > 0 and d > 0. Corresponding rows form one triplet. Entries are floating-point values.
margin is a nonnegative Python float, in squared-distance units.
Return a scalar tensor containing the mean loss over all N triplets.

Sum squared coordinate differences along the last axis. Do not take a square root or normalize the embeddings. Clamp each triplet's loss at zero before averaging. Include zero-loss triplets in the average. For per-triplet terms [2, -4], the answer is (2 + 0) / 2 = 1.

Use this squared-distance formula directly. A library triplet loss that uses unsquared Euclidean distance will produce different values.

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…

Triplet Loss

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