# EMC2: Efficient MCMC Negative Sampling for Contrastive Learning with Global
Convergence

Chung-Yiu Yau C.-Y. Yau and H.-T. Wai are with the Department of Systems
Engineering and Engineering Management, The Chinese University of Hong Kong,
Hong Kong SAR of China. Hoi-To Wai††footnotemark:  Parameswaran Raman P. Raman
and S. Sarkar are with Amazon Web Services, USA. Soumajyoti
Sarkar††footnotemark:  Mingyi Hong M. Hong is with Department of Electrical
and Computer Engineering, University of Minnesota, USA.

###### Abstract

A key challenge in contrastive learning is to generate negative samples from a
large sample set to contrast with positive samples, for learning better
encoding of the data. These negative samples often follow a softmax
distribution which are dynamically updated during the training process.
However, sampling from this distribution is non-trivial due to the high
computational costs in computing the partition function. In this paper, we
propose an Efficient Markov Chain Monte Carlo negative sampling method for
Contrastive learning (EMC2). We follow the global contrastive learning loss as
introduced in Yuan et al. (2022), and propose EMC2 which utilizes an adaptive
Metropolis-Hastings subroutine to generate hardness-aware negative samples in
an online fashion during the optimization. We prove that EMC2 finds an
$\mathcal{O}(1/\sqrt{T})$-stationary point of the global contrastive loss in
$T$ iterations. Compared to prior works, EMC2 is the first algorithm that
exhibits global convergence (to stationarity) regardless of the choice of
batch size while exhibiting low computation and memory cost. Numerical
experiments validate that EMC2 is effective with small batch training and
achieves comparable or better performance than baseline algorithms. We report
the results for pre-training image encoders on STL-10 and Imagenet-100.

§§footnotetext: The work of C.-Y. Yau was done while interning at Amazon Web
Services. M. Hong holds concurrent appointments as an Amazon Scholar and as a
faculty at University of Minnesota. This paper describes their work performed
at Amazon.

##  1 Introduction

Contrastive representation learning has been instrumental in self-supervised
learning for large-scale pre-training of foundation models Radford et al.
(2021); Cherti et al. (2023) as well as in the fine-tuning stage on downstream
tasks Xiong et al. (2020); Lindgren et al. (2021). It helps encode real-world
data into low-dimensional feature vectors that abstract the important
attributes about the data, and generalize well outside of the training
distribution. More recently, contrastive learning with multi-modal data has
helped embed different data modalities into the same feature space Li et al.
(2023), such as the studies with visual-language models Radford et al. (2021);
Alayrac et al. (2022); Cherti et al. (2023) and document understanding Xu et
al. (2020); Lee et al. (2023).

Contrastive learning uses pairwise comparison of representations in the
training objective, with the goal of learning representations of data where
positive pairs are drawn closer while negative pairs move apart in the
representation space. It is well known that generating a large dataset of
pairwise samples such as image-text pairs of the same semantics costs much
lower than manual labeling, e.g., the WebImageText dataset used for training
CLIP originates from Wikipedia articles Radford et al. (2021). While there
have been many studies about training objectives for contrastive learning,
optimizing the contrastive loss efficiently remains an open-problem as current
optimization methods critically rely on a large batch size to maintain the
quality of the negative sample distribution.

Formally, we aim to train models $\phi:\mathcal{X}\rightarrow\mathbb{R}^{d}$
and $\psi:\mathcal{Y}\rightarrow\mathbb{R}^{d}$ that are parameterized by
$\theta\in\mathbb{R}^{p}$. These models encode data from the input space
${\cal X},{\cal Y}$ of potentially two modalities into a vector space. Given
${\bf D}_{X}\subseteq{\cal X}$, ${\bf D}_{Y}\subseteq{\cal Y}$, we define a
distribution of positive pairs ${\cal D}_{\rm pos}$ such that ${\rm
supp}({\cal D}_{\rm pos})\subseteq{\bf D}_{X}\times{\bf D}_{Y}$ and a negative
set ${\bf D}_{\rm neg}(x)\subseteq{\bf D}_{Y}$ for every data point $x\in{\bf
D}_{X}$. We denote the dataset constructed in this manner by the tuple $({\cal
D}_{\rm pos},{\bf D}_{\rm neg}(\cdot))$.

This paper aims to minimize the following global contrastive loss Lindgren et
al. (2021); Yuan et al. (2022) with inverse temperature $\beta>0$:

$\displaystyle\min_{\theta\in\mathbb{R}^{p}}$
$\displaystyle\,\mathcal{L}(\theta)$ (1)
$\displaystyle:=\mathop{\mathbb{E}}_{(x,y)\sim{\cal D}_{\rm
pos}}\left[-\log\frac{\exp(\beta\phi(x;\theta)^{\top}\psi(y;\theta))}{{\displaystyle\text{\scalebox{0.8}{$\sum_{z\in{\bf
D}_{\rm neg}(x)}$}}}\exp(\beta\phi(x;\theta)^{\top}\psi(z;\theta))}\right]$

also see Mikolov et al. (2013) that proposed a similar loss to (1). Assume
$m:=|{\rm supp}({\cal D}_{X})|<\infty$ where ${\cal D}_{X}$ is the marginal
distribution of $x$ on ${\cal D}_{\rm pos}$. Moreover, we let $m_{\rm
neg}:=|{\bf D}_{\rm neg}(x)|<\infty$ for all $x\in{\bf D}_{X}$111The results
of this paper can be extended to the case where the negative sample size
$|{\bf D}_{\rm neg}(x)|$ is uneven..

The global contrastive loss (1) aims at finding the feature encoders
$\phi^{\star},\psi^{\star}$ that map data into $\mathbb{R}^{d}$ with maximized
similarity $\phi^{\star}(x)^{\top}\psi^{\star}(y)$ between positive data pair
$(x,y)\sim{\cal D}_{\rm pos}$ and minimized similarity
$\phi^{\star}(x)^{\top}\psi^{\star}(z)$ between negative data pair $(x,z)$,
$z\in{\bf D}_{\rm neg}(x)$. In other words, the loss maximizes the probability
of observing a positive sample $y$ under the context $x$ among the other
negative samples.

We notice that a highly related loss function design to (1) is the InfoNCE
loss Logeswaran & Lee (2018); Oord et al. (2018); Chen et al. (2020), which is
defined by:

$\displaystyle{\cal L}_{\rm NCE}(\theta;B)$ (2)
$\displaystyle=\mathop{\mathbb{E}}_{(x,y)\sim{\cal D}_{\rm
pos}}\mathop{\mathbb{E}}_{{\bf Z}\sim{\cal D}_{\rm
neg}(x;B)}\left[-\log\frac{\exp(\beta\phi(x;\theta)^{\top}\psi(y;\theta))}{{\displaystyle\text{\scalebox{0.8}{$\sum_{z\in{\bf
Z}}$}}\exp(\beta\phi(x;\theta)^{\top}\psi(z;\theta))}}\right]$

such that ${\cal D}_{\rm neg}(x;B)$ is a distribution of all subsets ${\bf
Z}\subseteq{\bf D}_{\rm neg}(x)$ with $|{\bf Z}|=B$. In fact, the global
contrastive loss (1) is an upper bound of the InfoNCE loss (2) since

$\displaystyle\mathop{\mathbb{E}}_{{\bf Z}\sim{\cal D}_{\rm
neg}(x;B)}\left[\log\sum_{z\in{\bf
Z}}\exp(\beta~{}\phi(x;\theta)^{\top}\psi(z;\theta))\right]$ (3)
$\displaystyle\leq\log\sum_{z\in{\bf D}_{\rm
neg}(x)}\exp(\beta~{}\phi(x;\theta)^{\top}\psi(z;\theta)),$

for any $x\in{\rm supp}({\cal D}_{\rm pos}),\theta\in\mathbb{R}^{p}$. As the
equality holds when $B=m_{\rm neg}$, we have ${\cal L}_{\rm NCE}(\theta;m_{\rm
neg})={\cal L}(\theta)$.

It was demonstrated that minimizing (2) with $B\gg 1$ produces encodings with
interpretable geometric and semantic properties Wang & Isola (2020); Robinson
et al. (2020); Zimmermann et al. (2021). At the same time, the convergence of
pre-training foundation models on large scale dataset happens only when using
a large enough batch size $B$ Radford et al. (2021). Such evidence points
towards the success of models that are trained to minimize ${\cal L}_{\rm
NCE}(\theta;B)$ of large negative batch size $B$. We believe that minimizing
${\cal L}(\theta)$, which yields the limiting upper bound to (2), can lead to
a better performance for contrastive learning.

Algorithms | Convergence Error | Memory | Computation
---|---|---|---
| SimCLR
---
Chen et al. (2020)
${\cal O}\left(\frac{\beta^{2}}{\sqrt{T}}+\frac{m_{\rm neg}^{2}\exp(4\beta)\sigma^{2}}{B}\right)$⋆ | ${\cal O}(BM_{\phi}+Bd)$ | ${\cal O}(BC_{\phi}+B^{2}d)$
| Negative Cache
---
Lindgren et al. (2021)
${\cal O}\left(\frac{\beta}{\sqrt{T}}+\frac{\beta^{6}}{\rho^{2}T}\right)$§ | ${\cal O}(M_{\phi}+m_{\rm neg}d+d)$ | ${\cal O}(C_{\phi}+m_{\rm neg}d+\rho mC_{\phi})$
| SogCLR
---
Yuan et al. (2022)
${\cal O}\left(\frac{\exp(6\beta)}{\sqrt{BT}}+\frac{m_{\rm neg}^{2}\exp(4\beta)\sqrt{m}\sigma^{2}}{B\sqrt{T}}+{m_{\rm neg}^{2}\exp(4\beta)\upsilon^{2}}\right)$† | ${\cal O}(BM_{\phi}+m+Bd)$ | ${\cal O}(BC_{\phi}+B^{2}d)$
EMC2 (Ours, $R=2$) | ${\cal O}\left({\frac{\beta}{\sqrt{T}}+\frac{m^{2}m_{\rm neg}^{2}\exp(6c^{2}\beta)\beta^{3}\sigma^{2}}{B\sqrt{T}}}\right)$ | ${\cal O}(BM_{\phi}+m+Bd)$ | ${\cal O}(BC_{\phi}+B^{2}d)$
EMC2 (Ours) | ${\cal O}\left({\frac{\beta}{\sqrt{T}}+\frac{2^{R}m^{2}m_{\rm neg}^{2}\exp(6c^{2}\beta)\beta^{3}\sigma^{2}}{BR^{2}\sqrt{T}}}\right)$ | ${\cal O}(BM_{\phi}+m+Bd)$ | ${\cal O}(BC_{\phi}+B^{2}d+BR)$

Table 1:  The second column shows the upper bound on convergence error
$T^{-1}\sum_{t=0}^{T-1}\mathbb{E}\left[\|\nabla_{\theta}{\cal
L}(\theta_{t})\|_{2}^{2}\right]$. The last two columns show the
memory/computation requirement per iteration of the algorithms. $C_{\phi}$
(resp. $M_{\phi}$) denotes the computational (resp. memory) cost to compute
the feature vector $\phi(x;\theta)$ or $\psi(x;\theta)$. ⋆From (Theorem 1,
Yuan et al. 2022). §The analysis shown in Lindgren et al. (2021) only consider
the case when batch size $B=1$. †SogCLR only converge to an ${\cal
O}(\upsilon^{2})$-stationary solution, where $\upsilon$ is the worst-case
feature heterogeneity error satisfying
$\sup_{\theta\in\mathbb{R}^{p}}\mathop{\mathbb{E}}_{(x,y)\sim{\cal D}_{\rm
pos}}\mathop{\mathbb{E}}_{z\sim{\rm Uniform}({\bf D}_{\rm
neg}(x))}\left[|\phi(x;\theta)^{\top}\psi(z;\theta)-\phi(y;\theta)^{\top}\psi(z;\theta)|\right]\leq\upsilon^{2}$.

