An Efficient Bayesian Data Augmentation Approach for Gradient-Bias Mitigation in Contrastive Learning
Anonymous Author(s)
AffiliationAddress email
Abstract
Contrastive learning (CL) has been the de facto technique for self-supervised representation learning (SSL), with impressive empirical success such as multi-modalrepresentation learning.
However, traditional CL loss only considers negativesamples from a minibatch, which could cause biased gradients due to the nondecomposibility of the loss.
For the first time, we consider optimizing a moregeneralized contrastive loss, where each data sample is associated with an infinitenumber of negative samples.
We show that directly using minibatch stochasticoptimization 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 appliedwithout gradient bias.
Specifically, our augmented loss defines a joint distributionover 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 foundationmodels, including SSL of ImageNet and SSL for vision-language representationlearning.
Experiment results indicate the existence of gradient bias in all cases, anddemonstrate the effectiveness of the proposed method on improving previous stateof the arts.
Remarkably, our method can outperform the strong MoCo-v3 under thesame hyper-parameter setting, with only around half of the minibatch size.
Introduction
Recent years have seen significant developments and applications of contrastive learning (CL), apromising self-supervised learning (SSL) technique, to a variety of striking research breakthroughssuch as learning of multi-modal foundation models [5] in text and image domains [45, 62, 66] fornatural-language and vision-language representation learning.
In addition, CL has also be appliedrecently to other scenarios such as supervised learning [32] and reinforcement learning [3, 36].
Themain idea of CL is to define a contrastive loss where each data sample is entangled with both itspositive and negative data samples, and optimize it to learn transformation-invariant features to betterdiscriminate positive samples from negative ones.
From a technical perspective, the contrastive lossis non-decomposable , in the sense that each data sample is coupled with all the other samples in theloss function, i.e. , to process one data sample, one has to get access to all its positive and negativesamples from the other data samples.
This is intractable when considering a more general CL settingwhere the number of negative samples are unbounded (potentially infinite).
In this case, if stochasticoptimization is applied, only negative samples in the current minibatch are available.
Thus, thiscould lead to biased gradients and consequently sub-optimal solutions (see Section 2.1 for a formalstatement), which is also evidenced in several existing works [21, 24].
A straightforward approach to overcoming the issue is to define some decomposable loss functionsfor 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 problemby completely dropping the negative samples.
These methods, however, are developed from quitedifferent perspective and are not equivalent to the standard contrastive learning, e.g. , SimCLRthat adopts the info-NCE loss [13].
Since standard contrastive learning has been shown to berelatively robust and easier to scale up to large foundation models [5] such as CLIP and its variants[44, 59, 62, 66], it is important to study its equivalent form in terms of a decomposable loss, andunderstand 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 contrastivelearning, by leveraging the well developed Bayesian data augmentation technique from the Bayesianlearning community [40].
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 randomvariables, also known as auxiliary data [38] 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 auxiliaryvariable/data * , allowing one to define a joint distribution over the model parameters and the auxiliarydata.
Conditioned on the auxiliary random variables, the potentially infinite negative samples willbe 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 newstochastic expectation-maximization (sEM) framework for maximal likelihood estimation (MLE) ofthe joint distribution.
Specifically, with the proposed augmentation technique, the auxiliary randomvariables enjoy simple Gamma posterior distributions, conditioned on which the model parametercan 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 probabilityperspective, which enjoys conditional decomposibility for unbiased stochastic optimization.
• We propose an efficient sEM algorithm for MLE of the joint distribution over modelparameters 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 aninput x , the feature extractor outputs its representation, written as z = enc ( x ; θ ) .
The basic ideaof CL is to maximize the representation similarities between positive data pairs while minimizingthose of negative pairs.
Formally, we are given an unlabeled dataset D ≜ { x i } .
Each x i ∈ Dwill 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 encoderand similarity scores will be calculated based on their representations.
For notation simplicity, wewill index positive similarity scores by ‘ j + ’ and negative similarity scores by k − , e.g. , we uses ij + ≜ sim ( enc ( x i ; θ ) , enc ( x ai ; θ )) to denote the similarity score between x i and its j -th positivedata point from { x ai } , where sim ( · , · ) denotes a positive similarity metric such as the exponentialcosine 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.
Intheory, 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 ofcontrastive 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 aslightly simpler sub-loss L x i ≜ s
[Equation]
k
[Equation]
s ik − , where the positive similarity score is written as s i +for notation simplicity.
We will consider this simpler case in our following development , althoughgeneralizing 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 infinitenumber of negative pairs for each data point.
Forpractical feasibility, a common trick is to replacethe infinite negative pairs with only those within aminibatch.
This gives us an approximation of L x ias ˜
L x i ≜
[Equation]
s
[Equation]
+ +
[Equation]
Bk =1 s ik − where B denotes the minibatch size.
We argue that this approximation is biased due to the non-decomposibility nature of thestandard CL loss L con , in the sense that the individualloss L con is associated with the negative samples thatdistribute across the whole data ˜ D .
In other words,the overall loss L con cannot be decomposed into asum over independent sub-losses.
Importantly, nondecomposibility indicates that L x i ̸ = E (cid:104)
˜
L x i (cid:105) , andconsequently, ∇ θ L x i ̸ = ∇ θ E (cid:104) ˜
L x i (cid:105) , where both the expectations are taken over the randomness ofminibaching.
Thus, using negative samples from a minibatch can cause gradient bias with minibatchstochastic optimization, leading to sub-optimal solutions.
We demonstrate the problem with a simplesynthetic experiment, where we try to optimize the simplest CL form of s 1 / ( s 1 + s 2 ) over s 1 > 0and s 2 > 0 .
We inject noise into the denominator to mimic stochastic optimization, resulting in astochastic objective of s 1 / ( s 1 + 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 curvesare plotted in Figure 1.
It is observed that direct optimization diverges on different runs due tothe 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 anduse 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 of0 .
1 × ( s 21 + s 22 ) .
At convergence, our method gives s 1 = 2 . 97 , s 2 = 0 .
34 , whereas direct optimizationgives s 1 = 0 .
15 , s 2 = 0 .
04 , indicating a sub-optimal solution from direct optimization.
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 somepreliminary works trying to mitigate the negative impact of minibatches [60, 4, 63], whether theycan 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
We first present our overall algorithm, then detail the derivations in the following sections.
Ouralgorithm is illustrated in Algorithm 1, which is very simple and intuitive.
The algorithm mainlyconsists of the following steps: 1) Load a minibatch, and calculate the similarity scores (same asstandard CL); 2) Estimate the so-called auxiliary data U , which will be used to adaptively weightthe negative data in the minibatch (see loss_1) to achieve gradient bias elimination; 3) Construct thetotal 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 thestandard minibatch info-NCE loss [13]; 4) Optimize the loss to update the encoder backbone.
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
The basic idea is to leverage Bayesian data augmentation techniques [40] to disentangle the infinite negative similarity scores in thedenominator of (1).
Originally, Bayesian data augmentation augments a complicated target distribution ‡p ( x ) ( e.g. , some posteriordistribution of interest) toa more well-behaved augmented distribution p ( x , u )with an auxiliary randomvariable u ; one can thendirectly deal with the augmented distribution p ( x , u )instead of the original p ( x ) .
For example, instead ofmaximizing p ( x ) , one canmaximize 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 theaugmented problem equivalent to the original one.
In other words, the augmented distribution mustsatisfy: 1) the marginal distribution should recover the original distribution by marginalization, i.e. , 163 (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 ; θ ) withauxiliary data U ,
such that E U
[ L dcon ( D , U ); θ ]
=
L con ( D ; θ ) , and L dcon ( D , U ; θ ) is expected todecouple negative samples given the auxiliary data U .
The main challenge then comes to dealingwith the entangling denominator of L x i .
This can be achieved by applying the Gamma identity, i.e. ,1 β α ∝ (cid:82) u α − 1 e − βu d u .
Specifically, we augment the denominator in the original contrastive loss intoan exponential form, and defines a joint distribution p ( θ , U |D ) over the model parameter θ and theaugmented 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 theaugmented 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 sampledfrom.
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 derivethe 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 equivalentto 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 negativelogarithm on both sides of p ( θ, U |D ) give us a new augmented loss:
[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 onthe single data point ˜ x i .
To prove our loss (3) is decomposable for a given U , we first construct anaugmented dataset ˜ D ′ ≜ { x i , x
+ i , x − i } from the original dataset D ≜ { x i } and the data augmentationfamily 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) Randomlysample another data sample x ′ ∼ D excluding x i ; 4) Randomly sample a transformation functiont ′ ∼ T and apply it to x ′ to generate a negative sample x − i =
t ′ ( x ′ ) .
From the construction, it is clearthat each data points ˜ x
′ i ≜ ( x i , x
+ i , x − i ) in the new augmented dataset ˜ D
′ are identical independentlydistributed, 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
underthe K → ∞ and infinite data limit.
Since it is easily verified that the sub-losses L ′
i ’s are independentfrom 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 issampled from ˜ D ′ (which can be constructed from a minibatch of the original data D ) and used toevaluate 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 thususing 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 theaugmented variables U are local random variables that are associated to the original data, it isunwise to directly optimize them in the same way as the global model parameter θ .
We advocatethat stochastic expectation-maximization (sEM) is a natural fit to this problem, as the conditionaldistribution 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 decomposableform of the contrastive loss, we define a joint distribution p ( θ , U | ˜ D ′ ) over ( θ , U ) based on theaugmented dataset ˜ D ′ .
Since ˜ D
′ is of infinite size, we will apply stochastic optimization to optimizelog 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 convergeunder some standard conditions [6].
The problem now reduces to optimizing log p ( θ , U | ˜ D ′ t ) in eachiteration.
To this end, we adopt the sEM algorithm proposed in [2], which is used to maximize a jointlikelihood of a latent variable model relying only on samples of the latent variables.
Specifically, withour 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) Stochasticapproxiation : update a stochastic objective as Q t +1 ( θ ) =
Q t ( θ ) + λ t (log p ( θ , U t | ˜ D ′ t ) − Q t ( θ )) ;3) Maximization : maximize Q t +1 ( θ ) .
Based on the general algorithm, we make a few adjustmentsof 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.
Afterscaling u
i by K and taking the K → ∞ limit on (2), it is known that the posterior distribution of u ifollows: p ( u i | θ , D ) =
Gamma (1 , E [ s ik − ]) .
Unfortunately, the exact value of the rate parameter 224 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 )) canalso gaurantee convergence of the algorithm.
To this end, we propose two methods to construct asequence of rate parameters to approximate E [ s ik − ] of the Gamma posterior p ( u i | θ , D ) : 1) Maintaina global buffer to store the average similarities score 1 K (cid:80) k s ik − for each u i , which is then updatedby moving averages over the minibatches; 2) Directly use the momentum network to calculate theaverage similarity scores from the current minibatch to approximate the grund-true expected similarityscores.
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 areexpected to converge to the expectation of the corresponding similarity scores.
Stochastic Approximation According to the description above, the stochastic objective containstwo terms: Q t ( θ ) and log p ( θ , U t | ˜ D ′ t ) .
The latter endows a convenient form, which has the sameexpression as (3) but with expectation evaluated on the current minibatch.
The term Q t ( θ ) hasthe 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 auxiliarydata U from previous iterations, giving us the same info-NCE loss evaluated on a minibatch (Loss_2in Algorithm 1), as in standard contrastive learning.
Maximization Once Q t +1 ( θ ) is constructed, which is in the form of a weighted average betweenthe info-NCE loss and log p ( θ , U t | ˜ D ′ t ) as explained above, optimization can be done in the sameway 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 estimatingthe auxiliary variables U .
This additional computational cost for U is negligible relative to thecomputation of other parts, e.g. , the back-propagation step.
Empirically, we also observe almost thesame 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 proposeddecomposable contrastive loss.
Let F be a hypothesis class containing feature extractors from thedata space X to feature space R k .
Similar to [27], we extend Rademacher complexity [49] to functionclasses with high-dimensional outputs, and define the Rademacher complexity of F on n data as(cid:98) R n ( F ) :
=
max x 1 , ··· ,x n ∈X E σ (cid:104) sup f ∈F ,i ∈
[ k ] 1 n (cid:16)(cid:80) nj
=1 σ j f i ( x j ) (cid:17)(cid:105) , where σ is a uniform randomvector in {− 1 , 1 } n and f i ( z ) is the i -th dimension of f ( z ) .
Note in CL, the function f is typicallynormalized, 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 aminimizer 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 populationloss 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 ∈
Fbe a minimizer of the population loss L con ( f ) .
Given a random dataset of size n , let ˆ f
emp ∈
Fbe 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 isbounded as
[Equation]
where constants c 1 ≲ kκ exp( 4 β ) , c 2 ≲ kκ 2 /β + exp( 2 β ) , k is the featue dimension and β is thecommonly 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 providedin Section A in the Appendix.
Spectral Contrastive Learning and Related
The recently proposed spectral contrastive learningloss [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 [65].
The main difference of ours and theirs lie on the scaling ofsimilarity scores: our method scales positive similarity scores with logarithm, and has the mechanismof adaptively weighting the negative data.
Moreover, our framework is theoretically equivalent to thestandard 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 forcontrastive learning.
Similar to the SimCLR case, we can also directly replace the original MoCocontrastive 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-freemethods, where only positive samples are considered, e.g. , BYOL [24] and SimSiam [16].
Ourframework 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 faraway), the adaptive weights will asymptotically approach zero and thus become negative-sample free.
Improved Contrastive Learning There have been many efforts trying to improve contrastivelearning from various perspectives, e.g. , [60] modifies the info-NCE loss to mitigate the robustness ofthe loss to the minibatch size; [53] shows the necessity of large batch size in contrastive training andproposed an improved method; [12] identifies the log - K curve of contrastive learning and proposesan improved FlatNCE loss for fix; Other efforts include the SimSiam [16], relative predictive coding[55], Wasserstein predictive coding [28], Barlow Twins [63], VICReg [4], gradient catching [20],and etc .
All these methods do not explicitly address the gradient bias problem.
There are also someprevious work considering better ways to sampling negative samples [22, 30, 47], but they still facethe gradient-bias problem.
Our method provides a principled way to decouple the negative sampleswith 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, andthen test it on large models for single-modal image and multi-modal vision-language representationlearning.
All the experiments are conducted on a NVIDIA A100 GPU Server.
All hyperparametersare adopted from the public codebases specified below without further tuning, although we believemore improvement can be achieved by further tuning the hyper-parameters, which is left as interestingfuture work.
Our goal of experiments is to demonstrate the effectiveness of our method to mitigategradient bias in contrastive learning, but not to achieve state-of-the-art results.
Thus, we do notcompare our method with other non-contrastive learning methods.
Code will be publicly available.
4.1 Robustness against Similarity-Score Noise
[Figure]
To illustrate the impact of similarity-score noise, weexplore two ways to add score noise: one manuallyperturb similarity scores and the other adopts theintrinsic noise induced from using minibatches.
Manual Similarity-Score Noise We gradually add
Gaussian noise to the similarity scores and monitorthe performance changes.
We compare our methodwith the popular SimCLR baseline [13], where ourmodel uses exactly the same network architecture as
SimCLR.
We conduct experiments on the popular
CIFAR-10 dataset, and implement our method basedon the publicly available codebase [46], where westrictly follow the default hyper-parameter settingwithout 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 SimCLR drops much more significantly than ours, demonstrating therobustness of our method to similarity-score noise.
Intrinsic Similarity-Score Noise
Next, we investigate the impact of similarity-score noise inheritedfrom minibatches.
In addition to the commonly used architecture for contrastive learning, where anormalization step is applied at the output of the backbone encoder network to project the featurerepresentation onto a hyper-sphere, we additionally test a variant by removing the normalizationstep.
The reason is that we find the normalization can make the impact of similarity-score noise lessprominent because the representations are constrained to a much more compact sub-space.
Whenremoving the feature normalization, we adopt exponential of negative Euclidean distance as thesimilarity metric.
We compare our method with SimCLR on the CIFAR-10 and STL datasets, andfollow the default hyper-parameter settings as detailed in the codebase [46].
We vary the minibatchsizes along with different pretraining epoches.
The results are shown in Table 1.
It is observed that ourmethod significantly improves over SimCLR in all cases.
Particularly, the performance gap withoutfeature normalization is more prominent, probably because the architecture is more vulnerable tosimilarity-score noise, as explained above.
Note the Github codebase [46] gives a top-1 accuracy of
[Figure]
[Table]
[Figure]
[Table]
89.1% with a minibatch of 512 and epochs of 500.
For directly comparison, we run our method withthe same setting, which gives 89.9%, further demonstrating gradient bias in the large-batch setting.
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 followthe default settings as in the publicly available codebase [15] without further tuning, except replacingthe original contrastive learning loss with our method.
We run the models on the public ImageNet-1Kdataset, and test it with different minibatch sizes, ranging among 128, 256 and 512.
Note the reportedresults in the original paper are based on very large batchsizes such as 1024 and 4096, which we werenot able to test them due to the time and computation constraint.
But we do fine our method can matchthose results with smaller batch sizes.
Similar to the original paper, we test our method with differentencoder backbones including ResNet50, ViT-Small and ViT-Base.
For evaluation, we finetune thepretrained models for downstream image classification on the ImageNet-1K validation data.
Wealso 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 linearprobing classification accuracy on different checkpoints.
For kNN evaluation, we follow previouswork [37] 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 thepretrained weights of the backbone network and supervised finetuning an added linear layer.
Weadopt the standard top-1 and top-5 accuracies on the validation data.
The final results are summarizedin 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 of68.9% and 76.8%, respectively, in terms of kNN and linear probing, but with around half of theminibatch size compared to previous art (67.8% and 76.7% of MoCo-v3 [14, 15]).
Note only the lastrow in the table is directly comparable to the numbers in the official Github.
All the other results usesmaller batch sizes.
Transfer Learning We transfer the pretrained checkpoints and prepare the datasets by strictlyfollowing the instructions from the MoCo-v3 website [15].
Similarly, we use the default settings forthe experiments.
The results are summarized in Table 3.
Consistent with previous findings [14], dueto the sizes of these downstream tasks, the performance gaps between different methods are not assignificant as other comparisons.
Nevertheless, with only around half of the batch size, our methodperform similarly with previously reported state of the arts [14].
4.3 Multi-Modal Vision-Language Representation Learning
Finally, we test our framework for vision-language representation learning by adopting the publiclyavailable Open-CLIP codebase [29].
Similarly to CLIP [45], Open-Clip uses contrastive learning tojointly embed images and the corresponding text.
We follow the default setting used in the codebaseto train the 3M images of the Conceptual Captions (CC) dataset [50], consisting of 2.89M trainingimages 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 pretrainedmodels, 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 andimage-to-text retrievals.
Following previous works, we adopt R@1, R@5, R@10 and validation lossas 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 aresummarized in Table 4.
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, byformulating it in a probability framework with Bayesian data augmentation.
An efficient andsimple stochastic EM algorithm is proposed, which constitutes a decomposable CL framework.
Generalization theory is also developed.
Experimental results on various settings, ranging from smallto large models and from single to multiple modality settings, demonstrate the effectiveness of ourframework.
Interesting future works include investigating more accurate methods to estimate theaugmented variables and improving on the multi-modal setting.
Furthermore, one potential advantageof our framework is its potential for efficient distributed training, especially in the asynchronizedsetting, 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 problemin 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 theauxiliary data U .
Thus, there is still room to improve the proposed method.
Finally, as this work istechnical in nature, we do not foresee potential negative societal impacts.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes]
(c) Did you discuss any potential negative societal impacts of your work?
[Yes]
(d) Have you read the ethics review guidelines and ensured that your paper conforms tothem?
[Yes]2.
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes](b) Did you include complete proofs of all theoretical results?
[Yes]
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[No](b) Did you specify all the training details (e.g., data splits, hyperparameters, how theywere chosen)?
[Yes]
(c) Did you report error bars (e
g., with respect to the random seed after running experiments multiple times)?
[N/A](d)
Did you include the total amount of compute and the type of resources used (e
g., typeof GPUs, internal cluster, or cloud provider)?
[Yes]4.
If you are using existing assets (e
g., code, data, models) or curating/releasing new assets...
(a)
If your work uses existing assets, did you cite the creators?
[Yes](b) Did you mention the license of the assets?
[N/A]
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes](d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating?
[N/A]
(e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content?
[N/A]
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]
A More Related Work
Self-Supervised Learning Earlier self-supervised learning approaches tries to generate pseudolabels to the unlabeled data through different transformations, such as solving jigsaw puzzles [42],colorization [64] and rotation prediction [23].
These methods, however, are typically not competitivewith supervised learning.
Recently, contrastive-learning based methods have been one of the mainstream methods for SSL [58, 56, 54, 28, 13], which can significantly close the gap to supervisedlearning without relying on ground-truth labels.
Our proposed framework serves as an importantaddition to the SSL literature, which addresses some limitation of the standard CL and is related toseveral 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 contrastiveloss [26], triple loss [57], multi-class N -pair loss [51], lifted structure loss [52], NCE loss [25],soft-nearest neighbors loss [48, 19], NT-Xent/infor NCE loss [13, 33, 56].
All of these losses aredefined 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. ,[60] modifies the info-NCE loss to mitigate the robustness of the loss to the minibatch size; [53]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 [55], Wasserstein predictive coding[28], Barlow Twins [63], VICReg [4], gradient catching [20], and etc .
All these methods do notexplicitly address the gradient bias problem.
There are also some previous work considering betterways to sampling negative samples [22, 30, 47], but they still face the gradient-bias problem.
Ourmethod 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 samplingby augmenting complex distributions with auxiliary random variables.
The techniques have foundapplications in many Bayesian model learning/inference problems, including but not limited tonegative Binomial processes [67], Gamma belief networks [68], and stochastic binary networks [61].
More cosely related to the proposed is the Bayesian data augmentation techniques to sampling thenormalized 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 1
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 samplex i ∈ D
[Equation] [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 [39] and stochastic expectation maximization [2].
According to our development in the main text, the marginal likelihood to be optimizedis 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 [39] corresponds to optimizing p ( θ | ˜ D ′ ) by alternating through thefollowing 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 infinitenumber 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 approximationof 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.3 for details
on how to get approximate samples.
• Stochastic approximation : Update a stochastic objective as:
[Equation]
Here log p ( θ , U t | ˜ D ′ t ) adopts the same form as in (6) but only includes data samples fromthe 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.
Toreduced variance, we first marginalize out the corresponding auxiliary data U t ′ <t , resultingin 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 historicalminibatches, which is not storage friendly.
To overcome the problem, we simply use thecurrent minibatch to evaluate Q k , giving us the standard minibatch info-NCE loss used inthe standard contrastive
learning
• Maximization : Maximize Q t
+1 ( θ ) .
We simply apply stochastic gradient descent for thisstep.
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 featureextractor f as
[Equation]
where u x ∼
[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.e.,
[Equation]
To prove our generalization theory, motivated by [27], 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 ] , thenwe sample tuples S =
[Equation]
z i , z
+ i , z
i , u i (cid:1)(cid:9) n/ 2 i =1 as follows:
[Equation] [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]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] [Equation]
[Table]
[Equation]
[Table]
[Equation]
[Table] [Equation] [Table] 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)
[Equation]
 where the first inequaltiy is by Talagrand’s lemma.
