Mini-Batch Optimization of Contrastive Loss
Authors: Jaewoong Cho
Kartik Sreenivasan
Keon Lee
Kyunghoo Mun
Soheun Yi
Jeong-Gwan Lee
Anna Lee
Jy-yong Sohn
Dimitris Papailiopoulos
Kangwook Lee
Venue: Accepted by TMLR ()

Abstract: Contrastive learning has gained significant attention as a pre-training method for self-supervised learning due to its ability to leverage large amounts of unlabeled data. A contrastive loss function ensures that embeddings of positive sample pairs (e.g., from the same class or different views of the same data) are similar, while embeddings of negative pairs are dissimilar. However, practical constraints such as large memory requirements make it infeasible to consider all possible positive and negative pairs, leading to the use of mini-batches. In this paper, we investigate the theoretical aspects of mini-batch optimization in contrastive learning with the InfoNCE loss. We show that mini-batch optimization is equivalent to full-batch optimization if and only if all $\binom{N}{B}$ mini-batches are selected, while sub-optimality may arise when examining only a subset. We then demonstrate that utilizing high-loss mini-batches can speed up SGD convergence and propose a spectral clustering-based approach for identifying these high-loss mini-batches. Our experimental results validate our theoretical findings and demonstrate that our proposed algorithm outperforms vanilla SGD, providing a better understanding of mini-batch optimization in contrastive learning.