Challenges in Optimizing (1). We notice that the contrastive loss gradient is
given by:

$\displaystyle\nabla{\cal L}(\theta)=\mathop{\mathbb{E}}_{(x,y)\sim{\cal
D}_{\rm
pos}}\Big{[}-\beta~{}\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(y;\theta))\Big{]}$
(4) $\displaystyle+\mathop{\mathbb{E}}_{(x,y)\sim{\cal D}_{\rm
pos}}\Big{[}\beta\sum_{z\in{\bf D}_{\rm
neg}(x)}p_{x,\theta}(z)\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(z;\theta))\Big{]}$
$\displaystyle\equiv\nabla{\cal L}_{\rm pos}(\theta)+\nabla{\cal L}_{\rm
neg}(\theta),$

with the softmax distribution:

$p_{x,\theta}(z)=\frac{\exp(\beta~{}\phi(x;\theta)^{\top}\psi(z;\theta))}{\sum_{z^{\prime}\in{\bf
D}_{\rm neg}(x)}\exp(\beta~{}\phi(x;\theta)^{\top}\psi(z^{\prime};\theta))}.$
(5)

The challenge of optimizing ${\cal L}(\theta)$ lies in the overwhelming
complexity to compute or approximate $\nabla{\cal L}_{\rm neg}(\theta)$ since
${\bf D}_{\rm neg}(x)$ often spans a large dataset, e.g., $m_{\rm
neg}=8.8\times 10^{6}$ for the MS MARCO dataset Bajaj et al. (2016).

To this end, SimCLR Chen et al. (2020) proposed to replace ${\bf D}_{\rm
neg}(x)$ in (4) by a randomly selected _negative batch_ ${\bf Z}$ of $B$
randomly augmented images, i.e., using the gradient of InfoNCE loss (2). While
achieving reasonable performance in certain scenarios, SimCLR requires a large
batch size $B$ on large-scale dataset training. For example, it requires up to
$B=32768$ negative samples per iteration in training CLIP Radford et al.
(2021), making it impossible to train such models in non-commercial data
center Cherti et al. (2023). We note that there exists a number of contrastive
learning tricks to avoid the problem of estimating $\nabla{\cal L}_{\rm
neg}(\theta)$. For example, Robinson et al. (2020) imposes a hardness-aware
distribution on ${\cal D}_{\rm neg}(x;B)$ of the InfoNCE loss. He et al.
(2020) utilizes a momentum mechanism to prevent feature collapse. Grill et al.
(2020); Zbontar et al. (2021); Bardes et al. (2021) use alternative loss
functions for contrastive learning, such as redundancy reduction loss and
covariance regularization. See Balestriero et al. (2023) for a comprehensive
overview on different contrastive learning methods.

An alternative approach is to consider the estimation of $\nabla{\cal L}_{\rm
neg}(\theta)$ through _negative sampling_. This approach is motivated from the
observation that (5) is a probability mass function of a softmax distribution
and thus the summation $\sum_{z\in{\bf D}_{\rm
neg}(x)}p_{x,\theta}(z)\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(z;\theta))$
in $\nabla{\cal L}_{\rm neg}(\theta)$ is equivalent to an expectation taken
w.r.t. the softmax distribution. Subsequently, it can be approximated through
_sampling_. In particular, if one draws $(x,y)\sim{\cal D}_{\rm pos}$,
$z\sim{\cal P}_{x,\theta}$ where

${\cal P}_{x,\theta}\equiv(p_{x,\theta}(z))_{z\in{\bf D}_{\rm neg}(x)},$ (6)

then $\beta\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(z;\theta))$ is an
unbiased estimate of $\nabla{\cal L}_{\rm neg}(\theta)$.

However, sampling from ${\cal P}_{x,\theta}$ remains a highly non-trivial task
due to the complexity in computing the partition function (i.e., denominator)
in the softmax distribution (5). Prior works have proposed remedies to the
sampling problem. The _negative cache_ algorithm Lindgren et al. (2021) stores
features $\\{\psi(z;\theta_{\tau})\\}_{z\in{\bf D}_{\rm neg}(x)}$ from
previous iterations and apply the Gumbel-max trick for sampling from an
approximate ${\cal P}_{x,\theta}$. Though it is proven to converge to a
stationary point of (1), the algorithm suffers from high computation and
memory complexity. Under a similar purpose, SogCLR Yuan et al. (2022) proposed
a running mean estimator for the normalization sum in ${\cal P}_{x,\theta}$.
While SogCLR corrects the gradient bias without a significant sampling nor
computation overhead, the algorithm is only guaranteed to converge to a
neighborhood of stationary points of (1) with a non-vanishing error. A
comparison of the above works is in Table 1.

Figure 1: Training 100 epochs on STL-10 with ResNet-18 using batch size
$b=32$. Horizontal axis is relative to the wall-clock training time in
seconds.

Our Contributions. We propose EMC2, an Efficient Markov Chain Monte Carlo
negative sampling method for Contrastive Learning. Our method departs from the
existing approaches since EMC2 directly tracks ${\cal P}_{x,\theta}$ to
generate negative samples for $\nabla{\cal L}_{\rm neg}(\theta)$. More
specifically,

  * •

The EMC2 utilizes a Metropolis-Hasting (M-H) algorithm specialized for
negative sampling from (5). Moreover, the samples from our M-H algorithm
_dynamically adjust_ to the evolution of the stochastic gradient (SGD)
iterates. This results in a state-dependent SGD scheme that enjoys low memory
and computation complexity per iteration.

  * •

We characterize the finite-time convergence rate of EMC2. We show that it
converges in expectation to a stationary point satisfying
$\mathbb{E}[\|\nabla{\cal L}(\theta_{t})\|^{2}]={\cal O}(1/\sqrt{T})$ for some
$t\in\\{1,...,T\\}$, where $T$ is the number of iterations. Moreover, the
latter is neither affected by small batch size, nor the burn-in period with
the M-H algorithm; see Table 1.

  * •

Our convergence analysis involves a non-trivial adaptation of the generic
result for biased stochastic approximation scheme in Karimi et al. (2019).
Specifically, we prove that the $\theta$ dependent kernel which induces the
MCMC’s Markov chain is ergodic, and is Lipschitz w.r.t. $\theta$.

Figure 1 previews the performance of EMC2 for pre-training the image encoder
on the STL-10 dataset. Observe that when training with a small batch size,
EMC2 is around 2x faster than SimCLR and SogCLR, and 3x faster than Negative
Caching. The rest of this paper is organized as follows. Section 2 develops
the EMC2 algorithm by showing how to combine M-H sampling in an online manner
with SGD updates. Section 3 presents the main convergence analysis results.
Finally, Section 4 shows the numerical experiments to corroborate our claims
on the efficacy of EMC2.

Related Works. We remark that existing works have considered using Markovian
samples that is not i.i.d. in SGD optimization. General convergence analysis
results have appeared in Sun et al. (2018); Doan (2022) for convex and non-
convex optimization with a homogeneous Markov chain, and Atchadé et al.
(2017); Karimi et al. (2019) with controlled Markov chain. Example
applications include policy evaluation in reinforcement learning Baxter &
Bartlett (2001); Bhandari et al. (2018); Srikant & Ying (2019), Bayesian
optimization for maximum likelihood De Bortoli et al. (2021), expectation
maximization with stochastic samples Kuhn & Lavielle (2004).

For tasks related to contrastive learning, prior works have proposed using
tree indexed structure to accelerate sampling Monath et al. (2023), and using
the graph structure for graph representation learning Yang et al. (2020).
However, they only proved that the bias of the gradient estimator is bounded
and lack convergence analysis for the overall learning algorithm. We remark
that there are earlier works on estimating the cross-entropy loss over a large
set of classes Vembu et al. (2012), e.g., by sub-sampling the classes based on
kernel methods Blanc & Rendle (2018) and random Fourier features Rawat et al.
(2019). In comparison, our work introduces the technique to contrastive
learning, and provide a comprehensive theoretical and empirical analysis on
its efficacy and convergence properties.

##  2 Our Proposed Method: EMC2

This section develops EMC2 for optimizing the global contrastive loss ${\cal
L}(\theta)$ in (1). To simplify notation, we define

$H(x,y;\theta):=\nabla_{\theta}\big{[}\phi(x;\theta)^{\top}\psi(y;\theta)\big{]}$
(7)

as the gradient of the sample pair $(x,y)\in{\cal X}\times{\cal Y}$. Recall
from (4) that the population gradient $\nabla{\cal L}(\theta)$ is composed of
two terms: (i) $\nabla{\cal L}_{\rm pos}(\theta)$ maximizes the correlation
between positive sample pairs, (ii) $\nabla{\cal L}_{\rm neg}(\theta)$
minimizes the correlation between _negative_ sample pairs.

To apply SGD on (1), with a positive sample pair drawn uniformly as
$(x,y)\sim{\cal D}_{\rm pos}$, the vector $-\beta H(x,y;\theta)$ yields an
_unbiased_ estimate for $\nabla{\cal L}_{\rm pos}(\theta)$. Our challenge lies
in obtaining an _unbiased_ estimate for the negative sample gradient
$\nabla{\cal L}_{\rm neg}(\theta)$. Observe that

$\nabla{\cal L}_{\rm
neg}(\theta)=\mathop{\mathbb{E}}_{\begin{subarray}{c}(x,y)\sim{\cal D}_{\rm
pos},~{}z\sim{\cal P}_{x,\theta}\end{subarray}}\Big{[}\beta
H(x,z;\theta)\Big{]}$ (8)

for ${\cal P}_{x,\theta}$ defined in (6). Compared to the case of $\nabla{\cal
L}_{\rm pos}(\theta)$, the distribution for the tuple $(x,z)$ in the
expectation above depends on $\theta$. It follows a softmax distribution (5)
with a large summation in the denominator that is difficult, if not
impossible, to evaluate when the number of negative samples is large.
Furthermore, the distribution also depends on $\theta$ and has to be updated
_dynamically_ as we optimize $\theta$.

We refer to the task of estimating (8) as the _negative sampling_ problem. To
improve memory consumption and computation cost, our idea is to develop an
MCMC negative sampling method that generates the desired samples in an online
manner. The resultant SGD method for (1) is then treated as a stochastic
approximation scheme with state-dependent samples coming from a _controlled
Markov chain_.

###  2.1 Negative Sampling via MCMC

The MCMC scheme (Robert & Casella, 1999) is a classical yet powerful method
for generating samples from an arbitrary distribution $\pi$. We focus on the
Metropolis-Hastings (M-H) algorithm Chib & Greenberg (1995) due to its
simplicity. In a nutshell, the algorithm generates new candidate samples from
a uniform distribution and adjusts the frequency of samples by a reject/accept
mechanism to match the target distribution $\pi$. It induces a Markov chain
whose stationary distribution is the target distribution.

For the negative sampling problem in (8), with $(x,y)\sim{\cal D}_{\rm pos}$,
we generate samples $z\sim{\cal P}_{x,\theta}$ by running the M-H algorithm.
In particular, we assign a state $Z_{i}\in[m_{\rm neg}]$ for each $i\in[m]$
that corresponds to a positive sample $x_{i}\in{\bf D}_{X}$. We generate a
candidate sample $Z_{i}^{\prime}\sim{\rm Uniform}([m_{\rm neg}])$, to be
accepted as the new sample $Z_{i}^{+}$ with probability:

