Weak Supervision

Deep metric learning seeks to organize images in an embedding space so that related images are close together and unrelated images are farther apart. In fully supervised settings, the training signal is usually derived from class labels. However, obtaining such labels at scale can be expensive, particularly when similarity is subjective or cannot be described by a fixed set of categories.

Weaker forms of supervision can be used to reduce the annotation burden. For example, instead of requiring every image to be labeled with a class, one could use image-level tags or similarity judgments to indicate that two images are related. These pairwise labels are often easier to collect and can still provide useful information for learning an embedding space. The tradeoff is that the relationships between images are not fully described. In particular, the absence of a pair does not necessarily mean that two images are dissimilar.

The Totally Looks Like dataset is an excellent demonstration of weak supervision. It contains image pairs that are perceived as visually similar by humans, even when the images belong to different semantic categories. The resemblances may arise from a shared shape, pose, texture, or composition. The dataset captures a nuanced understanding of visual similarity that would be difficult to express through ordinary class labels.

In this tutorial, we train an embedding model on the Totally-Looks-Like dataset using two objectives: symmetric triplet loss and canonical NT-Xent loss. We then evaluate the learned representations by using held-out images as queries and ranking images from a separate gallery.

Dataset

Totally-Looks-Like (TLL) is a dataset designed to reproduce human perception of image similarity. It is based on the similarly-named popular entertainment website. The dataset contains 6016 image pairs deemed to look alike by human annotators. The images are diverse and cover a wide range of categories, including objects, scenes, patterns, animals, and faces across various modalities (sketch, cartoon, natural images). We refer to the images in a pair as the “left view” and the “right view”.

Note

Unlike self-supervised learning, the left and right views are not two augmentations of one image; they are different images linked by a human similarity judgment.

Show code cell source

Hide code cell source

from pathlib import Path
from typing import Any, Callable
from tqdm import tqdm
import matplotlib.pyplot as plt
import torch
import torch.nn.functional as F
from PIL import Image
from sklearn.model_selection import train_test_split
from torch import nn, optim
from torch.utils.data import DataLoader, Dataset, Subset
from torchvision.models import MobileNetV3, MobileNet_V3_Large_Weights, mobilenet_v3_large

Download the images

The dataset is available for download from the Totally Looks Like website. Unzip the downloaded file and look for the left.zip and right.zip archives. Extract the contents of these archives into a directory created alongside this notebook. The directory structure should look like this.

Totally-Looks-Like/
├── left/
│   ├── 0001.jpg
│   ├── 0002.jpg
│   └── ...
└── right/
    ├── 0001.jpg
    ├── 0002.jpg
    └── ...

Define a custom dataset

We create a custom dataset class to load the TLL images. The dataset constructor requires the path to the directory containing the left and right folders. An optional transform argument can be passed to apply transformations to the images. The __getitem__ method loads the pair of left and right images corresponding to the given index and returns them as a tuple. The images are loaded using the Python Imaging Library and converted to RGB. If a transformation is provided, it is applied to both images.

class TotallyLookLike(Dataset):

    def __init__(self, root, transform=None):
        self.root = Path(root)
        self.transform = transform
        self.left_dir = self.root / "left"
        self.right_dir = self.root / "right"
        left_names  = TotallyLookLike.get_names(self.left_dir)
        right_names = TotallyLookLike.get_names(self.right_dir)
        assert left_names == right_names, "Directory names do not match"
        self.filenames = sorted(left_names)

    @staticmethod
    def get_names(path: Path) -> set[str]:
        return { p.name for p in path.glob('*.jpg') }

    def __len__(self):
        return len(self.filenames)

    def __getitem__(self, index):
        name = self.filenames[index]
        with Image.open(self.left_dir / name) as image:
            left = image.convert("RGB")
        with Image.open(self.right_dir / name) as image:
            right = image.convert("RGB")
        if self.transform:
            left = self.transform(left)
            right = self.transform(right)
        return left, right

We create the dataset and verify its size.

dataset = TotallyLookLike('.data/Totally-Looks-Like/')

print('Number of image pairs:', len(dataset))

Number of image pairs: 6016

