Metric Learning Losses

Status: Public API in netcl.autograd and netcl.nn

Metric learning trains embeddings so that semantically similar inputs are close together and dissimilar inputs are far apart. netcl provides a full suite of standard contrastive and metric losses: Triplet Margin Loss, InfoNCE (NT-Xent), Contrastive Loss (Siamese), Batch Hard Triplet Loss, Batch All Triplet Loss, and the l2_normalize utility.

Choosing a loss

Rule of thumb: start with Batch Hard Triplet Loss for labelled data and InfoNCE for self-supervised. Switch to Batch All early in training if gradients vanish. Use Contrastive or plain Triplet only when pair/triplet mining is already handled upstream.

Triplet Margin Loss

Formula

Given an anchor embedding a, a positive embedding p (same class as a), and a negative embedding n (different class):

L = mean_i( max(0,  d(a_i, p_i) − d(a_i, n_i) + margin) )

d is either squared Euclidean (default) or L2 norm.

The loss is zero when the anchor is already at least margin closer to the positive than to the negative — the "hard" triplets with a non-zero loss drive learning.

API

# With autograd (training):
from netcl import autograd as ag

loss = ag.triplet_loss(anchor_node, positive_node, negative_node,
                       margin=1.0,
                       distance="squared",   # "squared" | "euclidean"
                       reduction="mean")     # "mean"    | "sum"

# Forward only (inference):
import netcl.nn as nn
loss_val = nn.triplet_loss(anchor_tensor, positive_tensor, negative_tensor)

distance parameter

"squared" is the numerically stable default. Use "euclidean" when the embedding norms carry meaning (e.g. face verification).

Gradients

For distance="squared":

∂L/∂a_i = active_i / N · 2(n_i − p_i)
∂L/∂p_i = active_i / N · 2(p_i − a_i)
∂L/∂n_i = active_i / N · 2(a_i − n_i)

For distance="euclidean":

∂L/∂a_i = active_i / N · ((a_i−p_i)/d_pos − (a_i−n_i)/d_neg)
∂L/∂p_i = active_i / N · (p_i−a_i) / d_pos
∂L/∂n_i = active_i / N · (a_i−n_i) / d_neg

where active_i = 1 iff loss_i > 0.

Example training loop

import netcl as nc
import netcl.nn as nn
from netcl import autograd as ag

# Encoder maps inputs to D-dimensional embeddings
encoder = nn.Sequential(nn.Linear(784, 256), nn.ReLU(), nn.Linear(256, 64))
optimizer = nc.optim.Adam(encoder.parameters(), lr=1e-3)

for anchor, positive, negative in dataloader:
    tape = ag.Tape()
    ag.set_current_tape(tape)

    a_emb = encoder(anchor)
    p_emb = encoder(positive)
    n_emb = encoder(negative)

    loss = ag.triplet_loss(a_emb, p_emb, n_emb, margin=1.0)
    tape.backward(loss)
    ag.set_current_tape(None)

    optimizer.step()
    optimizer.zero_grad()

InfoNCE / NT-Xent Loss

Formula

InfoNCE (also known as NT-Xent or SimCLR loss) treats each sample's counterpart in the batch as its positive and all others as negatives:

L = mean_i( -sim(q_i, k_i)/τ + log Σ_j exp(sim(q_i, k_j)/τ) )

where sim(u, v) = u·v / (‖u‖ ‖v‖) (cosine similarity when normalize=True) and τ is the temperature. Positive pairs are on the diagonal of the N×N similarity matrix.

The loss equals the cross-entropy of a classifier that picks the correct positive among all N keys. A perfect model achieves log(N) → 0 as pairs become perfectly separated.

API

# With autograd (training):
from netcl import autograd as ag

loss = ag.info_nce_loss(query_node, key_node,
                        temperature=0.07,   # τ, typically 0.07–0.2
                        normalize=True,     # L2-normalise before similarity
                        reduction="mean")   # "mean" | "sum"

