Triplet Loss

Triplet Loss#

The triplet loss trains an embedding model by comparing three examples at a time: an anchor, a positive example of the same class, and a negative example of a different class. The goal is to place the positive closer to the anchor than the negative, usually by at least a fixed margin. The emphasis is not on achieving particular absolute distances, but on preserving the correct ordering of examples around each anchor.

The triplet loss is well suited to retrieval and recognition tasks, where the main requirement is that relevant examples rank ahead of irrelevant ones. Its effectiveness, however, depends strongly on triplet selection. Most possible triplets are either already correctly ordered or too easy to provide a useful learning signal, which motivates strategies for selecting informative triplets during training.

In this tutorial, we use the triplet loss to train an embedding model on MNIST dataset. We also explore two online triplet-mining strategies to select informative triplets during training.

Triplet Loss

Definition of triplet loss#

The triplet loss trains the model from groups of three examples. Each triplet contains an anchor, a positive example similar to the anchor, and a negative example dissimilar to it. The objective is to make the positive closer to the anchor than the negative by at least a margin. This leads to the following formulation.

\[ \mathcal{L}_{\rm triplet}(z_a, z_p, z_n) = \max(0, \|z_a - z_p\| - \|z_a - z_n\| + m) \]

Here, \(z_a\), \(z_p\), and \(z_n\) are the embeddings of the anchor, positive, and negative examples, respectively, and \(m\) is the required margin between the two distances.

Intuition#

The triplet loss compares the anchor-positive and anchor-negative distances directly.

Ordering. The positive should be closer to the anchor than the negative. If the negative is closer, the triplet contradicts the desired structure of the embedding space and produces a positive loss.
Margin. Correct ordering alone is not sufficient. The negative must be farther from the anchor by at least \(m\). A triplet in which the negative is only slightly farther away still contributes to the loss.

Training can reduce the loss by bringing the anchor and positive closer together, pushing the anchor and negative farther apart, or both. Once the distance ordering satisfies the margin requirement, the triplet contributes zero loss. The margin therefore defines sufficient relative separation. Unlike pair-based contrastive loss, the triplet loss does not impose an independent target on either distance. A positive pair need not be extremely close, and a negative pair need not be extremely far apart, provided that the positive is closer to the anchor by the required margin.

Triplet Loss

Triplet mining#

A dataset may contain an enormous number of possible triplets. Evaluating every possible combination would be computationally impractical. This would also be inefficient since many triplets already satisfy the margin and therefore contribute no loss. Triplet mining addresses this problem by selecting a smaller set of triplets that are informative for training. The idea is to focus on triplets that violate, or nearly violate, the desired ordering, as these provide meaningful gradients for learning.

The usefulness of a triplet depends on the current arrangement of the embeddings.

Easy triplet: The negative is already farther from the anchor than the positive by at least the margin. The constraint is satisfied, so the triplet contributes no loss.

\[ \mathcal{D}(a,n)\geq \mathcal{D}(a,p)+m. \]
Semi-hard triplet: The positive is correctly ranked closer to the anchor than the negative, but the separation is smaller than the required margin. These triplets remain informative because the ordering is correct but not yet sufficiently robust.

\[ \mathcal{D}(a,p)<\mathcal{D}(a,n)<\mathcal{D}(a,p)+m. \]

Hard triplet: The negative is at least as close to the anchor as the positive, so the desired ranking is violated. Hard triplets produce a strong training signal, although extremely difficult examples may correspond to outliers, ambiguous cases, or mislabeled data.

\[ \mathcal{D}(a,n)\leq \mathcal{D}(a,p). \]

Triplet Mining

Online mining#

Searching the entire dataset for triplets after every model update would be prohibitively expensive. In online mining, a batch is sampled first, and informative triplets are then constructed from the embeddings currently available in that batch. Because the embeddings change during training, the selected triplets adapt to the model’s current weaknesses. There are essentially two online strategies for triplet mining.