Visualize the image pairs

Let’s visualize a few image pairs from the dataset.

Show code cell source

Hide code cell source

indices = [26, 1635, 65, 115, 123, 322, 133, 625]

for i in range(0, len(indices), 2):
    plt.figure(tight_layout=True)
    for j in range(2):
        anchors, positives = dataset[indices[i+j]]
        plt.subplot(1, 4, 1 + j*2)
        plt.imshow(anchors)
        plt.title('Left')
        plt.axis('off')
        plt.subplot(1, 4, 2 + j*2)
        plt.imshow(positives)
        plt.title('Right')
        plt.axis('off')
    plt.show()

Loss functions

We consider two objectives for learning from annotated look-alike pairs: triplet loss and normalized temperature-scaled cross-entropy (NT-Xent) loss. Both encourage paired images to have similar embeddings, but they use negative examples differently. Triplet loss compares each positive pair with selected negatives, whereas NT-Xent contrasts each image with many candidates in the same batch.

Symmetric triplet loss

Triplet loss operates on triplets consisting of an anchor, a positive, and a negative. Instead, the TLL dataset only provides positive pairs. To construct triplets, each image serves as the anchor with its paired image as the positive, while images from other pairs in the same batch are treated as candidate negatives. Neither the left nor right images are privileged as anchors, so the loss is computed in both directions.

Given a batch of paired images, the implementation computes the matrix of distances between the two sides of the batch. The diagonal entries correspond to the annotated positive pairs. The off-diagonal entries compare images from different pairs and are used as candidate negatives. For each anchor, its positive distance is compared with all candidate-negative distances in the same row. The calculation is then repeated using the transposed distance matrix, which exchanges the roles of the two sides. Only comparisons that violate the margin are included in the average. If all margin constraints are already satisfied, the implementation returns a differentiable zero so that backpropagation remains valid.

class WeakTripletLoss(nn.Module):

    def __init__(self, margin=0.5):
        super().__init__()
        assert margin > 0, "Margin must be positive"
        self.margin = margin

    def forward(self, left, right):
        assert left.shape == right.shape, "Paired embedding batches must have equal shape"
        assert len(left) >= 2, "At least two pairs are required"
        return self._direction(left, right) + self._direction(right, left)

    def _direction(self, anchors, positives):

        # Compute distances between anchors and positives
        distances = torch.cdist(anchors, positives, p=2)

        # Get the anchor-positive distances (diagonal)
        positive_dist = distances.diag().unsqueeze(1)

        # Get the anchor-negative distances (off-diagonal)
        diagonal = torch.eye(len(anchors), dtype=torch.bool, device=anchors.device)
        negative_dist = distances.masked_fill(diagonal, torch.inf)

        # Compute the loss on hard and semi-hard triplets
        loss = positive_dist - negative_dist + self.margin
        loss = loss[loss > 0]

        if not loss.numel():
            return anchors.sum() * 0.0
        else:
            return loss.mean()

NT-Xent loss

The normalized temperature-scaled cross-entropy loss (NT-Xent) compares each embedding with all other embeddings in a batch. Its objective is to assign the highest similarity to the matching pair while reducing its similarity to the remaining candidates. Suppose a batch contains \(N\) annotated pairs, giving \(2N\) embeddings. Each embedding is treated as an anchor, its paired embedding is the positive, and the remaining \(2N-2\) embeddings serve as negatives. The loss for a pair \((z_i, z_j)\) is defined as follows.

\[ \mathcal{L}_\text{NT-Xent}(z_i, z_j) = -\log\left(\frac{e^{\text{sim}(z_i, z_j) / \tau}}{\sum_{k\neq i} e^{\text{sim}(z_i, z_k) / \tau}}\right), \qquad \text{sim}(z_i, z_j) = \frac{z_i^\top z_j}{\|z_i\| \|z_j\|}. \]

The NT-Xent loss employs the cosine similarity rather than the Euclidean distance. The numerator measures the similarity of the anchor to the positive. The denominator includes the similarity to the positive and to all candidate negatives, so the loss depends on how clearly the positive can be distinguished from the rest of the batch. The temperature parameter \(\tau\) controls how sharply these similarities are compared. A lower temperature places greater emphasis on the hardest negatives, whereas a higher temperature distributes the loss more evenly across the candidates.