$Q_{x_{i},\theta}(Z_{i}^{\prime},Z_{i})=\frac{p_{x_{i},\theta}(Z_{i}^{\prime})}{p_{x_{i},\theta}(Z_{i})}=\frac{\exp(\beta~{}\phi(x_{i};\theta)^{\top}\psi(Z_{i}^{\prime};\theta))}{\exp(\beta~{}\phi(x_{i};\theta)^{\top}\psi(Z_{i};\theta))},$
(9)

where $Z_{i}$ is the old sample; if the sample $Z_{i}^{\prime}$ is rejected,
then $Z_{i}^{+}=Z_{i}$. The above procedure induces a Markov chain $\cdots\to
Z_{i}\to Z_{i}^{+}\to\cdots$ whose stationary distribution coincides with
${\cal P}_{x_{i},\theta}$. Together with $(x_{i},y_{i})\equiv(x,y)\sim{\cal
D}_{\rm pos}$, this allows us to construct an unbiased estimate for
$\nabla{\cal L}_{\rm neg}(\theta)$ as $\beta
H(x_{i},{Z}_{i}^{\infty};\theta)$, where ${Z}_{i}^{\infty}$ is obtained after
running the M-H algorithm for a certain number of steps.

We note that evaluating (9) does not require us to compute the partition
function in the denominator of (5). As such, the computation complexity of
evaluating (9) is only $2C_{\phi}$. On the other hand, we require storing the
state $Z_{i}$ for each positive sample $x_{i}\in{\bf D}_{X}$, which results in
a memory cost of $m$ integers. The latter needs not be updated for every
iteration. Compared to the negative caching algorithm Lindgren et al. (2021)
which requires a cache memory consumption of $m_{\rm neg}d$ real numbers for
storing the feature vectors $\psi(z;\theta)$ for all $z\in{\bf D}_{\rm
neg}(x)$, and a computation cost of $\rho m_{\rm neg}C_{\phi}$ for updating a
$\rho$-fraction of cache, the MCMC scheme enjoys a lower memory cost and
computation complexity.

Single State MCMC vs Multi-State MCMC.  Alternatively, one can apply the M-H
algorithm to generate samples from the _joint distribution_ of $(x,z)$ in (8)
using a single state MCMC. However, the corresponding reject/accept ratio [cf.
(9)] involves the normalization constant $G(x,\theta)=\sum_{z\in{\bf D}_{\rm
neg}(x)}\phi(x;\theta)^{\top}\psi(z;\theta)$ from (5), for which $G(x,\theta)$
and $G(x^{\prime},\theta)$ have to be re-evaluated when transitioning from
$(x,\cdot)$ to $(x^{\prime},\cdot)$. Implementing such a scheme will require a
$2m_{\rm neg}C_{\phi}$ computation complexity, or require approximating the
normalization constant as a constant that is independent of $x$, both of which
are encouraging the use of a multi-state Markov chain as illustrated in (9).

###  2.2 State-dependent SGD Algorithm

As the target distribution in (8) depends on $\theta$, applying SGD with the
classical MCMC scheme or M-H algorithm requires first _freezing_ $\theta$ and
then simulating the Markov chain for a considerable amount of time _prior to_
forming the stochastic gradient. The latter is known as the _burn-in_ period
for MCMC and incurs additional complexity for every SGD iteration.

We propose to adjust the Markov chains in an online manner as $\theta$ is
updated, i.e., the method shall not maintain a long burn-in period. In
particular, at each SGD iteration, the M-H updates are executed with
initialization given by the state stored previously. Meanwhile, the Markov
chains with the reject/accept ratio (9) are controlled by the current
$\theta$.

Algorithm 1 Efficient MCMC Negative Sampling Method for Contrastive Learning
(EMC2)

1: input: Iteration number $T$, batch size $B$, negative batch size $R$, burn-
in period $P<R$, Markov chain state initialization $\\{Z_{j}\\}_{j=1}^{m}$ and
step size $\gamma$.

2: for $t=0,...,T-1$ do

3: Draw a mini-batch
$[(x_{i_{1}^{(t)}},y_{i_{1}^{(t)}}),...,(x_{i_{B}^{(t)}},y_{i_{B}^{(t)}})]\sim{\cal
D}_{\rm pos}$, with indices $\\{i_{1}^{(t)},...,i_{B}^{(t)}\\}\subseteq[m]$.

4: for $k=1,...,B$; $r=0,...,R-1$ do

5: Draw negative sample $Z^{\prime}_{i_{k}^{(t)}}\sim{\rm Unif}({\bf D}_{\rm
neg}(x_{i_{k}^{(t)}}))$.

6: Update the Markov chain state $Z_{i_{k}^{(t)}}$ by
$Z_{i_{k}^{(t)}}\leftarrow\begin{cases}Z^{\prime}_{i_{k}^{(t)}}&\text{w.p.}~{}\min\left\\{1,~{}Q_{x_{i_{k}^{(t)}},\theta_{t}}(Z^{\prime}_{i_{k}^{(t)}},Z_{i_{k}^{(t)}})\right\\},\\\
Z_{i_{k}^{(t)}}&\text{otherwise},\end{cases}$ where
$Q_{x,\theta}(Z^{\prime},Z)$ is defined in (9).

7: If $r\geq P$, store the sample
$\tilde{Z}_{i_{k}^{(t)}}^{(r)}=Z_{i_{k}^{(t)}}$.

8: end for

9: Update the model $\theta$ using (10).

10: end for

Let $t$ denotes the SGD iteration index and $\theta_{t}$ is the corresponding
model. Suppose that a mini-batch of $B$ positive samples are drawn
$\\{x_{i_{k}^{(t)}},y_{i_{k}^{(t)}}\\}_{k=1}^{B}$ and
$\\{\\{\tilde{Z}_{i_{k}^{(t)}}^{(r)}\\}_{r=0}^{R-1}\\}_{k=1}^{B}$ are the
sequences of samples generated by $R$ steps of the M-H algorithm, where the
latter is initialized by the states $Z_{i_{k}^{(t)}}$ from the previous SGD
iteration. Our idea is to update the model $\theta$ by the following
recursion:

$\theta_{t+1}=\theta_{t}-\gamma{\cal
H}\left(\bm{\xi}_{t+1};\theta_{t}\right),$ (10)

where
$\bm{\xi}_{t+1}:=\\{x_{i_{k}^{(t)}},y_{i_{k}^{(t)}},\\{\tilde{Z}_{i_{k}^{(t)}}^{(r)}\\}_{r=P}^{R-1}\\}_{k=1}^{B}$
collects the samples used, $\gamma>0$ is the SGD step size, and

$\displaystyle{\cal
H}\left(\bm{\xi}_{t+1};\theta_{t}\right):=-\frac{\beta}{B}\sum_{k=1}^{B}H(x_{i_{k}^{(t)}},y_{i_{k}^{(t)}};\theta_{t})$
(11)
$\displaystyle\qquad+\frac{\beta}{B(R-P)}\sum_{k=1}^{B}\sum_{r=P}^{R-1}H(x_{i_{k}^{(t)}},\tilde{Z}_{i_{k}^{(t)}}^{(r)};\theta_{t}),$

where $P$ is an adjustable burn-in parameter. The overall algorithm is
summarized in Algorithm 1.

We observe that (10) is different from a standard SGD algorithm as ${\cal
H}(\bm{\xi}_{t+1};\theta_{t})$ corresponds to a _biased_ estimate for the
gradient $\nabla{\cal L}(\theta_{t})$ of the contrastive learning loss (1) in
general. Instead, the algorithm belongs to the more general class of _biased
stochastic approximation_ scheme Karimi et al. (2019) whose stochastic updates
are driven by the Markov chain
$\cdots\to\bm{\xi}_{t}\to\bm{\xi}_{t+1}\to\cdots$. Moreover, the Markov
transition kernel is controlled by the current model $\theta_{t}$ that is
updated simultaneously.

To reduce bias, a traditional approach is to consider $P\to\infty$, $R=P+1$
which will ensure that the M-H algorithm can generate its samples from the
stationary distribution. However, as we shall demonstrate in Section 3, under
suitable and verifiable conditions, the $\theta$-dependent Markov transition
kernel is Lipschitz continuous and uniform geometrically ergodic.
Subsequently, it ensures the convergence of (10) towards a stationary point of
(1) for any $P,R$.

##  3 Convergence Analysis

This section shows that EMC2 converges to a stationary point of (1). Our
analysis is organized as follows: first we characterize the mixing rate of the
Markov chain $\cdots\to\bm{\xi}_{t}\to\bm{\xi}_{t+1}\to\cdots$ in (10) and its
smoothness property, then we analyze the convergence of the algorithm through
studying the latter as a special case of the biased stochastic approximation
scheme in Karimi et al. (2019). Finally, we present our main results in
Theorem 3.7 and Corollary 3.8.

###  3.1 Analysis of MCMC Negative Sampling

To enable the convergence of (10) utilizing the Markov chain
$\\{\bm{\xi}_{t}\\}_{t\geq 0}$, intuitively it requires the Markov chain to
(i) have a fast mixing time, and (ii) satisfy certain smoothness property as
$\theta_{t}$ is gradually updated. To fix notation, we denote the state-
dependent Markov transition kernel ${\tt
P}_{\theta}:\Xi\times\Xi\to\mathbb{R}_{+}$ such that $\bm{\xi}_{t+1}\sim{\tt
P}_{\theta_{t}}(\xi_{t},\cdot)$. For simplicity, we analyze the latter
transition probability of $\bm{\xi}_{t+1}$ when $P=R-1$, i.e., the state
dependent SGD update only takes the last sample in the M-H algorithm.
Moreover, with a slight abuse of notation, we consider the collection of all
hidden states $\bm{\xi}_{t+1}=\\{\tilde{Z}_{i}^{(R-1)}\\}_{i=1}^{m}\in\Xi$
which does not include $x,y$. As the latter is drawn i.i.d., it does not
affect the convergence analysis for (10) in Section 3.2.

For our first task in showing the fast mixing of $\\{\xi_{t}\\}_{t\geq 0}$, it
suffices to show that the Markov chain induced by ${\tt P}_{\theta}$ is
_geometrically ergodic_ for any fixed $\theta$. To this end, we note that as
the target distribution ${\cal P}_{x,\theta}$ belongs to the exponential
family, applying (Theorem 2.1, Roberts & Tweedie 1996) implies ergodicity. To
formally state the result, we impose the following assumptions:

######  Assumption 3.1.

There exists $c>0$ such that

$\max\\{\|\phi(x;\theta)\|,\|\psi(x;\theta)\|\\}\leq
c,~{}\forall~{}\theta\in\mathbb{R}^{p},~{}x\in{\cal X}.$ (12)

Assumption 3.1 is common with $c\leq 1$ for contrastive learning problems Chen
et al. (2020); Radford et al. (2021). This effectively controls the behavior
of $Q_{x,\theta}(z,Z)$ in the M-H algorithm (9) and thus the transition kernel
${\tt P}_{\theta}$. The following lemma shows the geometric ergodicity
property for the Markov chain:

######  Lemma 3.2.

Under Assumption 3.1. For any $\theta\in\mathbb{R}^{p}$ and any initialization
$\tilde{\bm{\xi}}_{0}\in\Xi$, the Markov chain
$\tilde{\bm{\xi}}_{0}\to\tilde{\bm{\xi}}_{1}\to\cdots$ induced by the
transition kernel ${\tt P}_{\theta}$ converges geometrically to the stationary
distribution
$\pi_{{\bm{x}},\theta}({\bm{z}})=\prod_{i=1}^{m}p_{x_{i},\theta}(z_{i})$,
where ${\bm{z}}=(z_{i})_{i=1}^{m}$. In particular, it holds

