Lifted Structured Loss

Triplets waste most of the batch

"We present a novel structured loss function designed to take full advantage of the training batches in the embedding learning problem."

Song et al. (2016) introduced a loss function that addresses a key limitation of triplet loss: each triplet update only considers one positive and one negative per anchor, wasting the vast majority of pairwise information available in a mini-batch.

To see why this matters, consider the progression of metric learning losses. Contrastive loss operates on pairs: it pulls matching pairs together and pushes non-matching pairs apart, one pair at a time. Triplet loss improves on this by considering triplets (anchor, positive, negative), giving the network a relative comparison. But the lifted structured loss goes further: it considers all pairwise distances in the entire mini-batch simultaneously.

From triplets to batch-wise comparisons

In triplet loss, each gradient update considers one anchor, one positive, and one negative. If we have a mini-batch of examples, there are pairwise relationships between them. Triplet loss only uses of these relationships per update. The rest of the information is simply discarded.

The lifted structured loss exploits all positive and negative pairs in the batch simultaneously. For each positive pair , it considers every negative that is close to either or , turning a sparse training signal into a dense one.

The loss function

Let denote the Euclidean distance between the embeddings of examples and . Let be the set of positive pairs (pairs with the same label) and the margin.

For each positive pair , the lifted structured loss first computes a per-pair cost:

Let us unpack each piece:

• is the distance between the positive pair. We want this to be small.
• finds the hardest negative for point : the negative that is closest to (largest margin violation).
• Similarly, finds the hardest negative for point .
• The outer takes the worst violation from either side. This is the key insight: we guard both members of the positive pair, not just the anchor.
• means: if a negative is closer than the margin , it contributes a positive value; if it is already far enough away, it contributes a negative value (which will be clipped).

The total loss over all positive pairs is then:

The squared hinge serves two purposes:

1. It ensures that satisfied pairs (where all negatives are far enough) contribute zero loss.
2. Compared to the standard hinge , the squared version produces smoother gradients near zero, which helps optimisation.

Smooth approximation

The nested in is not smooth: it selects a single hardest negative and ignores the rest. This creates discontinuous gradients that can make optimisation unstable. The authors replace it with a log-sum-exp approximation:

The log-sum-exp function is a classic smooth approximation to the operator. Recall that for any set of values :

So the log-sum-exp is an upper bound on the max, and it is tight when one value dominates. Crucially, it is everywhere differentiable, which gives gradient-based optimisers a smooth landscape to work with. It also means that near-hard negatives (not just the single hardest) can contribute to the gradient, providing a richer training signal.

Lifted Structured Loss in PyTorch

Here is a minimal implementation of the smooth version of the lifted structured loss. Given a batch of embeddings and their labels, we compute the full pairwise distance matrix and apply the loss over all positive pairs.

# lifted_structured_loss.py
import torch
import torch.nn.functional as F

def lifted_structured_loss(embeddings, labels, margin=1.0):
    """
    Lifted structured loss (smooth approximation).

    Args:
        embeddings: (B, D) tensor of L2-normalised embeddings.
        labels:     (B,) tensor of integer class labels.
        margin:     float, the margin alpha.

    Returns:
        Scalar loss.
    """
    # Pairwise Euclidean distance matrix: (B, B)
    dists = torch.cdist(embeddings, embeddings, p=2)

    # Masks for positive and negative pairs
    labels = labels.unsqueeze(0)
    pos_mask = (labels == labels.T).float()
    neg_mask = (labels != labels.T).float()

    # Remove self-pairs from positive mask
    pos_mask.fill_diagonal_(0)

    # For each point i, compute logsumexp over its negatives
    # Set non-negative entries to -inf so they don't contribute
    neg_exponents = margin - dists  # (B, B)
    neg_exponents = neg_exponents * neg_mask + (-1e9) * (1 - neg_mask)
    lse_per_row = torch.logsumexp(neg_exponents, dim=1)  # (B,)

    # For each positive pair (i, j): J_ij = lse_i + lse_j + D_ij
    # We build this for all pairs, then mask to positives
    # lse_per_row[i] + lse_per_row[j] approximates the nested max
    # But the paper sums inside one log, so we use:
    # log(exp(lse_i) + exp(lse_j)) + D_ij
    lse_i = lse_per_row.unsqueeze(1).expand_as(dists)
    lse_j = lse_per_row.unsqueeze(0).expand_as(dists)
    combined_lse = torch.logaddexp(lse_i, lse_j)

    J_ij = combined_lse + dists  # (B, B)

    # Squared hinge over positive pairs
    loss_mat = torch.clamp(J_ij, min=0).pow(2) * pos_mask

    num_pos = pos_mask.sum()
    if num_pos == 0:
        return torch.tensor(0.0, device=embeddings.device)

    return loss_mat.sum() / (2 * num_pos)

A few things to notice:

• torch.cdist computes the full pairwise distance matrix in one call, which is efficient on GPU.
• torch.logsumexp computes the smooth max over negatives for each point, and torch.logaddexp combines the two sides of the positive pair.
• We mask out self-pairs and non-positive pairs so that the loss is only computed over genuine positive pairs.
• The squared hinge clamp(..., min=0).pow(2) zeros out satisfied pairs and smooths the gradient near the boundary.