Batch-all mining. All valid anchor-positive-negative combinations within the batch are considered, and easy triplets with zero loss are excluded from the loss. For a class-balanced batch containing \(C\) classes and \(K\) examples per class, each anchor has \(K-1\) possible positives and \(CK-K\) possible negatives. The batch therefore contains \(CK(K-1)(CK-K)\) valid triplets before easy triplets are removed. Batch-all mining uses every informative combination, but many triplets may share the same examples and contribute highly redundant information.
Batch-hard mining. For each anchor, the farthest positive and the nearest negative within the batch are selected. This produces one challenging triplet per anchor and avoids enumerating every possible combination. Its effectiveness depends strongly on batch composition: each class must appear multiple times, and the batch must contain enough classes to provide meaningful negative candidates.

Batch-all uses a broad set of violated constraints, whereas batch-hard concentrates on the most difficult comparison available for each anchor. Neither strategy searches beyond the current batch, so class-balanced batches are important for both.

Note

An anchor is valid only if the batch contains another example of its class and at least one different class. This is why metric-learning samplers often place several examples from several classes in every batch. A large random batch only work for balanced datasets.

Best practices#

The effectiveness of triplet-loss training depends strongly on the triplets available in each batch. The following practices generally improve stability and efficiency.

Normalize embeddings. L2 normalization prevents the model from changing distances merely by scaling embedding norms. It also bounds Euclidean distances between normalized embeddings to the interval \([0,2]\), making the margin easier to interpret.
Use class-balanced batches. Each batch should contain multiple examples from several classes. Repeated examples per class provide positive candidates, while multiple classes provide negative candidates. Larger batches offer a broader selection of negatives, but they also increase computation and memory use. The batch size should therefore be chosen according to the model, dataset, and available hardware rather than fixed to a universal value.
Monitor triplet difficulty during training. Track the proportions of easy, semi-hard, and hard triplets, together with the anchor-positive and anchor-negative distance distributions. A high proportion of easy triplets may indicate that training is progressing, but it may also mean that the batches contain weak negative candidates. A persistent abundance of hard triplets may instead reveal poor representations, difficult batches, outliers, or noisy labels.
Tune the margin on validation data. A small margin imposes only weak separation, whereas a large margin keeps more triplets active and may enforce an unnecessarily difficult constraint. The appropriate value depends on the distance scale, the dataset, and the mining strategy. The training loss and the fraction of active triplets are useful diagnostics, but downstream retrieval or verification performance should guide the final choice.

PyTorch implementation#

We now implement the triplet loss with online mining in PyTorch. Since MNIST provides class labels, the loss receives a batch of embeddings and their corresponding labels, then constructs valid triplets directly within the batch. The implementation computes pairwise distances, identifies valid anchor-positive and anchor-negative relationships, and applies the selected mining strategy. This avoids generating triplets in advance and ensures that mining reflects the model’s current embedding geometry.

Note

All functions shown below are defined in triplet.py.

Helper function: `get_pairs()`#

The get_pairs() function identifies valid positive and negative pairs within a batch. Given a one-dimensional tensor of class labels, it returns two Boolean matrices.

Positive-pair mask. Entry \((i,j)\) is True when examples \(i\) and \(j\) have the same label and \(i \neq j\). The diagonal is excluded because an example should not be paired with itself.
Negative-pair mask. Entry \((i,j)\) is True when examples \(i\) and \(j\) have different labels.

Both masks are symmetric: whenever \((i,j)\) is selected, \((j,i)\) is selected as well.

def get_pairs(labels):
    """
    Get positive and negative pairs from labels and return them as boolean matrices.

    Args
        - labels: 1D tensor containing class labels

    Returns
        - positive: 2D boolean tensor where positive[i, j] = True if labels[i] == labels[j] and i != j
        - negative: 2D boolean tensor where negative[i, j] = True if labels[i] != labels[j]
    """
    positive = labels.unsqueeze(0) == labels.unsqueeze(1)
    positive.fill_diagonal_(False)
    negative = labels.unsqueeze(0) != labels.unsqueeze(1)
    return positive, negative

The output of this function is shown below for a simple test case.

