SimCLR NT-Xent Loss

Problem

Compute the SimCLR NT-Xent contrastive loss.

For 2N L2-normalized embeddings z (positive pairs are i and i + N), average -log(exp(sim(i, pos) / t) / sum over k != i of exp(sim(i, k) / t)) over all anchors, where t is the temperature. Return the loss rounded to 4 decimal places.

Hint

Subtract the max logit before exponentiating to keep the log-sum-exp stable.

Test cases

case 1

in: [[[1,0],[0,1],[1,0],[0,1]],0.1]

expected: 0.0001

case 2

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

expected: 1.0986

case 3

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

expected: 1.0598

Starter codePython

def nt_xent_loss(z, temperature):
    # Your code here
    pass

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