How do you write the CLIP loss in PyTorch?

In this article

1. Before you start
2. Why interviewers ask this
3. CLIP loss from first principles
4. Worked example
5. Trade-offs and edge cases
6. How to answer it in an interview
7. Common mistakes
8. Follow-up questions
9. Sources

Before you start

This page belongs to the PyTorch interview questions and machine learning (ML) coding round guide, and it builds on a few ideas.

Why interviewers ask this

Writing a contrastive loss by hand shows whether a candidate can turn a paper's pseudocode into tensor code without a library doing the thinking. InfoNCE is a contrastive loss (info noise-contrastive estimation; van den Oord et al., 2018) that scores one positive against a set of negatives with a softmax cross-entropy. The CLIP loss is InfoNCE applied in both directions, with in-batch negatives (the other pairs in the same batch, reused as wrong answers) and a learned temperature. Candidates report being asked to handwrite an InfoNCE-style contrastive loss in PyTorch, and to write a custom loss for a generative model; those are reports from the interview question bank, not tied here to any named company.

A weak answer says "it makes matching images and captions close" and cannot say what the labels are. A strong answer writes the matrix, names the labels as the integers from zero to N minus one, explains why both directions are averaged, and knows the scale is learned. A Staff-level answer adds what breaks: the matrix grows with the square of the batch, duplicate captions create wrong negatives, and the scale needs a cap to keep training stable.

CLIP loss from first principles

CLIP loss

CLIP loss is a symmetric cross-entropy loss that scores an N-by-N matrix of image-text similarities, treating the true pair in each row and each column as the correct class and every other entry as a wrong class, so it pulls matching pairs together and pushes all other pairs apart.

Logit scale

The logit scale is a number that multiplies every cosine similarity before the softmax and equals one divided by the temperature, so a larger scale makes the softmax sharper and a smaller scale makes it flatter.

Symmetric cross-entropy

In contrastive learning, a symmetric cross-entropy is the average of two cross-entropies computed over one score matrix, once along its rows and once along its columns, so that each side of a pair is trained to find the other side.

In plain words

Picture N photos and N captions laid out on a table. Each photo has exactly one matching caption. For each photo you must pick its caption from all N. For each caption you must pick its photo from all N. Your score is how confident you were in the right choice. The analogy breaks in one place. The model is not told which wrong captions are close calls, so a near-duplicate caption counts as fully wrong.

In one precise sentence: the loss turns an N-by-N table of scaled cosine similarities into two sets of N classification problems. The correct class for item i is position i. The loss then averages the two cross-entropies. Radford et al. (2021) state it as: "We optimize a symmetric cross entropy loss over these similarity scores."

Step 1: Normalize both sets of embeddings

The image encoder outputs one vector per image and the text encoder one vector per caption, so a batch gives two matrices with N rows each. The first step rescales every row to length one with F.normalize(x, dim=-1).

Input: two N-row matrices. Output: the same matrices with every row of length one. After this, the dot product of an image row and a text row is their cosine similarity, which is what the paper's objective maximizes for real pairs and minimizes for the rest, per Radford et al. (2021). Skip the step and the dot product also grows with vector length, so the score mixes direction with size and the scale in Step 2 loses its meaning. You now have comparable vectors; the next step compares all of them at once.

Step 2: Multiply into a matrix and apply the scale

One matrix product, image @ text.T, scores every image against every caption. Entry (i, j) is image i against caption j, so the diagonal holds the N real pairs and the other entries hold the wrong pairings. The paper counts them: N real pairs and $N^2-N$ wrong ones, which is 4 and 12 for a batch of 4.¹

Then every entry is multiplied by the logit scale. The temperature is the number the scores are divided by before the softmax, and the logit scale is its inverse. In the paper it is learned, stored as a log, starts at the equivalent of 0.07 and is clipped so the logits are never scaled by more than 100, a limit the authors found necessary for stable training (Radford et al., 2021). The starting scale is 1 / 0.07 = 14.2857.¹

Why it matters: on the toy matrix below, a scale of 1 gives a mean loss of 1.0358, while the starting scale of 14.2857 gives 0.0069.¹ A bigger scale is not simply better: it also magnifies a wrong ranking, which is why the scale is learned and capped.

Step 3: Use arange labels and cross-entropy in both directions

Row i must pick column i, so the labels are the integers torch.arange(N). The paper's pseudocode does exactly this: labels = np.arange(n), with one cross-entropy over each axis (Radford et al., 2021).