$\begin{split}\big{|}\mathbb{P}\big{(}\tilde{\bm{\xi}}_{\tau}={\bm{z}}\big{)}-\pi_{{\bm{x}},\theta}({\bm{z}})\big{|}\leq\left(1-\frac{BR}{2mm_{\rm
neg}\exp(2c^{2}\beta)}\right)^{\tau},\\\ \end{split}$ (13)

for any ${\bm{z}}$ and any $\tau\geq 0$.

See Appendix A for the proof, which relies on (Theorem 1.3, Mengersen &
Tweedie 1996). Moreover, from (13), we observe that through mini-batch
sampling among ${\cal D}_{\rm pos}$ and ${\cal P}_{x,\theta}$, the mixing rate
improves by $BR$ when consuming $BR$ samples in EMC2.

Our second task consists of showing that the Markov transition kernel is
smooth w.r.t. $\theta$. We require the following condition the models
$\phi,\psi$:

######  Assumption 3.3.

There exists $L_{P}\geq 0$ such that for any
$\theta,\theta^{\prime}\in\mathbb{R}^{p},(x,y)\in{\rm supp}({\cal D}_{\rm
pos}),z\in{\bf D}_{\rm neg}(x)$,

$\|\phi(x;\theta)^{\top}\psi(z;\theta)-\phi(x;\theta^{\prime})^{\top}\psi(z;\theta^{\prime})\|\leq
L_{P}\|\theta-\theta^{\prime}\|.$ (14)

Assumption 3.3 describes the Lipschitz condition of the similarity function
$\phi(x;\theta)^{\top}\psi(z;\theta)$, which is important in establishing the
Lipschitz condition of the transition kernel w.r.t. the model parameter
$\theta$. For example, it can be satisfied as $L_{P}=Lc$ for $L$-Lipschitz
functions $\phi,\psi$ w.r.t. $\theta$ with bounded norm as in Assumption 3.1.

The following lemma shows that the transition kernel ${\tt P}_{\theta}$ is
Lipschitz continuous with respect to $\theta$, as follows:

######  Lemma 3.4.

Under Assumption 3.1, 3.3. It holds:

$|{\tt P}_{\theta}(\bm{\xi},\bm{\xi}^{\prime})-{\tt
P}_{\theta^{\prime}}(\bm{\xi},\bm{\xi}^{\prime})|\leq
2^{R+1}BL_{P}\exp(2c^{2}\beta)\beta\|\theta-\theta^{\prime}\|,$ (15)

for any $\bm{\xi},\bm{\xi}^{\prime}\in\Xi$ and any
$\theta,\theta^{\prime}\in\mathbb{R}^{p}$.

The proof can be found in Appendix B. We remark that the dependence on
$2^{R+1}B$ in the Lipschitz constant is conservative. We anticipate the
effective Lipschitz constant for the transition kernel to be much smaller than
in (15). For example, (43) has applied a loose bound of $1$ to control the
transition probability.

###  3.2 Convergence to Stationary Point

Equipped with Lemma 3.2 & 3.4, we are ready to derive the convergence rate of
Algorithm 1 towards a stationary point of (1). Our idea is to treat (10) as a
biased stochastic approximation scheme and analyze the latter using the
general convergence theories given in Karimi et al. (2019).

We require the following conditions:

######  Assumption 3.5.

There exists $L_{H}\geq 0$ such that for any
$\theta,\theta^{\prime}\in\mathbb{R}^{p},x\in{\bf D}_{X},y\in{\bf D}_{Y}$,

$\displaystyle\|H(x,y;\theta)-H(x,y;\theta^{\prime})\|$ $\displaystyle\leq
L_{H}\|\theta-\theta^{\prime}\|.$ (16)

######  Assumption 3.6.

There exists $\sigma\geq 0$ such that for any
$\theta\in\mathbb{R}^{p},(x,y)\in{\rm supp}({\cal D}_{\rm pos}),z\in{\bf
D}_{\rm neg}(x)$,

$\displaystyle\|H(x,y;\theta)-\mathbb{E}_{{\cal D}_{\rm
pos}}[H(x,y;\theta)]\|$ $\displaystyle\leq\sigma,$ (17)
$\displaystyle\|H(x,z;\theta)-\mathbb{E}_{{\cal D}_{\rm pos},{\cal
P}_{x,\theta}}[H(x,z;\theta)]\|$ $\displaystyle\leq\sigma.$ (18)

Assumptions 3.5, 3.6 impose a uniform bound on the smoothness condition of the
similarity function and the variance of stochastic gradient. These assumptions
are standard in the literature, e.g., they are also used in Lindgren et al.
(2021); Yuan et al. (2022).

We present the main convergence result for EMC2:

######  Theorem 3.7.

Under Assumptions 3.1, 3.3, 3.5, 3.6. For any $T\geq 1$, there exists a
sufficiently small step size $\gamma>0$ such that the iterate $\theta_{t}$
generated by EMC2 satisfies $\displaystyle\textstyle
T^{-1}\sum_{t=0}^{T-1}\mathbb{E}[\|\nabla{\cal L}(\theta_{t})\|^{2}]$ (19)
$\displaystyle\leq\frac{2{\cal L}_{0,T}}{\gamma T}+\frac{24\beta\sigma mm_{\rm
neg}\exp(2c^{2}\beta)}{\gamma BRT}+\gamma\beta^{3}\sigma^{2}$
$\displaystyle\times\mathcal{O}\Big{(}\frac{L_{H}mm_{\rm
neg}\exp(2c^{2}\beta)}{BR}+\frac{L_{P}2^{R}m^{2}m_{\rm
neg}^{2}\exp(6c^{2}\beta)}{BR^{2}}\Big{)}$ where the above expectation is
taken with respect to the randomness in the algorithm, and ${\cal
L}_{0,T}=\mathbb{E}[{\cal L}(\theta_{0})-{\cal L}(\theta_{T})]\leq 4\beta$.

See Appendix C for the proof. Theorem 3.7 implies that in expectation, EMC2
finds an ${\cal O}(\frac{1}{\gamma T}+\gamma)$ stationary solution to (1) in
$T$ iterations, for any $T\geq 1$.

For a sufficiently large $T$, setting $\gamma=1/\sqrt{T}$ yields a global
convergence rate of ${\cal O}(1/\sqrt{T})$ regardless of the batch size $B$
nor the burn-in period controlled by $P,R$, as observed in the following
corollary:

######  Corollary 3.8.

Under Assumptions 3.1, 3.3, 3.5, 3.6. For a sufficiently large $T$, choosing
$\gamma=\frac{1}{\sqrt{T}}$, $R=2$, guarantees that the iterates generated by
EMC2 satisfy:

$\displaystyle\textstyle T^{-1}\sum_{t=0}^{T-1}\mathbb{E}[\|\nabla{\cal
L}(\theta_{t})\|^{2}]$ (20)
$\displaystyle\leq\frac{8\beta}{\sqrt{T}}+\frac{12\beta\sigma mm_{\rm
neg}\exp(2c^{2}\beta)}{B^{2}\sqrt{T}}+\frac{\beta^{3}\sigma^{2}}{\sqrt{T}}$
$\displaystyle\times\mathcal{O}\Big{(}\frac{L_{H}mm_{\rm
neg}\exp(2c^{2}\beta)}{B}+\frac{L_{P}m^{2}m_{\rm
neg}^{2}\exp(6c^{2}\beta)}{B}\Big{)}.$

We have presented a simplified version of the above result in Table 1.

##  4 Numerical Experiments: Unimodal Pre-Training of Image Encoder

This section examines the performance of EMC2 for the contrastive learning
task on training image encoders. We remark that although the analysis is only
shown for the SGD optimizer with EMC2, the stochastic gradient approximation
in (11) can be applied to other popular first-order optimizers such as Adam.

Our experiments consider the task of pre-training self-supervised image
encoder similar to Chen et al. (2020). This task is described by an instance
of (1) with unimodal encoders, i.e.,
$\phi(\cdot;\theta)\equiv\psi(\cdot;\theta)$ for any
$\theta\in\mathbb{R}^{p}$. Moreover, the self-supervised dataset is specified
with a set of augmented images. Let ${\cal A}$ be a set of image augmentation
operators $g:{\cal X}\to{\cal X}$, we specify $({\cal D}_{\rm pos},{\bf
D}_{\rm neg}(\cdot))$ as:

$\begin{cases}{\cal D}_{\rm pos}={\rm Uniform}(\\{(g(x),h(x)):x\in{\bf
D};~{}g,h\in\mathcal{A}\\}),\\\ {\bf D}_{\rm neg}(h(x))=\\{g(y):y\neq
x,y\in{\bf D};g,h\in\mathcal{A}\\}.\end{cases}$ (21)

For the dataset ${\bf D}\subseteq{\cal X}$, the positive pairs are set of the
random augmentations of the same image, while for $x\in{\bf D}$, its negative
pairs are set to be the whole dataset with augmentation except that of $x$.

Figure 2: Illustration of mini-batch MCMC sampling with 2 augmentations
$(x_{i}^{\prime},x_{i}^{\prime\prime})$ of each image $x_{i}$. Each horizontal
arrow represents a distribution tracked by one Markov chain and the direction
of M-H reject/accept step. Shaded area represents the samples used for burn-in
with burn-in period $P<2b-2$. Crossed-out diagonals are not regarded as
negative samples.

###  4.1 Implementation of EMC2

Image Augmentation. We follow previous studies on applying input augmentations
for contrastive learning on image data Xie et al. (2020) which would lead to
samples from a distribution ${\cal D}_{\rm pos}$ with infinite support.
Despite the fact that EMC2 is designed to track the Markov chain state for
each image in $x_{1},...,x_{m}$, in practice we approximate Algorithm 1 by
maintaining a Markov chain state for each set of $\\{g(x_{1}):g\in{\cal
A}\\},...,\\{g(x_{m}):g\in{\cal A}\\}$222When the state $Z_{i}$ is referred
outside the iteration it was sampled, it points to the image $x_{i}$ to avoid
storing the augmented image $g(x_{i})$ in memory.. This approximate algorithm
remains practical and performant as verified empirically in Section 4.3.

Mini-batch Sampling. As proposed by Chen et al. (2020), it is efficient to
draw a mini-batch and use the in-batch samples as negative samples. We can
implement M-H reject/accept steps on top of the same mini-batch sampling
scheme with $R=B-2$. In Figure 2, we illustrate a similarity matrix formed by
the matrix product ${\bf V}^{\top}{\bf V}$ such that ${\bf
V}=[\phi(x_{1}^{\prime}),...,\phi(x_{b}^{\prime}),\phi(x_{1}^{\prime\prime}),...,\phi(x_{b}^{\prime\prime})]\in\mathbb{R}^{d\times
2b}$ contains the feature vectors of a mini-batch of size $B=2b$ with two
random augmentations $(x_{i}^{\prime},x_{i}^{\prime\prime})$ of the same image
$x_{i}$, $i\in[b]$. EMC2 can regard the in-batch samples
$\\{x_{1}^{\prime},...,x_{b}^{\prime},x_{1}^{\prime\prime},...,x_{b}^{\prime\prime}\\}\backslash\\{x_{j}^{\prime},x_{j}^{\prime\prime}\\}$
as negative samples from the prior distribution ${\rm Uniform}({\bf D}_{\rm
neg}(x_{j}))$ and perform M-H acceptance/rejection steps over this negative
set. As illustrated in Figure 2, the M-H algorithm can run in parallel across
the multiple Markov chains. For the experiments, we consume burn-in samples
from the same mini-batch without requiring extra samples.

