Contrastive Loss (InfoNCE) kernel

[superposition]

Description

InfoNCE / NT-Xent loss for contrastive learning in latent space.

Operation

# Compute InfoNCE loss between anchor and positive/negative pairs
# anchor: (B, D), positives: (B, D), negatives: (B, N, D) or use in-batch
loss = info_nce_loss(anchor, positives, temperature=0.07)

Implementation

• Fused similarity computation + log-softmax
• Temperature scaling
• In-batch negatives (memory efficient)
• Optional: hard negative mining

Tests

• Matches manual implementation: -log(exp(sim_pos/t) / sum(exp(sim_all/t)))
• Gradient correct for anchor and positives
• Temperature scaling works
• Numerically stable (large batch, extreme similarities)
• gradcheck passes

Used In

CLIP, SimCLR, MoCo, JEPA, BYOL, contrastive world models