Input: the scaled N-by-N matrix and the labels. Output: two losses. Reading the rows scores image-to-text: given an image, find its caption. Reading the columns scores text-to-image, which in PyTorch is the same function applied to logits.T. Pass the raw scaled scores to F.cross_entropy, not softmax outputs: the function applies the softmax and the logarithm itself, and a second softmax would flatten the signal. The PyTorch documentation for undefined describes it as combining log-softmax and negative log-likelihood. The stable softmax page shows why the log-sum-exp form is used.

Step 4: Average the two losses

The final loss is the mean of the two cross-entropies: (loss_i + loss_t) / 2, as in the pseudocode of Radford et al. (2021). A 2025 paper, Structure-aware Language-Image Pretraining (SLIP), describes the same thing as the mean of the cross-entropy losses from both modalities (SLIP, 2025).

Averaging matters because a softmax over rows and a softmax over columns normalize differently, so one direction can be satisfied while the other is not. Here is the whole function:

import torch
import torch.nn.functional as F

def clip_loss(image_emb, text_emb, logit_scale):
    # Step 1: unit-length rows, so dot products are cosine similarities
    image_emb = F.normalize(image_emb, dim=-1)
    text_emb = F.normalize(text_emb, dim=-1)
    # Step 2: N x N matrix, entry (i, j) = image i vs text j, times the scale
    logits = logit_scale * image_emb @ text_emb.T
    # Step 3: row i must pick column i
    labels = torch.arange(logits.shape[0], device=logits.device)
    loss_i = F.cross_entropy(logits, labels)      # rows: image -> text
    loss_t = F.cross_entropy(logits.T, labels)    # columns: text -> image
    # Step 4: average both directions
    return (loss_i + loss_t) / 2

The function takes the scale as a number, so the same code works with a fixed or a learned temperature.

1 · Normalize2 · Matrix and scale3 · Two cross-entropies4 · AverageReplay

Step-by-step description

1. Image and text rows get length one.
2. One product scores every pair, times the scale.
3. Rows and columns are scored against arange.
4. The mean of both is the loss.

• Image and text vectors are normalized to length one first.
• One matrix product scaled by the logit scale gives the N by N logits.
• Cross-entropy reads the rows and the columns with arange labels.
• The loss is the mean of the two cross-entropies.

Worked example

Take a batch of four pairs. The similarity matrix below is illustrative: it holds cosine similarities of the kind a half-trained model might output, with the true pairs on the diagonal.¹

Work through row 1 at a scale of 1. The exponentials of 0.82, 0.31, 0.12 and 0.25 are 2.2705, 1.3634, 1.1275 and 1.2840, and they sum to 6.0454. The probability of the true caption is 2.2705 / 6.0454 = 0.3756, and the loss is minus the natural logarithm of that, 0.9793.¹

Text description of the diagram

• The exponentials of the four scores sum to 6.0454.
• The probability of the true caption is 0.3756 and the loss is 0.9793.

Repeating this for every row gives image-to-text losses of 0.9793, 1.0423, 1.1081 and 1.0139 at a scale of 1. Row 3 is the worst because its true score has the narrowest lead over the wrong ones.¹

A model that scores every pair equally gives each of N captions probability 1/N, so its loss is the natural logarithm of N. That is the sanity check for a new implementation: at random initialization the loss should start near ln N, which is 1.3863 for N = 4.¹ The paper trained with a very large batch of 32,768 (Radford et al., 2021), whose starting value is 10.3972.¹

Trade-offs and edge cases

The cost of the loss is the matrix. It has N times N entries, so a batch of 32,768 holds 1,073,741,824 scores, about 4.29 GB in 32-bit numbers if one device held it all.¹ The open_clip repository's ClipLoss class exposes options such as local_loss, gather_with_grad and cache_labels, which handle gathering features across graphics processing units (GPUs) and reusing labels (open_clip, read 2026-10-26). Large-batch training also needs sharded training; see fully sharded data parallel.

Larger batches give each real pair more negatives: N − 1 per row, so 3 at N = 4 and 32,767 at the paper's batch of 32,768 (N² − N = 12 and 1,073,709,056 wrong pairings in total).¹

Text description of the diagram

• A model that scores all pairs equally has a loss of ln N.
• That loss is 1.3863 at a batch of 4 and 10.3972 at a batch of 32,768.

The softmax is not the only choice. Zhai et al. (2023) propose a sigmoid loss for language-image pre-training (SigLIP) that scores each image-text pair as its own yes-or-no classification, with a learned temperature and a learned bias. It needs no normalization across a row or column, so it does not need a global softmax gathered across devices. The open_clip repository ships it as SigLipLoss in the same file (open_clip, read 2026-10-26). The paper notes the objective was first introduced in deep metric learning as the multi-class N-pair loss (Radford et al., 2021).

