Triplet Margin Loss

Problem

Implement triplet margin loss for metric learning. Given anchor, positive, and negative embeddings, compute the loss that encourages the anchor to be closer to the positive than the negative by at least a specified margin.

Solution

import numpy as np

def triplet_margin_loss(anchor: np.ndarray, positive: np.ndarray, negative: np.ndarray, margin: float = 1.0) -> float:
    # anchor, positive, negative: shape (batch_size, embedding_dim)
    dist_pos = np.sqrt(np.sum((anchor - positive) ** 2, axis=1))
    dist_neg = np.sqrt(np.sum((anchor - negative) ** 2, axis=1))

    losses = np.maximum(dist_pos - dist_neg + margin, 0.0)
    return float(np.mean(losses))

Explanation

1. Compute the Euclidean distance between each anchor and its positive example.
2. Compute the Euclidean distance between each anchor and its negative example.
3. The loss for each triplet is max(d_pos - d_neg + margin, 0). This is zero when the negative is already farther than the positive by at least the margin.
4. Average the per-triplet losses over the batch.

Complexity

• Time: O(B * d) where B is the batch size and d is the embedding dimension
• Space: O(B) for the distance arrays