###  4.2 Dataset and Metrics

We concentrate on two common datasets under this setup – STL-10 and
Imagenet-100. We apply the Adam optimizer for training with STL-10 and the
LARS optimizer with Imagenet-100. Details of the datasets and hyperparameters
used can be found in Appendix D.

To benchmark the performance of EMC2 and other algorithms, we report linear
probe (LP) accuracy and 1-nearest-neighbor (1-NN) accuracy of the embeddings
produced by the encoder, which are standard metrics for evaluating image
encoders Chen et al. (2020); Kukleva et al. (2023). They correspond to the
linear separability of features on a unit hypersphere Wang & Isola (2020) and
the local geometry of the learnt features.

Since the negative cache algorithm Lindgren et al. (2021) requires storing
feature vectors of negatives, we are restricted to implementing on an
alternative loss function with ${\bf D}_{\rm neg}(x)=\\{y:y\in{\bf
D}\backslash\\{x\\}\\}$, i.e., images without augmentation. Also, note that
negative cache algorithm uses $B$ extra samples per iteration as the Gumbel-
max negative sampling requires out-of-batch samples.

Figure 3: Comparison between different sizes of pre-augmented STL-10 with
ResNet-18 and batch size $b=256$. Horizontal axis is relative to the number of
samples accessed.

###  4.3 Effects of Image Augmentations

We first examine the effect of the number of pre-computed image augmentations,
$|{\cal A}|$, on the performance of EMC2. Note that as suggested in Chen et
al. (2020), increasing $|{\cal A}|$ can lead to improved performance. We also
compare a heuristic extension of EMC2 that effectively deploys an infinite
number of augmentations: the augmented images are generated on-the-fly at
every iteration and a single Markov chain state $Z_{i}$ is maintained for all
augmentations of the same image $x_{i}$.

From Figure 3, as the number of pre-computed augmentations increases, the
performance approaches to that of the infinite augmentation algorithm. This
illustrates that EMC2 adapts to random augmentations. In later experiments
with image augmentations, we implement EMC2 with the infinite augmentation
setting for the best performance.

###  4.4 Effects of Burn-in Period

For every positive pair $(x_{i},y_{i})$ in a mini-batch, EMC2 uses $P$ burn-in
negative samples to warm-up the Markov chain state $Z_{i}$ of sample $x_{i}$.
We study the effect of $P$ by comparing the performance for
$P\in\\{0,\frac{b}{2},b,\frac{3}{2b}\\}$. From Figure 4, we observe that the
performance of EMC2 improves by increasing $P$, and the improvement stops at a
certain threshold of around $P=b$. This amounts to the mechanism of Markov
chain convergence where the negative samples after the burn-in period are more
accurate to the true sample distribution in (5).

Figure 4: Comparison between different numbers of burn-in negative samples $P$
for each Markov chain state $Z_{i}$.

Figure 5: Comparison on STL-10 with ResNet-18 using batch size (top) $b=32$,
(bottom) $b=256$. Horizontal axis is relative to the number of samples
accessed.

###  4.5 Comparison to Baselines

We compare EMC2 to baseline algorithms: Negative Cache Lindgren et al. (2021),
SimCLR Chen et al. (2020), SogCLR Yuan et al. (2022). We first examine the
effects of batch sizes $b\in\\{32,256\\}$ for the experiments on STL-10, where
a standard setting for this dataset is $b=256$, yet we note that a small batch
size is often preferred for training large models due to limitations on GPU
memory. Observe from Figure 5 that EMC2 consistently outperforms the other
algorithms across all metrics regardless the choice of batch size. This is in
line with the intuition that SimCLR is bottle-necked by a poor loss
distribution on small batch size. For SogCLR, we suspect the performance drop
on linear probe accuracy is due to the violation of feature homogeneity
assumption on the setup of STL-10 with Resnet-18. For negative cache, its
performance on using one negative with stale cache error does not keep up with
in-batch negatives approaches.

Figure 6 shows a similar experiment but on the more difficult dataset
Imagenet-100. Note that the negative caching algorithm is not run due to the
excess computation complexity in cache refreshing and Gumbel-max sampling.
Observe that EMC2 shows performance gain over SimCLR when trained on the
standard batch size $b=256$, while the performance is on par with SogCLR.

From the above experiments, we observe that EMC2 delivers a consistent
performance over different datasets, batch size, that is on par with the best
compared baseline for the respective tasks. In Appendix D.2, additional time-
complexity comparison on STL-10 is presented to demonstrate the practical
benefits of EMC2.

Figure 6: Comparison on Imagenet-100 with ResNet-50 using batch size $b=256$.
Horizontal axis is relative to the number of samples accessed.

Convergence to Stationary Solution. To examine the accuracy of EMC2, we
construct a simple dataset (${\cal D}_{\rm pos},{\bf D}_{\rm neg}$) by taking
the first 500 images from STL-10 and using two pre-computed augmentations for
each image. Under this setting, the exact value of ${\cal L}(\theta)$ in (1)
and $\|\nabla{\cal L}(\theta)\|^{2}$ in (4) can be evaluated as shown in
Figure 7. We observe that the solution found by EMC2 is more accurate than the
other baselines by 2 orders of magnitude in terms of squared gradient norm.

Figure 7: Comparison on a subset of STL-10 using the first 500 images and pre-
computed two augmentations for each image. Trained using SGD with batch size
$b=4$.

##  5 Conclusion

This paper proposed a novel method EMC2 for optimizing the global contrastive
loss (1). The algorithm combines an adaptively adjusted MCMC scheme for
generating negative samples with a standard SGD update. We prove that EMC2
enjoys low memory and computation complexities, and admits a global
convergence rate of ${\cal O}(1/\sqrt{T})$ towards a stationary solution for
(1). Numerical experiments show that EMC2 enables small batch training for
contrastive learning which is due to its global convergence property. We hope
this work inspires future research into advanced sampling methods for
contrastive learning, such as the Langevin dynamics, and the convergence
property of optimization algorithms that rely on them.

## References

  * Alayrac et al. (2022) Alayrac, J.-B., Donahue, J., Luc, P., Miech, A., Barr, I., Hasson, Y., Lenc, K., Mensch, A., Millican, K., Reynolds, M., et al.  Flamingo: a visual language model for few-shot learning.  _Advances in Neural Information Processing Systems_ , 35:23716–23736, 2022.
  * Atchadé et al. (2017) Atchadé, Y. F., Fort, G., and Moulines, E.  On perturbed proximal gradient algorithms.  _The Journal of Machine Learning Research_ , 18(1):310–342, 2017.
  * Bajaj et al. (2016) Bajaj, P., Campos, D., Craswell, N., Deng, L., Gao, J., Liu, X., Majumder, R., McNamara, A., Mitra, B., Nguyen, T., et al.  Ms marco: A human generated machine reading comprehension dataset.  _arXiv preprint arXiv:1611.09268_ , 2016.
  * Balestriero et al. (2023) Balestriero, R., Ibrahim, M., Sobal, V., Morcos, A., Shekhar, S., Goldstein, T., Bordes, F., Bardes, A., Mialon, G., Tian, Y., et al.  A cookbook of self-supervised learning.  _arXiv preprint arXiv:2304.12210_ , 2023.
  * Bardes et al. (2021) Bardes, A., Ponce, J., and LeCun, Y.  Vicreg: Variance-invariance-covariance regularization for self-supervised learning.  _arXiv preprint arXiv:2105.04906_ , 2021.
  * Baxter & Bartlett (2001) Baxter, J. and Bartlett, P. L.  Infinite-horizon policy-gradient estimation.  _journal of artificial intelligence research_ , 15:319–350, 2001.
  * Bhandari et al. (2018) Bhandari, J., Russo, D., and Singal, R.  A finite time analysis of temporal difference learning with linear function approximation.  In _Conference on learning theory_ , pp. 1691–1692. PMLR, 2018.
  * Blanc & Rendle (2018) Blanc, G. and Rendle, S.  Adaptive sampled softmax with kernel based sampling.  In _International Conference on Machine Learning_ , pp. 590–599. PMLR, 2018.
  * Chen et al. (2020) Chen, T., Kornblith, S., Norouzi, M., and Hinton, G.  A simple framework for contrastive learning of visual representations.  In _International conference on machine learning_ , pp. 1597–1607. PMLR, 2020.
  * Cherti et al. (2023) Cherti, M., Beaumont, R., Wightman, R., Wortsman, M., Ilharco, G., Gordon, C., Schuhmann, C., Schmidt, L., and Jitsev, J.  Reproducible scaling laws for contrastive language-image learning.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 2818–2829, 2023.
  * Chib & Greenberg (1995) Chib, S. and Greenberg, E.  Understanding the metropolis-hastings algorithm.  _American statistician_ , pp. 327–335, 1995.
  * De Bortoli et al. (2021) De Bortoli, V., Durmus, A., Pereyra, M., and Vidal, A. F.  Efficient stochastic optimisation by unadjusted langevin monte carlo: Application to maximum marginal likelihood and empirical bayesian estimation.  _Statistics and Computing_ , 31:1–18, 2021.
  * Doan (2022) Doan, T. T.  Finite-time analysis of markov gradient descent.  _IEEE Transactions on Automatic Control_ , 68(4):2140–2153, 2022.
  * Grill et al. (2020) Grill, J.-B., Strub, F., Altché, F., Tallec, C., Richemond, P., Buchatskaya, E., Doersch, C., Avila Pires, B., Guo, Z., Gheshlaghi Azar, M., et al.  Bootstrap your own latent-a new approach to self-supervised learning.  _Advances in neural information processing systems_ , 33:21271–21284, 2020.
  * He et al. (2020) He, K., Fan, H., Wu, Y., Xie, S., and Girshick, R.  Momentum contrast for unsupervised visual representation learning.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pp. 9729–9738, 2020.
  * Karimi et al. (2019) Karimi, B., Miasojedow, B., Moulines, E., and Wai, H.-T.  Non-asymptotic analysis of biased stochastic approximation scheme.  In _Conference on Learning Theory_ , pp. 1944–1974. PMLR, 2019.
  * Kuhn & Lavielle (2004) Kuhn, E. and Lavielle, M.  Coupling a stochastic approximation version of em with an mcmc procedure.  _ESAIM: Probability and Statistics_ , 8:115–131, 2004.
  * Kukleva et al. (2023) Kukleva, A., Böhle, M., Schiele, B., Kuehne, H., and Rupprecht, C.  Temperature schedules for self-supervised contrastive methods on long-tail data.  _arXiv preprint arXiv:2303.13664_ , 2023.
  * Lee et al. (2023) Lee, C.-Y., Li, C.-L., Zhang, H., Dozat, T., Perot, V., Su, G., Zhang, X., Sohn, K., Glushnev, N., Wang, R., et al.  Formnetv2: Multimodal graph contrastive learning for form document information extraction.  _arXiv preprint arXiv:2305.02549_ , 2023.
  * Li et al. (2023) Li, C., Gan, Z., Yang, Z., Yang, J., Li, L., Wang, L., and Gao, J.  Multimodal foundation models: From specialists to general-purpose assistants.  _arXiv preprint arXiv:2309.10020_ , 1(2):2, 2023.
  * Lindgren et al. (2021) Lindgren, E., Reddi, S., Guo, R., and Kumar, S.  Efficient training of retrieval models using negative cache.  _Advances in Neural Information Processing Systems_ , 34:4134–4146, 2021.
  * Logeswaran & Lee (2018) Logeswaran, L. and Lee, H.  An efficient framework for learning sentence representations.  _arXiv preprint arXiv:1803.02893_ , 2018.
  * Mengersen & Tweedie (1996) Mengersen, K. L. and Tweedie, R. L.  Rates of convergence of the hastings and metropolis algorithms.  _The annals of Statistics_ , 24(1):101–121, 1996.
  * Mikolov et al. (2013) Mikolov, T., Sutskever, I., Chen, K., Corrado, G. S., and Dean, J.  Distributed representations of words and phrases and their compositionality.  _Advances in neural information processing systems_ , 26, 2013.
  * Monath et al. (2023) Monath, N., Zaheer, M., Allen, K., and McCallum, A.  Improving dual-encoder training through dynamic indexes for negative mining.  In _International Conference on Artificial Intelligence and Statistics_ , pp. 9308–9330. PMLR, 2023.
  * Oord et al. (2018) Oord, A. v. d., Li, Y., and Vinyals, O.  Representation learning with contrastive predictive coding.  _arXiv preprint arXiv:1807.03748_ , 2018.
  * Radford et al. (2021) Radford, A., Kim, J. W., Hallacy, C., Ramesh, A., Goh, G., Agarwal, S., Sastry, G., Askell, A., Mishkin, P., Clark, J., et al.  Learning transferable visual models from natural language supervision.  In _International conference on machine learning_ , pp. 8748–8763. PMLR, 2021.
  * Rawat et al. (2019) Rawat, A. S., Chen, J., Yu, F. X. X., Suresh, A. T., and Kumar, S.  Sampled softmax with random fourier features.  _Advances in Neural Information Processing Systems_ , 32, 2019.
  * Robert & Casella (1999) Robert, C. P. and Casella, G.  _Monte Carlo statistical methods_ , volume 2.  Springer, 1999.
  * Roberts & Tweedie (1996) Roberts, G. O. and Tweedie, R. L.  Geometric convergence and central limit theorems for multidimensional hastings and metropolis algorithms.  _Biometrika_ , 83(1):95–110, 1996.
  * Robinson et al. (2020) Robinson, J., Chuang, C.-Y., Sra, S., and Jegelka, S.  Contrastive learning with hard negative samples.  _arXiv preprint arXiv:2010.04592_ , 2020.
  * Srikant & Ying (2019) Srikant, R. and Ying, L.  Finite-time error bounds for linear stochastic approximation andtd learning.  In _Conference on Learning Theory_ , pp. 2803–2830. PMLR, 2019.
  * Sun et al. (2018) Sun, T., Sun, Y., and Yin, W.  On markov chain gradient descent.  _Advances in neural information processing systems_ , 31, 2018.
  * Vembu et al. (2012) Vembu, S., Gartner, T., and Boley, M.  Probabilistic structured predictors.  _arXiv preprint arXiv:1205.2610_ , 2012.
  * Wang & Isola (2020) Wang, T. and Isola, P.  Understanding contrastive representation learning through alignment and uniformity on the hypersphere.  In _International Conference on Machine Learning_ , pp. 9929–9939. PMLR, 2020.
  * Xie et al. (2020) Xie, Q., Dai, Z., Hovy, E., Luong, T., and Le, Q.  Unsupervised data augmentation for consistency training.  _Advances in neural information processing systems_ , 33:6256–6268, 2020.
  * Xiong et al. (2020) Xiong, L., Xiong, C., Li, Y., Tang, K.-F., Liu, J., Bennett, P., Ahmed, J., and Overwijk, A.  Approximate nearest neighbor negative contrastive learning for dense text retrieval.  _arXiv preprint arXiv:2007.00808_ , 2020.
  * Xu et al. (2020) Xu, Y., Xu, Y., Lv, T., Cui, L., Wei, F., Wang, G., Lu, Y., Florencio, D., Zhang, C., Che, W., et al.  Layoutlmv2: Multi-modal pre-training for visually-rich document understanding.  _arXiv preprint arXiv:2012.14740_ , 2020.
  * Yang et al. (2020) Yang, Z., Ding, M., Zhou, C., Yang, H., Zhou, J., and Tang, J.  Understanding negative sampling in graph representation learning.  In _Proceedings of the 26th ACM SIGKDD international conference on knowledge discovery & data mining_, pp. 1666–1676, 2020.
  * Yuan et al. (2022) Yuan, Z., Wu, Y., Qiu, Z.-H., Du, X., Zhang, L., Zhou, D., and Yang, T.  Provable stochastic optimization for global contrastive learning: Small batch does not harm performance.  In _International Conference on Machine Learning_ , pp. 25760–25782. PMLR, 2022.
  * Zbontar et al. (2021) Zbontar, J., Jing, L., Misra, I., LeCun, Y., and Deny, S.  Barlow twins: Self-supervised learning via redundancy reduction.  In _International Conference on Machine Learning_ , pp. 12310–12320. PMLR, 2021.
  * Zimmermann et al. (2021) Zimmermann, R. S., Sharma, Y., Schneider, S., Bethge, M., and Brendel, W.  Contrastive learning inverts the data generating process.  In _International Conference on Machine Learning_ , pp. 12979–12990. PMLR, 2021.

