Why do We Need Large Batchsizes in Contrastive Learning? A Gradient-Bias Perspective
Abstract
Contrastive learning (CL) has been the de facto technique for self-supervised representation learning (SSL), with impressive empirical success such as multi-modal representation learning.
However, traditional CL loss only considers negative samples from a minibatch, which could cause biased gradients due to the nondecomposibility of the loss.
For the ﬁrst time, we consider optimizing a more generalized contrastive loss, where each data sample is associated with an inﬁnite number of negative samples.
We show that directly using minibatch stochastic optimization could lead to gradient bias.
To remedy this, we propose an efﬁcient
Bayesian data augmentation technique to augment the contrastive loss into a decomposable one, where standard stochastic optimization can be directly applied without gradient bias.
Speciﬁcally, our augmented loss deﬁnes a joint distribution over the model parameters and the augmented parameters, which can be conveniently optimized by a proposed stochastic expectation-maximization algorithm.
Our framework is more general and is related to several popular SSL algorithms.
We verify our framework on both small scale models and several large foundation models, including SSL of ImageNet and SSL for vision-language representation learning.
Experiment results indicate the existence of gradient bias in all cases, and demonstrate the effectiveness of the proposed method on improving previous state of the arts.
Remarkably, our method can outperform the strong MoCo-v3 under the same hyper-parameter setting with only around half of the minibatch size; and also obtains strong results in the recent public benchmark ELEVATER for few-shot image classiﬁcation.
Introduction
Recent years have seen signiﬁcant developments and applications of contrastive learning (CL), a promising self-supervised learning (SSL) technique, to a variety of striking research breakthroughs such as learning of multi-modal foundation models [5] in text and image domains [46, 63, 67] for natural-language and vision-language representation learning.
In addition, CL has also be applied recently to other scenarios such as supervised learning [32] and reinforcement learning [3, 36].
The main idea of CL is to deﬁne a contrastive loss where each data sample is entangled with both its positive and negative data samples, and optimize it to learn transformation-invariant features to better discriminate positive samples from negative ones.
From a technical perspective, the contrastive loss is non-decomposable , in the sense that each data sample is coupled with all the other samples in the loss function, i.e. , to process one data sample, one has to get access to all its positive and negative samples from the other data samples.
This is intractable when considering a more general CL setting where the number of negative samples are unbounded (potentially inﬁnite).
In this case, if stochastic optimization is applied, only negative samples in the current minibatch are available.
Thus, this could lead to biased gradients and consequently sub-optimal solutions (see Section 2.1 for a formal statement), which is also evidenced in several existing works [21, 24].
A straightforward approach to overcoming the issue is to deﬁne some decomposable loss functions for learning.
For example, based on spectral theory, [27] proposes spectral contrastive learning, which deﬁnes the loss as the difference between positive and negative similarity scores.
Furthermore, negative-sample-free SSL methods such as BYOL [24] and SimSiam [16] deal with this problem by completely dropping the negative samples.
These methods, however, are developed from quite different perspective and are not equivalent to the standard contrastive learning, e.g. , SimCLR that adopts the info-NCE loss [13].
Since standard contrastive learning has been shown to be relatively robust and easier to scale up to large foundation models [5] such as CLIP and its variants [45, 60, 63, 67], it is important to study its equivalent form in terms of a decomposable loss, and understand what beneﬁts can be brought by this new form to further boost model performance.
To this end, we propose an equivalent conditional decomposable form of the standard contrastive learning, by leveraging the well developed Bayesian data augmentation technique from the Bayesian learning community [41].
The basic idea of Bayesian data augmentation is to augment a difﬁcultto-handle distribution to an easier one by augmenting the parameter space with auxiliary random variables, also known as auxiliary data [39] as each of them is associated with one data point.
When applying it to contrastive learning, each positive data pair will be associated with an auxiliary variable/data * , allowing one to deﬁne a joint distribution over the model parameters and the auxiliary data.
Conditioned on the auxiliary random variables, the potentially inﬁnite negative samples will be separated from each other without introducing bias, thus making a decomposable loss (details in Section 2.2.2).
We call our model Decomsable Contrastive Learning (DeCL) .
To learn from the proposed decomposable contrastive loss, we then propose a new stochastic expectation-maximization (sEM) framework for maximal likelihood estimation (MLE) of the joint distribution.
Speciﬁcally, with the proposed augmentation technique, the auxiliary random variables enjoy simple Gamma posterior distributions, conditioned on which the model parameter can be easily optimized by stochastic gradient descent.
To summarize, our contributions include:
• We propose an equivalent augmented form of the standard contrastive loss from a probability perspective, which enjoys conditional decomposibility for unbiased stochastic optimization.
• We propose an efﬁcient sEM algorithm for MLE of the joint distribution over model parameters and augmented data, which only needs minimal modiﬁcations based on various contrastive-loss optimization procedures such as SimCLR [13] and MoCo-v3 [14].
• We verify our framework with extensive experiments in large-scale foundation models, including the uni-modal ImageNet and multi-modal vision-language representation learning, as well as a public benchmark ELEVATER for few-shot image classiﬁcation [37].
Experimental results suggests that, with the ability of gradient-bias mitigation, our framework can improve the current CL-based state of the arts.
Decomposing Contrastive Learning 2.1 Setup & Gradient Bias in Contrastive Learning
In CL, a backbone network, parameterized by ✓ , is used for feature extraction/encoding.
Given an input x , the feature extractor outputs its representation, written as z = enc ( x ; ✓ ) .
The basic idea of CL is to maximize the representation similarities between positive data pairs while minimizing those of negative pairs.
Formally, we are given an unlabeled dataset D , { x i } .
Each x i 2 D will be applied a set of transformations to form an implicit augmented dataset ˜ D , { ( x i , { x ai } a ) } i , where x ai ⇠ t ( x i ) and t ⇠ T is a sampled transformation function from the transformation family
T .
A positive pair is represented as ( x i , x ai ) for 8 a ; and a negative pair is represented as ( x i , x j ) or ( x i , x aj ) for i 6 = j .
For both positive and negative pairs, the data points will be fed to the encoder and similarity scores will be calculated based on their representations.
For notation simplicity, we will index positive similarity scores by ‘ j + ’ and negative similarity scores by k  , e.g. , we use s ij + , sim ( enc ( x i ; ✓ ) , enc ( x ai ; ✓ )) to denote the similarity score between x i and its j -th positive data point from { x ai } , where sim ( · , · ) denotes a similarity metric (positive value) such as the exponential cosine similarity used in most contrastive learning methods (which we use unless explicitly stated); similarly, s ik  is used to denote the similarity score between x i and one of its negative data point.
Note the similarity scores depend on ✓ , but is omitted in the representation for notation simplicity.
In theory, each data point x i should be associated with an inﬁnite number of negative pairs due to the availability of arbitrary transformation functions.
Thus, we consider one of the most general form of contrastive loss with N i positive pairs and inﬁnite negative pairs for each x i , formulated as † :
[Equation]
A practical setting is to consider one single positive sample in each L x i , i.e. , N i = 1 , giving us a slightly simpler
sub-loss L x i ,
[Equation]
where the positive similarity score is written as s i + for notation simplicity.
We will consider this simpler case in our following development , although generalizing it to the more general case of N i positive pairs is straightforward.
Direct optimization (Direct-Opt, blue) fails to converge with heavy noise; whereas our method (red) can handle the noise and converges to the same optimal solution
Gradient Bias Note the general CL loss (1) is intractable for direct optimization due to the inﬁnite number of negative pairs for each data point.
For practical feasibility, a common trick is to replace the inﬁnite negative pairs with only those within a minibatch.
This gives us an approximation of L x i as ˜
L x i ,
[Equation]
Bk
[Equation]
s ik  where B denotes the minibatch size.
We argue that this approximation is biased due to the non-decomposibility nature of the standard CL loss L con , in the sense that the individual loss L con is associated with the negative samples that distribute across the whole data ˜ D .
In other words, the overall loss L con cannot be decomposed into a sum over independent sub-losses.
Importantly, nondecomposibility indicates that L x i 6 =
E h ˜
L x i i , and consequently, r ✓ L x i 6 =
r
✓ E h ˜
L x i i , where both the expectations are taken over the randomness of minibaching.
Thus, using negative samples from a minibatch can cause gradient bias with minibatch stochastic optimization, leading to sub-optimal solutions.
We demonstrate the problem with a simple synthetic experiment, where we try to optimize the simplest CL form of e s 1 / ( e s
+ e s 2 ) over s 1 and s 2 .
We inject noise into the denominator to mimic stochastic optimization, resulting in a stochastic objective of e s 1 / ( e s 1 + e s 2 +  t ) in each SGD iteration, where  is a random zero-mean Gaussian noise with variance set to 0.2.
We compare direct optimization with SGD and our method (introduced below) on 10 random runs, starting from different initializations.
The learning curves are plotted in
Figure 1.
It is observed that direct optimization diverges on different runs due to the high injected noise, whereas our method can always converge to the same optimum with loss = 1.
To further demonstrate the consequence of gradient bias, we inject noise of variance 0.1 and use a learning rate of 0.2 (other settings lead to similar conclusions) to run optimization until converged (around 500 steps).
To guarantee a unique optimal solution, we introduce a regularizer of 0 .
1 ⇥ ( s 21 + s 22 ) .
At convergence, our method gives s 1 = 2 . 97 , s 2 = 0 .
34 , whereas direct optimization gives s 1 = 0 .
15 , s 2 = 0 .
04 , indicating a sub-optimal solution from direct optimization.
We include more experimental results in
Appendix F.
In the literature, a common practice to mitigate the gradient bias problem is to use large batch sizes, which has been demonstrated to be necessary by existing works [13].
Although there are also some preliminary works trying to mitigate the negative impact of minibatches [61, 4, 64], whether they can mitigate gradient bias is not totally clear.
Our method deals with this problem from a principled
Bayesian perspective with a new efﬁcient Bayesian data augmentation technique.
2.2 Decomposable Contrastive Learning (DeCL) 2.2.1 The Overall Algorithm
1: # enc ( · ) : encoder backbone to extract features 2: # (  1 , · · · ,  t , · · · ) :
A decreasing sequence in [0 , 1] 3: t =4: for x in loader do .
load a minibatch x with B samples 5: x 1 , x 2 = aug ( x ) , aug ( x ) .
augmentation 6:
Calculate positive/negative similarity scores ( s + ) i , s i + (vector) and ( S  ) ik , s ik  (matrix) 7:
Estimate auxiliary variable/data U
(vector) as described in Section 2.2.
8: loss_1 =
( U ⇤ ( S  .
mean ( dim =  1))
log( s + )) .
mean () .
decomposable contrastive loss 9: loss_2 = (log( S  .
sum ( dim =  1))
log( s + )) .
mean () .
a smoothing loss induced from the stochastic EM principle 10: loss =
t ⇥ loss_1 + (1   t ) ⇥
loss_2 11: loss .
backward() .
update the encoder backbone 12: t = t +13: end for
We ﬁrst present our overall algorithm, then detail the derivations in the following sections.
Our algorithm is illustrated in Algorithm 1, which is very simple and intuitive.
The algorithm mainly consists of the following steps: 1)
Load a minibatch, and calculate the similarity scores (same as standard CL); 2)
Estimate the so-called auxiliary data U , which will be used to adaptively weight the negative data in the minibatch (see loss_1) to achieve gradient bias elimination; 3) Construct the total loss, which consists of a decomposable loss loss_1 and a smoothing loss loss_2 (coming from the stochastic EM algorithm described below), which we reformulate so that it is equivalent to the standard minibatch info-NCE loss [13]; 4) Optimize the loss to update the encoder backbone.
2.2.2 Decomposing Contrastive Loss with Bayesian Data Augmentation
The basic idea is to leverage Bayesian data augmentation techniques [41] to disentangle the inﬁnite negative similarity scores in the denominator of (1).
Originally, Bayesian data augmentation augments a complicated target distribution ‡ p ( x ) ( e.g. , some posterior distribution of interest) to a more wellbehaved augmented distribution p ( x , u ) with an auxiliary random variable u ; one can then directly deal with the augmented distribution p ( x , u ) instead of the original p ( x ) .
For example, instead of maximizing p ( x ) , one can maximize p ( x , u ) w.r.t.
x instead.
This will yield equivalent solutions while obtaining the convenience of dealing with an easier distribution p ( x , u ) .
One important requirement is to guarantee the augmented problem equivalent to the original one.
In other words, the augmented distribution must satisfy: 1) the marginal distribution should recover the original distribution by marginalization, i.e. ,
R u p ( x , u )d u = p ( x ) ; and 2) the augmented distribution is much easier to handle, e.g. , for MLE.
We adopt this general idea to augment the contrastive loss and make it conditional decomposable.
In other words, our goal is to construct an augmented objective, denoted as L dcon ( D , U ; ✓ ) with auxiliary data U ,
such that E U
[ L dcon ( D , U ); ✓ ] =
L con ( D ; ✓ ) , and L dcon ( D , U ; ✓ ) is expected to decouple negative samples given the auxiliary data U .
The main challenge then comes to dealing with the entangling denominator of L x i .
This can be achieved by applying the Gamma identity, i.e. , ↵ /
R u ↵  1 e   u d u .
Speciﬁcally, we augment the denominator in the original contrastive loss into an exponential form, and deﬁnes a joint distribution p ( ✓ , U |D ) over the model parameter ✓ and the augmented data U , { u i } i , where each u i is associated with a data sample x i 2 D :
[Equation]
Theorem 1.
The joint distribution (2) is a valid augmented distribution, i.e., marginalizing the augmented variable U recovers the original contrastive loss (1) up to some constant.
Constructing a Decomposable Loss With the probability interpretation of the contrastive loss in (2), we can apply maximal likelihood estimation to optimize the model parameter ✓ .
Note the augmented random variables have simple Gamma conditional distributions, which can be easily sampled from.
To this end, we propose an effective approximate method via stochastic expectationmaximization to optimize the joint distribution in Section 2.2.3.
In the following, we ﬁrst derive the decomposable loss from (2).
Suppose there are K negative samples for each data sample x i , where K = 1 represents the population-loss setting.
Note by scaling u i to K u i , (2) is equivalent to p ( ✓ , U |D ) /
Q
[Equation]
Q
[Equation]
e
[Equation]
.
Conditioning on U and taking negative logarithm on both sides of p ( ✓ , U |D ) give us a new augmented loss:
L dcon ,
[Equation]
where the expectation is taken over the distribution of negative samples of x i .
On Decomposibility Following the literature such as [31], we deﬁne decomposable loss as a loss
L ( ˜ D ) such that L ( ˜ D ) =
[Equation]
˜ x i 2 ˜ D L i ( ˜ x i ) for some dataset ˜ D , i.e. , the sub-loss L i only depends on the single data point ˜ x i .
To prove our loss (3) is decomposable for a given U , we ﬁrst construct an augmented dataset ˜ D 0 , { x i , x + i , x  i } from the original dataset D , { x i } and the data augmentation family
T as follows: 1) Randomly sample a data sample x i ⇠ D ; 2) Randomly sample a transfor- mation function t ⇠ T and apply it to x i to generate a positive sample x + i
=
t ( x i ) ; 3) Randomly sample another data sample x 0 ⇠ D excluding x i ; 4) Randomly sample a transformation function t 0 ⇠ T and apply it to x 0 to generate a negative sample x  i =
t 0 ( x 0 ) .
From the construction, it is clear that each data points ˜ x 0 i , ( x i , x
+ i , x  i ) in the new augmented dataset ˜ D 0 are identical independently distributed, and thus (3) is equivalent to L 0 dcon ,
P
[Equation]
P
[Equation] [Equation]
˜
[Equation]
L 0 i under the K ! 1 and inﬁnite data limit.
Since it is easily veriﬁed that the sub-losses L
i ’s are independent from each other, our new loss (3) is thus a valid decomposable loss.
Consequently, conditioned on U , one can safely apply stochastic optimization, where, at each iteration, a random minibatch of data is sampled from ˜ D 0 (which can be constructed from a minibatch of the original data D ) and used to evaluate the stochastic gradients from (3) for parameter updates.
To conclude, we have:
Corollary 2.
Given the auxiliary variable U , the new contrastive loss (3) is decomposable.
And thus, using minibatches for stochastic optimization maintains unbiased stochastic gradients.
2.2.3 Solving by Stochastic Expectation-Maximization
We adopt the MLE principle on the augmented joint distribution for model learning.
Because the augmented variables U are local random variables that are associated to the original data, it is unwise to directly optimize them in the same way as the global model parameter ✓ .
We advocate that stochastic expectation-maximization (sEM) is a natural ﬁt to this problem, as the conditional distribution of U follows a simple Gamma distribution, which allows one to easily get samples from.
There are different variants of sEM algorithms [2, 11, 17].
We adapt the method introduced in [2], and propose an efﬁcient sEM variant to solve our problem.
Speciﬁcally, based on our decomposable form of the contrastive loss, we deﬁne a joint distribution p ( ✓ , U | ˜ D 0 ) over ( ✓ , U ) based on the augmented dataset ˜ D 0 .
Since ˜ D 0 is of inﬁnite size, we will apply stochastic optimization to optimize log p ( ✓ , U | ˜ D 0 t ) in each iteration t , where ˜ D 0 t is a minibatch randomly sampled in this iteration.
Standard stochastic optimization theory implies such a stochastic approximation step would converge under some standard conditions [6].
The problem now reduces to optimizing log p ( ✓ , U | ˜ D 0 t ) in each iteration.
To this end, we adopt the sEM algorithm proposed in [2], which is used to maximize a joint likelihood of a latent variable model relying only on samples of the latent variables.
Speciﬁcally, with our notation and t indexing iteration, the algorithm consists of the following steps in each iteration: 1)
Simulation : sample the auxiliary data U t from their approximate posterior distribution; 2) Stochastic approxiation : update a stochastic objective as
[Equation]
The sequence {  t } is user-speciﬁed, and we recommend  t / 1 /t ;
3) Maximization : maximize Q t +1 ( ✓ ) .
Based on the general algorithm, we make a few adjustments of this basic framework to adopt to our setting, and elaborate on the three steps in the following:
Simulation This corresponds to simulating the auxiliary data U t for the current minibatch.
After scaling u i by K and taking the K ! 1 limit on (2), it is known that the posterior distribution of u i follows: p ( u i | ✓ , D ) =
Gamma (1 , E [ s ik  ]) .
Unfortunately, the exact value of the rate parameter
E [ s ik  ] is intractable, hindering exact simulation from its Gamma posterior.
Thanks to the theory in [2, Theorem 3], if one can construct a sequence of approximate distributions (˜ q 1 ( u i ) , · · · , ˜ q t ( u i ) , · · · ) such that the sequence converges in mean to p ( u i | ✓ , D ) , then sampling from these (˜ q t ( u i )) can also gaurantee convergence of the algorithm.
To this end, we propose two methods to construct a sequence of rate parameters to approximate E [ s ik  ] of the Gamma posterior p ( u i | ✓ , D ) : 1) Maintain a global buffer to store the average similarities score 1 K
P k s ik  for each u i , which is then updated by moving averages over the minibatches; 2) Directly use the momentum network to calculate the average similarity scores from the current minibatch to approximate the grund-true expected similarity scores.
Note this strategy can be directly applied to models with a momentum network such as MoCo [14] without any architecture changes.
In both ways, the rate parameters of the Gamma posteriors are expected to converge to the expectation of the corresponding similarity scores.
Stochastic Approximation According to the description above, the stochastic objective contains two terms: Q t ( ✓ ) and log p ( ✓ , U t | ˜ D 0 t ) .
The latter endows a convenient form, which has the same expression as (3) but with expectation evaluated on the current minibatch.
The term Q t ( ✓ ) has the same form as log p ( ✓ , U t | ˜ D 0 t ) , but with statistics taken from previous iterations (Loss_1 in
Algorithm 1).
For optimization convenience and algorithm stability, we marginalize out the auxiliary data U from previous iterations, giving us the same info-NCE loss evaluated on a minibatch (Loss_ in Algorithm 1), as in standard contrastive learning.
Although we did not investigate different versions of  t sequence, in practice, we ﬁnd performing alternative optimization between Loss_1 and Loss_ is beneﬁcial, e.g. , we alternate between one step optimization of Loss_1 and one step optimization of Loss_2 in the Open-CLIP experiment; and simply directly optimize over Loss_1 without the smoothing loss Loss_2 in the MoCo-v3 with ImageNet experiments.
Maximization Once Q t +1 ( ✓ ) is constructed, which is in the form of a weighted average between the info-NCE loss and log p ( ✓ , U t | ˜ D 0 t ) as explained above, optimization can be done in the same way as the standard CL, e.g. , by adopting the same optimization algorithms and hyperparameters.
More detailed explanation and derivations of the proposed sEM algorithm can be found in the
Appendix C.
Compared to standard contrastive learning, our method has an extra step of estimating the auxiliary variables U .
This additional computational cost for U is negligible relative to the computation of other parts, e.g. , the back-propagation step.
Empirically, we also observe almost the same computational speed of our method compared to the standard contrastive learning.
2.2.4 Generalization Bound
By extending the techniques in [27], we study the ﬁnite sample generalization bound of the proposed decomposable contrastive loss.
Let F be a hypothesis class containing feature extractors from the data space X to feature space R k .
Similar to [27], we extend Rademacher complexity [50] to function classes with high-dimensional outputs, and deﬁne the Rademacher complexity of F on n data as
[Equation]
⌘i , where  is a uniform random vector in {  1 , 1 } n and f i ( z ) is the i -th dimension of f ( z ) .
Note in CL, the function f is typically normalized, meaning b
R n ( F ) is controllable, e.g. , it is typically less than or equal to one.
For conciseness, we slightly overload some of the previously deﬁned notation.
Let f ⇤ pop 2 F be a minimizer of the loss L con ( f ) on inﬁnite data, called population loss; and let L dcon denote the loss of (3) on a ﬁnite dataset D , called empirical loss.
Since L dcon is an unbiased estimator of the population loss L con , we can derive generalization bounds via off-the-shelf concentration inequalities.
Assumption 1.
(Realizability)
At least one of the global minima of L con ( f ) belongs to F .
The following theorem bounds the population loss of a feature extractor trained on ﬁnite data.
Theorem 3.
Assume k
f ( x ) k 1   for some  > 0 and for all f 2 F , x 2 X .
Let f ⇤ pop 2 F be a minimizer of the population loss L con ( f ) .
Given a random dataset of size n , let ˆ f emp 2 F be a minimizer of the empirical loss L dcon .
Denote a .
b as a  Cb for some constant C .
Under
Assumption 1, with probability at least 1   over the randomness of data, the generalization loss is bounded as
[Equation]
where constants c 1 .
k  exp( 4  ) , c 2 .
k  2 /  + exp( 2  ) , k is the featue dimension and  is the commonly used temperature parameter in the similarity score [13, 14, 27].
Related Work
We discuss the most related research in this section.
A more complete list of related works is provided in Section A in the Appendix.
Spectral Contrastive Learning and Related
The recently proposed spectral contrastive learning loss [27], with our notation, can be written as: L =
[Equation]
CL also shares a similar objective [66].
The main difference of ours and theirs lie on the scaling of similarity scores: our method scales positive similarity scores with logarithm, and has the mechanism of adaptively weighting the negative data.
Moreover, our framework is theoretically equivalent to the standard contrastive loss, whereas the above two are not.
SimCLR [13] SimCLR is an instance of CL based on the standard info-NCE contrastive loss.
Since our method is independent of such speciﬁc implementation of CL, it can be directly applied to
SimCLR by replacing the loss.
We will demonstrate our method within SimCLR in the experiments.
MoCo The MoCo family, especially the recent MoCo-v3 [14], represents a strong baseline for contrastive learning.
Similar to the SimCLR case, we can also directly replace the original MoCo contrastive loss with our loss, which will also be comprehensively studied in the experiments.
Negative-Sample-Free SSL
Another line of research toward SSL is via negative-sample-free methods, where only positive samples are considered, e.g. , BYOL [24] and SimSiam [16].
Our framework can be considered as a bridge between contrastive and non-contrstive learning methods, as the contribution of nagative samples are adaptively weighted by the auxiliary random variables U .
Especially, when the expected negative similarly scores become smaller (negative samples are far away), the adaptive weights will asymptotically approach zero and thus become negative-sample free.
Improved Contrastive Learning There have been many efforts trying to improve contrastive learning from various perspectives, e.g. , [61] modiﬁes the info-NCE loss to mitigate the robustness of the loss to the minibatch size; [54] shows the necessity of large batch size in contrastive training and proposed an improved method; [12] identiﬁes the log - K curve of contrastive learning and proposes an improved FlatNCE loss for ﬁx; Other efforts include the SimSiam [16], relative predictive coding [56], Wasserstein predictive coding [28], Barlow Twins [64], VICReg [4], gradient catching [20], and etc .
All these methods do not explicitly address the gradient bias problem.
There are also some previous work considering better ways to sampling negative samples [22, 30, 48], but they still face the gradient-bias problem.
Our method provides a principled way to decouple the negative samples with Bayesian data augmentation, enabling an effective solution to mitigate gradient bias.
Experiments
We ﬁrst demonstrate the robustness of our method against gradient bias on small-scale problems, and then test it on large models for single-modal image and multi-modal vision-language representation learning.
All the experiments are conducted on a NVIDIA A100 GPU Server.
All hyperparameters are adopted from the public codebases speciﬁed below without further tuning, although we believe more improvement can be achieved by further tuning the hyper-parameters, which is left as interesting future work.
Our goal of experiments is to demonstrate the effectiveness of our method to mitigate gradient bias in contrastive learning, but not to achieve state-of-the-art results.
Thus, we do not compare our method with other non-contrastive learning methods, although we include some results in the Appendix.
Code will be publicly available after internal reviewing and approval.
4.1 Robustness against Similarity-Score Noise
To illustrate the impact of similarity-score noise, we explore two ways to add score noise: one manually perturb similarity scores and the other adopts the intrinsic noise induced from using minibatches.
Manual Similarity-Score Noise We gradually add Gaussian noise to the similarity scores and monitor the performance changes.
We compare our method with the popular SimCLR baseline [13], where our model uses exactly the same network architecture as SimCLR.
We conduct experiments on the popular CIFAR-10 dataset, and implement our method based on the publicly available codebase [47], where we strictly follow the default hyper-parameter setting without any tuning.
We ﬁx the minibatch size to 128 while varying the variance of the added Gaussian noise.
The results are plotted in Figure 2.
It is clear that with the noise variance increasing, the validation accuracy of
[Table]
[Figure]
[Table]
[Figure]
[Table]
SimCLR drops much more signiﬁcantly than ours, demonstrating the robustness of our method to similarity-score noise.
Intrinsic Similarity-Score Noise
Next, we investigate the impact of similarity-score noise inherited from minibatches.
In addition to the commonly used architecture for contrastive learning, where a normalization step is applied at the output of the backbone encoder network to project the feature representation onto a hyper-sphere, we additionally test a variant by removing the normalization step.
The reason is that we ﬁnd the normalization can make the impact of similarity-score noise less prominent because the representations are constrained to a much more compact sub-space.
When removing the feature normalization, we adopt exponential of negative Euclidean distance as the similarity metric.
We compare our method with SimCLR on the CIFAR-10 and STL datasets, and follow the default hyper-parameter settings as detailed in the codebase [47].
We vary the minibatch sizes along with different pretraining epoches.
The results are shown in Table 1.
It is observed that our method signiﬁcantly improves over SimCLR in all cases.
Particularly, the performance gap without feature normalization is more prominent, probably because the architecture is more vulnerable to similarity-score noise, as explained above.
Note the Github codebase [47] gives a topaccuracy of 89.1% with a minibatch of 512 and epochs of 500.
For directly comparison, we run our method with the same setting, which gives 89.9%, further demonstrating gradient bias in the large-batch setting.
We further compare our method with existing SSL methods including popular contrastive and noncontrastive learning methods.
The results are presented in Table 6 in the Appendix.
In general, our method outperforms existing methods in most case, especially when the batch sizes are small.
Please refer to Section F in the Appendix for more detailed experiments § .
4.2 Single-Modal Image Representation Learning
We then evaluate our framework on foundation models for large-scale image representation learning.
We adopt MoCo-V3 [14], one of the recent state-of-the-art method, as our baseline.
We strictly follow the default settings as in the publicly available codebase [15] without further tuning, except replacing the original contrastive learning loss with our method.
We run the models on the public ImageNet-1K dataset, and test it with different minibatch sizes, ranging among 128, 256 and 512.
Note the reported results in the original paper are based on very large batchsizes such as 1024 and 4096, which we were not able to test them due to the time and computation constraint.
But we do ﬁne our method can match those results with smaller batch sizes.
Similar to the original paper, we test our method with different encoder backbones including ResNet50, ViT-Small and ViT-Base.
For evaluation, we ﬁnetune the pretrained models for downstream image classiﬁcation on the ImageNet-1K validation data.
We
[Table]
also test the generalization ability of the pretrained models by conducting transfer learning on the
CIFAR-10, CIFAR-100, Flowers and Pets datasets, funetuning from the pretrained checkpoints.
Validation on ImageNet We evaluate the performance on ImageNet in terms of kNN and linear probing classiﬁcation accuracy on different checkpoints.
For kNN evaluation, we follow previous work [38] by simply performing k -nearest neighbor classiﬁcation in the embedding feature space, where k is set to 10, 20 and 100, respectively.
For linear probing, we following [14] by ﬁxing the pretrained weights of the backbone network and supervised ﬁnetuning an added linear layer.
We adopt the standard top-1 and top-5 accuracies on the validation data.
The ﬁnal results are summarized in Table 2.
It is observed that our method outperforms MoCo-v3 in all cases.
Remarkably, on the
ViT-Base/16 atchitecture, it slightly outperforms the CL-based art, achieving a top-1 accuracy of 68.9% and 76.8%, respectively, in terms of kNN and linear probing, but with around half of the minibatch size compared to previous art (67.8% and 76.7% of MoCo-v3 [14, 15]).
Note only the last row in the table is directly comparable to the numbers in the ofﬁcial Github.
All the other results use smaller batch sizes, thus are expected with lower accuracy.
Transfer Learning We transfer the pretrained checkpoints and prepare the datasets by strictly following the instructions from the MoCo-v3 website [15].
Similarly, we use the default settings for the experiments.
The results are summarized in Table 3.
Consistent with previous ﬁndings [14], due to the sizes of these downstream tasks, the performance gaps between different methods are not as signiﬁcant as other comparisons.
Nevertheless, with only around half of the batch size, our method perform similarly with previously reported state of the arts [14].
Few-Shot Image Classiﬁcation To further demonstrate the quality of our learned representations, we adopt the recently developed ELEVATER benchmark [37], which performs 5-shot transfer learning to 20 public image classiﬁcations datasets.
Each data set contains 5 randomly selected training samples in each class, and the model is trained for 50 epochs before the test score is reported,.
Three random seeds are considered for each data set.
We deploy the automatic hyper-parameter tuning pipeline implemented in ELEVATER to make a fair linear probing comparison of pre-trained models based on ViT-Base.
The original metrics of each dataset are used, with more details provided in [37].
The results are shown in Table 4.2.
Consistently, our method achieves an average accuracy of 54.7%, outperforming MoCo-v3 by 4.5% on average.
More signiﬁcantly, our results represents a new state of the art among the models pretrained on ImageNet-1K (to date Oct 7, 2022), outperforming the CACR model [66] with an average accuracy of 54.51%.
4.3 Multi-Modal Vision-Language Representation Learning
Finally, we test our framework for vision-language representation learning by adopting the publicly available Open-CLIP codebase [29].
Similarly to CLIP [46], Open-Clip uses contrastive learning to jointly embed images and the corresponding text.
We follow the default setting used in the codebase to train the 3M images of the Conceptual Captions (CC) dataset [51], consisting of 2.89M training images and 13K validation images.
We replace their contrastive learning algorithms (variant of
SimCLR) with our method.
The backbone architecture is the ResNet-50x4.
To evaluate the pretrained models, we test them on the CC validation dataset, as well as applying them for zero-shot evaluate on
[Table]
[Table] [Table]
the ImageNet-1K validation dataset.
The evaluation tasks on the CC dataset are text-to-image and image-to-text retrievals.
Following previous works, we adopt R@1, R@5, R@10 and validation loss as the evaluation metrics.
For zero-shot ImageNet classiﬁcation, we adopt top-1 and top-5 accuracies.
For computational convenience, following [29], we pretrain the model for 30 epochs.
The results are summarized in Table 5.
Our method gives consistently better results in all metrics, especially in the
R@10 score and top-5 accuracy, validating the effectiveness of our method in gradient bias reduction.
Conclusion
We investigate the problem of gradient bias in contrastive learning with respect to inﬁnite negativesample size, by formulating it in a probability framework with Bayesian data augmentation.
An efﬁcient and simple stochastic EM algorithm is proposed, which constitutes a decomposable CL framework.
Generalization bound on ﬁnite training samples is also developed.
Experimental results on various settings, ranging from small to large models and from single to multiple modality settings, demonstrate the effectiveness of our framework.
Interesting future works include investigating more accurate methods to estimate the augmented variables and improving on the multi-modal learning setting.
Furthermore, one potential advantage of our framework is its potential for efﬁcient distributed training, especially in the asynchronized setting, as one only needs to maintain the auxiliary data without communicating with local data.
Limitation & Potential Negative Societal Impacts
This paper identify the gradient bias problem in contrastive learning and proposes efﬁcient and effective solutions to mitigate the problem.
However, better approximations are still needed to tackle the problem, e.g. , better methods to estimate the auxiliary data.
Thus, there is still room to improve the proposed method.
Finally, as this work is technical in nature and does not include sensitive data, we do not foresee potential negative societal impacts.