Published in Transactions on Machine Learning Research (07 (07/2024)
Mini-Batch Optimization of Contrastive Loss
Jacwoong Cho KRAFTON
jucho Tuccho@kratonn.co com
Kartik Sreenivasan" Databricks
Kcon Lee KRAFTON
Strnlee@rrattonn com
Kyunghoo Mun Carnegie Mellon Universily
lamunnandrew.ccom cww. edu
Soheun Yi Carnegie Mellon University
soheuny www.aaawwwwwwww......mmmmmmmmm.. cmaa, eduu
Jeong-Gwan Lee KRAFTON
jiongguon.cee@krrattonccmmm
Anna Lee KRAFTON
ornnalee
Jy-yong Sohn Yonsei University
jysohn1 10sohn1108@@miee 10w@100@.@@mmmocooo
Dimitris Papailiopoulos Universily of Wisconsin-MMadison
dimitris@paapsii.i
Kangwook Lee
kangwook lee
Twisc. tee edu
University of Wisconsin-Madison
KRAFTON
Reviewed OpenReview: https: // opennevviee. net/ por/ff id= Nuz70VXpJJ
Abstract
Contrastive learning has gained significant attention as a pre-training method for self- supervised learning due to its ability to leverage large amounts of unlabeled data. A con- trastive loss function ensures that embeddings of positive sample pairs (e.g., from the same class or different views of the same data) are similar, while embeddings of negative pairs are dissimilar. However, practical constraints such as large memory requirements make infea- sible to consider all possible positive and negative pairs, leading to the use of mini-batches. In this paper, investigate the theoretical aspects of mini-batch optimization in contrastive learning with the InfoNCE loss. We show that mini-batch optimization is equivalent to full- batch optimization if and only if all (M) mini-batches are selected, while sub-optimality may arise when examining only a subset. We then demonstrate that utilizing high-loss mini-batches can speed up SGD convergence and propose spectral clustering- based ap- proach for identifying these high-loss mini-batches Our experimental results validate our theoretical findings and demonstrate that proposed algorithm outperforms vanilla SGD, providing better understanding of mini-batel optimization in contrastive learning.
"Equa Contributions.

---

Published in Transactions on Machine Learning Research (07/2024)
1 Introduction
Contrastive learning is commonly used as a pre-training method for self-supervised learning, due to its ability to leverage the vast amount of freely available unlabeled data (Jaiswal et al., 2020). The rostrastive loss function is designed to ensure that the embeddings of two samples are similar if they are considered a "positive" pair, in cases such as coming from the same class (Khosla et al., 2020); being an augmented version of one another (Chen et al., 2020a), or being two different modalities of the same data (Radford et al., 2021). Conversely, if two samples do not form a positive pair, they are considered a "negative" pair; and the contrastive loss encourages their embeddings to be dissimilar.
In practice, it is not feasible to consider all possible positive and negative pairs when implementing a contrastive learning algorithm due to the quadratic memory requirement. ( when working with N samples. To mitigate this issue of full-batch training, practitioners typically choose a set of mini-batches of size B = (1), and consider the loss computed for positive and negative pairs within each of the N/B batches (Chen et. al., 2022; 2020a; Hu et. al., 2021; Zeng el. al., 2021; Chen et al., 2021; Zolfaghari et al., 2021; Gadre et al., 2023). For instance, Gadre et al. (2023) train model B = 4096. This approach results in memory requirement of O(B²) (11) for each mini-batch, and total computational complexity linear in the number of chosen mini-batches batches. Despite the widespread practical use of mini-batch optimization in contrastive learning, there remains a lack of theoretical understanding as to whether this approach is truly reflective of the original goal of minimizing full-batch contrastive loss. This paper examines the theoretical aspects of optimizing mini-batches loaded for contrastive learning with the InfoNCE loss (Oord et al., 2018). Our study primarily focuses on this loss function due to its fundamental importance in contrastive learning and its analytical tractability.
dataset where N 1.28 x 10" with
Main Contributions. The main contributions of this paper are twofold. First, we show that under certain parameter settings, mini-batch optimization is equivalent to full-batch optimization if and only if all (%) mini-batches are selected. The results are based on an interesting connection between contrastive learning and the neural collapse (Lu & Steinerberger, 2022). From a computational complexity perspective, the identified equivalence condition may be seen as somewhat prohibitive, as it implics that all (+) O(NB) mini-batches must be considered.
Our second contribution, however, shows that Ordered SGD (OSGD) algorithm (Kawaguchi & Lu, 2020) can be effective in finding mini-batches that contain the most informative pairs and thereby speeding up convergence. We show that the convergence result from Kawaguchi & Lu (2020) can be applied directly to contrastive learning. We also show that OSGD can improve the convergence rate of SGD by a constant factor in certain scenarios. Furthermore, in a novel approach to address the challenge of applying OSGD to the (M) mini-batch optimization (which involves examining O(N") batches to select high-loss ones), we reinterpret the batch selection as a min-cut, problem in graph theory. This transformative shift. empowers to select high-loss batches efficiently via spectral clustering algorithm established in this field. The following informal theorems summarize the main findings.
Theorem 1 (informal). Under certain parameter settings, the mini-batch optimization is equivalent to full- batch oplimization if and only if all (+) mini-balches are selected. Allhough (H) mini-balch contrastive loss and full- batch loss are neither identical nor differ by a constant factor, the optimal solutions for both mini-batch and full-batch are identical (see Section 4).
Thcorem 2 (informal). In demonstrative toy example, OSGD operating on the principle of selecting high-loss batches, can polentially converge to the optimal solulion of botch contrastive loss oplimization faster by a conslant factor compared to SGD (see Section 5.1).
We validate our theoretical findings and the efficacy of the proposed spectral clustering- based batch selection method by conducting experiments on both synthetic and real data. On synthetic data, we show that our proposed batch-selection algorithms do indeed converge to the optimal solution of full-batch oplimization sig- nificantly faster than the baselines. We also conduct experiments by pre-training on CIFAR-100 (Krizhevsky et al., 2009) and Tiny ImageNet (Le & Yang, 2015) using the proposed method. We evaluate the performance of the subsequent models on retrieval tasks. Our experiments on real datasets demonstrate that our batch selection method outperforms vanilla. SGD in practically relevant settings.
2

---

Published in Transactions on Machine Learning Research (07/2024)
2 Related Work
Contrastive losses. Contrastive learning has been used for several decades to learn a similarity metric to be used later for verification or recognition applications (Misra & Maaten, 2020; Aberdam et al., 2021). Chopra et al. (2005) proposed one of the early versions of contrastive loss which has been updated and improved over the years (Sohn, 2016; Song & Ermon, 2020; Schroff et al., 2015; Khosla et al.; 2020; Oord et al., 2018). More recently, contrastive learning has been shown to rival and even surpass traditional supervised learning methods, particularly on image classification tasks (Chen et al., 2020b; Bachman et al., 2019). Further, its multi-modal adaptation leverages vast. unstructured data, extending its effectiveness beyond image and text modalities (Radford et al., 2021; Jia et al., 2021; Pham et al., 2021; Ma et al., 2021; Sachidananda et al., 2022; Elizalde et al., 2023; Goel et al., 2022; Lee et al., 2022; Ramesh et al., 2021; 2022). Unfortunately, these methods require extremely large batch sizes in order to perform effectively. Follow-up works showed that using momentum or carefully modifying the augmentation schemes can alleviate this issue to some extent (He et al., 2020; Chen et al., 2020b; Grill et al., 2020; Wang & Qi, 2022).
Effect of batch size. While most successful applications of contrastive learning use large batch sizes (e.g.. 32,768 for CLIP and 8,192 for SimCLR), recent efforts have focused reducing batch sizes and improving convergence rates (Yeh et al., 2022; Chen et al., 2022). Yuan et al. (2022) carefully studied the effect of the requirements on the convergence rate when a model is trained for minimizing SimCLR loss, and proved that the gradient of the solution is bounded by O ( They also propose SogCLR, an algorithm with a modified gradient update where the correction term allows for an improved convergence rate with better dependence on B. It is shown that the performance for small batches can be improved with the technique called hard negative mining (Robinson et al., 2021; Kalantidis et al., 2020; Zhang & Stratos, 2021).
Neural collapse. Neural collapse is phenomenon observed in Papyan et al. (2020) where the final classification layer of deep neural nets collapses to the simplex Equiangular Tight Frame (ETF) when trained well past the point of zero training error (Ji et al., 2022; Zhou et al., 2022). Lu & Steinerberger (2022) proved that this occurs when minimizing cross-entropy (CE) loss over the unit ball. Jiang et al. (2023) generalized neural collapse to practically relevant settings where the number of classes is larger than the dimension of embeddings. Extending beyond CE loss minimization, Graf et al. (2021) delved into the optimal embeddings by minimizing supervised contrastive (SC) loss in balanced datasets and compared them with CE loss minimization. Kini et al. (2023) extended this study by characterizing the optimal embeddings of SC loss minimization combined with ReLU in the last layer, which can learn symmetric embedding structures despite data imbalances, and the sufficient and necessary conditions for mini-batch to achieve the same optimal embeddings of full-batch optimization. Our work explores this problem for the unsupervised rastravvive learning setting. We characterize (1) the optimal embeddings that minimize the InfoNCE loss for any given temperature parameter under certain assumptions, and (2) conditions for mini-batch to achieve the same optimal embeddings of full-batch optimization.
Optimal permutations for SGD, The performance of SGD without replacement under different per- mutations of samples has been well studied in the literature (Bottou, 2009; Recht & Re, 2012; Recht & Ré, 2013; Nagaraj et al., 2019; Ying et al., 2020; Ahn et al., 2020; Rajput al., 2020; Mishchenko et al., 2020; Safran & Shamir, 2021b;a; Gürbüizaalaaan et al., 2021; Nguyen et al., 2021: Lu et al., 2021; Rajput et al. 2022; Tran et al., 2021; Lu et al., 2022; Cha et al., 2023; Cho & Yun, 2023). One can view batch selection in contrastive learning as method to choose specific permutation among the possible (7) mini-batche of size B. However, it is important to note that these bounds do not indicate an improved convergence rate for general non-convex functions and thus would not apply to the contrastive loss, particularly in the setting where the embeddings come from a shared embodding network. We show that in the case of OSGD (Kawaguchi & Lu, 2020), we can indeed prove that contrastive loss satisfies the necessary conditions in order to guarantee convergence.
3

---

Published in Transactions on Machine Learning Research (07/2024)
3 Problem Setting
Suppose we are given datasct {( yi) of N positive pairs (data sample pairs that are conceptually sim- ilar or related), where æ; and Yi are two different views of the same data. Note that this setup includes both the mulli-moda setting (e.g., CLIP (Radford et al., 2021)) and the uni-modal setting (e.g., SimCLR (Chen et al., 2020a)) as follows. For the multi-modal case, one can view (æ,, yi) as two different modalities of the same data, c.g., 2; is the image of scene while y, is the text description of the scene. For the uni-modal case, one can consider a; and yi different augmented images from the same image.
We consider the contrastive learning problem where the goal is to find embedding vectors for {x,) and (3ileen such that the embedding vectors of positive pairs (x; yt) are similar, while ensuring that the embeddings of other (negative) pairs are well separated. Tet ue E Ied be the embedding vector of a; and Vi E R be the embedding vector of yı- In practical settings, one typically considers parameterized encoders so that u, = fa( and v; - to We define embedding matrices U := [ui,...uy and := [I.....; uN] which are the collections of embedding vectors. Now, we focus on the setting of directly optimizing the embeddings instead of model parameters O and in order to gain theoretical insights into learning embeddings. Note that this setting is commonly utilized to achieve deeper understanding of the underlying principles and mechanisms (Graf et al., 2021; Kini et al., 2023). Consider the problem of directly optimizing the embeddings for N pairs which is given by
min U.V
Loon V)
s.t. ||uill = 1, || =1
[
(1)
where ||-I| denotes the & norm, the set [N] denotes the collection of all integers from 1 to N. In this work, we focus on the standard InfoNCE loss (Oord et al., 2018) defined as
ex/ma/ E e//
con (U.)) = N
Σ log
Σ
log
(2)
N
En= e"!
with т being a positive scalar known as temperature parameter. Note that (U, V) the full-batch version of the loss which contrasts all embeddings with each other. However, due the large computational complexity and memory requirements during optimization, practitioners often consider the following mini- batch version instead. Note that there exist (%) different mini-batches, each of which has B samples. For
k E
( let B& be the k-th mini-batch satisfying B& [N] and [Bxl B. Let UB. {u,)EB, and (1: JEB, Then, the contrastive loss for the k-th mini-batch is Loon(UB,) VB,.).
Vs.
4 Optimization for Full-Batch and Mini-Batch
In this section, we investigate the relationship between the problem of optimixing the full-batch loss C V) and the problem of optimizing the mini-batch loss (Uon Vs.). Towards this goal, we prove three main results, the proof which are in Appendix C.1.
We derive the optimal solution that minimizes the full-batch loss (Lemma 1, Theorem 3).
We show that the solution that minimizes the average of (A) mini-batch losses is identical to the one that mininizes the full-batch loss (Proposition 1, Theorem 4).
We show that minimizing the mini-batch loss summed over only strict subset of (M) mini-batches can lead to a sub-optimal solution that does not minimize the full-batch loss (Theorem 5).
4.1 Full-batch Contrastive Loss Optimzation
In this section, we characterize the optimal solution for the ball-batch loss minimization in Equation (1). We start by providing the definition of the simplex equiangular tight frame (ETF) which turns out to be the
4

---

Published in Transactions on Machine Learning Research (07/2024)
optimal solution in certain cases. The original definition of ETF (Sustik et al., 2007) is for N vectors in d-dimensional space where N 2 d Papyan et al. (2020) define the ETF for the case where N S dd+1 to characterize the phenomenon of neural collapse. In our work, we will introduce and use the latter definition of simplex ETFs which is stated below.
Definition 1 (Simplex ETF). We call a set of N vectors fu, form a simplex Equiangular Tight Frame (ETF) if ||ull = 1, Vi E [N] and u,u, = - 1/ 1), Vi + J.
In the following Lemma, we first prove that the optimal solution of full-batch contrastive learning is the simplex ETF for N d which follows almost directly from Lu & Steinerberger (2022).
Lemma 1 (Optimal solution when N S <d 1). Suppose N S 1. Then, the optimal solution (U*, V*) of the fall-batch contrastive learning problem in Equalion (1) salisfies two properties: (i) U V*, and (ti) the columns of U* form simpler ETF.
Recall that in this work, we explore the setting of directly optimizing embeddings an initial effort to gain thcoretical insight. However, Lemma 1 presents the result that can be extended to a more practically relevant setting where embeddings are generated through linear encoders parameterized by 0 and Φz u; = Waa; and v, = WWays. The details of this setting and the characteristics of the optimal encoders are stated below.
Corollary 1. Let us consider a setting where embedding vectors are generated via overparameterized linear encoders Wo. W E Raxdo; u; = Wgg v; Wyys where F, y: E Rdde We define X := [æ.....,, xx and Y = ........NN... Suppose d, do 2 N and XTX,YTY are invertible. Then, the optimal solution (Wo: ) given (i) W = U (XTX) TT; (6i) W LYY where (U+, (*) satisfies the two properties in Lemma 1.
Several practical scenarios satisfy N > d 1. However, the approach used in Lu & Steinerberger (2022) 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. However, the characterization for this case still challenging, so we employ a certain uonjecture; we conjecture that the optimal embeddings (U*, V*) for N 2d are symmetric and antipodal. These properties of embeddings are defined as follows.
Definition (Symmetric and Antipodal). Embedding matrices U and V are called symmetric and antipodal if (U,V) satisfies two properties: (i) Symmetric i.c., U V: (ii) Antipodal i.e., for cach [N], there exists j(i) such that U5(() uu;
Note that the symmetric property also holds for the optimal embeddings when N S d 1, and the antipodal property is a common assumption in geometric problems such as the sphere covering problem in Borodachov (2022). This provides a rationale for our conjecture.
Theorem shows that when N = =2d, the optimal solution for the full-batch loss minimization, under a symmetric and antipodal configuration, form cross-polytope which is defined as the following.
Definition 3 (Simplex pross-polytope). We call set of N vectors { form a simplex cross-polytope for all i, the following three conditions hold: |!u:|| 1; there exists a unique j such that u, u; -1; and u.uu uu for all k {i.j}.
Theorem (Optimal solution when N 2d). Let
(U* V*) == arg min
c (U, V)
(3)
(U.V)EA
s.t. |u.l| = 1, 1. = Vi e [N].
where A :== {U, V): : U, V are symmetric and antipodal). Then, the columns of U* form simplez cross- polytope for N = 2d.
Proof Outline. By the antipodality assumption, we can apply Jensen's 's incquality to N - indices without itself u; and antipodal point -u, for a given E [N]. Then we show that the simplex cross-polytope
See Def. 4 in Appendix for the full definition
5

---

Published in Transactions on Machine Learning Research (07 (2742)
1.9
2
full batch
1.8
mini batch
Train Data Config.
loss
1.7
{ (1,4)
Su 1, 2), [3, 4))
Sa (11,2,3,4)
(2.3): (2. 4), (3.4)
1.6
Some Mini-Batches (B=2)
Full Batch (B=4)
Al Combinations of
(B=2)
1.5
π 1/2
0 4/
Solution
M2
H2
arctan (11,2/ u [radian]
13
(a)
(b)
Figure I: (a) Comparing mini-batch loss and full-batch loss when N = 10, B = 2, and d = 2. We illustrate this by manipulating single embedding vector while maintaining all other embeddings (vı and (u,,))))2)) at their optimal solutions. Specifically, [u1,,, u1.2] is varied as (cos(0), sin( for 0 C [- v.. While the two loss functions are not identical, corroborating Propositionl, their minimizers align, providing empirical support for Theorem 4; (b) The relationship between full-batch and mini-batch optimization in contrastive learning. Consider optimizing N = pairs of d dimensional embedding vectors ((,,, where uy and os shown colored square and circle, respectively. The index written in the square/circle, The black rounded box represents batch. We compare three batch selection options:
full batch, f.c., B 4, (ii) all (G) = mini-batches with size B = 2, and (iii) some mini-batches. Here, Se is the set of mini-batches where each mini- batch is represented by the set of constituent samples' indices. Our theoretical/empirical findings are: the optimal embedding that minimizes full-batch loss and the one that minimizes the sum of (ž) mini-batch losses are identical, while the one that minimizes the mini-batch losses summed over only strict subset of (g) batches does not guarantee the negative correlation between u, and u, for i o j. This illustration supported by our mathematical results in Thms, 4 and 5.
minimizes this lower bound while satisfying the conditions that make the applications of Jensen's inequality tight. Seo Appendix for the full proof.
For the general case of N d d+11 1, excluding N =2d, we still leave it as an open problem.
4.2 Mini-batch Loss Optimization
Here we consider the mini-batch contrastive loss optirmization problem, where we first choose multiple mini- batches of size B and then find U, 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-bateh loss is intractable in practice. Let consider subset of all possible (š) mini-batches and
denote their indices by Sn C ( For fixed Sn, the mini-batch optimization problem formulated as:
min
Coon (U, V;SB)
(4)
U.V
s.t.
||2.||| 1, || 1 Vi [NE,
where the loss of given mini- batche is C ( s Eieso (Con VB:). To analyze the rela- tionship betweon the full-batch loss minimization in Equation (1) and the mini-batch loss minimization in Equation (4), we first compare the objective functions two problems as below.
Proposition 1. The mini-batch loss and full-batch loss are not identical, nor is one a simple scaling of the other by a constant factor. In other words, when SB (G) for all B 2 2, there erists no constant such that Coon (UV:SBB) = c. Loon (U.V) for all U
We visualize this proposition in Figure 1(a). Interestingly, the following result shows that the optimal solutions of both problems are identical.
Theorem 4 (Optimization with all possible
NN mini-batches). Suppose B 2 2. The set of minimizers of [
the (N) mini-batch problem in Equation (4) is the same that of the full-batch problem in Equation (1)

---

Published in Transactions on Machine Learning Research (07/2024)
for two coses: (t) N d 1, and (ti) N = 2d and the pairs (U, V) restricted to those satisfying the conditions stated in Def. 2. In such cases, the solutions (U, V) for the (N) mini-batch oplimization problem satisfies the following: Case (i) {u: forms a simplex ETF and u, = v; for all i E [N]] Case (ti): { forms a simplez eross-polytope.
(B)
Proof Outline. Similar to the proof of Lemma 1, we bound the objective function from below using Jensen's 's inequality. We then show that this lower bound is equivalent to a scaling of the bound from the proof of Lemma 1 by employing non-trivial counting arguments. Furthermore, we demonstrate that both the simplex ETF and the simplex cross-polytope minimize this lower bound for each case while satisfying the conditions that make the ioplications of Jensen's 's inequality tight. See Appendix C.1 for the detailed proof.
Now, we present mathematical results specifying the cases when the solutions of mini-batch optimization and butl-batch optimization differ. First, we show that when B 2, minimizing the mini-batch loss over any strict subset. of (g) batches, not equivalent to minimizing the full-batch loss.
Theorem (Optimization with fewer than (+) mini-batches). Suppose B and N <d++. Then, the minómizer of Equation (4) for SB & (() is not the minimizer of the full-batch optimization in Equation (1).
Proof Outline. Wc show that there exist embedding vectors that are not the simplex ETF, and have strictly lower objective value. This implies that the optimal solution of any set of mini-batches that does not contain all (2) mini-batches is not the same as that of the full-batch problem. See Appendix C.1 for the detailed proof.
The result of Theorem 5 is extended to the general case of B 2, under some mild assumption; please check Proposition 2 and 3 in Appendix C.I.
Figure 1 summarizes the main findings in this section. These findings suggest that mini-batch optimization can achieve the optimal embeddings of full-batch optimization, as long as all (g) possible mini-batches are considered. However, this poses a computational challenge, as involves loading all these mini-batches. In the following section, we introduce efficient algorithms for speeding up the (%) mini-batch optimization.
5 Ordered SGD for Mini-Batch Contrastive Learning
Recall that the optimal embeddings for the full-batch optimization problem in Equation (1) can be obtained by minimizing the sum of (A) mini-batch losses. An easy way of achieving the optimal embeddings is using gradient descent (GD) on the sum of losses for (N) mini-batches, or using stochastic approach that applies GD on the loss for randomly chosen mini-batch. Recent works found that applying GD on selective batches outperforms SGD in some cases (Kawaguchi & Tu,, 2020; Tu et al., 2021; Loshchilov & Hutter, 2015). A natural question is, whether this holds for mini-batch contrastive learning, i.e., (i) Is SGD enough to guarantee good convergence on selection method that performs better than vanilla SGD? To answer this question:
contrastive learning?, and (ii) Can we come up with a batch
We show that Ordered SGD (OSGD) (Kawaguchi & Lu, 2020) potentially accelerate convergence compared to vanilla SGD in demonstrative toy example (Section 5.1), and extend the convergence results from Kawaguchi & Lu (2020) to mini-batch contrastive loss optimization (Section 5.2).
We reformulate the batch selection problem into a min-cut problem in graph theory (Cormen et al., 2022), by considering a graph with nodes where each node is each positive pair and each edge represents a proxy to the contrastive loss between two nodes. This allows us to devise an efficient batch selection algorithm by leveraging a spectral clustering algorithm (Ng et al., 2001) (Section 5.3).
5.1 Convergence Comparison in a Toy Example: OSGD vs. SGD
This section investigates the convergence of two gradient- desprttdeeeeetttt-eee based methods, OSGD and SGD. The below lemma shows that the contrastive loss is geodesically non-quasi-convex, which implics the hardness of proving the convergence to the optimal embeddings of gradient-based methods for contrastive learning in Equation (1). The proof is provided in Appendix C.2.
7

---

Published in Transactions on Machine Learning Research (07 /024)
1.6
OSGD
= (- cos 0g, sin 04)
(cos 01, sin 01)
Contrastive
1.5
SCDD
choose (full)
04
1.4
0g
1.3
ug = (- cos 03. - sin 0z)
U2 (cos 02. - -sin 82)
25
50
73
100
125
150
175
200
Update steps
(a)
(b)
Figure 2: (a) 'oy example considered in Section 5.1; (b) The training loss curves of three algorithms (OSGD, SGD, and (#) full-batch gradient descent) applicd on the toy example when and B The x-axis represents the number of update steps, while the y-axis displays the loss Equation (2). OSGD converges the fastest among the three methods.
Lemma 2. Contrastive loss Loon V) a geodesically non-quasi-conwer function of U, V on T (U,V) ||ull |elll 1, E [N]]
In order to compare the convergence of OSGD and SGD to the optimal embeddings, we focus toy example where convergence to the optimal solution is achievable with appropriate initialization. Consider scenario where we have N = 4 embedding vectors {u with u, E R Each embedding vector is defined as u1 = (cos 1, sin 01); Uz = (cos 0g. sin 02); U3 = (- cos 0g. sin 0g); un = (- cos 0g, sin 04) for parameters (0.). For all i, the initial parameters are set as 0
0, and the other embedding vectors are initialized as vi u: This setting is illustrated in Figure 2(a). Over time step i, we consider updating the parameters (()) = [0() ( ( using the gradient descent based method on (UB, V8) with a learning rate yy ((++1) O(1) Vo (Uste, VB(.). At each time step 4, each learning algorithm begins by selecting mini-batch B(O c (1, 4} with batch size [B(*) 2. SGD randomly selects a mini-batch while OSGD selects a mini-batch as follows: B = arg max Leon (U)( (()( For a sufficiently small BES
margin p O, let TOSGD,TSca be the minimal time required for the algorithms to reach the condition that each component of the expected value B[0(2)] falls within (T/4 p,4/
TOSGD min T 2 0: E OSGD] (T/4 p, π //4
TSaD == min 2 0 E[0 saD. (T/4 0, /, J.
where (1) OSED and 0 SCDD denote the model parameters at step t when OSGD and SGD are used, respectively. 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.
Theorem 6. Consider the described setting where the parameters (( of embedding vectors are updated, shown in Pigure 2(b). Suppose there exist é, T such that for all i salisfying (()) {1,3} {2, 4), o Cron (Ugo, VB )l < and TOSGD: TSGD < T. Then, we have the following inequalities:
TOSGD T 7/4 - e O nE")
TSGD > 2/ -1 14----+0(ppe+
3(22/T +1)
m/c
T/e + Onje tzē)
Corollary 2. Suppose lower bounds of TOsGD, TSGD in Theorem 6 are tight, and the learning rate is small enough. Then, TOSGD/Tsan 11)3/- 2/ +1)<1. <.
In Figure 2(b), we present training loss curves of the full-batch contrastive loss in Equation (2) for various algorithms implemented on the toy example. One can observe that the loss of all algorithms eventually converges 1.253, the optimal loss achievable when the solution satisfics u, = v, and (u form simplex cross-polytope. Note that the optimality of this solution is given in Theorem 4. This empirical evidence corroborates our theoretical findings, indicating that OSGD converges faster than SGD.
8

---

Published in Transactions on Machine Learning Research (07 (204))
of batches aatche
epoch
epoch
epoch
epoch 101
200
200
200
200
random
spectral clustering
100
100
100
100
2
0.0
0.5
1.0
0.0
0.5
1.0
contrastive loss
Figure Histograms of batch counts for N/B batches, for the contrastive loss measured from ResNet-18 18 models trained CTFAR-100 using SGD, where N=50,000 and R=20. Each plot is derived from a distinct training epoch. Here compare two batch selection methods: (i) random shuffling N samples and partition them into N/B batches of size B, (ii) our SC method given in Algorithm 1. The histograms show that batches generated through the proposed spectral clustering method tend to contain a higher proportion of large loss values when compared to random batch selection. Similar results are observed in different scttings, details of which are given in Appendix E.1.
5.2 Convergence of OSGD in Mini-batch Contrastive Learning Setting
Recall that it is challenging to prove the convergence of gradkent-deeccen-baaed methods to the optimal em- beddings for contrastive learning problems in Equation (1) due to the non-quasi-coovveiity of the contrastive loss _eou, Instead, we show that OSGD minimizes modification of the empirical contrastive loss goon that considers the order of the batch losses, and is guaranteed to converge at a sublinear rate to a critical point the loss by extending the results in Kawaguchi & Lu (2020). We consider a proxy, the weighted contrastive
(an loss defined as con (U. E com (UB VB,,) with % ΣE (7) (7))))))) ( for two ar-
bitrary natural numbers ki, q S ( where Búy is mini-batch with j-th largest loss among batches of size B. Indeed, this is natural objective obtained by applying OSGD to our problem. OSGD updates the embeddings using the gradient averaged over q batches that have the largest losses among randomly chosen & batches (sce Algorithm 2 in Appendix C.2). Let U(0), v()) be the updated embedding matrices when applying OSGD for steps starting from U(O),, V using the learning rate rp- Then the following theorem, proven Appendix C.2, holds.
Theorem 7 (Convergence results). Consider sampling t" from [T] with probability proportional to (m TT that is, P(L" = 1) nr/ tnt Then up> > po v8/ B7 2/ /Br2, /c2/T/Br²
[Zcon(110) v(0)) Zoon ) + 8p FT of
(p po)2 pp) plp Po)
E v Foon (Ultt))
Σt= 7t
where x denotes the minimized value of ccon
Given sufficiently small learning rate T)t 0(121/2,, E/ coon decays at the rate of 1(T-1/2).
5.3 Spectral Clustering based Approach
Applying OSGD to mini-batch contrastive learning has a potential benefit as shown in Soction 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 (*) combinations. This process requires the computational complexity of O(N log log N which is computationally infcasible for large N. A naive solution to tackle this challenge is to first randomly choose k batches and then select high-loss batches among k batches. However, this naive random batch selection method does not guarantee that the chosen q batches have the highest loss among all (N) candidates.
Motivated by these issues of OSGD, we relax this problem by reducing our scarch space such that the q N/B chosen batches B.... BB, form a partition of N samples, í.e., B, B, = Ø and Use = [N]. We suggest an alternative batch selection method inspired by graph theory. The mini-batch contrastive loss
9

---

Published in Transactions on Machine Learning Research (07/2024)
Algorithm 1: Spectral Clustering Method
Input: the number of positive pairs N, mini-batch size B, embedding matrices: U, V
Output: selected mini-batches 1=,
N/B
j( Equation (5)
if #
Construct the affinity matrix A:
Aij
else
fi
Construct the degree matrix D from A: Dij
Ly Ag
else
L< D A; < N/B
Find the largest cigenvuctors of L, denoted as V& C NNN Normalize the rows of V& to have unit E2-norm
Apply the K-means clustering algorithm on the rows of the normalized Ve to get cluster centers Z C RAxk Construct bipartite graph Sesignn (i) the first partite set is V& and (ii) the second set is the collection of B copics of each center in Z
Compute distances between vectors of V2 and B copies of each center in Z, and assign these as edge weights in Fassign Solve the minimum weight matching problem in Gussign using a method such as the Hungarian algorithu
10 return {B,)
N/B
L (UOR VB) for given batch B is lower bounded using Jenson's incquality as follows:
2 (UB, VB)
1
log
wn//T
Σ
log
cn!uu/T
B
N E cu/os/T
N Σド=ı en/us/t
I
log
1+
Σ
Σ log
1+
Σ
(us--uuuil/
B
B
JEBi(i)
iEB
JEBI(()
Σ log 1 (B (ee wil/t)
+ log (1+ (B- 1)e (0e")(1)-uul) (u)))
B(B 1)
iEB (1)
We employ this lower bound due to nice property that can be expressed as a summation of losses over a pair (i, of samples within batch B. We consider graph 9 with N nodes and the weight wokk between node k and l is defined
w(k, :== ==
E (L)))(())))
log(1 (B (B 1)e" + log(1 + (B (T (w)-
(5)
Recall that our goal is to choose batches having the highest contrastive losses among (7) batches. In such a scenario, our larget problem is equivalent to the problem of clustering N nodes in graph G into clusters with equal size, where the objective to minimize the sum weights of inter-cluster edges. This problem is the min-cut problem (Cormen et al., 2022). To approximate the solution this problem basod on the graph G, we propose a variant of the well-known spectral clustering algorithm (Ng et al., 2001). Since the spectral clustering algorithm cannot guarantee to assign an equal number of nodes to each cluster, we introduce an additional step to ensure that each cluster (batch) has the equal number B of positive pairs. This step addresses a minimum weight matching problem in a bipartite graph (Crouse, 2016) involving N data points and replicated cluster centers from spectral clustering. One partite set is the collection of data points and the other set consists of B copies of each cluster center. The aim is to allocate exactly B data points to each center, minimizing overall cost, defined as the total distance from data points to their respective centers. This ensures each cluster receives an equal number of data points while minimizing the assignment cost. The pseudo-code of our batch selection 2 is provided in Algorithm 1. The computational complexity of our spectral clustering (SC) method is O(k2 O(k2B²N) while the complexity of OSGD is O(N log ND) due to the necessity of sorting (M) mini-batcher We provide a detailed complexity comparison of the algorithms in Appendix F. Figure 3 shows that the proposed SC method indeed favors batches with larger loss values.
?Our algorithm finds N/B clusters at once, instcad of only finding single best cluster. Compared with such an alternative approach, our method (1) more efficient when we update models for multiple iterations, and (ii) guaranteed load samples with N/B batches, thus expected have better convergence (Bottou, 2009; Haochen Sra, 2019; Gürbünba laban al., 2021).
10

---

Published in Transactions on Machine Learning Research (07 2224)
2.5
Gatens
4.D 23
1.5
180
N.D
stops
(a) optimal
(b) full-batch
(c) (2)
(d)
(g) suub
(e) norm difference
Figure The behavior of embedding matrices U, optimized by different batch selection methods for N 8 and (a)-(d): Heatmap of
matrix visualizing the pairwise inner products u uy where (a): ground-truth solution (ETF 2N, cross-polytope for d N/2), (b): optimized the full-batch loss with GD, (c): optimized the of (G) mini-batch losses with GD, (d): optimized partial sum of (R) mini-batch losses with GD. Note that both (b) and (c) reaches the ground-truth solution in (a), while (d) does not, supporting our theoretical results in Section 4.2. Further, (e) compares the convergence of three mini-batch selection algorithms: 1) SGD, 2) OSGD, and 3) our spectral clustering method, when updating embeddings for 500 steps. OSGD and method nearly converge to the optimal solution, while SGD does not. Here, y-axis represents the Frobenius norm of the difference between the heatmaps of the optimal solution and the updated embeddings. denoted | U..
6 Experiments
We validate our theoretical findings and the effectiveness of the proposed method by providing experimental results synthetic and real datasets. We first show that experimental results on synthetic dataset coincides with two main theoretical results: (i) the relationship between the full-batch contrastive loss and the mini- batch contrastive loss given in Section 4, (ii) the analysis on the convergence of OSGD and the proposed SC method given in Section 5. To demonstrate the practicality of our method, we provide experimental results on CIFAR-100 (Krizhovsky et al., 2009) and Tiny ImageNet (Lc & Yang, 2015). Details of the com/
experimental setting can be found in Appendix E, and our code is available at krafton Paafoon-aimiin htt.//iinibatchh-h...
6.1 Synthetic Dataset
Consider optimizing the embedding matrices U, V using GD, where each column of U, V is initialized as multivariate normal vector and then normalized as |uu ||m = 1, Vi. We use learning rate 0.5, and apply the normalization step at every iteration.
First, we compare the minimizers of three optimization problems: (i) full- batch optimization in Equation (1); (ii) mini-batch optimization in Equation (4) with SB = [G)]; (iii) mini-batch optimization with SB & [())]- We apply GD to each problem for N and B 2, obtain the updated embedding matrices, and then show the heatmap plot of N x N gram matrix containing all the pairwise inner products ulv, in Figure 4(b)- (d). One can observe that when cither full-batch or all (M) 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 Theorem 4. In contrast, when strict subset of (N) mini-batcher is used for training, these solutions are not achieved.
Second, we compare the convergence speed of three algorithms: (i) the proposed SC method; (ii) OSGD which updates the model using the top-1 batch that exhibits the highest loss from among the (N) possible batches each step (sce Algorithm 5); and (iii) SGD (sce Algorithm 4). Figure 4(c) shows that both OSGD and the proposed method nearly converge to the ground-truth solutions proved Theorem within 500 steps, while SGD does not. We obtain similar results for other values of N and d, given in Appendix F.2.
11

---

Published in Transactions on Machine Learning Research (07/2024)
Table 1: Top-1 retrieval accuracy on CTFAR- 100 (or Tiny ImageNet)), when each algorithm employs one of the following datasets to pretrain ResNet-18 with SimCLR and SogCLR objective: (1) CIFAR-100, (2) Tiny ImageNet, (3) CIFAR-100-C, or (4) Tiny ImageNet-C. SC algorithm proposed Section 5.3 outperforms existing baselines.
Image Retrieval
CIFAR-100
Tiny ImageNet
CIFAR-100-C
Tiny ImageNet-C
SimCLR
SogCLR
SimCLR
SogCLR
SimCLR
SogCLR
SimCLR
SogCLR
SGD
70.29
60.66
78.39
67.48
38.48
38.34
42.42
45.04
OSGD (approximated)
70.00
61.48
78.57
69.31
38.55
38.71
43.11
45.63
SC
74.71
73.32
82.95
79.73
38.99
41.51
43.45
46.81
6.2 Real Datasets
In this section, we show that the proposed SC method is effective in more practical settings where the em- bedding is learned by parameterized encoder, and can be easily applied to existing uni-modal frameworks, such as SimCLR (Chen et al., 2020a) and SogCLR (Yuan et al., 2022). We integrate different batch selection methods from (i) the proposed SC method; (ii) OSGD without replacement (see Algorithm 6), an approxi- mated OSGD, selected due to the intractability of evaluating all (+) batches on real datasets; and (iii) SGD (see Algorithm 4) into these frameworks. We compare the pre- trained models' performances in the retrieval downstream tasks.
We conduct mini-batch contrastive learning with the mini-batch size B = 32 using ResNet 18 babedd encoders on CIFAR-100 and Tiny ImageNet datasets. All learning is executed on a single NVIDIA A100 GPU. The training code and hyperparameters are based on the official codebase of SogCLR* (Yuan et al., 2022). We use LARS optimizer (You et al., 2017) with the momentum of 0.9 and the weight decay of 10-6, We utilize the learning rate scheduler which starts with warm-up phase in the initial 10 epochs, during which the learning rate increases linearly to the maximun value mmax 0.075 After this warm-up stage, we employ a cosine annealing shall-cycle) schedule for the remaining epochs. For the approximated OSGD, we employ k = 1500, = 150. To expedite batch selection in the proposed SC, we begin by randomly partitioning N positive pairs into kB-sized clusters, using A = 40. We then apply the SC method to each kB cluster to generate k mini-batches, resulting a lotal of x (N/kB) = N/B mini-batches. We train models for total of 100 epochs. We measure the performances of the models under the retrieval task which is defined as finding the positive pair image of given image among all pairs (the number of images of the validation dataset): (1) We extract normalized embedding features from augmented images by employing the pre-trained models; (2) We identify the pair image of the given augmented image among all of the validation images with the cosine similarity of embedding vectors.
We also measure performances of the model pre-trained on CIFAR-10 (or Tiny ImageNet) under a harder setting, where the various corruptions are applied per image so that we can consider a set of corrupted images as hard negative samples. We employ CIFAR-100-C (or Tiny ImageNet-CC which is the corrupted dataset (Hendrycks & Dieuterich, 2019) designed for robustness evaluation. CIFAR-100-C (Tiny ImageNet- C) has the same images as CIFAR-100 (Tiny ImageNet), but these images have been altered by 19 (15) different types of corruption (c.g., image noise, blur, etc.). Each type of corruption has five severity levels. We utilize images corrupted at severity level 1. These images tend to be more similar to each other than those corrupted at higher severity levels, which consequently makes more challenging to retrieve positive pairs among other images. To perform the retrieval task, we follow the following procedures: (i) We apply two distinct augmentations to each image to generate positive pairs; (ii) We extract normalized embedding features from augmented images by employing the pre-trained models; (iii) we identify the pair image of the given augmented image among augmentations of 19 (15) corrupted images with the cosine similarity of embedding vectors. This process is iterated across 10K CIFAR-100 images (10K Tiny-ImageNer images). The top-1 1 accuracy measures a percentage of retrieved images that match its positive pair image, where each pair contains two different modalitios stemming from a single image. Table presents the top-1 retrieval accuracy on CIFAR-100, Tiny ImageNet, and their corrupted counterparts. The proposed SC method
12

---

Published in Transactions on Machine Learning Research (07/2024)
surpasses SGD and the approximated OSGD. We show the efficacy of the SC method, yet it has limitation in that the proposed method requires more time for batch selection compared to SGD. To mitigate this issue in practice, we suggest parallel batch selection using the SC method while updating models through gradient descent. This approach involves using the same batches for multiple epochs as the next batch selection proceeds, based on the assumption that batches with higher loss values do not rapidly change. Exploring this approach is beyond the scope of this work so we leave it as a future work.
7 Conclusion
We theoretically analyzed mini-batch contrastive learning with the InfoNCE loss. First, we showed that the solution of mini-batch optimization and that of full-batch optimization are identical if and only if all (g) mini-batches are considered. Second, we analyzed the convergence of OSGD and devised spectral clustering (SC) method, a new batch selection method which handles the complexity issue of OSGD in mini-batch contrastive learning. Experimental results support our theoretical findings and the superiority of SC. Our future work includes (1) characterizing the optimal embeddings for N d + case, (2) generalizing findings to broader contrastive losses, and (3) exploring a complexity-efficien batch selection method.
References
Aviad Aberdam, Ron Litman, Shahar Tsiper, Oron Anschel, Ron Slossberg, Shai Mazor, R Manmatha, and Pietro Perona. Sequence-to-sequence contrastive learning for text recognition. In Proceedings of the IEEE/CVP Conference on Computer Vision and Pottern Recognition, pp. 15302 15312, 2021.
Kwangjun Ahn, Chulhoe Yun, and Suvrit Sra. Sgd with shuffling: optimal rates without component convexity and large epoch requirements. Advunces in Newral Information Processing Systems, 33:17526 14117526-77733 17535, 2020.
Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximixing mutual information across views. Advances in neural information processing systems, 32, 2019.
Sergiy Borodachov. Optimal antipodal configuration of 2d points on a sphere in Rd for covering. arXiv preprinl Xiv: 12472, 2022.
Léon Bottou. Curiously fast convergence of some stochastic gradient descent algorithms. In Proceedings of the symposium on learning and data science, Paris, volume 8, pp. 2624 2633, 2009.
Jacyoung Cha, Jacwook Lee, and Chulhec Yun. Tighter lower bounds for shuffling sgd: Random permutations and beyond, 2023.
Changyou Chen, Jianyi Zhang, Yi Xu, Liqun Chen, Jiali Duan, Yiran Chen, Son Dinh Tran, Belinda Zeng, and Trishul Chilimbi. Why do we need large batchsizes in contrastive learning? a gradient-bias perspective. In Alice H. Oh, Alckh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
Hao Chen, Benoit Lagadee, and Francois Bremond. Ice: Inter-instance contrastive encoding for unsupervised person re-identification. In Proceedings of the TEEE/CVF International Conference Computer Vision, pp. 14960 14969, 2021.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations, In International conference on machine learning, pp. 1597-1607. PMLR, 2020a.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXi::203..042 2020b.
Hanseul Cho and Chulhee Yun. Sgda with shuffling: faster convergence for nonconvex-pl minimax optimiza- tion, 2023.
13

---

Published in Transactions on Machine Learning Research (07/2024)
Sumit Chopra, Raia Hadsell, and Yann LeCun. Learning similarity metric discriminatively, with applica- tion to face verification. In 2005 JEEE Compuler Society Conference Compuler Vision and Pattern Recognition (CVPR'05) volume 1, pp. 539 546. LEEE, 2005.
Thomas Cormen, Charles E Leiserson, Ronald L Rivest, and Clifford Stein. Introduction to algorithms.
MIT press, 2022.
David F Crouse. On implementing 2d rectangular assignment algorithms. IEEE Transactions on Aerospace and Electronic Systems, 52(4):1679-15 1696, 2016.
Damek Davis and Dmitriy Drusvyatskiy. Stochastic model-based minimization of weakly convex functions. STAM Journal on Optimization, (9(1)2207-2338 2019.
Benjamin Elizalde, Soham Deshmuuh, Mahmoud Al Ismail, and Huaming Wang. Clap learning audio concepts from natural language supervision. In ICASSP 2023 2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1 5, 2023.
Samir Yitzhak Gadre, Gabriel Ilharco, Alex Fang, Jonathan Hayase, Georgios Smyrnis, Thao Nguyen, Ryan Marten, Mitchell Wortsman, Dhruba Ghosh, Jieyu Zhang, et al. Datacomp: In search of the next gener- ation of multimodal datasets. arXic preprint M.Xx:23004.44000 2023.
Shashank Goel, Hritik Bansal, Sumit Bhatia, Ryan Rossi, Vishwa Vinay, and Aditya Grover. Cyclip: Cyclic contrastive language-image pretraining. In Advances in Newral Information Processing Systems, volume 35, pp. 6704 6719. Curran Associates, Inc., 2022.
Florian Graf, Christoph Hofer, Mare Niethammer, and Roland Kwitt. Dissecting supervised contrastive learning. In International Conference on Machine Learning, pp. 3821-3834 PMI.R, 2021.
lean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Docrsch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al. Bootstrap your own latent-a new approach to self-supervised learning, Advances in neural information processing systems, 33:21271-2122 2020.
M. Gürbüzzaalban, A. Ozdaglar, and A. Parrilo. Why random reshullling beats stochastic gradient descent. Mathematical Programming, 186(1):49 84, 2021.
Jeff Haochen and Suvrit Sra. Random shuffling beats sgd after finite epochs. In International Conference on Machine Learning, pp. 2624 2633. PMLR, 2019.
Kaiming Hc, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Monentum contrast for unsupervisod visual representation learning. In Proceedings of the IEEE/CVP conference on computer vision and pallern recognition, pp. 9729 9738, 2020.
Dan Hendrycks and Thomas G. Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In 7th International Conference on Learning Representations, 1CLR 2019, 2019.
Qianjiang Hu, Xiao Wang, Wei Hu, and Guo-Jun Qi. Adco: Adversarial contrast for efficient learning of unsupervised representations from self-trained negative adversaries. In Proceedings of the IEEE/CVF Conference Compuler Vision and Pallern Recognilion, pp. 1083, 2021.
Ashish Jaiswal, Ashwin Ramesh Babu, Mohammad Zaki Zadeh, Debapriya Banerjee, and Fillia Makedon. A survey on contrastive self-supervised learning. Technologies, 9(1):2, 2020.
Wenlong Ji, Yiping Lu, Yiliang Zhang, Zhun Deng, and Weijie J Su. An unconstrained layer-peeled per- spective on neural collapse. In International Conference on Learning Representations, 2022.
Chao Jia, Yinfei Yang, Ye Xia, Yi-Ting Chen, Zarana Parckh, Hieu Pham, Quoe Le, Yun-Hsuan Sung, Zhen Li, and Tom Duerig. Scaling up visual and vision-language representation learning with noisy text supervision. In International Conference Machine Learning, pp. 4904-4916. PMI.R, 2021.
14

---

Published in Transactions on Machine Learning Research (07/2024)
Jiachen Jiang, Jinxin Zhou, Peng Wang, Qing Qu, Dustin Mixon, Chong You, and Zhihui Zhu. Generalized neural collapse for large number of classes. orv preprint arXiv:3310.05511, 2023.
Yannis Kalantidis, Mert Bulent Sariyildiz, Noe Pion, Philippe Weinzaepfel, and Diane Larlus. Hard negative mixing for contrastive learning. Advances in Neural Information Processing Systems, 33:21798-211 -2880,, 2020.
Kenji Kawaguchi and Haihao Lu. Ordered sgd: A new stochastic optimization framework for empirical risk minimization. In International Conference on Artificial Intelligence and Statistics, pp. 669-679. PMLR, 2020.
Prannay Khosla, Piotr Teterwak, Chen Wang, Aaron Sarna, Yonglong Tian, Phillip Isola, Aaron Maschinot, Ce Liu, and Dilip Krishnan. Supervised contrastive learning, Advances in Neural Information Processing Systems, 33:18661-18673, 2020.
Ganesh Ramachandra Kini, Vala Vakilian, Tina Behnia, Jaidev Gill, and Christos Thrampoulidis. Suppeviised-oonrrastive loss learns orthogonal frames and batching matters. arXiv preprint arXiv:2306.07960 2023.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning multiple layers of features from tiny images. 2009.
Ya Le and Xuan Yang. Tiny imagenet visual recognition challenge. CS 291N, 7(7):3, 2015.
Janghyeon Lee, Jongsuk Kim, Hyounguk Shon, Bumsoo Kim, Seung Hwan Kim, Honglak Lee, and Junmo Kim. UniCLIP: Unified framework for contrastive languago-image pre-training. In Advances in Neural Information. Processing Systems, 2022.
Nya Loshchilov and Frank Hutter. Online batch selection for faster training of neural networks. arXiv preprint 3XXi::511.003/33 2015.
Jianfeng Lu and Stefan Steinerberger. Neural collapse under cross-entropy loss. Applied and Computational Harmonic Analysis, 59:224-241 241, 2022. ISSN 1063-5203. Special Issue on Harmonic Analysis and Machine Learning.
Yucheng Lu, Si Yi Meng, and Christopher De Sa. A general analysis of example-selection for stochastic gradient descent. In International Conference on Learning Representations, 2021.
Yucheng Lu, Wentao Guo, and Christopher De Sa. Grab: Finding provably better data permutations than random reshuffling. In Advances in Neural Information Processing Systems, 2022.
Shuang Ma, Zhaoyang Zeng, Daniel McDuff, and Yale Song. Active contrastive lcarning of andio-visual vidco representations. In International Conference on Learning Representations, 2021.
Konslantin Mishchenko, Ahmed Khaled, and Peter Richtárík. Random reshullling: Simple analysis with
vast improvements. Advances in Neural Information Processing Systems, 33:17309 17320, 2020.
Ishan Misra and Laurens van der Maaten. Self-supervised learning of pretext-invariant representations. In Proceedings of the IEEE/CW Conference on Computer Vision and Pattern Recognition, pp. 6707-6717, 2020.
Dheeraj Nagaraj, Prateek Jain, and Pranceth Netrapalli. Sgd without replacement: Sharper rates for general smooth convex functions. In International Conference on Machine Learning, pp. 4703 4711. PMLR, 2019.
Andrew Ng, Michael Jordan, and Yair Weiss. On spectral clustering: Analysis and an algorithm. Advances in neural information processing systems, 14, 2001.
Lam M Nguyen, Quoc Tran-Dinh, Dzung T Phan, Phuong Ha Nguyen, and Marten Van Dijk. A unificd convergence analysis for shuffling-typ type gradient methods. The Journal of Machine Learning Research, 22 (1):9397-944) 2021.
15

---

Published in Transactions on Machine Learning Research (07/2024)
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arxie: 1tt:1807.00 180.003788 2018.
Vardan Papyan, XY Han, and David L Donoho. Prevalence of neural collapse during the terminal phase of doep learning training. Proceedings of the National Academy of Sciences, 117(40):246522225 246663, 2020.
Hicu Pham, Zihang Dai, Golnaz Ghiasi, Hanxiao Liu, Adams Wei Yu, Minh-Thang Luong, Mingxing Tan, and Quoc V Le. Combined scaling for zero-shot transfer learning. arXiv prepriut arXiv:2111.10050 2021.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference Machine Learning, pp. 8748-8763. PMLR, 2021.
Shashank Rajput, Anant Gupta, and Dimitris Papailiopoulos. Closing the convergence gap of sgd without replacement. In International Conference on Machine Learning, pp. 7964-7973. PMLR, 2020.
Shashank Rajput, Kangwook Lee, and Dimitris Papailioponlos. Pormutation-bass SGD: Is random optimal? International Conference on Learning Represenlations, 2022.
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Πya Sutskever. Zero-shot text-to-image generation. In International Conference on Machine Learning, pp. 8821-8831. PMLR, 2021.
Adilya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. Hierarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204006195 2022.
Benjamin Recht and Christopher Re. Toward a noncommutative arithmetic-geomeeric mean inequality: Conjectures, casc-studics, and consoquences. In Proceedings of the 25th Annual Conference on Learning Theory, volume of Proceedings of Machine Learning Research, pp. 11.1 11.24. PMLR, 2012.
Benjamin Recht and Christopher Ré. Parallel stochastic gradient algorithms for large-scale matrix comple- tion. Mathematical Programming Computation, 5(2):201-22) 226, 2013.
Joshua David Robinson, Ching-Yat Yao Chuang, Suvrit Sra, and Stefanic Jegelka. Contrastive learning with hard negative samples. In International Conference Learning Represenlations, 2021.
Vin Sachidananda, Shao-Yor Yen Tseng, Erik Marchi, Sachin Kajarekar, and Panayiotis Georgiou. Calm: Contrastive aligned audio-language multirate and multimodal representations. arXiv preprint arXiv:2202.03587 2022.
Itay Safran and Ohad Shamir. How good is sgd with random shuffling?, 2021a.
Itay Safran and Ohad Shamir. Random shuffling beats sgd only after many epochs on ill-conditioned prob- lems. Advonces in Neural Information Processing Systems, 34:15151-15111, 2021b.
Florian Schroff, Dmitry Kalenichenko, and James Philbin. Facenet: A unified embedding for face recognition and clustering. In Proceedings of the IEEE conference computer vision and pattern recognition, pp. 815-823, 2015.
Kihyuk Sohn. Improved deep metric learning with multi-class n-pair loss objective. Advances in neural information processing systems, 29, 2016.
Jiaming Song and Stefano Ermon. Understanding the limitations of variational mutual information estima- tors. In International Conference on Learning Representations, 2020.
Mátyás A Sustik, Joel A Tropp, Inderjit Dhillon, and Robert W Heath Jr. On the existence of equiangular tight frames. Linear Algebra and its applications, 426(2-3):619 635, 2007.
16

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Trang H Tran, Lam M Nguyen, and Quoc Tran-Dinh. Smg: A shuffling gradient-based method with mo- mentum. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 10879-103330 PMLR, 2021.
Xiao Wang and Guo-Jun Qi. Contrastive learning with stronger augmentations. IEEE Transactions Pullern Analysis and Machine Intelligence, 2022.
Chun- Ych, Cheng- Yao Hong, Yen-Chi Hsu, Tyng-Luh Liu, Yubei Chen, and Yann LeCun. Docoupled contrastive learning. In European Conference on Compuler Vision, pp- 668 684. Springer, 2022.
Bicheng Ying, Kun Yuan, and Ali H. Sayed. Variance-reduced stochastic learning under random reshullling. IEEE Transactions on Signal Processing, 1013390-1408 2020. doi: N..109/TSSP2020.2222
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprinl 1rXiv:1708.000 2017.
Zhuoning Yuan, Yuexin Wu, Zi-Hao Qiu, Xianzhi Du, Lijun Zhang, Denny Zhou, and Tianbao Yang. Prov- able stochastic optimization for global contrastive learning: Small batch does not harm performance. In International Conference on Machine Learning, pp. 25760 25782. PMLR, 2022.
Dewen Zeng, Yawen Wu, Xinrong Hu, Xiaowei Xu, Haiyun Yuan, Meiping Huang, Jian Zhuang, Jingtong Hu, and Yiyu Shi. Positional contrastive learning for volumetric medical image segmentation. In Medical Image Computing and Computer Assisted Intervention-MICCAA 2021: 24th International Conference, Strasbourg, France, September 27 October 1, 2021, Proceedings, Part II 24, pp. 221 230. Springer, 2021.
Wenzheng Zhang and Karl Stratos. Understanding hard negatives in noise contrastive estimation. North American Chapter of the Association for Computational Linguistics, 2021.
Jinxin Zhou, Xiao Li, Tianyu Ding, Chong You, Qing Qu, and Zhihui Zhu. On the optimization landscape of neural collapse under mse loss: Global optimality with unconstrained features. In International Conference on Machine Learning, pp. 27179-27200 27202. PMLR, 2022.
Mohammadreza Zolfaghari, Yi Zhu, Peter Gehler, and Thomas Brox Crosselr: Cross-modal contrastive learning for multi-modal video representations. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 1450-1459, 2021.
17

---

Published in Transactions on Machine Learning Research (07/2024)
Organization of the Appendix
1. Appendix A outlines the limitations of this paper.
2. In Appendix B introduce some additional definitions for posterity.
In Appendix C we provide detailed proofs of the results from the main paper as well as any intermediate results/lemmas that found useful.
(a) Specifically, in Appendix C.1 we provide proofs of the results from Section 4 which focuses on the relationship between the optimal solutions for minimizing the mini-batch and full-batch constrastive loss.
(b) Appendix 5 contains the proofs of results from Section C.2 which concern the application of Ordered SGD to mini-batch contrastive learning.
(c) Appendix C.3 is intended to supplement Appendix 5 in that it contains auxiliary notation and proofs required in the proof of Theorem 7.
4. Appendix D specifics the pscudo-code and details for the two algorithms: (i) Stochastic Gradient Descent and (ii) Ordered SGD.
5. Appendix F. provide some additional results.
6. In Appendix F we provide computational complexities for the three algorithms.
A Limitations
We note that theoretical results have two major limitations:
1. While we would like to extend our results to the general case of N > d 1+1, we were only able to characterize the optimal solution for the specific case of N = 2d. Furthermore, our result for the case of = 2d in Theorem 4 requires the use of the conjecture that the optimal solution is symmetric and antipodal. However, as mentioned by Lu Steinerberger (2022), the general case of d + 1 seems quite challenging in the non-asymptotic regime.
2. In practice, the embeddings are usually the output of a shared neural network encoder. However, our results are for the case when the embeddings only have a norm constraint. Thus, our results do not readily indicate any generalization to unseen data. We expect however, that it is possible extend our results to the shared encoder setting by assuming sufficient overparamelerization..
B Additional Definitions
Definition 4 (Sustik et al. (2007)). A set of N vectors [N in the Id form an equiangular tight frame (ETF) if (i) they are all unit norm: |w,|l = 1 for every E [N], (ii) they are equiangular: [|uTu,l =a 2 0 for all i * and some 2 O, and (iii) they form a tight frame: UUT= = (N/d)La where U is a d x N matrix whose columns are u1 u2...., uN, and La is the x d identity matrix.
C Proofs
C.1 Proofs of results from Section 4
Lemma 1 (Optimal solution when S +d+ 1). Suppose < 1 1. Then, the optimal solution (U*, V*) of the full-batch contrastive learning problem in Equation (1) satisfies two properties: (i) U* V*, and (ti) the columns of U form simplex ETF.
18

---

Published in Transactions on Machine Learning Research (07 (22/))
Proof. First, define the contrastive loss the sum of two symmetric one-sided contrastive loss terms to simplify the notation. We denote the following term as the one-sided contrastive loss
eu/u//T Σ j== E n//T
C(U, N
log
(6)
cut
Then, the overall contrastive loss given by the sum of the two one-sided contrastive losses:
n (U. = LUU V) + +VV, U).
(7)
Since fron is symmetric in its arguments, results pertaining the optimum of C(U, V) readily extend C Now, let us consider the simpler problem of minimizing the one-sided contrastive loss from Equation (6)
which reduces the problem to exactly the same sctting as Lu & Steinerberger (2022):
eu/v//
C)
log
N
N 2j=1
E4T
1
log
1+
(;)-
j=1,jşi
Note that, we have for any fixed 1 1 < N,
N
N Σ
Σ
(o)=v;)Tu::+
= (17u//T
j=
-N 1)e (v) u:/T)
Σ clult
N
j=1,zæi
(a)
(N 1)e
(1) (1/T) exp
1 Σ uuulT
N 1ipe
(
(N 1)e - (vT/:/)
vTu - viuu 4
exp
T(N 1)
vTu; N(uT.u u:) T(N
(N - exp
(8)
where (a) follows by applying Jensen inequality for e' and (b) follows from v = Em U. Since log(-) is monotonic, we have that > + log(x) > log(y) and therefore,
N
utu;
NoOT u:) T(N 1) )
L(U. V) 2
log 1 + (N 1)
N
(c)
oTu;
N(u. u.)
2 log
1 1 (N - exp
N
7(N
-
T(N 1)
1=1
1
TTu
N
log
1 (N - exp
Σ (o) mi)
(9)
N
-1
r(N 1)
where (c) follows by applying Jensen incquality to the convex function <(t) log(1 + at for a, > and
(d) follow from u = Σmm u,
Note that for equalities to hold in Equation (8) and (9), we need constants 6, e such that
1. u; = G3
Vj + i v vTu: Nov. u.)
(10)
Vi [N].
(11)
N
19

---

Published in Transactions on Machine Learning Research (07/2024)
Since log(-) and exp(:) are both monotonic, minimizing the lower bound in Equation (8) is equivalent to
N
TTu
u
min
Ui
N
N
x=1
N
N
max
N
u.uu
Vi
T
(12)
==1
1=1
All that remains is show that the solution that maximizes Eq 12 also satisfies the conditions in Equa- tion (10) and (11). To see this, first note that the maximization problem can be written
max
v Scack((NNT - 1x1]) @ Ua)
where Vatack = (V1, 02 ......, Un) a vector in RNd formed by stacking the vectors v; together. Ustack is similarly defined. In denotes the N N identity matrix, 1, denotes the all-one vector IK",, and & denotes /N since each ||u/|| = ||o:|| 1. Since
the Kronecker product. It is easy to see that ||tssack|| || |||stack the eigenvalues of A are the product of the eigenvalues of A and B, in order to analyze the spectrum of the middle term in the above maximization problem, it suffices to just consider the eigenvalues of (Nix - 1 As shown by the elegant analysis in Lu & Steinerberger (2022), (NIN - 1y 1 ) Np for any p E R. such for any q E E such that q = k1v for some R. Therefore it follows that its eigenvalues are N with multiplicity (N - 1) and 0. Since its largest eigenvalue is N and since || || = |sttckl| vN, applying cauchy schwarz inequality, we have that
that Σi= P; and (NIx 1v 1x13)q
sack (NIN 1, 1 @ Ig) stack
max
||sstckl [I(NI, 1, @ Ia)l [stack
VN(NNVN N².
Moreover, that setting u, v; and setting {u: to be the simplex ETF attains the maximum above while also satisfying the conditions in Equation (10) and (11) with = -1/(N 1) and e = -N/(N - 1). Therefore, the inequalities in Equation (8) and (9) are actually equalities for u; v; when they are chosen to be the simplex ETF in Id which attainable since d Z N - 1. Therefore, we have shown that if U== = (ut), 1 is the simplex ETF and u = of Vi [N], then U arg miny, L(U, V) over the U
unit sphere in R". All that remains is to show that this also the minimizer for gron
First note that U*. V is also the minimizer for L(V, U) through symmetry. One repeat the proof exactly by simply exchanging u, and v, to see that this is indeed true. Now recalling Equation (7), we have
min coon min (C(U, V) + ((U,
min (C(U, V)) min (C(U, V))
(13)
C(U', VV) + C(V*,U")
However, since the minimizer of both terms in Equation (13) is the same, the inequality becomes an equality. Therefore, we have shown that (U+, V*) is the minimizer of _con completing the proof.
Remark 1. In the proof the above Lemma, we only show that the simplex ETF attains the minimum loss in Equation (1), but not that is the only minimizer. The proof of Lu Steinerberger (2022) can be extended to show that this is indeed true well. We omit it here for ease of exposition.
Theorem (Optimal solution when N 2d). Let
UU+ V*) == arg min
con (U, V)
(3)
(U.V)
s.t. |u.ll = 1,11 =1 Vi [N].
where A:= := {(U, V) : U. V are symmetric and antipodal). Then, the columns of U* form simplez cross- polytope for N = 2d.
20

---

Published in Transactions on Machine Learning Research (07/2024)
Proof. By applying the logarithmic property that allows division to be represented as subtraction,
1
v//T
C(U
log
N
b eu/Us/T
u. uTvu/T r- log
Σ (ulos/T)
Since U V (symmetric property), the contrastive loss satisfies
L (U, V) 2C(U,U)
N
N
2
log Σ cufus(T)
N
1=1
2
2
log Σee culus/t).
(14)
N
Since
|ul| = for any E [M], we can derive the following relations:
uu u
uu 1 2uJu, Uj, u, =
2
We incorporate these relations into Equation (14) as follows:
N log
2
2
- uu 11/11277
coom (U.)) v) =
N
j=1
2
log
N
The antipodal property of U indicates that for each E [N], there exists j(i) such that 2y()) m By
applying this property, we can manipulate the summation of 1151//11/22 over j as the following:
N Σ
e | 1uu
1//2 = I/u -M:1//2///2 + 1--1111 u,,;; +
Σ
- I|u 1llu:uy113/4 12/1//2
1=1
= =1 +
c -1,111,11//2-
($1((1
Therefore,
N
con (U.) V)
Σ log 1 -2/7 2/7 + Σ c elllu,u,,11/2
N
km
j+1,(()
(a)
2
N
log (1 +e 2/T
2// + (N - 2)e -1/u 1u,,0/11//77
N(N 2)
1=1 jfi)((()
N
2
Σ
log (1 te 2/T + (N - 2)e 1lui 11l:- 11:1////
2
log 1 + (N 1)e 1)e-2/1
N(N 2)
N 2
1=1 jæY
N
2
Σ
Σ
log (1 + 22/ + (N 2)- 1lue -u;l1/22))
2
log + (N -1)e -2r))
N(N- - 2)
N -
1=1
21

---

Published in Transactions on Machine Learning Research (07/)024)
e- + t); and (b) follows by Lemma 3, and the fact that function s()) log[1 + e e-2/T + (N 2)2- is convex and monotonically decreasing, {ui,..,,, , denotes a set of vectors which forms cross-polytope
where (a) follows by applying Jensen's inequality to the concave function f(t) log(1 + e 2/7
Both inequalities in (a) and (b) are equalities only when the columns of U form a pross-polytope. Therefore, the columns of U form a pross-polytope.
Lemma 3. Given a function g(t) is convex and monotomically decreasing, let
U² arg min UEA
E
all - -nll3) s.l. ||ui|| 1, [|v/l1 =1 Vi [N],
(15)
where A = (U UU U is antipodal). Then, the columns of U* form simplez cross-polytope for N 2d.
Proof. Suppose N 2d and U E A. Given a function g(l) is convex and monotonically decreasing. j(i) denotes the corresponding index for such that (y(()) Ui and l|u uyca)l 4. Under these conditions, we derive the following:
s(Ilu: =
Σ
g(1lee - -2,113)
(a)
Ng(4) + N(N 2)g
1
Σ
[uu - u,)1²)
N(N - 2)
i=1 jtti((t
N N
= =No(4) + -NNN - -2)g( NNNN - 2)
1
(-4N +
Σ
[u; - 2, 112)
/=1 j=1
NN
Ng((4) N(N 2)99 ((NN -2)
1
-4N Σ
Σ
(2 2n (p))
/=1 j=1
No(44) + N(N - -200
1
-4N + 2N - ||
Σ
u.
119)
N(N -2)
1=1
(6) 2 Ng(4) + -N(N- ((N-2)9 2)9 N(N - 2) (44N+2N2
Ng(4) + N(N 2)g(2).
where (a) follows by Jensen's inequality; and (b) follows from the fact that | u112 2 0 and the function g(t) is monotonically decreasing. The equality conditions for (a) and (b) only hold when the columns of U form a cross-polytope. We can conclude that the columns of U form a cross polytope.
Proposition 1. The mini-bated loss and full-batch loss are not identical, nor is one a simple scaling of the other by a constant factor. In other words, when SB ( for all B 2 2, there erists no constont e such
that fcon (U, V: Sn) Loon .V) for all U.V.
Proof. Consider Ü, V defined such that u; = 0, = e, Vi E [N], where e; i-th unit vector in R., First note that a, = Vi E [N] and ūā 63 = Vi # j. Then,
1/T 1/T +N-1) 1//r c(Ü v- log(e
(16)
1
C(ÜB, Vo.) = log(e +B 1) 1/T.
(17)
N
B
22

---

Published in Transactions on Machine Learning Research (07 (07/2024)
We now consider the second part of the statement. For contradiction, assume that there exists some e e R such that milll (U, V; Loon (U, VV for all U.V. Let U, V be defined such that á; 0; e vi E [N], where e = (1, (1,0,... (1,0,--,,0) ,). Note that O = Vi,j E [N]. Then,
C(Ù, V) = log(IV),
(18)
1
C(ÜB. VB,) log (B).
(19)
From Equation (16) and 17, we have that logfe N-10-1--11)1--- Whereas from Equation (18) and (19),
we have that c ())(B) which is contradiction. Therefore, there exists no e E satisfying the given condition.
(N)
Theorem (Optimization with all possible (x) mini-batches). Suppose B 2 2. The set of minimizers of the (g) mini-batch problem in Equation (4) is the same as that of the full-batch problem in Equation (1) for two cases: (i) N S d 1, and (ii) N = 2d and the pairs (U, V) are restricted to those satisfying the conditions stated in Def. In such cases, the solutions (U, V) for the (ø) miniibutch oplimization problem satisfies the following: Case (i) fuili forms a simplex ETF and u, v; for all E [N]; Case (ti): {ui) forms a simplex cross-potylope.
Proof. Case (i): Suppose N <d+1.
For simplicity, first consider just one of the two terms in the two-sided loss. Therefore, the optimization problem becomes
min U.V
L(UB, VE,) s.l. |uu||| 1, || Vi E [N].
Similar to the proof of Lemma 1, we have that
C(UB. V8,)
log
1+ Σ 4 (vL
JEB.
kEB, kfj
(a)
Σ EkEB, u1 Uÿ)
log
1 (B - -xx
B
T(B 1)
1=1 JEB
1
EKEB, (ujen Un Buy
log
1+ I+ - (B-1) exp
exp
B
T(B 11-1)
JEB:
(b)
N
Di
EjeB, EREB, u.us Σ Ei=1 E Ejes, Buju,
log
1 (B - 1) exp
B
(N) B
where (a) and (b) follows by applying Jensen's inequality to et and log(1 (1 + ae for a, b O, respectively. Note that for equalities to hold in Jensen's 's inequalities, we need constants Cjs such that
u uJr p4 = Vk #
(20)
u uTv. Nu,
=c Vi [N].
(21)
N N
23

---

Published in Transactions on Machine Learning Research (07 2024)
Now, carefully consider the two terms in the numerator:
4 = A
A2 =
Buy
X JEB REB;
JEB,
N-2 batches that contain To simplify A1, first note that for any fixed 1,m [N] such that 1 m, there are B-2 and m. And for / = m, there are (m=1) batches that contain that pair. Since these terms all occur in A we have that
N -
N
(B-1) (M--1
N -W-2) 2)
-2
A1
Σ
u Um +
Σ u
B 2)
2
N 2
N N
(N-2)
(N B
N
Σ
Σ
u
um Vm +
Σ lml
u 1
B 2)
B -
B - 1-10
1ml m=1
Similarly, we have that
(
N -
A
B BB-17 B
u v1
Plugging these back into the above inequality, we have that
N
Σ E1=1 ΣM umun N NE ut ufd )
C(UB,, Vo.)
log 1 (B--) exp
B
+N(N -
N
uTv N N 1=1 u.Vv
log
1 [B (B-1)exxp
NN( 1
Observe that the term inside the exponential is identical to Equation (9) and therefore, we can reuse the same spectral analysis argument to show that the simplex ETF also minimizes Σ L(UB. Vs.). Once again, since the proof is symmetric the simplex ETF also minimizes Σ L(VB,, UB.).
Case (ii): Suppose N 2d, and U, V are symmetric and antipodal. Next, consider the following optimization problem
min (U.))EA
con (UB,, VB.) s.t. || = ||/l = Vi E [N],
(22)
where A (UUV)) U, V symmetric and antipodal). Since U
V
(symmetric property) the con-
trastive loss salies
Con (UB, VB.) = 2C(UB, UB.)
= =
Σ
uJu//T log
Σ
jEB,
KEB.
2
+
2
Σ
log Σ
(23)
T
B
jeB,
REB;
24

---

Published in Transactions on Machine Learning Research (07/2024)
Therefore, the solution of the optimization problem in Equation (22) is identical to the minimizer of the following optimization problem:
U = arg min
log
Σ
jEB
REB
The objective of the optimization problem can be rewritten by reorganizing summations as
Σ
log
Σ eajue/t)
(24)
j=I
REB.
where I; : j :i:
Bi} represents the set of batch indices containing j. We then divide the summation term in Equation (24) into two terms:
N
N
N
log Σ
eufuu/t
log Σ eufus/t)
log Σ eufuu/t).
(25)
j=1 SEZ,
REB.
j=1 iEA,
kEB
J=1
j=1
REB.
by partitioning the set Ij for each j E [N] into the following with k()) being the index for which (k()
Uÿ
= = fi:je B, and k() E BJ:: = [i:je B, and k(j) & B.).
We will prove that the columns of U form a cross polystoo by showing that the minimizer each term of the RHS Equation (25) also forms a cross-polytope. Let us start with the first term of the RHS Equation (25). Starting with applying Jensen's 's inequality the concave function f(x) le 11/ 2), get:
log Σ e cupultt
log 21/ +e 41/1
Σ
enJui/T)
IEA,
KEB;
5=1 EEA,
kEBAGGRG)
Σ
log (6'/T 11/= + (B- 2)c23ue/1
B-2 2
j=1 iEA, KEB, (AAAD)
N -3 log (e +e-1/1 + +BB -2)cmJux//
Σ
Σ
B-2
B-3
1=1 (E(A)()
(N-9)
log (el/T 11/1 + +(B -2)e4Jul/7) - N log (B (BB- 1/)-
+
B-2
J=1 k
(N-9)
log (et/T 11- + +B - 2))e//r . 10
Nlog (6t/7 + + - 118-11e-10 -1)))
B 2
(a)
NN
N
llu 11gl- 27
log 2/ +e -1/7 + (B- (B-2)e2/1
+ +(B - 1)e N Nlog (
1//)]
B
where (a) follows by Lemma and the fact that g(t) log(a+be - for a, is convex and monotonically decreasing. [ui.--. uxy denotes of vectors which forms cross-polytope. All equalities hold only when the columns of U form cross-polytope.
25

---

Published in Transactions on Machine Learning Research (07 (07/2024)
Next consider the second term of the RHS in Equation (25). By following similar procedure above, we get:
N
N
log
Σ
e
2 B
log (e"/+ +(B (B 1)e"3uu/1)
iE
BEB
j=l iEA, REB:
Σ
E
N N-33
log (e"/T + 1)c")u//7)
B
B B
j=1 (E(iA(S))
3N-3
B-2
log (e +(B 1)e4/uz/t) - N log (e ( -B-1)e- 1)-
B
J=1 ドマデ
NN-3
log (e!/T + (B-1)et/ 11ug -111- 25
B-2
N log (E!/r + +(B -B 118-11e-111 1)e-1/1)
B
where {ui,-.. ) denotes set of vectors which forms cross-polytope.
Both terms of RHS in Equation (25) have the minimum value when U forms a cross-polytope. Therefore, we can conclude that the columns of U form a cross-polytope.
Theorem (Oplimization with lewer than (g) mini-batchess). Suppose B and S dd dd+1. Then, the minimizer of Equation (4) for SB & [x) is not the minimizer of the full-batch optimization in Equation (1).
Proof. Consider a set of batches SB c that (1, & USes [B.]. For contradiction, assume that the simplex ETF (U*, V*) is indeed the optimal solution of the loss over these Sn batches. Then, by definition, have that for any (U, # (U*, V*).
(*) with the batch size 2. Without loss of generality, assume
Σ C(UB, VB,)
Σ L(UB,, VB,) JESn
SB
SB
iESp
C(U VE, S L(UB. V8.),
(26)
icSB
iESp
where (U*, V*) is defined such that ut of for all e [N] and u To -1/(N - 1) for all i + j. Also recall that |u,|| ||v,|| 1 for all i e [N]. Therefore, we also have
Σ C(UB, Vg,)= 1ESg
ΣΣ log Σ exp -
JESo jEB,
KEB
Σ log
1
Σ exp (-T(N-1))
-1) +)
iESp JEB,
KEB
1 ( )
Σ
Σ
log
1 1
-
(27)
JEB
where the last equality is due to the fact that |B,|
Now, let us consider (Ü, V) defined such that ú; o; for all i E [N], and u. = - 2) for all
* j. (i,j) (21,2) (2, 1)}. Intuitively, this equivalent to placing ūz.... ūN on simplex ETF of N -1
points and setting û U2 This is clearly possible because d > - d d> N 22, which is the condition
26

---

Published in Transactions on Machine Learning Research (07 (224))
required to place N points simplex ETF in R Therefore,
Σ
L(UB, VB,)
Σ log
1+ Σ exp (ūj -0;1/7
iESn
JEB,
kEB
Σ exp (TT(N- 2) +)
1
log
1
iESp jEB,
KEB
Σ SESp JEB,
Σ log 1 + cxp
T(N-2) (-
1
(28)
where the last equality follows since (1,2) & USes (B:). It is easy to from Equation (27) and (28) that Eies, L(UB;, Vs,) < Eies, B,U V&,) which contradicts Equation (26). Therefore, the optimal solution of minimizing the contrastive loss over any SB proof.
[+) batches is not the simplex ETF completing the
7
Proposition 2. Suppose B 2 2, and let SB S ( be set of mini- indices. If there erist two data points that never belong logether in any mini-balch, i.e., Fi,j E [N] s.l. {i,j} & BR for all E SB, then the optimal solution of Equation (4) is not the minimizer of the full-batch problem in Equation (1).
Proof. The proof follows in fairly similar manner to that of Theorem 5. Consider set of batches of size B 2, SB C [.)]. Without loss of generality, assume that {1,2} & Bą for any k E SB. For contradiction, assume that the simplex ETF (U* V*) is the optimal solution of the loss over these S3 batches. Then, by definition, we have that for any (U, V) + (U*,V*)
Once again, for contradiction assume that the simplex ETF (U*, V*) is indeed the optimal solution of the loss over these Sn batches. Then, by definition for any (U, V) + (U", VV),
1
Σ L(U
Vg,)
1
Σ L(UB,, VB,)
SB
SB
iESs
L(U Vig,) S Σ LUU
VB.)
(29)
iESB
iSSp
where (U+, V*) is defined such that u² = of for all i [N] and u T = -1/(N - 1) for all + j. Also
rocall that [ull = ||0;|| for all E [N]. Therefore, we also have
Σ C(UB, V$,) RESB
ΣΣ
log
1+ Σ exp [njj(o--)))1 wz)/T)
B
iESn JEB;
kEB,,
T(N 11 (- +1
B Σ
Σ
log
Σ exp
1E8n JEB,
KEB
Σ
log
1 (B - 1)
T(N - 1) -
(30)
B
iESp JEB;
Now, let us consider (Ü, V) defined such that म; = 0, for all i e [M] ū = 8, and ūā - 1/(N -2) for all i (t,j)) & ((12),, (2, 1)}. Once again, note that this is equivalent to placing ū2,..., uN simplex
27

---

Published in Transactions on Machine Learning Research (07 (07/2024)
ETF of N 1 points and setting ú₁ dg. Hence,
Σ iESn
C(UB, VB,)
1 Σ log
Σ iESA jEB,
1+
Σ exp (ū) -0,)88)
B
kEB,Jys
T(N T(N--2 -2) -1
Σ log
1
Σ exp
B
iESp JEB;
kEB
E
log
1+ (B - -B
(31)
B
T(N -)
IESB JEB;
where for the final equality note that following. The only pair for which új n2 + - 1/N is (j, (1, 2).
Since there is no E SB such that {1,2] E B;, this term never appears in our loss. From Equation (30)
and Equation (31), we have that Σiese [(UB, VB,) < Eiesa L(U Vg. which contradicts Equation (29).
Therefore, we conclude that. the optimal solution of the contrastive loss over any Sn the simplex ETF.
[>) batches is not
Proposition 3. Suppose 2 2, and let SB S [(x) be a set of mini-batch inidices satisfying Ba Bj 2, Vi.j E SB and Ules B, = [N], i.e., (Bi)ieSa forms non-overlapping mini-batches that cover all data samples. Then, the minimizer of the mini-batch loss optimization problem in Equation (4) different from the minimizer of the full-batch loss oplimizalion problem in Equation (1).
Proof. Applying Lemma specifically to single batch B₁ gives that the optimal solution for just the loss over this batch is the simplex ETF over B points. In the case of non-overlapping batches, the objective function can be separated across batches and therefore the optimal solution for the sum of the losses is equal to the solution of minimizing each term independently. More precisely, we have
N/B Σ
N/B Σ min Leon (CD, VB.),
min U.V
ccnn (UB. VB.)
Us Ve.
where UB. = {uy : E B,} and VB. = {u, j E B.}. respectively, and the equality follows from the fact that Bis are disjoint.
C.2 Proofs of Results From Section 5
Lemma 2. Contrastive loss Loon (U, V) is a geodesically non-quasi-conwer function of U,V on T =
(U,V) : [|u:]] = -[vi|| = Vi [N]].
Proof. Without loss of generality, we focus on the case where 7 = 1. The contrastive loss function coon is geodesically quasi-convex for any two points (U, V) and (U", V') in the domain and for all t in [0, 1]:
_ (tU V) + (1 -1 UU v')) < max( (U, V), L (U', V')).
We provide a counter-cxample for geodesically quasi-convexity, which is triplet of points (UL, V!), (U2 (U2, (U's V2) where (U's v3) is on the geodesic between other two points and satisfies ( v3) max foon (U1 V1), Loon (U2,, V2) Let N = and
1-= = V
U2
V
V
Now, define U normalizo (UL + U2)/2) and V = normalizo(( + V2) which is the "midpoint" of the geodesic between (U', v\) and (U², V2). By direct calculation, we obtain f CU3, V x 2.798 > 2.773 max(Ccon(UU) Vt), (U2, V2)), which indicates Leon is geodesically non-quasi-convvx.
28

---

Published in Transactions on Machine Learning Research (07/2024)
Theorem 8 (Theorem 6 restated). Consider N 4 samples and their embedding vectors {u,) fo with dimension d = 2. Suppose u; 's are parametrized by O(t) = [0(t) 0 -00h) ,04' as in the setting described in Section 5.1 (see Figure 2(a)). Consider initializing u; (0) of (0) and 0 = > 0 for all i, then updoting
o( via OSGD and SGD with the batch size B = 2 described in Section 5.1. Let TOSGD, TSGD be the minimal time required for OSGD, SGD algorithm to have E[O(T)] E (T/4- p. p/4/4)". Suppose there exist E, T such that for all satisfying B(0) = {1,3} or (2, 4), Vee (UBB, BBD) ) S ?, and TOSGD, TSGD < T Then,
TOSGD& 2T x/4 -/4-p-10 c O11 --
n
3(e2/T +1), x/4-p- p- c+ 0 722)
TsaD 2 22/T -1
m/ + One) + +1²)
Proof. We begin with the proof of
TOSGD 2 7/4- - p + (17 + n)"))
7/6
Assume that the parameters are initialized at (01) 0 0 with the batch size (r 1):
(,, Then, there are six batches 2, and we can categorize the batches according to the mini batch contrastive loss
1. B= [1.2] or [3, 4}: Lon VB) -2/7 2ogg(e cns 2r/T);
2. (1.3) or 42: L Coon (UB, VB) = -2/7 + + - 11/T);
3. B = [1. or (2. c (UB, VB) = -2/ loggee + -e- c0002///1)
Without loss of generality, we assume that OSGD algorithm described in Algorithm 5 chooses the mini-batch B = [1,2] corresponding to the highest loss at time = 0, and updates the parameter as
0f = Vo, Loon (UB, V8), 0 = V L (UB, VB).
Then, for the next update, OSGD choose uz un which is now closer than updated U1 u2 And u3 us would be updated as same what previously u1, un have changed. Thus, O1 updates only at the even time, and stays at the odd time, i.e.
100) V Leon (UB, VB) if is even, (
o((+1)
if t is odd.
Iterating this procedure, we can view OSGD algorithm as onramaraettrrize algorithm of parameter g( e(2t) as:
1 =, 610) = 1(t-1) + g(o(t-1)), (TTnau) E c -
π 7
where ((Ø) = 2 (in())) + ell c(1-0a((0))))), and Thatt = TOSGD/2. In the procedure of updates, we may assume that ((() E (0, 3) for all t. To analyze the drift of (!)), we firstly study smoothness of
g'()) 4 cos 20(1 + c(l-coo20)/T) 20)/T) sin? 20 ell co-cos 20)/r
(1+ co1-coo20)//73
We can observe that max [g((") = 2/T, hence g( has Lipschitz constant 2/T, i.e. 6
(((dit-11) - - s (t(-1) -000
29

---

Published in Transactions on Machine Learning Research (07 2024)
Therefore,
(( 0(--1) [1
|]lo(c)] - (0
2n ) + O(nc²),
where the first inequality is from Tipschitz-continiity of g( and the second equality is from Taylor expan-
sion of g at e 0 as;
4r 6
s(c) c 3r
c
Hence, 5( S (1 20/T) + One)) indicates that
(lthar)
< (1 (hi/t) T O(1)²)
S +
+ O1 + mee)
for some constant T >> TOSGD. Moreover $ - p < (Taat))) implies that
Thaif T 7/4-p- p- O(ne + n2²c) 2
n/c
So, we obtain the lower bound of TOSGD by doubling Thats We estimate of TOSGD-
Now, we study convergence rate of SGD algorithm. We claim that
3(62/T + +1), s/A-p--+0gp(e- TSGD 2 e2/r -
+ 0
Without loss of gencrality, we firstly focus on the drift of 0j. Since batch selection is random, given ( (0(2 (()) (()) )
1. B = (1,2) with probability 1/6. Then, coo (UB. V8) = 227 r + 2log(e +
implics
(0) + +0
e(++1) = o(9) + 1) + (1 coss(( +0 1)/T
2. = (1,3} with probability 1/6. At t = 0, the initial batch selection can be primarily categorized into three distinct sets; closely positioned vectors {u), U2} or {u3, u4}, vectors that form obtuse angles {u1,20 41} or {uz, us}; and vectors diametrically opposed
180", {u1 uz} or (uz, Given that 6 is substantially small, the possibility of consistently selecting batches from the same category for subsequent updates is relatively low. As such, it is reasonable to infer that each batch is likely to maintain ils position within the initially assigned calegories. From this, one can deduce that vector sets such as {uI, u3] or (ug,us) continue to sustain an angle close to 180", Given these conditions, it is fcasible to postulate that the selected batch B encompasses either (1,3) or {2,4], the magnitude of the gradient of the loss function (oon VB), denoted by ||V Vis]l,, would be less than a particular threshold ē, i.e.
[ VB)il <
Then,
0((+1)) = 0(() + r(()).
3. B = {1,4) with probability 1/6. Then, coo (UB. V8) = -2/7 + 2log(e +e- c00(0)+0))/)) implics
2 (0) + +0(2)/T
((++1) 0(t)) " 1 (0 1//" D)/
30

---

Published in Transactions on Machine Learning Research (07/2024)
Since there is no update on 81 for the other cases, taking expectation yields
E E (+)
8/01 (
? (i O(zē)
6
where Fi( is defined as:
OO 2 sin + 02)/T
2 ssin (1 + +A)/TT
+ e(t-coo(f,++2))1
1 + elL+ccoo(( +0,))//+
We study smoothness of F1 by setting Fi(0)= = (- + - [ + +0,); where
2sint/T
S+(1) = 2 2sin 1/T
S-(1)= = 1 1 ell-cuut)/+
1 e(L+oot))/
Note that
max [5-(0) = 1/TT,
max [0
[1) = CC/T,
1E (2 /2)
for some constant C' E (0, Then for 0 = (01, 03 01), 0- = (0%, 02 O333 01),
(F'(0') (O) s [ff + 82) - f0 + 02)1 + +If+ (0' + 01) 1f f+1 + 04)l
< (1/T) 10 + 02- 0, 02 + (C/T) 10 02 0 04
< (2(1 + C)/T) 01l
In the same way, we can define the functions F2, F3, F4 all having Lipschitz constant 2(1 + C)/T. As we define F(0)- (Fi (0), (2)) Fa (0), (o)) it has Lipschitz constant 4(1 + C)/7 satisfying that
e 0 " F(0) O(1²),
where Big O() is applied elementwise to the vector, denoting that each element follows O(-) independently. From Tipschitzness of F, for any 2 1,
E|| (1) - (((-1) 0 11 - |((((()))) + O()²)
3 | ||| + (1( - - O(r))
2n/(1 + C3 110(0-1) - -000) 0(1))
3r
By taking expecations for both sides,
El ( ||| " F(O()))
(1 + C)
El ( 1 + O(n²).
37
Applying the triangle inequality, ||0(0) llo) < (0(()) 0((11)|| 10| + 10 - ((0) |l. we further deduce that
217(1 + C)
B1 - + //| [ + O(rr)
El|l@(()
< [0)
Ell (
37
6
( ||| + (nr)) we can write
Setting r [n[++C]
E 0 و l7 < (I
2m(1 + +00
[1110(-1)- 0000 + 17,.
37
Thus, with constant T > TSGD:
211(1 + C)
TSOD
B||l[(sen)) - | + +T < (1
r
37
2m(1+ + C)
<
1 +
TSGD)
r T O
3r
31

---

Published in Transactions on Machine Learning Research (07 (224))
By Taylor expansion of F1 near = 0:
2(62/T - 1) c O(e), |F(0")| 4(62/T c 0(2), T(62/T +1)
F((GGG(()
7(62/T +
we get
2//T - + O(e² + = O(e +() I (1+ +1)
Since ||] - - 2 2 e),
2m(1 + C)T
TSGD 2 R|] - 0 (112(c+
3r
2( π - - t) +0(7((++2))
4
Therefore,
3(62/T +1)
x/4 m/A-p-cc0((p(--0
TSGD
2
1/E n 0 +
Remark 2. To simply compare the convergence rates of two algorithms, we assumed that there is some constant T such that TSGD, TOSCD < T Theorem 8. However, without this assumption, we could still obtain lower bounds of both algorithms at == as;
2
TOSGD log(1 2m)
log
c+e)
11-p-(1-()++0p +8)
TSGD
log
log (1+ 2110011) 1)
Ö
0 +
where δ = (e2 1))2(C2 + 1)(2 1), = max [2sin + el+00s2)) (oboo>)]] ")] and their approximations are ze 4)
Ö ~ 0.265, C R 0.436. For small enough 7), 6, we can observe OSGD algorithm converges faster than SGD algorithm, the inequalities are tight.
Direct Application of OSGD and its Convergence We now focus exclusively on the convergence of OSGD. We prove Theorem 7, which establishes the convergence of an application of OSGD to the mini-batch contrastive learning problem, with respect to the loss function coon
Algorithm 2: The direct application of OSGD our problem
1: Parameters: l: the number of batches to be randomly chosen at each iteration; q: the number of batches of the largest losses to be chosen among & batches at each iteration; T: the number of iterations.
Inputs: an initial feature vector (U(1), V(()), the set of learning rates {ne 12-1
for t = 0 to
do Randomly choose SC with =
Choose in.....; S having the largest losses, i.e., Loon (U V..
g+ y Eies Loom (T(() v(t))
(U(++1, v(++1) + (T(1)) ((1) nh
(U(+1)) v((t1)) normalize (t++1) ((++)) ((+1))
32

---

Published in Transactions on Machine Learning Research (07 (024))
For case of reference, repeat the following definition:
Coon (U, = coon (UB()) VBgD); Y El (77)
(32)
where B)) represents the batch with the j-th largest loss among all possible batches, and q, ko are parameters for the OSGD.
Theorem (Convergence results). Consider sampling 1" from [T) with probabilily proportional to {n ) Then Vp > po = B BT² + 2/7 /Br", that is, P( = t) = / (/
(p+pu) (Zoon(U v _oon ) + +88p Et ni p(p po)
E | 2 (U
El n
where Coom+ denoles the minimized value of Lcon.
Proof. Define (01)) D()) argmin (Zron(U', V') + (") ven))a].
We begin by refer-
U'.V"
ring to Lemma 2.2. in Davis & Drusvyatskiy (2019), which provides the following equations:
[(U!))) v(t) - (0(0) ((()) - i (Ult)) v())),, | Zo ( (()))) < Vi cron (ult), V(O)
Furthermore, we have that V con is Pu-Lipschitz in ((Ba(0, 1)) )2 by Theorem 11. This gives
(u(tt) v(t))) v -
- p (U(t) v(t)) - -0 p()) (0() p(')) (
Therefore,
I con (V('m)((l)) v(()))
< | V coo (((1)) |(t)) + + (U(t)) v(t)) - - ( p(ce)
< V con (()) 1(C + pol(((()) v(()) - -00() b(1)))
+ Pa
v Ccon (U(1', v('')]1
p
As a consequence of Thm 9,
E V om (U(t), |o0))1]
(p+ pb)2
voo roon (Uit") v(ns)/
E
(Coon(U v(O)) - -zcom
+ 8p Σt= vi
(p+ pu)2 pp)
p(p ((P-p) Po)
Σt 1t
(p+ + pu)² po)²
(Zon)(r))), v(1) Zcon * ))
+ +8p Σt= nh
p(p Pa)
Σ Ttt
Note that Zoon+ is the minimized value of Zrom, and the last inequality is due to (gum( v E <
Ccon (U()), v(0)) - < because
Eron (U))) v(0) = min
f (U', V') + I0' V') - (v)1)) [(0)
U.V'
< coo (U(1), ((0)
33

---

Published in Transactions on Machine Learning Research (07 (224))
by putting (U', V' (Ul0), ) and
[eon+
min
{
(U', (U'
U.
{ +') 2 v') (U. -u)
I min U..
for any U, V, implying that Ccon+ < Eponn
We provide details, including proof of theorems and lemmas in the sequel.
Theorem 9. Consider sampling 1* from (T] with probubility PP(L" = = 1// 1e/ 9i). 2 2/2/B72 4e2/T /BT2, have
Then vp Po
Ccon (U()), v( c 8Pp 8pT=-22 "Ei=
v coon (Ult)) vitt)/]
E
s p Pa
El 1/t
where Zoon (U,)) = min U.V
coo (U', V") GI" (') - (U, v)133 and Zgon denoles the minimized value
of Eron.
Proof. eon is Po-Lipschitz in ((Ba(O, 1) by Theorem 11. Hence, is Po-weakly convex by Lemma 5. Furthermore, the gradient norm of a mini-batch loss, or V.u. VB.)Il is bounded by L 4/T. Finally, (Kawaguchi & Lu, 2020, Theorem 1) states that the expected value of gradients of the OSGD algorithm is Vu.v (onn(U v(0) at each iteration t. Therefore, we can apply (Davis & Drusvyatskiy, 2019, Theorem 3.1) to the OSGD algorithm to obtain the desired result.
Roughly speaking, Theorem 7 shows that (Ule", v())) are close to stationary point of Lgon. We refer readers to Davis & Drusvyatskiy (2019) which illustrates the role of the norm of the gradient of the Morean envelope, |/ [g (U)(),) V(')],, being small in the context of stochastic optimization.
We leave the results of some auxiliary theorems and lemmas to Subsection C.3.
C.3 Auxiliaries for the Proof of Theorem 7
For square matrix A, we denote its trace by tr(A). If matrices A and C are of the same shape, we define the canonical inner product (4, by
(A = E
i
Following a pythonic notation, we write Ad; and A. for the i-th row and j-th column of matrix A, respectively. The Mauchy-Schware inequality for matrices given by
(A,C) S ||A||||),
1/2
where a norm ||-|| Frobenius norm in matrix i.e. |||||
Σ Ab,
Lemma 4. Let A pmxn, C' E gank, Then, ||AC||| < ||||||||||
Proof. By basic calculation, we have
||AC|| tr(CT = (TCCC (CCT, AT < ||CC [TAA]]].
34

---

Published in Transactions on Machine Learning Research (07/2024)
Meanwhile, for any positive semidefinite matrix D, let D UAUT be spectral decomposition of D. Then, we have
= ((D'))tt(UAAPUT)) ((AAAA)))))(AAUUUT) = tr((A) < (tr(A))² = (tr(D)))
where (D) denotes the i-th eigenvalue of a matrix D. Invoking this fact, we have
& (tr(CCT))² ||C1
or equivalently, ||CCT]| - ||C]? Similarly, we have ||ATA|| = |||||?. Therefore, we obtain ||AC1|2 = ||CCT|||| A|| || S ||A|||||C|??
which means ||A||| < ||||||||]
If L: Rmxa + is a function a matrix X E Rmxx, we write a gradient of with respect to as a matrix-valuee function defined by
ac (VxL)y = ax 0Xij
0L
Then, the chain rule gives
d
dx VxC) di
di [())
for a scalar variable t. If C(U, V) function of matrices U, V R=x", define Vu,v a horizontal stack of two gradient matrices, i.c., VU,vC= (VuL, VyL.
Now, we brielly review some necessary facts about Lipschitz functions.
Lemma 5 (Rendering of weak convexity by Lipschitz gradient). Let f: R - R be p-smooth function, i..., ff is p-Tipschitz function. Then, f is p-weakly convex.
Proof. For the sake of simplicity, assume f is twice differentiable. We claim that V 2 -pla, where Ta is the x d identity matrix and A > B means A - B is positive semidefinite matrix. It is clear that this claim renders f + 111-112 to be convex.
Let us assume, contrary to our claim, that there exists ay Rdd with f²)())) Y -pla- Therefore, ())()))) has an cigenvalue X < -p. Denote corresponding cigenvector by u, so we have V fxx)) Au, and consider g(e) V + eu); the (elementwise) Taylor expansion of g at 0 gives
f(zo + eu) Vf(xx) + f(zo) o(c).
which gives
[V cu) - (/(o)) = v2 f o(c)
[
Taking e 0, we obtain [V + eu) Vf(zo)1/E [X| > p. which is contradictory to p-Lipschitzness of Vf.
For X E BxB let us define
B
B
B
B
M (X) (M)))=
B -2tr(X)/7+ +
Σ log exp X + log Σ
Σ
exp((Xii/T)
Using this function, we can write the loss corresponding mini-batch B of size B by
M (ULVs) crn (UB VB).
We now claim the following:
35

---

Published in Transactions on Machine Learning Research (07/2024)
Lemma 6. Consider X E R where [Xayl $ 1 for all 1 1 i.j S B. Then, (x is bounded by v8/Br² and 202/r B Lipschitz.
Proof. With basic calculus rules, we obtain
xX) -21l + Px + +Qx,
(33)
where In is the B x identity matrix and
B
exp(Xx//T)/ (xx)x=ex((X,//1)) (Px)a = exp(X1)/ /T) Σ exp((x//T), (Qx)
Σ
exp(Xx;) (Xz)/T).
k=1
From Ej Psj = 1 for all i, it is casy to see that [[[TB - P)i,,l? < 2. This gives [ - Pxl1 - 2B, and
similarly [[g Qxl1² 2 < 2B. Therefore, we have
XB-VxLM(XX) & [IB Pxll + Il - Qxll s v8B.
(34)
or equivalently
Vx - V8/Bt2
(35)
We now show that Vx
2-2/ -ipscciiz. Define p: R + B by
(p(x)), cxp(x;/T) Ek=1 exp((x/T)
Then, we have
dr (2x) = - [diag(p(x)) - p(=)p(x)1)
For [-), have p(x); v
2/r
217
for any i. Thus,
170777
B
p(x) 1 diag(p(x)) = 2/7
0 dr
IB.
Br
which means p(x) is 2/-
Br Lipschitz, i... llp(z) p(y)l
< - || yll for any x, y 1|., Using this fact,
can bound Ppx Pyll for X, Y E [-1, BxB follows:
B
2/r
B
22/T 2
|Px - Px
Σ ((()..) - P(V..)11
X
Yr.11
||
-
BT
Br
Similarly, we have |[xx - - < - XX - Y1I-- Summing up,
202/7
(x) C (Y) (X) < |Px - Pr - Qrl1
Br 1|X YI1-
which renders
207/7 -
B272 || YII.
Recall that L0nn (Ug, VB) = LM (U)V)) for Ug. Vg d (They correspond embeddings corresponding to mini-batch B). Using this relation, calculate the gradient of goon with respect to Ug. Denote
36

---

Published in Transactions on Machine Learning Research (07/2024)
E.j R one-hot matrix, which is matrix of zero entries except for (i.j) indices being 1, and write
G = VxL* (Ug)s. Then,
a
V(U
Lo (UB, VB) =
VxCM(U
((B))
@UBij
ET,
EnG(VV))
()V
(VBGT)g
This elementwise relation means
con (UB. V8) = VBG VB(Vx M (UTVE)))
(36)
0UB
and similarly,
C (UB. V8 = UBV (M (UgV)).
(37)
V 8
We introduce simple lemma for bounding the difference between two multiplication of matrices. Lemma 7. For A1, A pm and B1, B2 E g have
|A - AAball S ||A - A2ll Bill || - B2ll.
Proof. This follows from a direct calculation and Lemma 4
|A - A2Ball || - B2 + A₁ - A2B2ll
||A[ B2)l + [[(A| B2)
|| A2lll B |2|||| B2ll.
7
Theorem 10. For any U, V E (Ba(0, 1)) and any batch B of size B, we have Vou (UB. V8)l <
Proof. Suppose UB, VB E (Ba(0, 1)) have
conn VUs. C (UB, VB) (Va(V (MMVV) (
from Equation (36) and (37). By following the fact that UB|]. |&||| S BB and VxLM S /8/B72 (see Lemma 6), we get
VVs (UgVE)) V (ULV V8 7
and
(UTV3)) < V8 (
UBV M (UIVE)
U (UgVE)
Then,
V LoT V8)l =
(U.V < =
37

---

Published in Transactions on Machine Learning Research (07/2024)
Since L (UB, V8) is independent of UNN and VIJJS we have
VV (UB, VB3il Vy (UB, VB)l S 4
Theorem 11. V (U, V) is Po-Lipschitz for U, V (Ba(0, or to clarify,
V V\) (U v²)1 = pol v\)
for any U², Vl, U² V (Ba(0,1))² where Pa /8/ B+ + 462/T /BT².
Proof. Denoting Ug. V as parts of U V that correspond to a mini-batch B, we first show UV L (U Vig) VU Conn (UB, VBll I poll(U Vil) (U2, Vi)Il holds. For any UB, VB e (Ba(0, 1)) have
VUs Vg Loon (UB, VB)) VV CM (U)V V CM (U)VE))
from Equation 36 and Equation 37. Recall Lemma 6; for any Ug. Vi E (Ba 1)) (i= 1.2,, we have
((U)) S /8/B22
and
2
X
Beae B 3 VV3
We invoke Lemma 7 and obtain
|UAV M ((U)) (
s | - V (1 - VxC*
((
202/T
S / UBl + B3/2,2 (U
TV - V331
20
S /8/BB>2 U B3/2,2 (1UU)
| | V)
-
S ((V8/BB++22 (115555--22000000100------811113 - -/ll.
-
and similarly
< (2e2TT || U2 +V8/ 1
/ V - V/21.
Using the fact that
(ax + by) (bx + az)² = (a² + (1² + + 4abay S (a S () x²)
holds for any a, 6 2 0 and H, y E R, we obtain
I L (U. Vi)) VL VL
= (U
+
=
38

---

Published in Transactions on Machine Learning Research (07 (224))
Restating this with Pu 18/BT2 4c2 /Br?, we have
|V (UB ViA) (U3 v poll(Ug Vgg (Ug, v3)
(38)
Recall the definition of g
conn Con = Σ Ji (aLoom(Un (UB VB,,),
( and Ej 2= =q. For any U, V E (8") we can find neighborhood of (U, V)
where 7i
that value rank of Leon (UB, V&,) over i E (I...., (+) does not change, since _ po-Lipschitz. More precisely speaking, find rank that be accepted by all points in the neighborhood. Therefore, we have
VU.v Zoon roon (U, V)
Σ 7o Vu.v. con (UB VBgg):
and since ||UB V ll 5 IIU VlI, VU (U, (U,V) is locally po-Lipschitz. Since Leon is smooth, (U, V) 3 Poll,, where Ny is the N x N identity matrix. such property is equivalent to -PoIN I V Vv (U, Therefore, foon is Po-Lipschitz on (( ((Ba(0,1)) ?)
Coon
D Algorithm Details
D.1 Stochastic Gradient Descent (SGD)
We consider two SGD algorithms:
1. SGD with replacement (Algorithm 3) with & 1 for the theoretical analysis Section 5.1.
SGD without replacement (Algorithm 4) for experimental results in Section 6, which is widely employed practical settings.
In the more practical setting where u, fo(x;) and (v,),, SGD updates the model pa- rameters 8, using the gradients 1/2 Eieso Vo.sCro (UB, VB.) instead of explicitly updating U and V.
Algorithm 3: SGD with replacement
Input:
the number of positive pairs N, mini-batch size B, the number of mini-batches k, the number of iterations T, the learning rate initial embedding matrices: U, V
for t = 1 to 7' do
Randomly select mini-batch indices SE C [(>) ((Sel = k)
Compute the gradient: g + Exes Vu.v foom (Us VB,))
Update the weights: (U.V) + (),t)
Normalize column vectors of embedding matricus (U, V)
Algorithm 4: SGD without replacement
Input: the number of positive pairs N, mini-batch size B, the number of mini-batches k the number of epochs E, the learning rate n, initial embedding matrices: U, V
for to E do
Randomly partition the N positive pairs into N/ B mini-batchhe batches: (B:)
N/R
for j to N/Bk do
Select mini-batch indices Sg [k[j 1) 1) ....,, kj}
Compute the gradient: g Σiese Von (UB Vi,))
6
Update the weights: (U, V) + (U, -
Normalize column vectors of embedding matrices (U, V)
39

---

Published in Transactions on Machine Learning Research (07 (07/2024) 2024)
D.2 Ordered SGD (OSGD)
We consider two OSGD algorithms:
1. OSGD (Algorithm 5) with == (m) for the theoretical analysis in Section 5.1 and experimental results on synthotic datasets in Scction 6.1.
2. OSGD without replacement (Algorithm 6) for experimental results on real datasets Section 6.2, which is implemented for practical settings.
In the more practical setting where u, Jo(w);) and o, (o OSGD updates the model pa- rameters 0 using the gradients # Eiese Ve.sCCon(UB, VB,) instead of explicitly updating U and V.
Algorithm 5: OSGD
Input: the number of positive pairs N, mini-batch size the number of mini-batches k, the number of iterations T, the of learning rates (n 17-1- initial embedding matrices: U, V
for 1 to do
Randomly select & mini-batch indices SB S. [[²)] (1Sel k)
3 Choose q mini-batch indixes Sq == f61, 2 C Sg having the largest losses i.c., [Con(UB, VB,))
Compute the gradient: y / EiEs, Vu.ve VE,))
Update the weights: v (U, V) n/()
Normalize column vectors of embedding matrices (U, V)
Algorithm 6: OSGD without replacement
Input: the number of positive pairs N, mini-batch size B, the numbers of mini-batches & and the number of epochs E, the set of learning v, initial embedding matrices: U, V
for E do
Randomly partition the N positive pains into N/B mini-batches N/ for == to N/Bk do
B/B Select & mini-batch indices Sg {k[[ 1) 1,Kj 1) 2...,,.
Choose g mini-batch indices Sq {6,, i4 Sg having the largest losses ie., LUB, VB,)
Compute the gradient: 4 S, Vu.v Vs.)
Update the weights: (V) < (U, V) 7.3
Normalize column vectors of embedding matrices (U,
E Additional Experimental Results
In this section, we provide additional exporimental results. First, we present histograms of mini-batch counts for different loss values from models trained with different batch selection methods. Next, we provide the results for E 16) on the synthetic dataset.
E.1 Batch Counts: SC method vs. Random Batch Selection
We provide additional results comparing the mini-batch counts of two batch selection algorithms: the pro- posed SC method and random batch selection. The mini-batee counts are based the mini-batch contrastive loss (UB. V8) with r = 1. We measure mini-batch losses from ResNet-18 models trained on CIFAR-100 using the gradient descent algorithm with different. batch selection methods: (i) SGD (Algorithm 4), (ii) OSGD (Algorithm 6), and (iii) the SC method (Algorithm 1). Figure illustrates histograms of mini-batch counts for N/B mini-batches, where 50000 and B 20. The results show that mini-batches generated through the proposed spectral clustering method tend contain higher proportion of large loss values when compared to the random batch selection, regardless of the pre-trained models used.
40

---

Published in Transactions on Machine Learning Research (07 2224)
epoch
epoch
epoch
epoch 101
of batches
200
random
spectral clustering
100
3
5
0.0
0.5
1.0
0.0
0.5
1.0
batches
200
100
2
4
0.0
0.5
1.0
0.0
0.5
1.0
# batches
200
100
0
2
6
1
3
4
0.0 0.5
1.0
0.0
0.5
1.0
contrastive loss
Figure 5: Histograms of mini-batch counts for N/B mini-batches, for the contrastive loss measured from ResNet-18 models trained on CIFAR-100 using different batch selection methods: (i) SGD (Top), (ii) OSGD (Middle), (ii) SC method (Bottom), where 5=50,000 and B=20. Each column of plots is derived from distinct training epoch. Here we compare two batch selection methods: (i) randomly shuffling N samples and partition them into N/B mini-batches of size B, (ti) the proposed SC method given in Algorithm 1. The histograms show that minihbatches generated through the proposed spectral clustering method tend contain higher proportion of large loss values when compared to random batch selection, regardless of the pre-trained models used.
10
2.0
DGGD
Spentol Mehwv
2N
0.0
0.0
-10
3.0
N222
0.5
0.01
100
200
300
400
500
Update stage
(a) solutions
(b) full-batch
(c) (?) -all
(d)
(%) -sub
(c) norm differences
Figure 6: Heatmap of N x N matrix visualizing the resulting values from the same settings with Fig 4 except N = 4.
E.2 Synthetic Dataset
With the settings from Scction 6.1, where each column of embedding matrices U, V is initialized as a multivariate normal vector and then normalized as ||u=||| = || = 1, for all i, we provide the results for N (4, 16} and 2N or d = N/2. Figure and show the results for We additionally present the results for theoretically unproven cases, specifically for N and d E (3. 5) (sce Figure 8). The results provide empirical evidence that all combinations of mini-batches leads to the optimal solution of full-batch minimization for the theoretically unproven cases.
and N 16, respectively.
F Computational Complexity
We compare the complexities of three batch selection algorithms: (i) SGD (Algorithm 4), (ii) OSGD (AI- gorithm 6), and (iii) the SC mothod (Algorithm 1). We first present the time complexity for N/B batch selection for each algorithm using big-O notation:
For SGD, batches selected through just random shuffling resulting in O(N).
41

---

Published in Transactions on Machine Learning Research (07 (272))
0500
5.5
Spierr Cumening Wehee
3.0
560
4-22
1.0
1.0
5.0
4.5
12 2/2
0.0
1.0
Update stope
solutions
(b) full-batch
(c)
(7) -all
(d)
(c)
differences
Figure 7: Heatmap of x matrix visualizing the resulting values from the same settings with Fig except 16.
(a) full-batch
(b) (3)
-all
(c)
(7) sub
Figure Thcoretically unproven setting.
Heatmap of N N matrix when and d N 1.
For OSGD, we randomly select batches, cach of size B, and then measure the loss for each of these batches. After measuring, we sort the results and select the lop-g batches based on their losses. The computational complexity of evaluating the loss and sorting to generate q batches is O + log k). Consequently, the complexity for processing N/B batches becomes O(kB²² klogg k) N/ O((kB/q + & logg k/ N).
For the SC method, the primary computational bottleneck arises from the Hungarian algorithm, which operates at a complexity of O(V3). This complexity is the predominant factor influencing the overall performance of our algorithm. In order to accelerate batch selection in the proposed SC approach, randomly partition N positive pairs into clusters, each of size kB. We then subsequently apply the SC method every k--sized cluster instead of applying it directly on the entire N pairs. This process generates k mini- batchs. In this context, the computational complexity for producing & batches is (I)3 B3). Extending this to generate N/B batches, the complexity becomes (133 B3 x N/kB) O BBN)
We further report the wall-clock time required for the selection and update of N/B batches using each algorithm on the CIFAR100 dataset, where N = 50, 000 and B = 32 (see Section 6.2). The SC method requires an additional 11 minutes for batch selection compared to SGD (2 minutes) per epoch. As mentionod in Section 6.2, we suggest parallel batch selection using the SC method while updating models through mitigate this issue in practice. This approach involves using the same batches for multiple epochs as the next batch selection proceeds, based the assumption that batches with higher loss values do not rapidly change.
gradient descent
42


Reviews:
[{"review_id": "mIskGpOekS", "reviewer": "Reviewer_geuj", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- The analysis seems solid and correct.\n- Studying the effective of mini-batching in contrastive loss is important due to its prevalent usage.\n\nWeaknesses:\n- For the first part of the contribution, i.e., proving that the optimal embeddings are identical if and only if all the possible mini-batches are involved in the objective function, I am not so sure about the significance of this result. It seems straight forward to me when considering an objective function with a finite-sum structure. That is, comparing the optimal solutions of full-batch loss and the average of all possible mini-batch losses. Probably the interesting part is that the mini-batch loss is not a simple scaling of the full-batch loss by a constant factor in the case of constrative loss as stated in Proposition 1? I think the authors need to at least comment on the effect of randomness. Does the analysis imply that solving mini-batch constrative loss with SGD using uniformly sampled mini-batches would converge to the same optimal solution as GD solving the full-batch loss? Does the analysis imply that SGD using shuffling without replacement is a nice idea for the contrastive loss since it scans every possible mini-batch in an epoch? I am not sure about the significance or insight of the theoretical results in this part.\n- For the second part of the contribition, i.e., showing that OSGD is better than SGD in a toy example with some theoretic supports, establishing the convergence of OSGD in contrstive loss, and proposing a spectral clustering-based OSGD variant with empirical advantages, I feel that the contribution is not closely connected to the first part, i.e., the story line is not so clear to me. The authors show that the problem is geodesically non-quasi-convex, so it is expected to be hard. This part confused me a lot. It seems to me that (mini-batch) SGD for solving general non-convex finite-sum problems has already been thoroughly studied (e.g., \"Garrigos, Guillaume, and Robert M. Gower. \"Handbook of convergence theorems for (stochastic) gradient methods.\" arXiv preprint arXiv:2301.11235 (2023).\"). What is special in the mini-batch contrastive loss setting, that makes the analysis hard? Is it the unit norm constraint? For handling unit norm constraint, I remember there is a line of work proving convergence for Riemannian SGD in the non-convex case (e.g., \"Hosseini, Reshad, and Suvrit Sra. \"Recent advances in stochastic Riemannian optimization.\" Handbook of Variational Methods for Nonlinear Geometric Data (2020): 527-554.\"). I think it would be great if the authors could provide some discussion on the setup of contrastive loss, comparing to the standard finite-sum one and the riemannian finite-sum one.\n- Does the Spectral Clustering-based OSGD enjoy the same convergence guarantee as the one in Section 5.2?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the effect of mini-batching in the optimization of contrastive loss. This paper first compares the optimal embeddings of minimizing the full-batch and mini-batch contrastive losses, in which the authors prove that the optimal embeddings are identical if and only if all the possible mini-batches are involved in the objective function. Then, the authors studies the convergences of OSGD and SGD, which shows that OSGD can be faster than SGD in a toy example, and proposes a spectral clustering method to improve the minibatch selection problem of OSGD.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "See above\n\nOther Typos and Questions:\n- The first paragraph in sec. 5.1: geodesically non-quasi-convex\n- The second paragraph in the Intro: $O(\\mathcal{B}^2) = O(1)$, in the case of $\\mathcal{B}=4096$, $\\mathcal{B}^2$ is at the scale of $10^7=\\Theta(N)$, how can we simply treat it as a constant?\n- Theorem 6 is studying the convergences of two methods containing randomness, but there is no probablity or expectation in the statement, which has to be wrong.", "broader_impact_concerns": "NA", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "8HGRowqQhT", "reviewer": "Reviewer_QzrE", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n\n* This paper establishes the convergence guarantee for the stochastic version of InfoNCE loss, which is close to the practice setting.\n* This paper proves the convergence of OrderedSGD and shows its faster convergence than SGD for optimizing InfoNCE loss.\n\nWeaknesses:\n\n* Proving the convergence of optimizing $U$ and $V$ in (1) is not that interesting to me. In practice, one may try to train the model parameters in the function $f_{\\theta}(x)$ and $g_\\phi(y)$. The training dynamic will be very different from training over $U$ and $V$.\n\n* Besides, if only consider training $U$ and $V$, it is also not surprising that the optimal solution is ETF, as all vectors $u_i$'s and $v_i$'s are symmetric in the InfoNCE loss, as well as its mini-batch version.\n\n* More discussions regarding Theorem 7 need to be added. The authors need to compare the convergence rates of SGD and OSGD in the setting of Theorem 7.\n\n* The authors mention that finding the best $q$ loss among k batches is challenging, and can be kind of reduced to the min-cut problem. However, the authors provide the neither approximation error, nor the complexity for solving the min-cut problem, making it difficult to precisely evaluate the computational complexity of OSGD.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper studies the optimization for contrastive loss when using stochastic mini-batches. The contributions include\n* The paper examines the theoretical aspects of optimizing mini-batches loaded for contrastive learning with the InfoNCE loss and shows that under certain parameter settings, mini-batch optimization is equivalent to full-batch optimization if and only if all N/B mini-batches are selected.\n* The paper proposes a batch selection method based on the Ordered SGD algorithm and spectral clustering to efficiently select high-loss mini-batches, which can potentially improve the convergence rate of mini-batch contrastive loss optimization compared to vanilla SGD.\n* Numerical experiments are conducted to support the theoretical arguments.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. Add discussion on the difference between optimizing $U, V$ and model parameters.\n2. Add discussion about the convergence rate comparisons between SGD and OSGD in the setting of Theorem 7.\n3. Add the approximation error between min-cut problem and the best $q$ mini-batches problem, and also provide a complete complexity analysis for OSGD.", "broader_impact_concerns": "No concern", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "AXt45sWP0w", "reviewer": "Reviewer_oR4M", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strengths** \n\n1.\tA well-written study of an important open problem of mini-batch convergence when optimizing InfoNCE loss\n2.\tInteresting theoretical results regarding the soundness of mini-batch approach and its limits\n3.\tA practical approach for improving the mini-batch convergence with spectral clustering\n\n**Weaknesses**\n\n1.\tThe assumption about U=V seems rather strange in Definition 2. Is there any error here? I wonder how meaningful it is to use it for the two matrices and not for the minimizers as done in Lemma 1. It feels that this assumption is too strong to actually hold in practice, otherwise why not directly optimizing over U or V in isolation? Figure 4e confirms that U and V do not immediately behave as U* and V* and are quite far from them in the beginning of training.\n\tGiven that several main contributions of this work rely on this assumption, it is hard to not to \tsee it as a very limiting, unless the authors can provide more background on why it is not the \tcase. \n2.\tI do not have a strong background in this field so maybe the results and the techniques of this paper are completely new but I would like the authors to elaborate more on why the proofs of mini-batch convergence in the contrastive learning setting are harder to obtain that similar such results in general SGD setting? The sentence just before the beginning of Section 3 hints on this but could be elaborated a bit more. \n3.\tFrom Figure 2, Figure 4e and Table 1 it is not immediately clear why SC outperformance OSGD in the end? It seems that OSGD converges to the oracle in Figure 2, in Figure 4e it convergence faster to the optimal U and V and yet SC improves over it in Table 1. Is there some explanation to it or a visualization showing where exactly SC improves compared to OSGD?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "In this paper the authors provide a theoretical study of the mini-batch optimization of the popular InfoNCE loss often used in contrastive learning. They start by carefully explaining the intrinsic difficulties of the mini-batch optimization in the contrastive learning setting consisting in the need to draw a prohibitively large number of batches. The paper further presents several theoretical results showing that 1) when all possible mini-batches are selected then the convergence to the full-batch solution is achieved, yet the losses of the two approaches differ; 2) the convergence is never achieved when some mini-batches are omitted. The authors further study a variation of the SGD called OSGD that selects the mini-batches having the highest contrastive loss and show its benefits over vanilla SGD in some cases. The experimental results on a toy examples and several real-world datasets confirm the utility of the provided theory.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1.\tProvide explanation regarding the assumption U=V in Definition 2.\n2.\tProvide a more comprehestive comparison to the approaches used to prove mini-batch convergence in our learning setups. Explain the differences in the techniques employed and the challenges that need to be overcome.\n3.\tExplain the improved empirical performance of SC vs. OSGD with a clear example or some sound intutions.", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: The paper studies aspects of a very important problem: mini-batch optimization for the widely popular contrastive losses. The paper proves the relationship between full-batch and the mini-batch optimization under some assumptions and provides efficient batch sampling algorithms for improved performance. The reviewers agree that the claims made in the paper are sufficiently supported with theoretical and some empirical evidence.
Decision: Accept (poster)