Note

The condition \(k \neq i\) ensures that the anchor is not compared with itself in the denominator.

The implementation first concatenates the embeddings from all annotated pairs. It then computes the complete cosine-similarity matrix and divides it by the temperature. Each row corresponds to one anchor. The diagonal is excluded because it compares the anchor with itself, while the index of its paired view defines the correct target. Cross-entropy then encourages the positive to receive the largest similarity among the remaining \(2N-1\) candidates. Because every embedding serves as an anchor, the loss is inherently symmetric. It also uses all available negatives rather than selecting individual triplets. These negatives are assumed to be dissimilar, so incomplete pair annotations may introduce false negatives.

class NTXentLoss(nn.Module):
    """
    NT-Xent loss for contrastive learning. 
    The input embeddings are expected to be normalized.
    """

    def __init__(self, temperature=0.1):
        super().__init__()
        assert temperature > 0, "Temperature must be positive"
        self.temperature = temperature

    def forward(self, left, right):
        assert left.shape == right.shape, "Paired embedding batches must have equal shape"
        assert len(left) >= 2, "At least two pairs are required"

        # Concatenate left and right views
        batch_size = len(left)
        embeddings = torch.cat([left, right], dim=0)

        # Compute the cosine similarity matrix (dot product of normalized embeddings)
        logits = embeddings @ embeddings.T / self.temperature

        # Mask out self-similarities (diagonal)
        self_mask = torch.eye(2 * batch_size, dtype=torch.bool, device=logits.device)
        logits = logits.masked_fill(self_mask, -torch.inf)

        # Indices of the positive pairs in the concatenated batch
        #   (0, N), (1, N + 1), ..., (N - 1, 2 * N - 1), 
        #   (N, 0), (N + 1, 1), ..., (2 * N - 1, N - 1)
        targets = torch.arange(2 * batch_size, device=logits.device).roll(batch_size)
        
        # Compute the cross entropy loss, where each positive is treated as a class.
        return F.cross_entropy(logits, targets)

Preparation

Instead of training the embedding model from scratch, we will use a pretrained backbone to extract features from the images, and a separate embedding head to transform these features into the final embeddings.

MobileNet

We begin by loading a pretrained MobileNetV3 model from the torchvision library.

weights = MobileNet_V3_Large_Weights.DEFAULT
preprocess = weights.transforms()
mobilenet = mobilenet_v3_large(weights=weights)

Data split

The pretrained weights come with their own preprocessing pipeline, so we need to ensure that the data is transformed accordingly.

dataset = TotallyLookLike('.data/Totally-Looks-Like/', transform=preprocess)

The pairs are divided into training, validation, and test subsets before any features are extracted.

SPLIT_SEED = 42

# Training: 70%
train_idx, holdout_idx = train_test_split(
    range(len(dataset)),
    test_size=0.30,
    random_state=SPLIT_SEED,
)

# Validation: 15%, Test: 15%
valid_idx, test_idx = train_test_split(
    holdout_idx,
    test_size=0.50,
    random_state=SPLIT_SEED,
)

train_images = Subset(dataset, train_idx)
valid_images = Subset(dataset, valid_idx)
test_images = Subset(dataset, test_idx)

print(f"Train: {len(train_idx)}\nValid: {len(valid_idx)}\nTest:  {len(test_idx)}")

Train: 4211
Valid: 902
Test:  903

Feature extraction

Features are computed and cached separately for each subset, avoiding unnecessary repeated passes through the backbone. Caching makes training considerably faster, but it also fixes the backbone features.

Show code cell source

Hide code cell source

def extract_pair_features(model: MobileNetV3, dataset, batch_size=128):
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=False)

    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
    old_device = next(model.parameters()).device

    was_training = model.training
    model.to(device).eval()

    left_features, right_features = [], []
    with torch.inference_mode():
        for left, right in tqdm(loader, desc="Extracting features"):
            left, right = left.to(device), right.to(device)
            left = model.avgpool(model.features(left))
            right = model.avgpool(model.features(right))
            
            left_features.append(torch.flatten(left, 1).cpu())
            right_features.append(torch.flatten(right, 1).cpu())

    model.to(old_device).train(was_training)

    return torch.cat(left_features), torch.cat(right_features)

