Mini-Batch Optimization of Contrastive Loss

In this section, we characterize the optimal solution for the full-batch loss minimization in Eq. (1). We start by providing the definition of the simplex equiangular tight frame (ETF) which turns out to be the optimal solution in certain cases. The original definition of ETF [62] is for $N$ vectors in a $d$-dimensional space where $N\geq d+1$ ¹¹1See Def. 4 in Appendix A for the full definition. Papyan et al. [46] defines the ETF for the case where $N\leq d+1$ to characterize the phenomenon of neural collapse. In our work, we use the latter definition of simplex ETFs which is stated below.

In the following Lemma, we first prove that the optimal solution of full-batch contrastive learning is the simplex ETF for $N\leq d+1$ which follows almost directly from Lu & Steinerberger [36].

Actually, many practical scenarios satisfy $N>d+1$. However, the approach used in Lu & Steinerberger [36] cannot be directly applied for $N>d+1$, leaving it as an open problem. While solving the open problem for the general case seems difficult, we characterize the optimal solution for the specific case of $N=2d$, subject to the conditions stated below.

We conjecture that the optimal solutions for $N=2d$ are symmetric and antipodal. Note that the symmetric property holds for $N\leq d+1$ case, and the antipodality is frequently assumed in geometric problems such as the sphere covering problem in [4].

Thm. 3 shows that when $N=2d$, the optimal solution for the full-batch loss minimization, under a symmetric and antipodal configuration, form a cross-polytope which is defined as the following.

Proof Outline. By the antipodality assumption, we can apply Jensen’s inequality to $N-2$ indices without itself ${{\bm{u}}}_{i}$ and antipodal point $-{\bm{u}}_{i}$ for a given $i\in[N]$. Then we show that the simplex cross-polytope also minimizes this lower bound while satisfying the conditions that make the applications of Jensen’s inequality tight.

For the general case of $N>d+1$, excluding $N=2d,$ we still leave it as an open problem.

Here we consider the mini-batch contrastive loss optimization problem, where we first choose multiple mini-batches of size $B$ and then find ${\bm{U}},{\bm{V}}$ that minimize the sum of contrastive losses computed for the chosen mini-batches. Note that this is the loss that is typically considered in the contrastive learning since computing the full-batch loss is intractable in practice. Let us consider a subset of all possible $\binom{N}{B}$ mini-batches and denote their indices by ${\mathcal{S}}_{B}\subseteq\left[\binom{N}{B}\right]$. For a fixed ${\mathcal{S}}_{B}$, the mini-batch loss optimization problem is formulated as:

where the loss of given mini-batches is ${\mathcal{L}}^{\operatorname{con}}_{\operatorname{mini}}({\bm{U}},{\bm{V}};{\mathcal{S}}_{B}):=\frac{1}{|{\mathcal{S}}_{B}|}\sum_{i\in{\mathcal{S}}_{B}}{\mathcal{L}}^{\operatorname{con}}({\bm{U}}_{{\mathcal{B}}_{i}},{\bm{V}}_{{\mathcal{B}}_{i}}).$ To analyze the relationship between the full-batch loss minimization in Eq. (1) and the mini-batch loss minimization in Eq. (4), we first compare the objective functions of two problems as below.

We illustrate this proposition by visualizing the two loss functions in Fig. 1 when $N=10,B=2$, and $d=2$. We visualize it along a single embedding vector ${\bm{u}}_{1}$ by freezing all other embeddings (${\bm{v}}_{1}$ and $\{{\bm{u}}_{i},{\bm{v}}_{i}\}_{i=2}^{10}$) at the optimal solution and varying ${\bm{u}}_{1}=[u_{1,1},u_{1,2}]$ as $[\cos(\theta),\sin(\theta)]$ for $\theta\in[-\pi,\pi]$. One can confirm that two losses are not identical (even up to scaling).

Interestingly, the following result shows that the optimal solutions of both problems are identical.

Proof Outline. Similar to the proof of Lem. 1, we bound the objective function from below using Jensen’s inequality. Then, we show that this lower bound is equivalent to a scaling of the bound from the proof of Lem. 1, by using careful counting arguments. Then, we can simply repeat the rest of the proof to show that the simplex ETF also minimizes this lower bound while satisfying the conditions that make the applications of Jensen’s inequality tight.