Helper function: `get_triplets()`#

The get_triplets() function enumerates all valid triplets within a batch. Given a one-dimensional tensor of class labels, the function obtains the positive and negative masks from get_pairs(). It then combines them through broadcasting to construct a 3D Boolean tensor. Entry \((i,j,k)\) is True when \(i\) can serve as the anchor, \(j\) as its positive, and \(k\) as its negative. This requires that the anchor and positive are different examples sharing the same label, while the negative has a different label. Finally, the indices of all True entries are extracted to produce the anchor, positive, and negative index tensors. The size of each tensor is equal to the number of valid triplets in the batch. Entries at the same position define one triplet.

def get_triplets(labels):
    """
    Get all valid triplets (i, j, k) where i != j and labels[i] == labels[j] and labels[i] != labels[k]

    Args
        - labels: 1D tensor containing class labels

    Returns
        - anchor:   1D tensor containing indices of anchor samples
        - positive: 1D tensor containing indices of positive samples
        - negative: 1D tensor containing indices of negative samples
    """
    pos, neg = get_pairs(labels) # (N, N)
    pos = pos.unsqueeze(2) # (N, N, 1)
    neg = neg.unsqueeze(1) # (N, 1, N)
    triplets = pos & neg   # (N, N, N)
    return torch.nonzero(triplets, as_tuple=True)

The output of this function is shown below for a simple test case.

Batch-all mining#

Batch-all mining evaluates every valid triplet in the batch, but averages the loss only over triplets that violate the margin. These active triplets are either hard or semi-hard. The implementation generates all valid triplets, computes their unthresholded loss, remove easy triplets, and average the remaining values. All-easy batches return a differentiable zero connected to the embeddings, so backward() remains valid.

def triplet_loss_batch_all(embeddings: torch.FloatTensor, labels: torch.IntTensor, margin: float):

    # Get all triplets in the batch
    anchor, positive, negative = get_triplets(labels)

    # Compute all pairwise distances
    distances = torch.cdist(embeddings, embeddings, p=2)

    # Evaluate the triplet loss
    loss = distances[anchor, positive] - distances[anchor, negative] + margin
    
    # Remove easy triplets
    loss = loss[loss > 0]

    if loss.numel() == 0:
        return embeddings.sum() * 0.0
    else:
        return loss.mean()

Batch-hard mining#

Batch-hard mining constructs as many triplets as there are examples in the batch. Each example serves as an anchor and is paired with the most distant example of the same class (hardest positive) and the closest example of a different class (hardest negative). This produces one challenging triplet per anchor and avoids enumerating every valid triplet. The implementation makes sure to exclude from the loss any anchors that do not have at least one positive and one negative candidate.

def triplet_loss_batch_hard(embeddings, labels, margin):

    # Get all pairs in the batch
    positive, negative = get_pairs(labels)

    # Identify valid anchors
    valid_anchor = positive.any(dim=1) & negative.any(dim=1)

    # Ensure there is at least one valid anchor
    if not valid_anchor.any():
        return embeddings.sum() * 0.0

    # Compute all pairwise distances
    distances = torch.cdist(embeddings, embeddings, p=2)

    # Find hardest positive and negative for each anchor
    hardest_positive = distances.masked_fill(~positive, -torch.inf).amax(dim=1)
    hardest_negative = distances.masked_fill(~negative,  torch.inf).amin(dim=1)

    # Evaluate the triplet loss
    loss = torch.relu(hardest_positive - hardest_negative + margin)

    return loss[valid_anchor].mean()

Module `TripletLoss`#

The TripletLoss module provides a common interface for the mining strategies. Its constructor receives the margin and the mining strategy. During the forward pass, the module takes a batch of embeddings and class labels, applies the selected strategy to construct informative triplets, and returns their average loss.