train_left, train_right = extract_pair_features(mobilenet, train_images)
valid_left, valid_right = extract_pair_features(mobilenet, valid_images)
test_left,  test_right  = extract_pair_features(mobilenet, test_images)

print("Train features:", tuple(train_left.shape))
print("Valid features:", tuple(valid_left.shape))
print("Test  features:", tuple(test_left.shape))

Extracting features: 100%|██████████| 33/33 [03:29<00:00,  6.35s/it]
Extracting features: 100%|██████████| 8/8 [00:42<00:00,  5.29s/it]
Extracting features: 100%|██████████| 8/8 [00:40<00:00,  5.10s/it]

Train features: (4211, 960)
Valid features: (902, 960)
Test  features: (903, 960)

Feature datasets

The extracted features must be organized into datasets that support efficient data loading. We construct a custom PyTorch dataset that handles the pairing of precomputed features corresponding to look-alike images. This dataset is designed to supply the necessary inputs for the triplet loss or NT-Xent loss.

class FeatureDataset(Dataset):

    def __init__(self, left_features, right_features):
        self.left_features = left_features
        self.right_features = right_features

    def __len__(self):
        return len(self.left_features)

    def __getitem__(self, idx):
        return self.left_features[idx], self.right_features[idx]

We construct separate training, validation, and test feature datasets.

train_ds = FeatureDataset(train_left, train_right)
valid_ds = FeatureDataset(valid_left, valid_right)
test_ds  = FeatureDataset(test_left, test_right)

Embedding head

The embedding head is a neural network that projects the features extracted by the pretrained model into a lower-dimensional embedding space suitable for metric learning. This module typically consists of one or more fully connected layers, optionally followed by batch normalization and a non-linearity, with a final L2 normalization to enforce unit norm. Normalizing the embeddings is crucial for ensuring that the distance between embeddings is scale-invariant. It is also necessary for computing the cosine similarity in the NT-Xent loss.

class EmbeddingHead(nn.Module):

    def __init__(self, in_dim=960, out_dim=128):
        super().__init__()
        self.head = nn.Sequential(
            nn.Linear(in_dim, 512),
            nn.ReLU(),
            nn.Linear(512, out_dim, bias=False),
            nn.BatchNorm1d(out_dim),
        )

    def forward(self, x):
        x = self.head(x)
        x = F.normalize(x, p=2, dim=1)
        return x

Training

We now train the embedding head using the cached features. The training process involves optimizing the parameters of the embedding head to minimize the selected loss function (triplet loss or NT-Xent loss) on the training dataset, while monitoring performance on the validation dataset.

from training import Trainer

Custom adapter

The Trainer class is designed to be used across different training scenarios. Under the hood, it depends on a user-defined function that specifies how to compute the loss function for a single batch of data. The default implementation is tailored for supervised learning and makes various assumptions about the batch structure, the model architecture, and the loss computation. The code of this function is shown below.

def supervised_adapter(model: nn.Module, batch: Any, func: Callable) -> torch.Tensor:
    inputs, labels = batch       # Batch contains inputs and labels
    outputs = model(inputs)      # Model takes inputs and returns predictions
    return func(outputs, labels) # Loss compares predictions to labels

The default adapter does not fit the weakly-supervised learning scenario. We need a custom adapter that takes in a batch containing pairs of inputs, forwards them separately through the model to obtain their embeddings, and then uses these embeddings to compute the triplet loss or NT-Xent loss. The code of this custom adapter is shown below.

def weakly_supervised_adapter(model: nn.Module, batch: Any, func: Callable) -> torch.Tensor:
    anchors, positives = batch      # Batch contains pairs of inputs
    anchors   = model(anchors)      # Model is applied to both inputs separately
    positives = model(positives) 
    return func(anchors, positives) # Loss compares embeddings of both inputs

Training setup