Now, we present mathematical results specifying the cases when the solutions of mini-batch optimization and full-batch optimization differ. First, we show that when $B=2$, minimizing the mini-batch loss over any strict subset of $\binom{N}{B}$ batches, is not equivalent to minimizing the full-batch loss.

Proof Outline. We show that there exist embedding vectors that are not the simplex ETF, and have a strictly lower objective value. This implies that the optimal solution of any set of mini-batches that does not contain all $\binom{N}{2}$ mini-batches is not the same as that of the full-batch problem.

The result of Thm. 5 is extended to the general case of $B\geq 2$, under some mild assumption; please check Prop. 2 and 3 in Appendix B.1. Fig. 1 summarizes the main findings in this section.

This section investigates the convergence of two gradient-descent-based methods, OSGD and SGD. The below lemma shows that the contrastive loss is geodesic non-quasi-convex, which implies the hardness of proving the convergence of gradient-based methods for contrastive learning in Eq. (1).

We provide the proof in Appendix B.2.

In order to compare the convergence of OSGD and SGD, we focus on a toy example where convergence to the optimal solution is achievable with appropriate initialization. Consider a scenario where we have $N=4$ embedding vectors $\{{\bm{u}}_{i}\}_{i=1}^{N}$ with ${\bm{u}}_{i}\in{\mathbb{R}}^{2}$. Each embedding vector is defined as ${\bm{u}}_{1}=(\cos\theta_{1},\sin\theta_{1});{\bm{u}}_{2}=(\cos\theta_{2},-\sin\theta_{2});{\bm{u}}_{3}=(-\cos\theta_{3},-\sin\theta_{3});{\bm{u}}_{4}=(-\cos\theta_{4},\sin\theta_{4})$ for parameters $\{\theta_{i}\}_{i=1}^{n}$. Over time step $t$, we consider updating the parameters ${\bm{\theta}}^{(t)}:=[\theta_{1}^{(t)},\theta_{2}^{(t)},\theta_{3}^{(t)},\theta_{4}^{(t)}]$ using gradient descent based methods. For all $i$, the initial parameters are set as $\theta_{i}^{(0)}=\epsilon>0$, and the other embedding vectors are initialized as ${\bm{v}}_{i}^{(0)}={\bm{u}}_{i}^{(0)}$. This setting is illustrated in Fig. 2.

At each time step $t$, each learning algorithm begins by selecting a mini-batch ${\mathcal{B}}^{(t)}\subset\left\{1,2,3,4\right\}$ with batch size $|{\mathcal{B}}^{(t)}|=2$. SGD randomly selects a mini-batch, while OSGD selects a mini-batch as follows: ${\mathcal{B}}^{(t)}=\arg\max\limits_{{\mathcal{B}}\in{\mathcal{S}}}{\mathcal{L}}^{\operatorname{con}}({\bm{U}}_{{\mathcal{B}}}({\bm{\theta}}^{(t)}),{\bm{V}}_{{\mathcal{B}}}({\bm{\theta}}^{(t)}))$. Then, the algorithms update ${\bm{\theta}}^{(t)}$ using gradient descent on ${\mathcal{L}}^{\operatorname{con}}({\bm{U}}_{{\mathcal{B}}},{\bm{V}}_{{\mathcal{B}}})$ with a learning rate $\eta$: ${\bm{\theta}}^{(t+1)}={\bm{\theta}}^{(t)}-\eta\nabla_{{\bm{\theta}}}{\mathcal{L}}^{\operatorname{con}}({\bm{U}}_{{\mathcal{B}}^{(t)}},{\bm{V}}_{{\mathcal{B}}^{(t)}})$. For a sufficiently small margin $\rho>0$, let $T_{\textnormal{OSGD}},T_{\textnormal{SGD}}$ be the minimal time required for the algorithms to reach the condition $\mathbb{E}[{\bm{\theta}}^{(T)}]\in(\pi/4-\rho,\pi/4)^{N}$. Under this setting, the following theorem compares OSGD and SGD, in terms of the lower bound on the time required for the convergence to the optimal solution.