##  Appendix A Proof of Lemma 3.2

Notice that our target softmax distribution ${\cal P}_{x,\theta}$ has a finite
support with each state having non-zero mass. For any
$\theta\in\mathbb{R}^{p}$,

$\min_{x,z,Z}Q_{x,\theta}(z,Z)=\min_{x,z,Z}\frac{\exp(\beta~{}\phi(x;\theta)^{\top}\psi(z;\theta))}{\exp(\beta~{}\phi(x;\theta)^{\top}\psi(Z;\theta))}\stackrel{{\scriptstyle({\rm
Assm.}~{}\ref{assm:bounded_embd})}}{{\geq}}\exp(-2c^{2}\beta)$ (22)

Now recall that $m_{\rm neg}=|{\bf D}_{\rm neg}(x)|$ for any $x\in{\bf
D}_{X}$, i.e., the state space consists of $m_{\rm neg}$ states on each Markov
chain $Z_{i}$. Then the smallest probability to transition to state $z$ from
any initial state is lower bounded by $m_{\rm neg}^{-1}\exp(-2c^{2}\beta)$,
where the factor $m_{\rm neg}^{-1}$ comes from using the uniform prior ${\rm
Uniform}({\bf D}_{\rm neg}(x))$ in Algorithm 1.

By Theorem 1.3 of Mengersen & Tweedie (1996), the transition kernel of each
Markov chain of $Z_{i}$ converges to its target distribution ${\cal
P}_{x,\theta}$ at a geometric rate of $1-m_{\rm neg}^{-1}\exp(-2c^{2}\beta)$.
Therefore, the joint Markov chain of states $(Z_{1},...,Z_{m})$ in Algorithm 1
has a mixing rate of $1-m^{-1}m_{\rm neg}^{-1}\exp(-2c^{2}\beta)$ when the
batch size $B=1$. As we perform mini-batch sampling of batch size $B$ on the
expectation over $(x,y)\sim{\cal D}_{\rm pos}$ and $R$ steps of M-H algorithm
over $z\sim{\cal P}_{x,\theta}$ _in parallel_ , the mixing rate can be
improved to $[1-Bm^{-1}m_{\rm neg}^{-1}\exp(-2c^{2}\beta)]^{R}$. We can
simplify the rate by the following upper bound, as

$\displaystyle\left[1-\frac{B}{mm_{\rm
neg}\exp(2c^{2}\beta)}\right]^{R}\leq\exp\left(-\frac{BR}{mm_{\rm
neg}\exp(2c^{2}\beta)}\right)\leq 1-\frac{BR}{2mm_{\rm neg}\exp(2c^{2}\beta)}$
(23)

where the first inequality uses $1-x\leq\exp(-x)~{}\forall x\in\mathbb{R}$,
the second inequality uses $\exp(-x)\leq 1-x/2$ for $x\in[0,1]$ and the
assumption that $BR\leq mm_{\rm neg}\exp(2c^{2}\beta)$. $\square$

##  Appendix B Proof of Lemma 3.4

For each Markov chain corresponding to $Z_{i}$ in the M-H algorithm, the
acceptance probability is given by
$\min\\{1,Q_{x_{i},\theta}(Z,Z^{\prime})\\}$. To establish Lemma 3.4, we first
observe that for any $Z,Z^{\prime}\in{\bf D}_{\rm neg}(x_{i})$ and
$\theta,\theta^{\prime}\in\mathbb{R}^{p}$, it holds

$\displaystyle|Q_{x,\theta}(Z,Z^{\prime})-Q_{x,\theta^{\prime}}(Z,Z^{\prime})|$
(24)
$\displaystyle=|\exp\left(\beta\phi(x;\theta)^{\top}\psi(Z;\theta)-\beta\phi(x;\theta)^{\top}\psi(Z^{\prime};\theta)\right)-\exp\left(\beta\phi(x;\theta^{\prime})^{\top}\psi(Z;\theta^{\prime})-\beta\phi(x;\theta^{\prime})^{\top}\psi(Z^{\prime};\theta^{\prime})\right)|$
(25)
$\displaystyle\leq\exp(2c^{2}\beta)\cdot\beta\cdot|\phi(x;\theta)^{\top}\psi(Z;\theta)-\phi(x;\theta)^{\top}\psi(Z^{\prime};\theta)-(\phi(x;\theta^{\prime})^{\top}\psi(Z;\theta^{\prime})-\phi(x;\theta^{\prime})^{\top}\psi(Z^{\prime};\theta^{\prime}))|$
(26) $\displaystyle\stackrel{{\scriptstyle({\rm
Assm.}~{}\ref{assm:smooth_kernel})}}{{\leq}}2L_{P}\cdot\exp(2c^{2}\beta)\cdot\beta\cdot\|\theta-\theta^{\prime}\|$
(27)

where (26) uses the fact that $\exp(\cdot)$ is $\exp(2c^{2}\beta)$-Lipschitz
when restricted on the domain $[-2c^{2}\beta,2c^{2}\beta]$ and the restriction
is due to Assumption 3.1.

Now denote ${\bf K}_{i,\theta}\in\mathbb{R}^{m_{\rm neg}\times m_{\rm neg}}$
as the transition matrix of $Z_{i}$ and

$\displaystyle[{\bf K}_{i,\theta}]_{z^{\prime},z}=\begin{cases}\frac{1}{m_{\rm
neg}}\cdot\min\\{1,Q_{x_{i},\theta}(z,z^{\prime})\\}&\text{if}~{}z\neq
z^{\prime},\\\ \frac{1}{m_{\rm neg}}+\frac{1}{m_{\rm neg}}\sum_{j\in[m_{\rm
neg}]\backslash\\{z\\}}(1-\min\\{1,Q_{x_{i},\theta}(z,j)\\})&\text{if}~{}z=z^{\prime}.\end{cases}$
(28)

Then for $z\neq z^{\prime}$,

$\displaystyle|[{\bf K}_{i,\theta}]_{z^{\prime},z}-[{\bf
K}_{i,\theta^{\prime}}]_{z^{\prime},z}|$ $\displaystyle=\frac{1}{m_{\rm
neg}}|\min\\{1,Q_{x_{i},\theta}(z,z^{\prime})\\}-\min\\{1,Q_{x_{i},\theta^{\prime}}(z,z^{\prime})\\}|$
(29) $\displaystyle\leq\frac{1}{m_{\rm
neg}}|Q_{x_{i},\theta}(z,z^{\prime})-Q_{x_{i},\theta^{\prime}}(z,z^{\prime})|\stackrel{{\scriptstyle\eqref{proof:lip_q}}}{{\leq}}\frac{1}{m_{\rm
neg}}\cdot
2L_{P}\cdot\exp(2c^{2}\beta)\cdot\beta\cdot\|\theta-\theta^{\prime}\|$ (30)