# Forward only (inference):
import netcl.nn as nn
loss_val = nn.info_nce_loss(query_tensor, key_tensor)

Inputs

Both can be raw network outputs; normalize=True (default) applies L2 normalisation internally, including its Jacobian in the backward pass.

Temperature

Lower temperature (e.g. 0.07) sharpens the distribution and pushes the model to discriminate more finely. Values around 0.1–0.2 are common for self-supervised learning; very small values (< 0.05) can cause training instability.

Gradients

The gradient of the loss w.r.t. the pre-normalisation query embeddings:

grad_S   = (softmax(S) − I) · scale            # N×N, scale=1/N for mean
grad_q_n = (grad_S / τ) @ k_n                  # N×D  (matmul)
grad_k_n = (grad_S.T / τ) @ q_n                # N×D
grad_q   = (grad_q_n − (grad_q_n·q_n) q_n) / ‖q‖   # L2-norm Jacobian

The backward runs host-side (numpy) after syncing GPU tensors, which is fast because the N×N similarity matrix is small for typical batch sizes (N ≤ 512).

Example: SimCLR-style training

import netcl as nc
from netcl import autograd as ag

encoder = build_encoder()
projection = nc.nn.Sequential(nc.nn.Linear(512, 128), nc.nn.ReLU(), nc.nn.Linear(128, 64))
optimizer = nc.optim.Adam(list(encoder.parameters()) + list(projection.parameters()), lr=1e-3)

for (view1, view2) in dataloader:          # two augmented views of same images
    tape = ag.Tape()
    ag.set_current_tape(tape)

    q = projection(encoder(view1))          # (N, D)
    k = projection(encoder(view2))          # (N, D)

    loss = ag.info_nce_loss(q, k, temperature=0.1)
    tape.backward(loss)
    ag.set_current_tape(None)

    optimizer.step()
    optimizer.zero_grad()

Contrastive Loss (Siamese Networks)

Formula

For each pair (x1_i, x2_i) with binary label y_i (1 = similar, 0 = dissimilar):

d_i  = ‖x1_i − x2_i‖₂
L_i  = y_i · d_i²  +  (1 − y_i) · max(0, margin − d_i)²

Similar pairs are pulled together; dissimilar pairs are pushed apart until their distance exceeds margin. Pairs where the dissimilar distance already exceeds margin contribute zero loss (no wasted gradient on easy negatives).

API

from netcl import autograd as ag
import netcl.nn as nn

# Autograd (training):
loss = ag.contrastive_loss(x1_node, x2_node, labels,   # labels: (N,) 0/1 array
                           margin=1.0, reduction="mean")

# Module:
module = nn.ContrastiveLoss(margin=1.0)
loss = module(x1_node, x2_node, labels, tape=tape)

Margin

margin controls how far apart dissimilar pairs must be pushed. Typical values are 0.5–2.0 depending on the scale of the embedding space. With L2-normalised embeddings (unit sphere) the maximum distance is 2.0, so margin=1.0 is a sensible default.

Gradients

Let d_i = ‖x1_i − x2_i‖₂, δ_i = x1_i − x2_i, and h_i = max(0, margin − d_i).

∂L/∂x1_i = 1/N · [ y_i · 2δ_i  −  (1 − y_i) · 2h_i · δ_i/d_i ]
∂L/∂x2_i = −∂L/∂x1_i

For similar pairs (y=1) the gradient pulls x1 and x2 together proportional to their current distance. For dissimilar pairs (y=0) it pushes them apart only when d_i < margin; once h_i = 0 the gradient vanishes.

Example

import netcl as nc
from netcl import autograd as ag

siamese = build_siamese_encoder()   # shared-weight encoder
optimizer = nc.optim.Adam(siamese.parameters(), lr=1e-3)

for (x1, x2, labels) in dataloader:   # labels: 1 = same person, 0 = different
    tape = ag.Tape()
    ag.set_current_tape(tape)

    emb1 = siamese(x1)
    emb2 = siamese(x2)

    loss = ag.contrastive_loss(emb1, emb2, labels, margin=1.0)
    tape.backward(loss)
    ag.set_current_tape(None)

    optimizer.step()
    optimizer.zero_grad()

