Similarity search generalisation in contrastive learning with InfoNCE loss

2026-07-10

AI Summary

This work addresses the lack of theoretical characterization regarding the deviation of the InfoNCE loss from the ideal similarity search objective under a finite number of negative samples in contrastive learning. The authors introduce a novel continuity bound based on Gâteaux differentiation that preserves the mean structure of negative samples and aligns the algorithmic temperature via an adjustable “inverse temperature” parameter. This framework establishes a precise relationship between the population risk of InfoNCE and the ideal softmax cross-entropy loss. Theoretical analysis reveals that the deviation is of order O(1/k), where k denotes the number of negative samples, and that the generalization error stabilizes as k increases due to an averaging effect. These findings elucidate the critical role of negative sample size in determining generalization performance.

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.

Problem

• similarity search
• contrastive learning
• InfoNCE loss
• generalisation
• embedding models

Innovation

• InfoNCE loss
• similarity search generalisation
• Gâteaux differentiation
• continuity bound
• contrastive learning