1 / 1

Implement InfoNCE Loss

Loss Functions

Hard

Compute one-directional InfoNCE loss for two aligned embedding batches. Build the similarity logits:

S = \frac{Z_1Z_2^{\mathsf T}}{\tau}

Treat entries $S_{ii}$ as positive pairs and every entry in row $i$ as a candidate:

L = -\frac{1}{N}\sum_{i=1}^{N}\log\left(\frac{e^{S_{ii}}}{\sum_{j=1}^{N}e^{S_{ij}}}\right)

Here, $N$ is batch size, $\tau$ is temperature, and row $i$ of $Z_1$ is paired with row $i$ of $Z_2$ . Subtract each row maximum before exponentiation and return the mean loss as a Python float.

Loading visualization...

Examples

Input: Z1 = [[1, 0], [0, 1]], Z2 = [[1, 0], [0, 1]], temperature = 0.1

Output: 0.000045

Explanation: Each diagonal similarity is much larger than the competing similarity in its row.

Input: Z1 = [[1, 0], [0, 1]], Z2 = [[0, 1], [1, 0]], temperature = 0.1

Output: 10.000045

Input: Z1 = [[1, 0], [0, 1]], Z2 = [[1, 0], [0, 1]], temperature = 1.0

Output: 0.313262

Hint 1

Use Z1 @ Z2.T / temperature to build the logits.

Hint 2

Subtract np.max(logits, axis=1, keepdims=True) before exponentiation.

Hint 3

Use np.diag(shifted) for the positive-pair logits.

Requirements

Compute all pairwise dot-product logits
Use diagonal entries as positive pairs
Stabilize each row before exponentiation
Return the mean row loss as a Python float

Constraints

Z1 and Z2 have the same numeric shape (N,D)
N is between 1 and 256
D is between 1 and 512
temperature is positive
Use NumPy only

Prev Next