Batch Hard Triplet Loss

For each anchor in the batch, mines the hardest positive (furthest same-class embedding) and the hardest negative (closest different-class embedding), then applies the standard triplet margin loss on that selected pair.

hard_pos_i = argmax_{j: same class, j≠i} d(i, j)
hard_neg_i = argmin_{j: diff class}      d(i, j)
L_i        = max(0, d(i, hard_pos_i) − d(i, hard_neg_i) + margin)

Anchors with no valid positive or negative in the batch contribute 0 to the loss.

Hard mining within the batch is more efficient than offline mining — every forward pass automatically uses the most informative triplets for the current model state.

API

from netcl import autograd as ag
import netcl.nn as nn

# Autograd:
loss = ag.batch_hard_triplet_loss(embeddings_node, labels,
                                   margin=1.0,
                                   distance="squared",   # "squared" | "euclidean"
                                   reduction="mean")

# Module:
module = nn.BatchHardTripletLoss(margin=1.0, distance="squared")
loss = module(embeddings_node, labels, tape=tape)

labels accepts a numpy array or a Tensor of integers.

Batch composition

For effective mining the batch should contain at least 2 classes and at least 2 samples per class. A common strategy is to sample P classes × K samples each (e.g. P=16, K=4 → batch of 64). Batches with a single class per sample produce no valid triplets and a loss of zero.

distance parameter

Gradients

Let j = hard_pos_i, k = hard_neg_i, active_i = 1 iff L_i > 0.

For distance="squared":

∂L/∂E[i] += active_i / N · 2(E[k] − E[j])
∂L/∂E[j] += active_i / N · 2(E[j] − E[i])   # summed over all i that selected j
∂L/∂E[k] += active_i / N · 2(E[i] − E[k])   # summed over all i that selected k

For distance="euclidean" (letting dp = d(i,j), dn = d(i,k)):

∂L/∂E[i] += active_i / N · ((E[i]−E[j])/dp − (E[i]−E[k])/dn)
∂L/∂E[j] += active_i / N · (E[j]−E[i]) / dp
∂L/∂E[k] += active_i / N · (E[i]−E[k]) / dn

Contributions from all anchors that selected the same embedding as their positive or negative are accumulated via scatter-add.

Example

import numpy as np
import netcl as nc
from netcl import autograd as ag

encoder = build_encoder()
optimizer = nc.optim.Adam(encoder.parameters(), lr=1e-3)

# Batch: P=8 classes × K=8 samples = 64 embeddings
for images, class_ids in dataloader:       # class_ids: (N,) int array
    tape = ag.Tape()
    ag.set_current_tape(tape)

    embeddings = encoder(images)           # (N, D)

    loss = ag.batch_hard_triplet_loss(embeddings, class_ids, margin=0.5)
    tape.backward(loss)
    ag.set_current_tape(None)

    optimizer.step()
    optimizer.zero_grad()

Batch All Triplet Loss

Uses every valid (anchor, positive, negative) triplet in the batch:

L = mean over {(i,j,k) : same[i,j], diff[i,k], loss>0}( max(0, d(i,j) − d(i,k) + margin) )

Compared to batch-hard mining, this provides denser gradient signal early in training when most triplets are still active. As the model improves, the fraction of active triplets shrinks and the effective signal approaches that of hard mining. Both the forward pass and the backward pass are fully vectorised via numpy broadcasting — no per-triplet Python loops. The O(N³) active-triplet mask is computed once and reused for gradient accumulation.

API

from netcl import autograd as ag
import netcl.nn as nn

# Autograd:
loss = ag.batch_all_triplet_loss(embeddings_node, labels,
                                  margin=1.0,
                                  distance="squared",   # "squared" | "euclidean"
                                  reduction="mean")

# Module:
module = nn.BatchAllTripletLoss(margin=0.5)
loss = module(embeddings_node, labels, tape=tape)