The training process begins by assembling the model. In our pipeline, the extracted features are passed through the embedding head to obtain the final embedding representations. A DataLoader is used to handle batching with a large batch size. An optimizer is then defined to update the parameters of the embedding head. The loss function can be either Triplet Loss or NT‑Xent Loss.

Show code cell source

Hide code cell source

seed = 42
torch.manual_seed(seed)
if torch.cuda.is_available():
    torch.cuda.manual_seed_all(seed)

OBJECTIVE = "nt_xent"  # or "triplet"

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

model = EmbeddingHead(in_dim=960, out_dim=128)
optimizer = optim.Adam(model.parameters(), lr=1e-3, amsgrad=True)

if OBJECTIVE == "nt_xent":
    loss_fn = NTXentLoss(temperature=0.1)
elif OBJECTIVE == "triplet":
    loss_fn = WeakTripletLoss(margin=0.5)
else:
    raise ValueError(f"Unknown objective: {OBJECTIVE}")

trainer = Trainer()
trainer.set_adapter(weakly_supervised_adapter)

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

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

Epoch  1/10: 100%|██████████| 16/16 [00:00<00:00, 59.88it/s, train_loss=5.4978, valid_loss=5.1681]
Epoch  2/10: 100%|██████████| 16/16 [00:00<00:00, 84.69it/s, train_loss=4.0349, valid_loss=5.0095]
Epoch  3/10: 100%|██████████| 16/16 [00:00<00:00, 84.80it/s, train_loss=3.2104, valid_loss=5.0071]
Epoch  4/10: 100%|██████████| 16/16 [00:00<00:00, 67.70it/s, train_loss=2.5632, valid_loss=5.0231]
Epoch  5/10: 100%|██████████| 16/16 [00:00<00:00, 86.61it/s, train_loss=2.0432, valid_loss=5.0738]
Epoch  6/10: 100%|██████████| 16/16 [00:00<00:00, 85.47it/s, train_loss=1.6006, valid_loss=5.1004]
Epoch  7/10: 100%|██████████| 16/16 [00:00<00:00, 74.04it/s, train_loss=1.2746, valid_loss=5.1686]
Epoch  8/10: 100%|██████████| 16/16 [00:00<00:00, 89.99it/s, train_loss=0.9840, valid_loss=5.1880]
Epoch  9/10: 100%|██████████| 16/16 [00:00<00:00, 63.67it/s, train_loss=0.7866, valid_loss=5.2521] 
Epoch 10/10: 100%|██████████| 16/16 [00:00<00:00, 77.99it/s, train_loss=0.6242, valid_loss=5.2728]

Show code cell source

Hide code cell source

