An Efficient Bayesian Data Augmentation Approach for Gradient-Bias Mitigation in Contrastive Learning
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 first time, we consider optimizing a more generalized contrastive loss, where each data sample is associated with an infinite number of negative samples.
We show that directly using minibatch stochastic optimization could lead to gradient bias.
To remedy this, we propose an efficient
Bayesian data augmentation technique to augment the contrastive loss into a decomposable one, where standard stochastic optimization can be directly applied without gradient bias.
Specifically, our augmented loss defines 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 obtains strong results in the recent public benchmark ELEVATER for few-shot image classification, and can outperform the strong MoCo-v3 under the same hyper-parameter setting, with only around half of the minibatch size.
Introduction
Recent years have seen significant 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 define 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 infinite).
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 define some decomposable loss functions for learning.
For example, based on spectral theory, [27] proposes spectral contrastive learning, which defines 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 benefits 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 difficultto-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 define a joint distribution over the model parameters and the auxiliary data.
Conditioned on the auxiliary random variables, the potentially infinite negative samples will be separated from each other without introducing bias, thus making a decomposable loss (details in
Section 2.2.2).
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.
Specifically, 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 efficient sEM algorithm for MLE of the joint distribution over model parameters and augmented data, which only needs minimal modifications 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.
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 ∈ 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 ∀ a ; and a negative pair is represented as ( x i , x j ) or ( x i , x aj ) for i ̸
= 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 a i ; θ )) 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 infinite 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 infinite 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.
[Figure]
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 infinite number of negative pairs for each data point.
For practical feasibility, a common trick is to replace the infinite negative pairs with only those within a minibatch.
This gives us an approximation of L x i as ˜
L x i ≜
[Equation]
s
[Equation]
− 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 ̸ = E (cid:104)
˜
L x i (cid:105) , and consequently, ∇ θ L x i ̸ = ∇ θ E (cid:104) ˜
L x i (cid:105) , 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 zeromean 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. 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 experiments 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 efficient Bayesian data augmentation technique.
2.2 Decomposable Contrastive Learning 2.2.1 The Overall Algorithm
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_ 11: loss .
backward() ▷ update the encoder backbone 12: t = t +13: end for
We first 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.
The basic idea is to leverage Bayesian data augmentation techniques [41] to disentangle the infinite 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. , (cid:82) 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. ,β α ∝ (cid:82) u α − 1 e − βu d u .
Specifically, we augment the denominator in the original contrastive loss into an exponential form, and defines 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 ∈ 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 will propose an effective approximate method via stochastic expectationmaximization to optimize the joint distribution in Section 2.2.3.
In the following, we first derive the decomposable loss from (2).
Suppose there are K negative samples for each data sample x i , where K = ∞ represents the population-loss setting.
Note by scaling u i to K u i , (2) is equivalent to p ( θ, U |D )
∝ (cid:81) i : x i ∈D s i + e − u
i
K s i + (cid:81)
k
e − u
i 1 K s ik − .
Conditioning on U and taking negative logarithm on both sides of p ( θ, U |D ) give us a new augmented loss: [Equation] dcon ≜ (cid:88) [Equation]
i
[Equation]
where the expectation is taken over the distribution of negative samples of x i .
On Decomposibility Following the literature such as [31], we define decomposable loss as a loss
L ( ˜ D ) such that L ( ˜ D ) =
(cid:80) ˜ x i ∈ ˜
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 first construct an augmented dataset ˜ D ′ ≜ { 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 transformation 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 ′ ∼ D excluding x i ; 4) Randomly sample a transformation function t ′ ∼ T and apply it to x ′ to generate a negative sample x − i =
t ′ ( x ′ ) .
From the construction, it is clear that each data points ˜ x ′ i ≜
( x i , x + i , x − i ) in the new augmented dataset ˜ D
′ are identical independently distributed, and thus (3) is equivalent to L ′ dcon ≜ (cid:80) i : x ′
i ∈ ˜
D ( u i s ik − − log( s i + )) ≜ (cid:80) i : x ′
i ∈ ˜
D L ′ i under the K → ∞ and infinite data limit.
Since it is easily verified 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 ′ (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 fit 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 efficient sEM variant to solve our problem.
Specifically, based on our decomposable form of the contrastive loss, we define a joint distribution p ( θ , U | ˜ D ′ ) over ( θ , U ) based on the augmented dataset ˜ D ′ .
Since ˜ D
′ is of infinite size, we will apply stochastic optimization to optimize log p ( θ , U | ˜ D ′ t ) in each iteration t , where ˜ D ′ 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 ′ 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.
Specifically, 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 Q t +1 ( θ ) =
Q t ( θ ) + λ t (log p ( θ , U t | ˜ D ′ t ) − Q t ( θ )) ; The sequence { λ t } is user-specified, 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 → ∞ 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 (cid:80) 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 ′ 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 ′ 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.
In practice, we find performing alternative optimization between Loss_1 and Loss_2 is more beneficial than the joint optimization (we alternate between one step optimization of Loss_1 and one step optimization of Loss_2 in the Open-CLIP experiment); And sometimes, we simply directly optimize over Loss_1 without the smoothing loss
Loss_2 (this is adopted in the 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 ′ 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 finite 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 define the Rademacher complexity of F on n data as (cid:98)
[Equation]
, 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 (cid:98) 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 defined notation.
Let f ∗ pop ∈ F be a minimizer of the loss L con ( f ) on infinite data, called population loss; and let L dcon denote the loss of (3) on a finite 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 finite data.
Theorem 3.
Assume ∥
f ( x ) ∥ ∞ ≤ κ for some κ > 0 and for all f ∈ F , x ∈ X .
Let f ∗ pop ∈ F be a minimizer of the population loss L con ( f ) .
Given a random dataset of size n , let ˆ f emp ∈ 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 = (cid:80) i : x ′ i ∈
˜ D ( s ik − − 2 s i + ) .
The attraction-repulsion
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 specific 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] modifies 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] identifies the log - K curve of contrastive learning and proposes an improved FlatNCE loss for fix; 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 first 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 specified 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.
Code will be publicly available.
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 fix 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 [Figure] [Table] SimCLR drops much more significantly than ours
[Figure]
demonstrating the robustness of our method to similarity-score noise.
[Figure] 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 find 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 significantly 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 fine 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 finetune the [Table] pretrained models for downstream image classification on the ImageNet-1K validation data.
We 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 classification accuracy on different checkpoints.
For kNN evaluation, we follow previous work [38] by simply performing k -nearest neighbor classification in the embedding feature space, where k is set to 10, 20 and 100, respectively.
For linear probing, we following [14] by fixing the pretrained weights of the backbone network and supervised finetuning an added linear layer.
We adopt the standard top-1 and top-5 accuracies on the validation data.
The final 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 official 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 findings [14], due to the sizes of these downstream tasks, the performance gaps between different methods are not as significant 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 Classification 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 classifications 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 significantly, 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
[Table]
[Table] [Table]
models, we test them on the CC validation dataset, as well as applying them for zero-shot evaluate on 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 classification, 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 infinite-negative-sample contrastive learning, by formulating it in a probability framework with Bayesian data augmentation.
An efficient and simple stochastic EM algorithm is proposed, which constitutes a decomposable CL framework.
Generalization theory 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 setting.
Furthermore, one potential advantage of our framework is its potential for efficient 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 efficient 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 U .
Thus, there is still room to improve the proposed method.
Finally, as this work is technical in nature, we do not foresee potential negative societal impacts.
Acknowledgments and Disclosure of Funding
C.Chen is partial supported by NSF CISE-2223292 and an Amazon research award.
A More Related Work
Self-Supervised Learning Earlier self-supervised learning approaches tries to generate pseudo labels to the unlabeled data through different transformations, such as solving jigsaw puzzles [43], colorization [65] and rotation prediction [23].
These methods, however, are typically not competitive with supervised learning.
Recently, contrastive-learning based methods have been one of the main stream methods for SSL [59, 57, 55, 28, 13], which can significantly close the gap to supervised learning without relying on ground-truth labels.
Our proposed framework serves as an important addition to the SSL literature, which addresses some limitation of the standard CL and is related to several state-of-the-art SSL methods.
Contrastive Learning Contrastive loss functions were probably first invented for metric learning, which intends to learn similarity functions that measure the similarity between a pair of objects.
Researchers have explored various contrastive loss functions, including the max-margin contrastive loss [26], triple loss [58], multi-class N -pair loss [52], lifted structure loss [53], NCE loss [25], soft-nearest neighbors loss [49, 19], NT-Xent/infor NCE loss [13, 33, 57].
All of these losses are defined in a non-decomposable way that can lead to gradient bias ¶ .
There have been some efforts trying to improve contrastive learning from various perspectives, e.g. , [61] modifies 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] identifies the log − K curve of contrastive learning and proposes an improved FlatNCE loss for fix;
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.
Bayesian Data Augmentation Bayesian data augmentation is a useful tool to facilitate sampling by augmenting complex distributions with auxiliary random variables.
The techniques have found applications in many Bayesian model learning/inference problems, including but not limited to negative Binomial processes [68], Gamma belief networks [69], and stochastic binary networks [62].
More cosely related to the proposed is the Bayesian data augmentation techniques to sampling the normalized random measure, a general family of normalized random probability measures [10, 9].
Our method adopts a similar idea to the CL setting to address one missing piece of the framework .
B Proof of Theorem
Proof.
For completeness, we rewrite the joint distribution p ( θ , U |D ) over the model parameter θ and the augmented data U ≜ { u i } i to be optimize, where each u i is associated with a data sample x i ∈ D :
[Equation]
Obviously, u i ’s are independent with each other given the similarity scores.
Thus, we can apply
Gamma identity to each u i , resulting in
[Equation]
which concludes the proof.
C Stochastic Expectation Maximization
Using our notation, we introduce expectation maximization [40] and stochastic expectation maximization [2].
According to our development in the main text, the marginal likelihood to be optimized is defined as:
[Equation]
which is intractable.
To solve the problem, we augment the marginal distribution with auxiliary data
U ≜ { u i } i as
[Equation]
The standard EM algorithm [40] corresponds to optimizing p ( θ | ˜ D ′ ) by alternating through the following two steps:
• Expectation : Compute the conditional expected log-likelihood:
[Equation]
• Maximization : Maximize θ → Q ( θ | θ k ) in the parameter space to find θ k +1 ∈ arg max Q ( θ | θ k ) .
Unfortunately, directly applying the EM algorithm is infeasible in our setting as U contains an infinite number of auxiliary data, i.e. , Q ( θ | θ k ) in the expectation step is infeasible.
To mitigate the problem, we adapt the stochastic EM algorithm proposed in [2], which constructs a stochastic approximation of Q ( θ | θ k ) in each iteration.
Specifically, the algorithm contains the following steps:
• Simulation : Sample the auxiliary data U t for the current minibatch from their approximate posterior distribution.
Please refer to the “Simulation” paragraph in Section 2.2. for
details on how to get approximate samples.
• Stochastic approximation : Update a stochastic objective as:
[Equation]
θ , U t | ˜ D ′ t ) Here log p ( θ , U t | ˜ D ′ t ) adopts the same form as in (6) but only includes data samples from the current minibatch.
Thus, it is feasible to evaluate.
According to the recursive definition, Q t ( θ ) has the same form as log p ( θ , U t | ˜ D ′ t ) but is evaluated on previous minibatches.
To reduced variance, we first marginalize out the corresponding auxiliary data U t ′ < t , resulting in the same marginal form as (5) but the expectations are only evaluated on minibatches, thus it is feasible.
One problem in evaluating Q t ( θ ) is that we need to store the historical minibatches, which is not storage friendly.
To overcome the problem, we simply use the current minibatch to evaluate Q k , giving us the standard minibatch info-NCE loss used in the standard contrastive learning.
• Maximization : Maximize Q t
+1 ( θ ) .
We simply apply stochastic gradient descent for this step.
D Generalization
For clearness, we redefine some terms used in our proof.
Consider a dataset (cid:98) X
= { ¯ x 1 , ¯ x 2 , · · · , ¯ x n } containing n data points i.i.d.
sampled from P X .
Let ˆ P X be the uniform distribution over (cid:98) X .
Let ˆ P ¯ x, ¯ x ′ be the uniform distribution over data pairs (¯ x i , ¯ x j ) where i ̸
= j .
Based on our algorithm development in the main text, we define the empirical decomposable contrastive loss of a feature extractor f as
[Equation]
where
[Equation]
Our argument in the “On Decomposibility” paragraph in Section 2.2.2 has proved that (cid:98) L n ( f ) is an unbiased estimator of population spectral contrastive loss, i [Equation] e.,
[Equation]
To prove our generalization theory
[Equation] motivated by [27] [Equation] we construct tuples of the original dataset (cid:98) X .
Specifically, we sample a subset of tuples as follows: first sample a permutation π : [ n ] → [ n ] , then we sample tuples S =
[Equation]
i [Equation] i
[Equation] i , u
[Equation]
n/ 2 i =1 as follows:
[Equation]
We define the following loss on S :
[Equation]
It is easy to see that (cid:98) L S ( f ) is an unbiased estimator of (cid:98) L n ( f ) .
We prove the following lemma.
Lemma 4.
Let F be a hypothesis class of feature extractors from X to R k .
Assume ∥
f ( x ) ∥ ∞ ≤ κ for all x ∈ X .
For i ∈ [ k ] , define f i : X → R be the function such that f i ( x ) is the i -th dimension of f ( x ) .
Let F i be the hypothesis containing f i for all f ∈ F .
For m ∈ Z + , let (cid:98) R m ( F i ) be the maximal possible empirical Rademacher complexity of F i over m data:
[Equation]
where x 1 , x 2 , · · · , x m are in X , and σ is a uniform random vector in {− 1 , 1 } m .
Then, the empirical
Rademacher complexity on any m tuples (cid:8)(cid:0) z i , z
+ i , z ′ i (cid:1)(cid:9) m i
=1 can be bounded by
[Equation]
Proof.
[Table]
[Equation]
+
[Equation]
≤
[Equation]
Notice that for any z 1 , z 2 · · · z m and z ′ 1 , z ′ 2 , · · · , z ′ m in X and any i ∈ [ k ] we have
E σ   sup f i ∈F i   1 m m (cid:88)
j = σ j f i ( z j ) f i (cid:0)
z ′ j (cid:1) ≤2 E σ   sup f i ∈F i   1 m m (cid:88)
j = σ j (cid:0) f i ( z j ) + f i (cid:0)
z ′ j (cid:1)(cid:1)  
 + 1 2 E σ   sup f i ∈F i   1 m m (cid:88)
