Similarity search generalisation in contrastive learning with InfoNCE loss

2026-07-10Machine Learning

Machine Learning
AI summary

The authors study a popular method called InfoNCE used in training embedding models that find similar items. They show mathematically how using more negative examples in the training loss makes it closer to an ideal similarity search. They also provide a new way to measure how well the method generalizes to new data by introducing a continuity bound that considers the number of negative samples. Their analysis suggests that increasing negative samples helps reduce the chance of errors when applying the model beyond the training data. This work provides deeper theoretical insights into how InfoNCE works beyond existing ideas based on mutual information.

Similarity searchEmbedding modelsContrastive learningInfoNCE lossCross-entropySoftmaxGeneralisationLipschitz continuityGâteaux differentiationNegative sampling
Authors
Nick Whiteley
Abstract
Similarity search is a primary application of embedding models trained by contrastive learning. For one of the most popular contrastive learning loss functions, InfoNCE, we show that the population risk with $k$ negative samples is $O(1/k)$ close to an expected cross-entropy which quantifies deviation between i) a softmax similarity search over unseen data using the learned embedding function, and ii) an idealised softmax search over the same data but using similarity implicitly represented in the positive sample generator. This complements existing interpretations of InfoNCE in the $k\to\infty$ limit which are phrased in terms of mutual information, and alignment versus uniformity in embeddings. To quantify generalisation performance, we introduce a new continuity bound for the InfoNCE loss, obtained via Gâteaux differentiation. The bound preserves the structure of averaging over negative samples present in the loss function and features an ``inverse temperature'' parameter which can be tuned to account for the algorithmic temperature. For embedding functions which are Lipschitz in a parameter, this yields a simple demonstration that the averaging effect of $k$ negative samples in the InfoNCE loss carries over to stabilisation of the generalisation error as $k$ grows.