FlatNCE: Is Floating-Point Error the Real Reason Contrastive Learning Struggles with Small Batches?
Ever since SimCLR made a splash in unsupervised visual representation learning, contrastive learning has steadily gained popularity in CV and even NLP, with more and more related studies and works appearing. A well-known drawback of standard contrastive learning is that it requires fairly large batch sizes (SimCLR performs best at batch_size=4096), and performance drops noticeably with small batch sizes. Consequently, one direction of follow-up work has been to reduce this dependence on large batch sizes. A natural question then arises: what exactly causes standard contrastive learning to perform poorly with small batch sizes?
Recently, a paper titled Simpler, Faster, Stronger: Breaking The log-K Curse On Contrastive Learners With FlatNCE has offered an answer to this question: floating-point error. This might sound hard to believe, but the paper's analysis is actually quite reasonable, and the proposed improvement, FlatNCE, indeed works better in practice—enough to be convincing.
The Devil in the Details
In what follows, I'll present the main content of the original paper using my own understanding and notation. I won't go through the full details of contrastive learning here—broadly speaking, for a given sample $x$, we construct $K$ paired samples $y_1,y_2,\cdots,y_K$, where $y_t$ is the positive sample and the rest are negatives. We then score each sample pair $(x, y_i)$, denoting the scores as $s_1,s_2,\cdots,s_K$. Contrastive learning aims to widen the score gap between positive and negative pairs, typically using cross-entropy as the loss directly:
\begin{equation}-\log \frac{e^{s_t}}{\sum\limits_i e^{s_i}} = \log \left(\sum_i e^{s_i}\right) - s_t = \log \left(1 + \sum_{i\neq t} e^{s_i - s_t}\right)\end{equation}
For simplicity, we'll denote it as $\xi=\sum\limits_{i\neq t} e^{s_i - s_t}$ from here on. In practice, positive samples are usually highly similar samples obtained via data augmentation, while negative samples are all the other samples within the batch—so roughly speaking, the negatives can be regarded as $K-1$ randomly selected samples. This implies that the gap between positive and negative pairs tends to be quite pronounced, so the model can easily achieve $s_t \gg s_i(i\neq t)$, i.e., $e^{s_i - s_t}\approx 0$. As a result, when the batch size is small (equivalent to $K$ being small), $\xi$ will also get quite close to 0, meaning the loss function above will also approach 0 quite closely.
A loss function close to 0 usually implies that the gradient is also close to 0. However, this doesn't mean the update step of the model is small. This is because the optimizers currently used in contrastive learning are adaptive ones like Adam, whose update magnitude roughly takes the form $\frac{\text{gradient}}{\sqrt{\text{gradient}\otimes\text{gradient}}}\times\text{learning rate}$. This means that no matter how small the gradient is, as long as it's stable, the update will still remain on the order of $\text{learning rate}$. Contrastive learning is exactly such a scenario: to achieve $e^{s_i - s_t}\to 0$, one would need $s_i - s_t\to -\infty$, but since the scores in contrastive learning are usually cosine similarities divided by a temperature parameter, they are bounded, so $s_i - s_t\to -\infty$ is unattainable. Consequently, after a certain number of training steps, the loss function will stay close to—but greater than—0 for a long stretch.
However, the computation of $\xi$ itself carries floating-point error, and when $\xi$ is very close to 0, this floating-point error can actually be larger than the exact value itself. Then the computation of $\log(1+\xi)$ will also carry floating-point error, and subsequently the gradient computation will too. As these errors accumulate, the gradients that are ultimately computed can end up close to pure random noise, providing no useful guidance for updates. This, according to the original paper, is the reason contrastive learning degrades noticeably at small batch sizes.
Turning the Subtle into the Substantial
Once we understand the cause, it isn't hard to come up with a targeted fix. Taking a first-order expansion of the loss function, we have:
\begin{equation}\log \left(1 + \sum_{i\neq t} e^{s_i - s_t}\right)\approx \sum_{i\neq t} e^{s_i - s_t}\end{equation}
In other words, after a certain number of training steps, the model is effectively being trained with $\xi$ as the loss function. Of course, since $\log(1+\xi)\leq \xi$—i.e., $\xi$ is an upper bound of $\log(1+\xi)$—the outcome wouldn't be much different even if we had used $\xi$ as the loss function from the very start. The main issue to address now is that $\xi$ approaching 0 leads to floating-point error problems. As mentioned earlier, the update magnitude of adaptive optimizers roughly takes the form $\frac{\text{gradient}}{\sqrt{\text{gradient}\otimes\text{gradient}}}\times\text{learning rate}$, which means that if we simply multiply the loss function by a constant, the update in theory won't change. So since $\xi$ is too small, we can just multiply it by a suitable constant to scale it up.
What should we multiply it by? A fairly direct idea is that the loss function shouldn't be too small or too large—keeping it around the order of $\mathcal{O}(1)$ would be ideal. So we might as well multiply by the reciprocal of $\xi$, which gives us
\begin{equation}\frac{\xi}{\text{sg}(\xi)} = \frac{\sum\limits_{i\neq t} e^{s_i - s_t}}{\text{sg}\left(\sum\limits_{i\neq t} e^{s_i - s_t}\right)}\label{eq:flatnce-1}\end{equation}
as the loss function. Here $\text{sg}$ denotes stop_gradient (called "detach" in the original paper), meaning the denominator is treated purely as a constant, and gradients are only taken with respect to the numerator. This is the alternative proposed in the original paper, called FlatNCE.
That said, a loss function written with the $\text{sg}$ operator isn't exactly the form we're used to, so let's convert it. Notice that:
\begin{equation}\nabla_{\theta}\left(\frac{\xi}{\text{sg}(\xi)}\right) = \frac{\nabla_{\theta}\xi}{\xi} = \nabla_{\theta}\log \xi\end{equation}
In other words, the gradient provided by using $\frac{\xi}{\text{sg}(\xi)}$ as the loss function is exactly identical to the gradient from using $\log \xi$ as the loss function. So we can replace the loss with $\log \xi$, which no longer contains the $\text{sg}$ operator:
\begin{equation}\log\left(\sum\limits_{i\neq t} e^{s_i - s_t}\right) = \log\left(\sum\limits_{i\neq t} e^{s_i}\right) - s_t\label{eq:flatnce-2}\end{equation}
Compared to cross-entropy, this loss is simply the $\text{logsumexp}$ computation with the score of the positive pair $s_t$ removed. Note that $\text{logsumexp}$ can typically be computed efficiently, without floating-point error dominating. So we can replace cross-entropy with this loss function—theoretically equivalent to cross-entropy, but in practice performing better at small batch sizes. It's also worth pointing out that the result of the expression above isn't necessarily non-negative, so there's no need to be surprised if negative loss values show up during training when using this loss—that's normal behavior.
Truth in Practice
The analysis sounds plausible enough—but does it actually hold up? Naturally, we need experiments to settle this. Unsurprisingly, FlatNCE does indeed work remarkably well.
All the experiments in the original paper are on CV tasks, mainly replacing SimCLR's loss with FlatNCE—the resulting method is called FlatCLR. What we probably care about most is whether FlatNCE really solves the dependence on large batch sizes, and the figure below gives an affirmative answer:
Comparison of SimCLR and FlatCLR across different batch sizes
Below is a comparison of SimCLR and FlatCLR results across various tasks, showing FlatCLR's superior performance:
Comparison of SimCLR and FlatCLR across various tasks
Nitpicking
Overall, the original paper's contribution is quite creative—the "floating-point error" angle is unconventional but also remarkably precise, which deserves credit.
Intuitively, the original goal of cross-entropy is "maximize the gap between positive and negative scores," which works fine for standard classification problems, but isn't quite enough for contrastive learning. That's because the goal of contrastive learning is to learn features: beyond the "coarse" feature of positive samples scoring higher than negatives, negatives also need to keep being compared against each other to learn finer-grained features. FlatNCE's objective, by contrast, is "make the positive score as large as possible, and the negative scores as small as possible"—shifting from learning relative values to learning absolute values. This allows optimization to continue even after positives and negatives have been pulled apart by some margin, rather than stopping prematurely (for non-adaptive optimizers) or having the update dominated by floating-point-error noise (for adaptive optimizers).
That said, some aspects of how the original paper is presented invite criticism. For instance, the paper spends a considerable amount of space discussing mutual information estimation, which has no substantive connection to the paper's main thread and adds to the reader's difficulty in following it. Of course, a paper isn't the same as a popular-science article, and adding extra theoretical derivations to flesh things out is reasonable enough—it would just have been better if the analysis of floating-point error had been given more prominence. What I find hardest to understand, though, is that the paper directly presents $\eqref{eq:flatnce-1}$ as the final result. This kind of formulation involving "stop_gradient" isn't exactly hard, but it's not reader-friendly either—normally this approach is only "forced" upon you when it's difficult to find the antiderivative, which clearly isn't the case for FlatNCE.
Summary
This post has introduced a new piece of work on contrastive learning, which analyzes the floating-point error problem in cross-entropy under small-batch contrastive learning, pointing out that this may well be the main reason contrastive learning underperforms at small batch sizes. It also proposes a targeted improved loss function, FlatNCE. Experiments show that contrastive learning based on FlatNCE indeed alleviates the dependence on large batch sizes and achieves better results.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.