InfoNCE Loss

Problem

Compute the InfoNCE contrastive loss for one anchor. Similarity is the dot product; logits are sim(anchor, positive) / temperature followed by each negative's similarity / temperature. The loss is the negative log softmax probability of the positive entry.

The signature is infonce_loss(anchor, positive, negatives, temperature). Return a scalar.

Hint

Lower temperature sharpens the distribution and increases the penalty from hard negatives.

Test cases

case 1

in: [[1,0],[1,0],[[0,1],[-1,0]],1]

expected: 0.4076059644

case 2

in: [[1,0],[1,0],[[1,0]],2]

expected: 0.6931471806

case 3

in: [[0,1],[0,1],[[1,0],[0,-1]],0.5]

expected: 0.1429316285

case 4

in: [[2,0],[1,1],[[0,0]],1]

expected: 0.126928011

Starter codePython

import math
def infonce_loss(anchor, positive, negatives, temperature):
    # Your code here
    pass

4 test cases · run your code against them in the editor