and for $z=z^{\prime}$,

$\displaystyle|[{\bf K}_{i,\theta}]_{z^{\prime},z}-[{\bf
K}_{i,\theta^{\prime}}]_{z^{\prime},z}|$ $\displaystyle=\left|\frac{1}{m_{\rm
neg}}\sum_{j\in[m_{\rm
neg}]\backslash\\{z\\}}(-\min\\{1,Q_{x_{i},\theta}(z,j)\\}+\min\\{1,Q_{x_{i},\theta^{\prime}}(z,j)\\})\right|$
(31) $\displaystyle\leq\frac{1}{m_{\rm neg}}\sum_{j\in[m_{\rm
neg}]\backslash\\{z\\}}|Q_{x_{i},\theta}(z,j)-Q_{x_{i},\theta^{\prime}}(z,j)|\stackrel{{\scriptstyle\eqref{proof:lip_q}}}{{\leq}}2L_{P}\cdot\exp(2c^{2}\beta)\cdot\beta\cdot\|\theta-\theta^{\prime}\|$
(32)

To quantify the smoothness of the transition kernel with $R$ Metropolis-
Hastings steps, we observe that by the notation $[{\bf K}]_{z,z^{\prime}}={\bf
e}_{z}^{\top}{\bf K}{\bf e}_{z^{\prime}}$ for some basis vectors ${\bf
e}_{z},{\bf e}_{z^{\prime}}$,

$\displaystyle\max_{z,z^{\prime}}|{\tt
P}_{\theta}(Z_{i}^{(0)}=z;Z_{i}^{(R)}=z^{\prime})-{\tt
P}_{\theta^{\prime}}(Z_{i}^{(0)}=z;Z_{i}^{(R)}=z^{\prime})|$ (33)
$\displaystyle=\max_{z,z^{\prime}}|[{\bf
K}_{i,\theta}^{R}]_{z^{\prime},z}-[{\bf
K}_{i,\theta^{\prime}}^{R}]_{z^{\prime},z}|$ (34)
$\displaystyle=\max_{z,z^{\prime}}|{\bf e}_{z^{\prime}}^{\top}({\bf
K}_{i,\theta}^{R}-{\bf K}_{i,\theta^{\prime}}{\bf K}_{i,\theta}^{R-1}+{\bf
K}_{i,\theta^{\prime}}{\bf K}_{i,\theta}^{R-1}-{\bf
K}_{i,\theta^{\prime}}^{R}){\bf e}_{z}|$ (35)
$\displaystyle\leq\max_{z,z^{\prime}}|{\bf e}_{z^{\prime}}^{\top}({\bf
K}_{i,\theta}-{\bf K}_{i,\theta^{\prime}}){\bf K}_{i,\theta}^{R-1}{\bf
e}_{z}|+|{\bf e}_{z^{\prime}}^{\top}{\bf K}_{i,\theta^{\prime}}({\bf
K}_{i,\theta}^{R-1}-{\bf K}_{i,\theta^{\prime}}^{R-1}){\bf e}_{z}|$ (36)
$\displaystyle\stackrel{{\scriptstyle(i)}}{{\leq}}\max_{z,z^{\prime}}\|{\bf
e}_{z^{\prime}}^{\top}({\bf K}_{i,\theta}-{\bf
K}_{i,\theta^{\prime}})\|_{\infty}\cdot\|{\bf K}_{i,\theta}^{R-1}{\bf
e}_{z}\|_{1}+\|{\bf e}_{z^{\prime}}^{\top}{\bf
K}_{i,\theta^{\prime}}\|_{1}\cdot\|({\bf K}_{i,\theta}^{R-1}-{\bf
K}_{i,\theta^{\prime}}^{R-1}){\bf e}_{z}\|_{\infty}$ (37)
$\displaystyle\stackrel{{\scriptstyle(ii)}}{{\leq}}\max_{z,z^{\prime}}\|{\bf
e}_{z^{\prime}}^{\top}({\bf K}_{i,\theta}-{\bf
K}_{i,\theta^{\prime}})\|_{\infty}\stackrel{{\scriptstyle(iii)}}{{+}}2\cdot|[{\bf
K}_{i,\theta}^{R-1}]_{z^{\prime},z}-[{\bf
K}_{i,\theta^{\prime}}^{R-1}]_{z^{\prime},z}|$ (38)
$\displaystyle\stackrel{{\scriptstyle\eqref{proof:lip_q}}}{{\leq}}2L_{P}\cdot\exp(2c^{2}\beta)\cdot\beta\cdot\|\theta-\theta^{\prime}\|+2\cdot|[{\bf
K}_{i,\theta}^{R-1}]_{z^{\prime},z}-[{\bf
K}_{i,\theta^{\prime}}^{R-1}]_{z^{\prime},z}|$ (39)
$\displaystyle\leq(1+2+...+2^{R-1})\cdot
2L_{P}\cdot\exp(2c^{2}\beta)\cdot\beta\cdot\|\theta-\theta^{\prime}\|$ (40)
$\displaystyle=(2^{R}-1)\cdot
2L_{P}\cdot\exp(2c^{2}\beta)\cdot\beta\cdot\|\theta-\theta^{\prime}\|$ (41)

where $(i)$ uses Hölder’s inequality for the norm pair
$(\|\cdot\|_{1},\|\cdot\|_{\infty})$, $(ii)$ uses the fact that ${\bf
K}_{i,\theta}^{R-1}$ is a column stochastic matrix and $(iii)$ uses the
inequality $\|{\bf e}_{z^{\prime}}^{\top}{\bf K}_{i,\theta^{\prime}}\|_{1}\leq
2$ because

$\displaystyle\|{\bf e}_{z^{\prime}}^{\top}{\bf
K}_{i,\theta^{\prime}}\|_{1}=\sum_{z\in[m_{\rm neg}]}[{\bf
K}_{i,\theta}]_{z^{\prime},z}$ (42)
$\displaystyle\stackrel{{\scriptstyle\eqref{eq:kernel_value}}}{{=}}\sum_{z\in[m_{\rm
neg}]\backslash\\{z^{\prime}\\}}\frac{1}{m_{\rm
neg}}\min\\{1,Q_{x_{i},\theta}(z,z^{\prime})\\}$ (43)
$\displaystyle\quad+\frac{1}{m_{\rm neg}}+\frac{1}{m_{\rm
neg}}\sum_{j\in[m_{\rm
neg}]\backslash\\{z^{\prime}\\}}(1-\min\\{1,Q_{x_{i},\theta}(z^{\prime},j)\\})$
(44) $\displaystyle\leq\frac{m_{\rm neg}-1}{m_{\rm neg}}+\frac{1}{m_{\rm
neg}}+\frac{m_{\rm neg}-1}{m_{\rm neg}}\leq 2$ (45)

Finally, in the Markov chain
$\tilde{\bm{\xi}}_{0}\to\tilde{\bm{\xi}}_{1}\to\cdots$, since only $B$ state
variables $\\{Z_{i_{k}^{(t)}}\\}_{k=1}^{B}$ are active or updated in each
transition, we have

${\tt P}_{\theta}(\bm{\xi};\bm{\xi}^{\prime})=\begin{cases}\quad
0&\text{if}~{}\Delta_{\xi}=\|\xi-\xi^{\prime}\|_{0}>B,\\\
\left(\begin{matrix}m\\\
B\end{matrix}\right)^{-1}\cdot\left(\begin{matrix}m-\Delta_{\xi}\\\
B-\Delta_{\xi}\end{matrix}\right)\cdot\prod_{k=1}^{B}{\tt
P}_{\theta}(Z_{i_{k}^{(t)}}^{(0)};Z_{i_{k}^{(t)}}^{(R)})&\text{otherwise}.\end{cases}$
(46)

where $\left(\begin{matrix}m\\\
B\end{matrix}\right)^{-1}\left(\begin{matrix}m-\Delta_{\xi}\\\
B-\Delta_{\xi}\end{matrix}\right)$ is the probability that
$j\in\\{i_{k}^{(t)}\\}_{k=1}^{B}$ for $j$ s.t. $\xi(j)\neq\xi^{\prime}(j)$. To
simplify notations, suppose $Z_{i_{k}^{(t)}}^{(0)}=z_{t,k}$ and
$Z_{i_{k}^{(t)}}^{(R)}=z_{t,k}^{\prime}$, then

$\displaystyle|{\tt P}_{\theta}(\xi;\xi^{\prime})-{\tt
P}_{\theta^{\prime}}(\xi;\xi^{\prime})|=\left(\begin{matrix}m\\\
B\end{matrix}\right)^{-1}\left(\begin{matrix}m-\Delta_{\xi}\\\
B-\Delta_{\xi}\end{matrix}\right)\left|\prod_{k=1}^{B}{\tt
P}_{\theta}(Z_{i_{k}^{(t)}}^{(0)};Z_{i_{k}^{(t)}}^{(R)})-\prod_{k=1}^{B}{\tt
P}_{\theta^{\prime}}(Z_{i_{k}^{(t)}}^{(0)};Z_{i_{k}^{(t)}}^{(R)})\right|$ (47)
$\displaystyle=\left(\begin{matrix}m\\\
B\end{matrix}\right)^{-1}\left(\begin{matrix}m-\Delta_{\xi}\\\
B-\Delta_{\xi}\end{matrix}\right)\left|\prod_{k=1}^{B}[{\bf
K}_{i,\theta}^{R}]_{z_{t,k},z_{t,k}^{\prime}}-\prod_{k=1}^{B}[{\bf
K}_{i,\theta^{\prime}}^{R}]_{z_{t,k},z_{t,k}^{\prime}}\right|$ (48)
$\displaystyle=\left(\begin{matrix}m\\\
B\end{matrix}\right)^{-1}\left(\begin{matrix}m-\Delta_{\xi}\\\
B-\Delta_{\xi}\end{matrix}\right)\Big{|}\prod_{k=1}^{B}[{\bf
K}_{i,\theta}^{R}]_{z_{t,k},z_{t,k}^{\prime}}-[{\bf
K}_{i,\theta^{\prime}}^{R}]_{z_{t,1},z_{t,1}^{\prime}}\prod_{k=2}^{B}[{\bf
K}_{i,\theta}^{R}]_{z_{t,k},z_{t,k}^{\prime}}$ (49) $\displaystyle\quad+[{\bf
K}_{i,\theta^{\prime}}^{R}]_{z_{t,1},z_{t,1}^{\prime}}\prod_{k=2}^{B}[{\bf
K}_{i,\theta}^{R}]_{z_{t,k},z_{t,k}^{\prime}}-\prod_{k=1}^{B}[{\bf
K}_{i,\theta^{\prime}}^{R}]_{z_{t,k},z_{t,k}^{\prime}}\Big{|}$ (50)
$\displaystyle\stackrel{{\scriptstyle\eqref{proof:smooth_r_step_kernel}}}{{\leq}}\left(\begin{matrix}m\\\
B\end{matrix}\right)^{-1}\left(\begin{matrix}m-\Delta_{\xi}\\\
B-\Delta_{\xi}\end{matrix}\right)B\cdot 2^{R+1}\cdot
L_{P}\cdot\exp(2c^{2}\beta)\cdot\beta\cdot\|\theta-\theta^{\prime}\|$ (51)