Hard vs All: when to switch

A common schedule: start with BATL for the first few epochs, then switch to BHTL once the loss stabilises.

Gradients

Let A[i,j,k] = 1 iff triplet (i,j,k) is active (valid and loss > 0), s = 1/|A|.

For distance="squared":

∂L/∂E[i] = s · (A.sum(axis=1) − A.sum(axis=2)) @ E   # anchor contributions
∂L/∂E[j] = s · (A.sum(axis=0).T − A.sum(axis=2).T) @ E    # positive
∂L/∂E[k] = s · (A.sum(axis=1).T − A.sum(axis=0).T) @ E    # negative

Computed entirely with numpy matmuls — no Python loop over triplets.

Example

import netcl as nc
from netcl import autograd as ag

encoder = build_encoder()
optimizer = nc.optim.Adam(encoder.parameters(), lr=1e-3)

for images, class_ids in dataloader:
    tape = ag.Tape()
    ag.set_current_tape(tape)

    embeddings = encoder(images)   # (N, D)

    loss = ag.batch_all_triplet_loss(embeddings, class_ids, margin=0.5)
    tape.backward(loss)
    ag.set_current_tape(None)

    optimizer.step()
    optimizer.zero_grad()

l2_normalize

Normalises a tensor to unit length along a given dimension — needed as a pre-processing step before InfoNCE and contrastive losses when normalize=False.

from netcl import autograd as ag
import netcl.nn as nn

# With gradient:
normed = ag.l2_normalize(x_node, dim=-1, eps=1e-8)

# Forward only:
normed_tensor = nn.l2_normalize(x_tensor, dim=-1)

eps is added to the norm before dividing to avoid division by zero for zero vectors. Only dim=-1 (last axis, i.e. row-wise normalisation) is supported for 2-D tensors; other shapes always use the CPU path.

The backward applies the exact Jacobian of the projection onto the unit sphere:

grad_x = (grad_out − (grad_out · x̂) x̂) / ‖x‖

Note: ag.info_nce_loss with normalize=True (default) calls l2_normalize internally — no need to normalise manually before passing to InfoNCE.

GPU acceleration

For 2-D inputs with dim=-1, the GPU kernel path (l2_norm_rows / l2_jac) activates when N·D ≥ 8 192 (e.g. N=64, D=128 or N=128, D=64). Below that threshold numpy host-side is faster due to kernel-launch overhead. Measured speedup at threshold: 26–31× over the numpy round-trip path.

Implementation Details

Where it lives

Backend support

All ops accept GPU tensors (backend="opencl"). Whether computation runs on GPU or CPU depends on the op:

For ops with a GPU condition, GPU tensors below the threshold are synced to host, computed in numpy, and placed back on the source backend. The threshold for l2_normalize was determined empirically: at N·D < 8 192 kernel-launch overhead dominates; at N·D ≥ 8 192 the GPU kernel is 9–31× faster than the numpy round-trip path (e.g. N=64, D=128: 31×; N=128, D=128: 29×).

batch_hard_triplet_loss uses the GPU at all batch sizes. Int32 index buffers (pos, neg, labels) are allocated once per queue+N and reused across iterations; grad_E is zero-filled on-device via clEnqueueFillBuffer. Measured speedups: N=32: 26×, N=64: 23×, N=128: 18×, N=256: 9×.

Numerical stability

• Triplet / BHTL / BATL: epsilon of 1e-8 added under the sqrt for distance="euclidean", preventing division by zero for identical embeddings.
• InfoNCE: uses the log-sum-exp trick (subtract row max before exp). Both forward kernel and numpy reference apply this trick identically. Intermediate computations in the backward use float64 for precision.
• Contrastive: gradient is numerically undefined when d_i = 0 for a similar pair; eps=1e-8 is added to d_i in the backward to prevent NaN.

See Also

• Tape — the autograd graph losses register on.
• mse_loss — regression loss.
• CrossEntropyLoss — classification loss.
• nn API — full API reference.