class TripletLoss(torch.nn.Module):

    def __init__(self, margin: float, batch_all: bool = False):
        super().__init__()
        assert margin > 0, "Margin must be positive"
        self.margin = margin
        self.loss_fn = triplet_loss_batch_all if batch_all else triplet_loss_batch_hard

    def forward(self, embeddings, labels):
        assert embeddings.ndim == 2, "Embeddings must be a 2D tensor"
        assert labels.ndim == 1, "Labels must be a 1D tensor"
        assert len(embeddings) == len(labels), "Number of embeddings and labels must match"
        return self.loss_fn(embeddings, labels, self.margin)

Mining diagnostics#

The average training loss does not reveal how many triplets contribute to it or how difficult those triplets are. A separate helper therefore reports statistics such as the proportions of easy, semi-hard, and hard triplets, together with the positive and negative distance distributions. These quantities help interpret the training process and diagnose weak batch composition, unsuitable margins, or unstable mining. They are computed for monitoring only and do not alter the loss function or introduce additional optimization objectives.

Training with triplet loss#

We now train the embedding model on MNIST data using triplet loss. Because the custom loss accepts a batch of model outputs and class labels, it is compatible with the Trainer class defined in training.py and requires no changes to the training loop.

from training import Trainer

Dataset#

First, we load the MNIST dataset with a minimal preprocessing pipeline that converts images to tensors.

preprocess = v2.Compose([v2.ToImage(), v2.ToDtype(torch.float32, scale=True)])

full_ds = MNIST('.data', train=True, download=True, transform=preprocess)
test_ds = MNIST('.data', train=False, download=True, transform=preprocess)

train_idx, valid_idx = train_test_split(
    range(len(full_ds)),
    test_size=0.10,
    stratify=full_ds.targets,
    random_state=42,
)

train_ds = Subset(full_ds, train_idx)
valid_ds = Subset(full_ds, valid_idx)

Model#

Next, we define the embedding network as a convolutional backbone followed by a fully-connected neck, with an optional L2-normalization step at the output. Normalization scales each embedding to unit length, placing all embeddings on the unit hypersphere. This gives distances a consistent scale and prevents the model from changing them simply by increasing or decreasing vector norms. An explicit epsilon is included to ensure numerical stability when an embedding has a norm close to zero.