j = σ j (cid:0) f i ( z [Equation]   

[Equation]
   where the first inequaltiy is by Talagrand’s lemma.
Combine these two equations and we get:
[Equation]

[Equation]

[Equation]
D.1 Proof of Theorem
Proof.
We know that E S (cid:104) (cid:98) L S ( f ) (cid:105) =
L con ( f ) , where S is sampled by first sampling (cid:98) X then sample S according to (7).
Notice that when (cid:98) X contains n i.i.d.
natural data samples, the set of random tuples S contains n i.i.d tuples.
Therefore, we can apply generalization bound with
Rademacher complexity to get a uniform convergence bound.
In particular, by Lemma 4 and noticing the fact that E u i (cid:104) u
i exp (cid:16) f ( z i ) ⊤ f
( z ′ i ) /β
(cid:17) − f
( z i )
⊤ f (cid:0) z
+ i (cid:1) /β
(cid:105) always take values in range (cid:104) − kκ 2 /β, kκ 2 /β
+ exp( 2 β ) (cid:105) , we apply standard generalization analysis based on Rademacher complexity.
With similar techniques as those in the proof of Theorem 4.1 in [27], we can get the following result: with probability at least 1 − δ 2 / 4 over the randomness of (cid:98) X and S , we have for any f ∈ F ,
L con ( f )
[Equation]
n
This conclusion translates to the random tuple case as follows:
With probability at least 1 − δ/ 2 over random tuples S conditioned on (cid:98) X , the above equation holds.
Since both L ( f ) and (cid:98) L n ( f ) take value in range (cid:104) − kκ 2 /β, kκ 2 /β + exp( 2 β ) (cid:105) , we have that with probability at least 1 − δ/ 2 over random (cid:98) X , we conclude that for any f ∈ F ,
[Equation]
Since negating the functions in a function class doesn’t change its Rademacher complexity, we can also get results from the other direction: With probability at least 1 − δ/ 2 over random (cid:98) X , we have for any f ∈ F ,
[Equation]
Combine them together we get the excess risk bound, stated as: With probability at least 1 − δ , we have
L con ( ˆ f ) ≤ L ( f ∗F ) + (cid:18) 32 kκ exp( 4 β ) (cid:19) max i ∈
[ k ] (cid:98) R n/ 2 ( F i ) + (cid:18) 4 kκ 2 /β + 2 exp( 2 β ) (cid:19) · (cid:32)(cid:114) 4 log 2 /δ n + δ(cid:33) where ˆ f is minimizer of (cid:98) L n ( f ) in F and f ∗F is minimizer of L ( f ) in F .
Set c 1 = 32 kκ exp( 4 β ) and c 2 = 8 kκ 2 /β + 4 exp( 2 β ) and notice that max i ∈
[ k ] (cid:98) R n/ 2 ( F i ) =
(cid:98) R n/ 2 ( F ) finishes the proof.
E An Alternative Approximate Loss
We propose an alternative loss to approximate the original info-NCE loss, which is potentially more expressive.
Note we did not apply this in our current experiments.
However, we are expecting to add it in to our future version.
Suppose there are K negative samples in a minibatch.
Denote s i, : K ≜ (cid:80)
Kk
=1 s ik − and s i,K : ≜ (cid:80) ∞
k = K +1 s ik − .
The standard contrastive learning with infinite negative samples corresponds to optimizing
[Equation]
where the equality applies due to the assumed unbiased estimation of the expected negative similarity score with a minibatch.
We propose to replace the arithmetic mean of ( s i, : K , s i, : K , s i,K : ) with their geometric mean, and define a geometric contrastive learning objective as
[Equation]
We call this loss function the geometric contrastive loss.
We optimize the loss by maximizing its lower bound (cid:80)
i L i with our Bayesian data augmentation technique.
Specifically, the stochastic EM iterates through:
• Estimate u
i : This is done by sampling from a Gamma distribution as: u i ∼ Gamma (2 , E [ s ik − ]) , where we can use the same technique described in the main text to approximate E [ s ik − ] .
• Maximizing p ( θ |D ) : This can be done by maximizing the above lower bound of the geometric contrastive loss as: L = (cid:80) i log s i
+ s i, : K + 2 log
s i + − u i ˜ s ik − , where ˜ s ik − is a stochastic version of E [ s ik − ] estimated from a minibatch.
F Additional Details on the Experiments
Synthetic experiments We run the experiment for Figure 1 with a different setting to exam the convergence solution, where we mimic the noisy contrastive learning with a simplest synthetic experiment.
Specifically, we define our objective to be maximized as F = e s 1 / ( e s 1 + e s 2 ) w.r.t.
( s 1 , s 2 ) .
In this example, e s 1 can be considered as the positive similarity score, and e sas the sum of all the negative similarity score for one data point.
Thus, this is a simplification of the standard contrastive loss.
To introduce noise into the negative similarity score, we explicitly inject random Gaussian noise into the denominator, 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.09.
To guarantee a unique optimal solution, we introduce a regularizer of 0 .
1 × ( s 21 + s 22 ) .
We compare direct optimization with SGD and our method over 100 independent random runs, starting from zero initializations.
The learning curves of the objective values F w.r.t. the iteration numbers are plotted in Figure 3.
It is observed that the two methods indeed converge to different solutions, with a better objective value from our method.
Comparison with non-contrastive learning methods To demonstrate the effectiveness of our proposed method, we compare it with more popular self-supervised learning methods, including 1) the contrastive-learning methods SimCLR [13], DCL [61], NNCLR [18] and SwaV [7]; 2) the non-contrastive learning methods BYOL [24], DINO [8] and BarlowTwins [64].
For all these methods except SimCLR and ours, we use the public implement in the LIGHTLY benchmark [1], downloaded
[Table]
from https://docs.lightly.ai/_downloads/b99fe89a7fc2b4740cb9f1e34d3229ad/cifar10_benchmark.py, and use the default settings on the optimizer with 300 epochs.
Table 6 shows the comparison of our method with SimCLR on the CIFAR-10 dataset under different batch sizes.
It is clear that the proposed method performs better in general.
The best non-contrastive learning method, BYOL, can perform better than the standard contrastive learning method SimCLR, especially in the smallbatchsize settings.
When correcting gradient bias in contrastive learning with our method, we can outperform BYOL when batch sizes are large enough ( e.g. , ≥ 256), and slightly worse when batch sizes are too small.
The results also suggest there is still room the design better mechanism to correct the gradient bias in contrastive learning.
Furthermore, it is observed that contrastive learning methods can generally perform better with increasing batchsizes, whereas non-contrastive methods do not seem to be improving, suggesting that contrastive learning can be scaled up by using larger batch sizes.
This is part of the reasons why recent large foundation models such as CLIP adopt contrastive learning in the training.