Combine these two equations and we get:
[Equation]

[Equation]

[Equation]
[Table]
[Equation]
[Table]
[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 (
[Equation]
n
This conclusion translates to the random tuple case as follows:
With probability at least 1 − δ/ 2 overrandom tuples S conditioned on (cid:98) X , the above equation holds.
Since both L ( f ) and (cid:98) L n ( f ) take valuein 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 732 have
L con ( ˆ f ) ≤ L ( f
[Equation]
+
[Equation]
[Equation]
[Equation]
max
[Equation]
∈
[Equation]
[Equation]
F
[Equation]
+
[Equation]
+
[Equation]
(cid:19) · (cid:32)(cid:114) 4 log
[Equation]
[Equation]
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 β ) andc 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 moreexpressive.
Note we did not apply this in our current experiments.
However, we are expecting to addit 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 : ≜ 741 (cid:80) ∞
k = K +1 s ik − .
The standard contrastive learning with infinite negative samples corresponds to 742 optimizing
[Equation]
where the equality applies due to the assumed unbiased estimation of the expected negative similarityscore with a minibatch.
We propose to replace the arithmetic mean of ( s i, : K , s i, : K , s i,K : ) with theirgeometric 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 itslower bound (cid:80)
i L i with our Bayesian data augmentation technique.
Specifically, the stochastic EMiterates 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 astochastic version of E [ s ik − ] estimated from a minibatch.
F Additional Details on the Experiments
Synthetic experiments We re-run the exper- iment for Figure 1, where we mimic the noisycontrastive learning with a simplest syntheticexperiment.
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 s 2as the sum of all the negative similarity scorefor one data point.
Thus, this is a simplificationof the standard contrastive loss.
To introducenoise into the negative similarity score, we explicitly inject random Gaussian noise into thedenominator, resulting in a stochastic objectiveof e s 1 / ( e s 1 + e s 2 + δ t ) in each SGD iteration,where δ is a random zero-mean Gaussian noisewith variance set to 0.09.
To guarantee a uniqueoptimal solution, we introduce a regularizer of0 .
1 × ( s 21 + s 22 ) .
We compare direct optimizationwith SGD and our method over 100 independentrandom 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 todifferent solutions, with a better objective value from our method.
Comparison with non-contrastive learning methods To demonstrate the effectiveness of ourproposed method, we compare it with more popular self-supervised learning methods, including1) the contrastive-learning methods SimCLR [13], DCL [60], NNCLR [18] and SwaV [7]; 2) thenon-contrastive learning methods BYOL [24], DINO [8] and BarlowTwins [63].
For all these methodsexcept SimCLR and ours, we use the public implement in the LIGHTLY benchmark [1], downloadedfrom https://docs.lightly.ai/_downloads/ [Table] and use the default settings on the optimizer with 300 epochs.
Table 5 shows the comparison ofour method with SimCLR on the CIFAR-10 dataset under different batch sizes.
It is clear that theproposed 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 canoutperform BYOL when batch sizes are large enough ( e.g. , ≥ 256), and slightly worse when batchsizes are too small.
The results also suggest there is still room the design better mechanism to correctthe gradient bias in contrastive learning.
Furthermore, it is observed that contrastive learning methodscan generally perform better with increasing batchsizes, whereas non-contrastive methods do notseem to be improving, suggesting that contrastive learning can be scaled up by using larger batchsizes.
This is part of the reasons why recent large foundation models such as CLIP adopt contrastivelearning in the training.