class EmbeddingNet(nn.Module):

    def __init__(self, out_dim, normalize=True):
        super().__init__()
        self.backbone = nn.Sequential(
            nn.Conv2d(1, 64, 5, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(64, 128, 5, padding=2),
            nn.ReLU(),
            nn.MaxPool2d(2),
            nn.Conv2d(128, 256, 3, padding=1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.neck = nn.Sequential(
            nn.Flatten(),
            nn.Linear(256 * 3 * 3, 1024),
            nn.ReLU(),
            nn.Linear(1024, 256),
            nn.ReLU(),
            nn.Linear(256, out_dim),
        )
        self.normalize = normalize

    def forward(self, x):
        x = self.neck(self.backbone(x))
        if self.normalize:
            return F.normalize(x, p=2, dim=1, eps=1e-12)
        else:
            return x

Setup#

Then, we instantiate the embedding model, triplet-loss module, optimizer, and data loaders. Training uses batch-all mining with a margin of \(0.2\) and a batch size of \(512\). This large batch is likely to contain several examples of every MNIST class, providing many positive and negative candidates for online mining. For imbalanced datasets or datasets with rare classes, random batching is less reliable; a class-balanced sampler should instead construct batches with multiple examples from several classes.

torch.manual_seed(42)

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model  = EmbeddingNet(50, normalize=True).to(device)

optimizer = optim.Adam(model.parameters(), lr=1e-4)

loss_fn = TripletLoss(margin=0.2, batch_all=True)

train_loader = DataLoader(train_ds, batch_size=512, shuffle=True)
valid_loader = DataLoader(valid_ds, batch_size=512, shuffle=False)

Training loop#

Finally, we train the model for a few epochs.

epochs = 5

trainer = Trainer()
history = trainer.fit(model, train_loader, loss_fn, optimizer, epochs, valid_loader)

===== Training on cpu device =====

Epoch  1/5: 100%|██████████| 106/106 [02:17<00:00,  1.29s/it, train_loss=0.1502, valid_loss=0.1528]
Epoch  2/5: 100%|██████████| 106/106 [02:18<00:00,  1.30s/it, train_loss=0.1411, valid_loss=0.1616]
Epoch  3/5: 100%|██████████| 106/106 [02:28<00:00,  1.40s/it, train_loss=0.1378, valid_loss=0.1621]
Epoch  4/5: 100%|██████████| 106/106 [02:42<00:00,  1.53s/it, train_loss=0.1385, valid_loss=0.1726]
Epoch  5/5: 100%|██████████| 106/106 [02:48<00:00,  1.59s/it, train_loss=0.1312, valid_loss=0.1749]

The training and validation loss are plotted below.

../../_images/0254387e76ca62dde2230b49feb900bc794a51f3d77c3c24677a56d21ac5640e.png

Evaluation#

The model evaluation follows the same procedure as in the previous tutorial.

Embedding generation#

First of all, we generate embeddings for the training gallery, validation query, and test query sets. Each dataset is passed through the trained embedding model, and the resulting embeddings and labels are collected into tensors. Since this step does not involve training, the model is placed in evaluation mode and gradient computation is disabled.

gallery_embeddings, gallery_labels = get_embeddings(model, train_ds)
valid_embeddings, valid_labels = get_embeddings(model, valid_ds)
test_embeddings, test_labels = get_embeddings(model, test_ds)

k-NN classification#

For each test embedding, we find the \(k\) closest embeddings in the training gallery and predict the label by majority vote among their class labels. The proportion of correct predictions provides a simple measure of how well the embedding space organizes examples by class.

knn = KNeighborsClassifier(n_neighbors=5, metric="euclidean")
knn.fit(gallery_embeddings.numpy(), gallery_labels.numpy())

predictions = knn.predict(test_embeddings.numpy())
score = accuracy_score(test_labels.numpy(), predictions)

print(f"Test k-NN accuracy: {score:.2%}")

Test k-NN accuracy: 99.26%

Retrieval#

For each query embedding, the gallery examples are ranked from nearest to farthest. In MNIST, a gallery image is considered relevant when it has the same digit label as the query. A good embedding model places relevant images near the top of this ranking.

test_retrieval = retrieval_metrics(test_embeddings, test_labels, gallery_embeddings, gallery_labels)

for name, value in test_retrieval.items():
    print(f"{name}: {value:.2%}")

recall@1: 99.18%
recall@5: 99.56%
mAP: 98.93%

Qualitative visualization#

t-SNE projects high-dimensional embeddings into two dimensions while attempting to preserve local neighborhoods. It can therefore help reveal whether examples from the same class tend to appear near one another. To obtain a coherent view, we fit t-SNE to 2’000 held-out embeddings.

../../_images/6edebf7ae6451ad539c8cc6689da68b933831e48b3c3fa4a4f931ada80324da5.png

Summary#

In this tutorial, we trained an embedding model on MNIST data using triplet loss. Unlike pair-based objectives, triplet loss learns a relative ordering: a positive example should be closer to its anchor than a negative example by at least a specified margin. Because most possible triplets provide little or no training signal, triplet selection is a central part of the method. We implemented two online mining strategies. Batch-all mining averages over all active triplets in a minibatch, whereas batch-hard mining selects the most difficult positive and negative for each valid anchor. Their behavior depends strongly on batch composition, since each anchor requires suitable positive and negative candidates.

Triplet Loss

Contents

Triplet Loss#

Definition of triplet loss#

Intuition#

Triplet mining#

Online mining#

Best practices#

PyTorch implementation#

Helper function: get_pairs()#

Helper function: get_triplets()#

Batch-all mining#

Batch-hard mining#

Module TripletLoss#

Mining diagnostics#

Training with triplet loss#

Dataset#

Model#

Setup#

Training loop#

Evaluation#

Embedding generation#

k-NN classification#

Retrieval#

Qualitative visualization#

Summary#

Helper function: `get_pairs()`#

Helper function: `get_triplets()`#

Module `TripletLoss`#