The proof is concluded by taking maximum over $0\leq\Delta_{\xi}\leq B$.
$\square$

##  Appendix C Proof of Theorem 3.7

Our idea is to apply Theorem 2 of Karimi et al. (2019) which shows the
convergence of biased stochastic approximation scheme with Markovian noise
such as (10). In particular, we shall take the Lyapunov function therein as
our global contrastive loss, i.e., $V(\theta)={\cal L}(\theta)$. We proceed by
verifying the required assumptions.

#### Verifying A1-A2 of Karimi et al. (2019)

Using Lemma 3.2, we observe that the mean field of the stochastic update in
(10), i.e., the expected value of ${\cal H}(\xi;\theta)$ when $\xi$ is drawn
from the stationary distribution of the Markov chain induced by ${\tt
P}_{\theta}$, coincides with the gradient of ${\cal L}(\theta)$, i.e.,
$h(\theta)=\nabla{\cal L}(\theta)$. As such, A1-A2 of Karimi et al. (2019) are
satisfied with $c_{0}=d_{0}=0$, $c_{1}=d_{1}=1$.

#### Verifying A3 of Karimi et al. (2019)

For this, we need to show that $\nabla{\cal L}(\theta)$ is Lipschitz
continuous. We observe that a stronger condition holds as the stochastic
gradient map is Lipschitz w.r.t. $\theta$. For any sample $x,y,z$ and any
$\theta,\theta^{\prime}\in\mathbb{R}^{p}$, we have

$\displaystyle\|-\beta\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(y;\theta))+\beta\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(z;\theta))-[-\beta\nabla_{\theta}(\phi(x;\theta^{\prime})^{\top}\psi(y;\theta^{\prime}))+\beta\nabla_{\theta}(\phi(x;\theta^{\prime})^{\top}\psi(z;\theta^{\prime}))]\|$
(52)
$\displaystyle\leq\beta\|\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(y;\theta))-\nabla_{\theta}(\phi(x;\theta^{\prime})^{\top}\psi(y;\theta^{\prime}))\|+\beta\|\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(z;\theta))-\nabla_{\theta}(\phi(x;\theta^{\prime})^{\top}\psi(z;\theta^{\prime}))\|$
$\displaystyle\stackrel{{\scriptstyle({\rm
Assm.}~{}\ref{assm:smooth_sgrad})}}{{\leq}}2\beta
L_{H}\|\theta-\theta^{\prime}\|$

#### Verifying A5-6 of Karimi et al. (2019)

For these assumptions, we observe that A12-A14 of Karimi et al. (2019) can be
satisfied with the constants

$\bar{L}_{P}=2^{R+1}BL_{P}\cdot\exp(2c^{2}\beta)\cdot\beta,\quad\bar{L}_{H}=2\beta
L_{H},\quad\bar{\rho}=1-\frac{BR}{2mm_{\rm neg}\exp(2c^{2}\beta)},\quad
K_{P}=1,\quad\bar{\sigma}\stackrel{{\scriptstyle\eqref{proof:sigma_bar}}}{{=}}2\beta\sigma$
(53)

Applying Lemma 7 of Karimi et al. (2019) shows that A5-A6 can be satisfied
with

$L_{PH}^{(0)}=\frac{2\beta\sigma\bar{\rho}}{BRm^{-1}m_{\rm
neg}^{-1}\exp(-2c^{2}\beta)/2},\quad L_{PH}^{(1)}=\frac{6\cdot
2^{R+1}B\exp(2c^{2}\beta)\beta^{2}\sigma L_{P}}{(BRm^{-1}m_{\rm
neg}^{-1}\exp(-2c^{2}\beta)/2)^{2}}+\frac{2\beta L_{H}}{BRm^{-1}m_{\rm
neg}^{-1}\exp(-2c^{2}\beta)/2}$ (54)

#### Verifying A7 of Karimi et al. (2019)

The stochastic gradient used in (10) has a uniformly bounded error from its
mean field as

$\displaystyle\left\|{\cal H}(\xi_{t+1};\theta_{t})-\nabla{\cal
L}(\theta_{t})\right\|$ (55)
$\displaystyle\leq\beta\left\|\frac{1}{B}\sum_{\ell=1}^{B}H(x^{(t)}_{i_{\ell}},y^{(t)}_{i_{\ell}};\theta_{t})-\mathbb{E}_{(x,y)\sim{\cal
D}_{\rm pos}}[\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(y;\theta))]\right\|$
$\displaystyle\quad+\beta\left\|\frac{1}{B(B-P)}\sum_{k=1}^{B}\sum_{r=P}^{R-1}H(x_{i_{k}^{(t)}},\tilde{Z}_{i_{k}^{(t)}}^{(r)};\theta_{t})-\mathbb{E}_{(x,y)\sim{\cal
D}_{\rm pos}}\mathbb{E}_{z\sim{\cal
P}_{x,\theta}}[\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(z;\theta))]\right\|$
$\displaystyle\leq\frac{\beta}{B}\sum_{\ell=1}^{B}\|H(x^{(t)}_{i_{\ell}},y^{(t)}_{i_{\ell}};\theta_{t})-\mathbb{E}_{(x,y)\sim{\cal
D}_{\rm pos}}[\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(y;\theta))]\|$
$\displaystyle\quad+\frac{\beta}{B(B-P)}\sum_{k=1}^{B}\sum_{r=P}^{B-1}\|H(x_{i_{k}^{(t)}},\tilde{Z}_{i_{k}^{(t)}}^{(r)};\theta_{t})-\mathbb{E}_{(x,y)\sim{\cal
D}_{\rm pos}}\mathbb{E}_{z\sim{\cal
P}_{x,\theta}}[\nabla_{\theta}(\phi(x;\theta)^{\top}\psi(z;\theta))]\|$
$\displaystyle\stackrel{{\scriptstyle({\rm
Assm.}~{}\ref{assm:bounded_var})}}{{\leq}}2\beta\sigma$ (56)

#### Convergence of (10)

Upon verifying A1-A3, A5-A7 of Karimi et al. (2019), we can apply Theorem 2
therein to analyze the convergence of (10). In particular, we choose a
constant step size $\gamma$ satisfying $\gamma\leq\frac{1}{2(2\beta
L_{H}+C_{h})}$. For any $T\geq 1$ and the randomly drawn $t\sim{\rm
Uniform}([0,...,T-1])$, we have

$\displaystyle\mathbb{E}[\|\nabla{\cal L}(\theta_{t})\|^{2}]$
$\displaystyle=\frac{1}{T}\sum_{t=0}^{T-1}\mathbb{E}[\|\nabla{\cal
L}(\theta_{t})\|^{2}]$ (57) $\displaystyle\leq\frac{\mathbb{E}[{\cal
L}(\theta_{0})-{\cal
L}(\theta_{T})]+3L_{PH}^{(0)}+[8\beta^{3}L_{H}\sigma^{2}+2L_{PH}^{(1)}\beta\sigma+2\beta
L_{H}L_{PH}^{(0)}(1+2\beta\sigma)](\gamma^{2}T)}{\gamma T/2}$ (58)

where $C_{h}=L_{PH}^{(1)}(1+\sigma)+L_{PH}^{(0)}(\beta L_{H}+1)$. Simplifying
and rearranging terms leads to the conclusion.

#### Estimating $\mathbb{E}[{\cal L}(\theta_{0})-{\cal L}(\theta_{T})]$

We observe that by Assumption 3.1, for any $\theta\in\mathbb{R}^{p}$,

$\displaystyle{\cal L}(\theta)$
$\displaystyle=\mathop{\mathbb{E}}_{(x,y)\sim{\cal D}_{\rm
pos}}\left[-\beta~{}\phi(x;\theta)^{\top}\psi(y;\theta)\right]+\mathop{\mathbb{E}}_{(x,y)\sim{\cal
D}_{\rm pos}}\left[\log\sum_{z\in{\bf D}_{\rm
neg}(x)}\exp(\beta~{}\phi(x;\theta)^{\top}\psi(z;\theta))\right]$ (59)
$\displaystyle\leq\beta+\mathop{\mathbb{E}}_{(x,y)\sim{\cal D}_{\rm
pos}}\left[\log\sum_{z\in{\bf D}_{\rm neg}(x)}\exp(\beta)\right]$

By a similar argument,

${\cal L}(\theta)\geq-\beta+\mathop{\mathbb{E}}_{(x,y)\sim{\cal D}_{\rm
pos}}\left[\log\sum_{z\in{\bf D}_{\rm neg}(x)}\exp(-\beta)\right]$ (60)

Therefore,

$\displaystyle\mathbb{E}[{\cal L}(\theta_{0})-{\cal L}(\theta_{T})]$
$\displaystyle\leq 2\beta+\mathop{\mathbb{E}}_{(x,y)\sim{\cal D}_{\rm
pos}}\left[\log\sum_{z\in{\bf D}_{\rm neg}(x)}\exp(\beta)-\log\sum_{z\in{\bf
D}_{\rm neg}(x)}\exp(-\beta)\right]$ (61)
$\displaystyle=2\beta+\mathop{\mathbb{E}}_{(x,y)\sim{\cal D}_{\rm
pos}}\left[\log\left(\frac{\sum_{z\in{\bf D}_{\rm
neg}(x)}\exp(\beta)}{\sum_{z\in{\bf D}_{\rm
neg}(x)}\exp(-\beta)}\right)\right]$
$\displaystyle=2\beta+\mathop{\mathbb{E}}_{(x,y)\sim{\cal D}_{\rm
pos}}\left[\log\exp(2\beta)\right]=4\beta.$

$\square$

##  Appendix D Experiment Details

###  D.1 Details of Dataset and Hyperparameters

In Table 2, we list the important attributes of datasets we used in the
experiment section.

Dataset | # of Pos. Pairs $m$ | # of Neg. Samples $m_{\rm neg}$ | Image Crop Size
---|---|---|---
STL-10 | 100,000 (inf. aug) | 100,000 (inf. aug) | 96 $\times$ 96
Imagenet-100 | 124,689 (inf. aug) | 124,689 (inf. aug) | 224 $\times$ 224
Table 2:  Datasets attributes.

In Table 3, we list the hyperparameter values adopted in our experiments.

Dataset | Model |  | Inverse
---
Temp. $\beta$
| Batch
---
Size $b$
| Learning
---
Rate $\gamma$
| Feature
---
Dim. $d$
| Weight
---
Decay
|  Cache Refresh $\rho$
---
(Negative Cache)
|  Burn-in Steps $P$
---
(EMC2)
STL-10 | Resnet-18 | 14.28 | 32 | $10^{-4}$ | 512 | $10^{-4}$ | 0.01024 | 31
STL-10 | Resnet-18 | 5 | 256 | $10^{-3}$ | 512 | $10^{-4}$ | 0.1 | 255
Imagenet-100 | Resnet-50 | 14.28 | 256 | 1.2 | 128 | $10^{-6}$ | - | 255
STL-10 subset | Resnet-18 | 5 | 4 | $10^{-3}$ | 512 | $10^{-4}$ | 0.1 | 3

Table 3:  Hyperparameter values of the experiments.

###  D.2 Run-Time Complexity

Since the computational complexity differs among algorithms, we provide a
performance-time plot in Figure 8 to compare the time complexity. Note that in
this setup, negative cache algorithm uses four Tesla T4 GPUs for training and
refreshing the negative cache while the other algorithms run on one Tesla T4
GPU.

Figure 8: Training 100 epochs on STL-10 with ResNet-18 using batch size
$b=32$. Horizontal axis is relative to the wall-clock training time in
seconds.