In Fig. 2, we present training loss curves of the full-batch contrastive loss in Eq. (2) for various algorithms implemented on the toy example. One can observe that the losses of all algorithms eventually converge to 1.253, the optimal loss achievable when the solution satisfies ${\bm{u}}_{i}={\bm{v}}_{i}$ and $\{{\bm{u}}_{i}\}_{i=1}^{N}$ form simplex cross-polytope. As shown in the figure, OSGD converges faster than SGD to the optimal loss. This empirical evidence corroborates our theoretical findings in Corollary 1.

Recall that it is challenging to prove the convergence of gradient-descent-based methods for contrastive learning problem in Eq. (1) due to the non-quasi-convexity of the contrastive loss $\mathcal{L}^{\text{con}}$. Instead of focusing on the contrastive loss, we consider a proxy, the weighted contrastive loss defined as $\widetilde{{\mathcal{L}}}^{\operatorname{con}}({\bm{U}},{\bm{V}})\coloneqq\frac{1}{q}\sum_{j=1}^{{N\choose B}}\gamma_{j}{\mathcal{L}}^{\operatorname{con}}({\bm{U}}_{{\mathcal{B}}_{(j)}},{\bm{V}}_{{\mathcal{B}}_{(j)}})$ with $\gamma_{j}={\sum_{l=0}^{q-1}{j-1\choose l}{{N\choose B}-j\choose k-l-1}}/{{{N\choose B}\choose k}}$ for two arbitrary natural numbers $k,q\leq\binom{N}{B}$ where ${\mathcal{B}}_{(j)}$ is a mini-batch with $j$-th largest loss among batches of size $B$. Indeed, this is a natural objective obtained by applying OSGD to our problem, and we show the convergence of such an algorithm by extending the results in Kawaguchi & Lu [29]. OSGD updates the embedding vectors using the gradient averaged over $q$ batches that have the largest losses among randomly chosen $k$ batches (see Algo. 2 in Appendix B.2). Let ${\bm{U}}^{(t)}$, ${\bm{V}}^{(t)}$ be the updated embedding matrices when applying OSGD for $t$ steps starting from ${\bm{U}}^{(0)}$, ${\bm{V}}^{(0)}$, using the learning rate $\eta_{t}$. Then the following theorem, proven in Appendix B.2, holds.

Given sufficiently small learning rate $\eta_{t}\sim O(t^{-1/2}),$ $\mathbb{E}\|\nabla\widetilde{{\mathcal{L}}}^{\operatorname{con}}\|^{2}$ decays at the rate of $\widetilde{O}(T^{-1/2}).$ Therefore, this theorem guarantees the convergence of OSGD for mini-batch contrastive learning.

Applying OSGD to mini-batch contrastive learning has a potential benefit as shown in Sec. 5.1, but it also has some challenges. Choosing the best $q$ batches with high loss in OSGD is only doable after we evaluate losses of all $\binom{N}{B}$ combinations, which is computationally infeasible for large $N$. A naive solution to tackle this challenge is to first randomly choose $k$ batches and then select $q$ high-loss batches among $k$ batches. However, this naive random batch selection method does not guarantee that the chosen $q$ batches are having the highest loss among all $\binom{N}{B}$ candidates. Motivated by these issues of OSGD, we suggest an alternative batch selection method inspired by graph theory. Note that the contrastive loss $\mathcal{L}^{\sf{con}}(U_{{\mathcal{B}}},V_{{\mathcal{B}}})$ for a given batch ${\mathcal{B}}$ is lower bounded as follows:

This lower bound is derived using Jensen’s inequality. Detailed derivation is provided in Appendix C.1. A nice property of this lower bound is that it can be expressed as a summation of terms over a pair $(i,j)$ of samples within batch ${\mathcal{B}}$. Consider a graph ${\mathcal{G}}$ with $N$ nodes, where the weight between node $k$ and $l$ is defined as $w(k,l):=\sum_{(i,j)\in\{(k,l),(l,k)\}}\log\left(1+(B-1)e^{{\bm{u}}_{i}^{\intercal}({\bm{v}}_{j}-{\bm{v}}_{i})}\right)+\log\left(1+(B-1)e^{{\bm{v}}_{i}^{\intercal}({\bm{u}}_{j}-{\bm{u}}_{i})}\right)$. Recall that our goal is to choose $q$ batches having the highest contrastive loss among $N\choose B$ batches. We relax this problem by reducing our search space such that the $q=N/B$ chosen batches ${\mathcal{B}}_{1},\cdots,{\mathcal{B}}_{q}$ form a partition of $N$ samples, i.e., ${\mathcal{B}}_{i}\cap{\mathcal{B}}_{j}=\varnothing$ and $\cup_{i\in[q]}{\mathcal{B}}_{i}=[N]$. In such scenario, our target problem is equivalent to the problem of clustering $N$ nodes in graph ${\mathcal{G}}$ into $q$ clusters with equal size, where the objective is to minimize the sum of weights of inter-cluster edges. This problem is nothing but the min-cut problem [13], and we can employ even-sized spectral clustering algorithm which solves it efficiently. The pseudo-code of our batch selection method²²2Our algorithm finds $N/B$ good clusters at once, instead of only finding a single best cluster. Compared with such alternative approach, our method is (i) more efficient when we update models for multiple iterations, and (ii) guaranteed to load all samples with $N/B$ batches, thus expected to have better convergence [5, 21, 20]. is provided in Algo. 1, and further details of the algorithm are provided in Appendix C. Fig. 3 shows the histogram of contrastive loss for $N/B$ batches chosen by the random batch selection method and the proposed spectral clustering (SC) method. One can observe that the SC method favors batches with larger loss values.

Consider the problem of optimizing the embedding matrices ${\bm{U}},{\bm{V}}$ using GD, where each column of ${\bm{U}},{\bm{V}}$ is initialized as a multivariate normal vector and then normalized as $\lVert{\bm{u}}_{i}\rVert=\lVert{\bm{v}}_{i}\rVert=1$, $\forall i$. We use learning rate $\eta=0.5$, and apply the normalization step at every iteration.

First, we compare the minimizers of three optimization problems: (i) full-batch optimization in Eq.(1); (ii) mini-batch optimization in Eq. (4) with ${\mathcal{S}}_{B}=\left[\binom{N}{B}\right]$; (iii) mini-batch optimization with ${\mathcal{S}}_{B}\subsetneq\left[\binom{N}{B}\right]$. We apply GD algorithm to each problem for $N=8$ and $B=2$, obtain the updated embedding matrices, and then show the heatmap plot of $N\times N$ gram matrix containing all the pairwise inner products ${\bm{u}}_{i}^{\intercal}{\bm{v}}_{j}$ in Fig. 4(b)-(d). Here, we plot for two regimes: $d=2N$ for the top row, and $d=N/2$ for the bottom row. In Fig. 4(a), we plot the gram matrix for the optimal solution obtained in Sec. 4.2. One can observe that when either full-batch or all $\binom{N}{B}$ mini-batches are used for training, the trained embedding vectors reach a simplex ETF and simplex cross-polytope solutions for $d=2N$ and $d=N/2$, respectively, as proved in Thm 4. In contrast, when a strict subset of $\binom{N}{B}$ mini-batches are used for training, these solutions are not achieved.

Second, we compare the convergence speed of three algorithms in mini-batch optimization: (i) OSGD; (ii) the proposed SC method; and (iii) SGD (see details of the algorithms in Appendix C). Fig. 4(e) shows the $\|{\bm{U}}^{\star\intercal}{\bm{V}}^{\star}-{\bm{U}}^{(t)\intercal}{\bm{V}}^{(t)}\|_{F}$ which is the Frobenius norm of the difference between heatmaps of the ground-truth solution (${\bm{U}}^{\star},{\bm{V}}^{\star}$) and the embeddings at each step $t$. We restrict the number of updates for all algorithms, specifically 500 steps. We observe that both OSGD and the proposed method nearly converge to the ground-truth solutions proved in Thm. 4 within 500 steps, while SGD does not. We obtain similar results for other values of $N$ and $d$, given in Appendix D.2.

Here we show that the proposed SC method is effective in more practical settings where the embedding is learned by a parameterized encoder, and can be easily applied to existing uni-modal frameworks, such as SimCLR [9] and SogCLR [68]. We conduct mini-batch contrastive learning on CIFAR-100 and Tiny ImageNet datasets and report the performances in the image retrieval downstream task on corrupted datasets, the results of which are in Table 1. Due to the page limit, we provide detailed experimental information in the Appendix D.3.