Edge cases follow from the setup. Two identical or near-identical captions in one batch make a correct match count as wrong; this is an inference from how the labels work, not a finding from the paper, so deduplicate or mask them. A learned scale needs a cap. A minimal version:

import math
import torch
import torch.nn as nn

class LearnedScale(nn.Module):
    def __init__(self):
        super().__init__()
        # stored as a log so the scale stays positive
        self.log_scale = nn.Parameter(torch.tensor(math.log(1 / 0.07)))

    def forward(self):
        return self.log_scale.exp()

    @torch.no_grad()
    def clamp(self):
        # call after each optimizer step: the scale never exceeds 100
        self.log_scale.clamp_(max=math.log(100))

How to answer it in an interview

The spoken structure: say the matrix and the labels, then the two directions, then the sanity check, then the failure modes.

Common mistakes

• "I just take the dot product of the raw embeddings." (wrong) Without normalization the scores depend on vector length. (right) "I normalize first, so the dot product is cosine similarity."
• "I apply softmax, then pass the probabilities to cross-entropy." (wrong) The function expects raw scores and applies the softmax itself. (right) "I pass scaled logits and let cross-entropy do the softmax."
• "I compute the loss over rows only." (wrong) That trains images to find captions but not the reverse. (right) "I average the row loss and the column loss."
• "Temperature is always a fixed hyperparameter." (wrong) Fixed-temperature contrastive setups exist, but in CLIP it is learned and capped at a scale of 100. (right) "I learn it as a log value and clamp it after each step."
• "Only the diagonal matters." (wrong) The diagonal sets the labels, while every other entry acts as a negative in the softmax denominator. (right) "The diagonal is the target; the rest of each row is the competition."

Follow-up questions

How do you implement a custom loss function in PyTorch for a generative model?

Write it as a function of tensors that returns one scalar, built from differentiable operations, then check it against a reference on a tiny input. The generative pre-trained transformer (GPT) from scratch walkthrough shows the same habit for a next-token loss.

How do you implement masked cross-entropy with label smoothing?

In PyTorch call F.cross_entropy(logits, targets, ignore_index=-100, label_smoothing=ε): positions whose target equals ignore_index are left out of the loss, and label_smoothing mixes the one-hot target with a uniform distribution (PyTorch documentation, read 2026-10-26). By hand, take the per-position negative log-probabilities, multiply by a mask that is zero at padding, and divide by the number of unmasked positions.

How do you implement custom losses, attention mechanisms and training loops from scratch?

Build each from tensors and test it on a tiny case with a hand-computed answer, as the worked example above does. For the attention part, see multi-head attention without undefined and what the causal mask does.

Why is the CLIP loss symmetric?

Each direction normalizes over a different set: the row softmax compares captions for one image, and the column softmax compares images for one caption. Averaging both makes each side of a pair learn to find the other.

What goes wrong with a small batch or duplicate captions?

A small batch gives each real pair few negatives, so the task is easy and the loss is a weak signal; at N = 4 each image competes with only 3 wrong captions (12 wrong pairings in total).¹ Duplicate captions make a correct match count as wrong, so deduplicate or mask them.

How is the CLIP loss related to InfoNCE?

It is the same family: the CLIP loss is InfoNCE applied in both directions, with the other pairs in the batch as negatives and a learned temperature. The paper traces the batch construction to the multi-class N-pair loss, and a 2025 paper notes that InfoNCE can be rewritten as a cross-entropy over similarity logits (SLIP, 2025).

Sources

1. Learning Transferable Visual Models From Natural Language Supervision (arXiv 2103.00020) — arXiv (Radford et al., OpenAI), 26 Feb 2021
2. SLIP: Structure-aware Language-Image Pretraining (arXiv 2511.03019) — arXiv, 1 Nov 2025
3. open_clip: loss.py (ClipLoss), read 2026-10-26 — GitHub (mlfoundations/open_clip), 26 Oct 2026
4. Sigmoid Loss for Language Image Pre-Training (arXiv 2303.15343) — arXiv (Zhai et al.), 27 Mar 2023
5. Representation Learning with Contrastive Predictive Coding (arXiv 1807.03748) — arXiv (van den Oord et al.), 10 Jul 2018
6. torch.nn.functional.cross_entropy, PyTorch documentation, read 2026-10-26 — PyTorch, 26 Oct 2026