InfoNCE Contrastive Loss From Scratch in NumPy

Given queries Q and keys K (both n × d, no zero rows) and a temperature tau > 0: normalise every row to unit length, form the logits L_ij = q̂_i·k̂_j / τ, and return the mean over i of log Σ_j exp(L_ij) − L_ii. Key i is the positive for query i; the other keys are its negatives.

Examples

Input Q = [[1, 0], [0, 1]], K = [[1, 0], [0, 1]], tau = 1

Output 0.3133

Each query matches its key exactly: loss log(e + 1) − 1.

Input Q = [[1, 0], [0, 1]], K = [[1, 0], [0, 1]], tau = 0.001

Output 0

A tiny temperature: logits of 1000. The naive exp overflows.

Constraints

1 ≤ n ≤ 512, 1 ≤ d ≤ 256
10⁻³ ≤ τ ≤ 10
Checked to a relative tolerance of 10⁻⁹

Hints

Hint 1

Normalise with Q / np.linalg.norm(Q, axis=1, keepdims=True).

Hint 2

The loss is the cross-entropy of each row of the logits with the diagonal as the target.

Hint 3

Subtract each row's maximum before exponentiating, and add it back after the log.

Where candidates go wrong

Forgetting to normalise: then the loss rewards long vectors.
np.log(np.exp(L).sum(axis=1)) with L in the thousands: inf, then nan.
Averaging over all n² entries instead of over the n rows.

The idea

Normalise both matrices, build the n × n logits with one matrix product, and compute a row-wise log-sum-exp with the maximum subtracted. The loss is the mean of the log-sum-exp minus the diagonal.

Time O(n²d) · Space O(n²)

Reference solution

import numpy as np

def info_nce(Q, K, tau):
    Q = np.asarray(Q, dtype=float)
    K = np.asarray(K, dtype=float)
    Qn = Q / np.linalg.norm(Q, axis=1, keepdims=True)
    Kn = K / np.linalg.norm(K, axis=1, keepdims=True)
    L = (Qn @ Kn.T) / tau
    m = L.max(axis=1, keepdims=True)
    lse = m[:, 0] + np.log(np.exp(L - m).sum(axis=1))       # stable log-sum-exp per row
    return float(np.mean(lse - np.diag(L)))

Console⌘/Ctrl + Enter runs

Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.

Why interviewers ask it

Contrastive pre-training is how most modern embedding models are trained. Writing its loss tests whether a candidate can combine normalisation, a temperature and a numerically stable cross-entropy without a library.

The theory behind it is in the lesson Cosine search, contrastive loss and negative sampling.

Keep going

Easy12 min

The InfoNCE contrastive loss