plt.plot(history['train_loss'], label='Train loss')
plt.plot(history['valid_loss'], label='Validation loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.legend()
plt.show()

Retrieval evaluation

After training the model, we evaluate its ability to retrieve images that look alike. The general idea is to use the learned embeddings to perform a retrieval task, where we query the model with a left image and expect it to retrieve the corresponding right image from a gallery of candidates.

Pairwise distances in the embedding space

The first step in evaluating the model is to compute the pairwise distances between the left and right images. To this end, we take the features extracted by the pretrained model, pass them through the trained embedding head, and compute the pairwise distances between them. Normalized embeddings make cosine and Euclidean rankings equivalent.

model.eval()
with torch.inference_mode():
    left_embeddings = model(test_left)
    right_embeddings = model(test_right)

distances = 1 - left_embeddings @ right_embeddings.T

Collapse diagnostics

To ensure that the learned embeddings are meaningful, we inspect the norms of the embeddings, the per-dimension standard deviation, and the average off-diagonal similarity. Near-zero variation or uniformly high similarity can indicate that the embeddings have collapsed to a trivial solution, where all images are mapped to the same point in the embedding space.

combined = torch.cat([left_embeddings, right_embeddings])
similarity = combined @ combined.T
off_diagonal = ~torch.eye(len(combined), dtype=torch.bool)

print("Mean embedding norm:", combined.norm(dim=1).mean().item())
print("Mean per-dimension std:", combined.std(dim=0).mean().item())
print("Mean off-diagonal cosine similarity:", similarity[off_diagonal].mean().item())

Mean embedding norm: 1.0
Mean per-dimension std: 0.08796964585781097
Mean off-diagonal cosine similarity: 0.007899009622633457

Retrieval metrics

The metric function locates each diagonal match in the sorted gallery. It also supports tiny synthetic tests because its inputs are only a distance matrix and matching indices.

Show code cell source

Hide code cell source

def paired_retrieval_metrics(distances, ks=(1, 5, 10)):

    order = distances.argsort(dim=1)
    target = torch.arange(len(distances))[:, None]
    ranks = (order == target).nonzero()[:, 1] + 1

    result = {
        f"recall@{k}": (ranks <= k).float().mean().item()
        for k in ks
    }
    reciprocal_rank = 1.0 / ranks.float()
    result["median_rank"] = ranks.float().median().item()
    result["mAP"] = reciprocal_rank.mean().item()
    result["MRR"] = reciprocal_rank.mean().item()
    return result

results = paired_retrieval_metrics(distances)

for name, value in results.items():
    if name == "median_rank":
        print(f"{name}: {value:.1f}")
    else:
        print(f"{name}: {value:.2%}")

recall@1: 7.86%
recall@5: 14.40%
recall@10: 18.72%
median_rank: 116.0
mAP: 11.91%
MRR: 11.91%

Retrieval examples

To illustrate the retrieval process, we select left images from the test set and show the top-5 retrieved right images from the whole dataset. The retrieved images are ranked by their distance to the query image, and we visually inspect the results to assess whether the model has learned to retrieve images that look alike.

Show code cell source

Hide code cell source

def visualize_retrieval_results(anchor_index, distances, k=5):
    
    original_images = TotallyLookLike('.data/Totally-Looks-Like/')

    # Get the distance row for the selected anchor.
    anchor_distances = distances[anchor_index]
    
    # Since lower distances are better, sort in ascending order.
    topk = torch.topk(anchor_distances, k, largest=False)
    topk_indices = topk.indices.cpu().numpy()
    topk_values = topk.values.cpu().numpy()
    
    # Setup the plot: one column for the anchor and k columns for retrievals.
    fig, axes = plt.subplots(1, k+1, figsize=(3*(k+1), 4))
    
    # Plot the anchor image.
    left_image, _ = original_images[anchor_index]
    axes[0].imshow(left_image)
    axes[0].set_title("Anchor")
    axes[0].axis("off")
    
    # Plot the top-k retrieval results from the right folder.
    for i, (idx, dist_val) in enumerate(zip(topk_indices, topk_values), start=1):
        _, right_image = original_images[idx]
        axes[i].imshow(right_image)
        # You can also indicate if this is the correct match.
        label = f"Rank {i}\nDist: {dist_val:.2f}"
        if idx == anchor_index:
            label += "\n(Correct)"
        axes[i].set_title(label)
        axes[i].axis("off")
    
    plt.tight_layout()
    plt.show()

left_features, right_features = extract_pair_features(mobilenet, dataset)

with torch.inference_mode():
    left_embeddings = model(left_features)
    right_embeddings = model(right_features)

all_distances = 1 - left_embeddings @ right_embeddings.T

Extracting features: 100%|██████████| 47/47 [04:46<00:00,  6.10s/it]

for index in [28, 115, 207, 809, 1743, 3043, 3799]:
    visualize_retrieval_results(index, all_distances)

Summary

In this tutorial, we trained an embedding model on the Totally-Looks-Like dataset using either symmetric triplet loss or NT-Xent loss. The dataset provides pairs of images judged to look alike, allowing the model to learn from pairwise similarity annotations rather than class labels. We then inspected the learned representation by retrieving the five nearest gallery images for several held-out queries.

The choice of loss function is only one factor that determines retrieval quality. Batch construction, negative selection, normalization, feature extraction, and evaluation protocols can be equally important. Musgrave et al. (2020) showed that, under more consistent experimental conditions, several sophisticated metric-learning losses offered only modest improvements over established baselines such as triplet loss. This result emphasizes the importance of careful implementation and controlled evaluation when comparing methods.