Implement InfoNCE Loss

InfoNCE is what contrastive learning methods like SimCLR and CLIP optimize. For each query, treat its matching key (same index) as the one positive among a batch full of negatives, and turn similarity scores into a classification problem.

L_i = -\log \frac{\exp(\text{sim}(q_i, k_i) / \tau)}{\sum_j \exp(\text{sim}(q_i, k_j) / \tau)}

Task: write infonce_loss(queries, keys, temperature) using the dot product as sim, returning the mean of L_i over the batch, rounded to 4 decimal places.

queries[i] is only ever compared against keys[i] as its positive — every other key in the batch counts as a negative for it.
A lower temperature sharpens the distribution, making the model pay more for confusing a negative with the positive.

Implement InfoNCE Loss

Examples

Hints