Your contrastive learning problem is secretly a distribution alignment problem
Authors: Zihao Chen
Chi-Heng Lin
Ran Liu
Jingyun Xiao
Eva L Dyer
Venue: NeurIPS 2024 poster (2024)
Keywords: Optimal transport
Distribution alignment
Noise contrastive estimation
Primary area: probabilistic_methods
TL;DR: In this work, we introduce a novel framework for representation learning that recasts contrastive estimation as a distribution alignment problem.

Abstract: Despite the success of contrastive learning (CL) in vision and language, its theoretical foundations and mechanisms for building representations remain poorly understood. In this work, we build connections between noise contrastive estimation losses widely used in CL and distribution alignment with entropic optimal transport (OT). This connection allows us to develop a family of different losses and multistep iterative variants for existing CL methods. Intuitively, by using more information from the distribution of latents, our approach allows a more distribution-aware  manipulation of the relationships within augmented sample sets.
We provide theoretical insights and experimental evidence demonstrating the benefits of our approach for generalized contrastive alignment. Through this framework, it is possible to leverage tools in OT to build unbalanced losses to handle noisy views and customize the representation space by changing the constraints on alignment.
By reframing contrastive learning as an alignment problem and leveraging existing optimization tools for OT, our work provides new insights and connections between different self-supervised learning models in addition to new tools that can be more easily adapted to incorporate domain knowledge into learning.

Your contrastive learning problem is secretly a distribution alignment problem
Zihao Chen', Chi-Heng Lin, Ran Liu, Jingyun Xiao, Eva L. Dyer* School of Electrical & Computer Engineering Georgia Tech, Atlanta, GA
Abstract
Despite the success of contrastive learning (CL) in vision and language, its the- oretical foundations and mechanisms for building representations remain poorly understood. this work, we build connections between noise contrastive estima- tion losses widely used in CL and distribution alignment with entropic optimal transport (OT). This connection allows us to develop family of different losses and multistep iterative variants for existing CL methods. Intuitively, by using more information from the distribution latents, our approach allows a more distribution-awall manipulation of the relationships within augmented sample sets. We provide theoretical insights and experimental evidence demonstrating the benefits of our approach for generalized contrastive alignment. Through this framework, it is possible to leverage tools in OT to build unbalanced losses to handle noisy views and customize the representation space by changing the con- straints on alignment. By reframing contrastive learning as an alignment problem and leveraging existing optimization tools for OT, our work provides new insights and connections between different self-supervised learning models in addition to new tools that can be more easily adapted to incorporate domain knowledge into learning.
Introduction
In machine learning, the availability of vast amounts of unlabeled data has created opportunity to learn meaningful representations without relying on costly labeled datasets [26] [22 27). Self- supervised learning has emerged as a powerful solution to this problem, allowing models leverage the inherent structure in data to build useful representations. Among self-supervised methods, con- trastive learning (CL) is widely adopted for its ability create robust representations by distinguishing between similar (positive) and dissimilar (negative) data pairs. With success in fields like image and language processing g 46), contrastive learning now also shows promise in domains where cross-modal, noisy, or structurally complex data make labeling especially challenging 34 56 [0].
Traditional contrastive learning methods primarily aim bring positive pairs-often often augmentations of the same sample-closed together representation space. While effective, this approach often struggles with real-world challenges such noise in views, variations in data quality, shifts introduced by complex transformations, where positive pairs may not perfectly align. Additionally, in tasks requiring domain generalization, aligning representations across diverse domains (e.g.. variations in style or sensor type) is critical but difficult to achieve with standard contrastive learning, which typically lacks mechanisms for incorporating domain-speciffc relationships. These limitations highlight the need for more flexible approach that can adapt alignment strategies based on the data structure, allowing for finer control over similarity and dissimilarity among samples.
*Contact: (zchen959, evadyer) @ @ggatech.edt
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

To address this challenge, we introduce a novel generalized contrastive alignment (GCA) framework, which reinterprets contrastive learning as a distributional alignment problem. Our method allows flexible control over the alignment of samples by defining target transport plan, Pigt Pitt. that serves as a customizable alignment guide. For example, setting Pigt to resemble a diagonal matrix encourages each positive to align primarily with itself or its augmentations, thereby reducing the effect of noise between views. Alternatively, we can incorporate more complex constraints, such as weighting alignments based on view quality or enforcing partial alignment structures where noise data heterogeneity is prevalent. This flexibility enables GCA to adapt effectively to a wide range tasks, from simple twin view alignments scenarios with noisy or variably aligned data.
Our approach also bridges connections between GCA and established methods, such as InfoNCE (INCE) [38], Robust InfoNCE (RINCE) [[2)] and BYOL (22] demonstrating that these can be viewed as iterative alignment objectives with Bregman projections 6 21]. This perspective allows us to systematically analyze and improve uniformity within the latent space, property that enhances representation quality and ultimately boosts downstream classification performance.
We validate our method through extensive experiments on both image classification and noisy data tasks, demonstrating that GCA's unbalanced OT (UOT) formulations improve classification performance by relaxing our constraints on alignment. Our results show that GCA offers a robust and versatile framework for contrastive learning, providing flexibility and performance gains over existing methods and presenting promising approach to addressing different sources of variability in sulf-super learning.
The contributions of this work include:
A new framework called generalized contrastive alignment (GCA); which reinterprets standard contrastive learning as a distributional alignment problem, using optimal transport to provide flexible control over alignment objectives. This approach allows us to derive novel class of contrastive losses and algorithms that adapt effectively to varied data structures and build customizable transport plans.
We present GCA-UOT, contrastive learning method that achieves strong performance on standard augmentation regimes and excels in scenarios with more extreme augmentations or data corrupted by transformations. GCA-UOT leverages unbalanced transport to adaptively weight positive alignments, enhancing robustness against view noise and cross-domain variations.
We provide theoretical guarantees for the convergence of our GCA-based methods and show that our alignment objectives improve representation quality by enhancing the uni- formity of negatives and strengthening alignment within positive pairs. This leads to more discriminative and resilient representations, even in challenging data conditions.
Empirically, we demonstrate the effectiveness of GCA in both image classification and domain generalization tasks. Through flexible, unbalanced OT-based losses, GCA achieves superior classification performance and adapts alignment to include domain-speciff informa- tion where relevant, without compromising classification accuracy in domain generalization.
2 Background
2.1 Contrastive learning
Contrastive learning (CL) is representation learning methodology that uses positive and negative pairs to define similarity in the latent space. Let D = {x; J denote our dataset. For each sample x; in a batch of training data with size B, we create two augmented copies x", and x" independently, i.e., x = (xx)) where & is a randomly drawn augmentation function from some augmentation class A and likewise for x". The (x/, x"" is called a positive pair of x; while (x), xx) is treated negative pair for any # i. One of the most widely used formulations of the CL problem, InfoNCE (INCE) 8 seeks to maximize the negative log probability that a sample is correctly classified as
Bii
LINCE - llog
e
Eite enis
where sty =====f((()) -====f((x]]T66((2/1155666 [e(x))//116(())()111( is the score between augmented samples.
2

---

Building upon the principles of INCE, SimCLR g and MoCo 24 are two representative works that form the foundation of contrastive learning methods for visual representation tasks. Alternatively, BYOL [22] and SimSiam [ discard the use of negative samples to avoid large batch size and instead use exponential moving average-based updates to avoid representational collapse. Recent contrastive methods have focused on improving the tolerance to noise in samples to enhance robustness in diverse scenarios [3]. Among them, Robust INCE (RINCE) is a robust contrastive loss function characterized by its symmetric properties and theoretical resistance noisy labels [47 [2]. Specifically, RINCE provides robustness to noisy views by introducing adjustable parameters and g [1] which rebalance the cost of positive and negative views, resulting in the following loss:
CRINCE = ( - + X9(e +
Σ i9gt eau)))
By optimizing the above loss functions, the encoder ʃ is trained to construct a semantically coherent representation space where positive pairs of samples are positioned nearby, while those negative pairs with divergent semantic attributes are separated [57].
2.2 Proximal Operators and Projections
To make the connections between different CL losses clearer later, use the notion of proximal operators. In words, the proximal operator will provide a way to find the closest point in some closed convex set. Formally, we define the proximal operator follows.
Definition 1 (Proximal Operator). Let (x v) r(x) r(v) (VT(v), x v) be a Bregman divergence with a convex function r. The proximal operator of h: X' - R {+oo} is defined for a point v E X with closed convex set B X
Prox h
(v) arg min (h(x) dr(x, xeB
Moreover, we can define the concept of projection as a special case of the proximal operator when we let hxx) be an indicator function hz(r)) = {0, if x E 00, if x & B} on constraint set B. See Appendix A.2 A.2 for more details.
2.3 Solving Optimal Transport Through Proximal Point Methods
Optimal transport (OT) is widely used in characterizing the distance between two collections of samples {x:) and {y; ) with associated measures μ = - 6x and -Ø= 8y, with Dirac delta function 5x and by on finite support 433. Here, p and are vertices of the IR B simplex defined AB {u RB U4 2 0, Ei=
1}. OT aims learn joint coupling matrix, or transport plan P RBxB that minimizes the cost of transporting mass encoded by cost matrix C E BxB from one distribution to another. In practice, entropy regularization is used to solve the OT objective, resulting in the following entropy-regularize OT (EOT) objective:
min PEB
(P.C) - (P), where -(P)=-E P log (og(Pss)
=
(3)
where is user specified parameter that controls the amount of smoothing in the transport plan, and C(x,y) = - / |y| is often set to encode the cosine similarity between pairs of samples.
The Sinkhorn Algorithm and its Interpretation as a Bregman Projection. Solving Equation (3) could be interpreted as iterative alignment problem a Hilbert space generated from the kernel Kij = exp( (C,,,,/)).. This alignment problem can be solved through iterative Bregman projections onto the two constraints sets that encode the marginals along the rows and columns [ 5
C = [P Pl = c = (P P
1 =v}
The first step of Bregman projection to find the minimizer P(1) = arg (in(cKL(P| P1 4} by the proximal operator Prox (K) with Lagrange multiplier f on the row constraint set CP
and compute its derivatives with respect to with u = ef/= > 0:
elog(P!!) /) - f1 ( uK, (P 1) = μ - (uK, 1) H. - P
(5)
K1
3

---

Next, we project P()) onto the column constraint set C2 resulting in P = Prox KL (P()))
dia) pud 12 The iterative updates can succinctly expressed as the Sinkhorn iterations:
p(2+1) diag( u (++1)) K diag(v()),
p = diag( (++1) K K diag(v!(+!))
(6)
with the scaling vectors u and v(O) updated according to:
(+1) def
M
((+1)
del
Kv
K [T)
Here, iterations converge to stable transport plan ( as the optimal solution of Equation (3) which provides the minimum cost matching between two distributions. The convergence and dynamics of OT and its dual formulation have been studied extensively in 4 that the iterates will converge the optimal solution of the EOT objective, that p(t) + p(x)) with t + co. See Appendix A.3 for more details on both the continuous and discrete formulations of OT.
19 1. Thus, these results guarantee
2.4 Wasserstein Dependency Measure
The Wasserstein Dependency Measure (WDMM) is a measure of deviation between two probability measures. We will use this later and thus provide the formal definition here [39].
Definition (Wasserstein Dependency Measure). Define the WDM as the Wasserstein distance (W1) between the joint distribution x(I y) and the product of marginal distributions & v(x, v) of two random variables x and Wi (T, p. & v) supfec(xx>) (Bn(as)))(a y)l 16 [p v))), where C(X x y) denotes the set of all 1-Lipschitz functions from X x y to R.
2.5 Optimal Transport and Alignment in Representation Learning
Distribution alignment and OT have been widely used for domain adaptation
14
30
59
and in generative modeling 2 55 49 58). The connections between distribution alignment and contrastive learning. however, are still nascent. In [51]. the authors explore the connection between inverse OT (IOT) [3) [8] and INCE. Our work builds on this connection to OT to build robust divergences (RINCE) and to build novel unbalanced optimal transport (UOT) method (Section (33)) Additionally, we show how our framework can be used to build flexible methods for encouraging contrast multiple levels. We this concept of hierarchical contrast and show that it be used in domain generalization settings (Section 6.2). It is note that GCA-UOT focuses on relaxing the hard constraints on the row and columns into the soft penalties, which is different with the idea of "unbalanced matching" in [5]] which considers the case where the encoders may not have the same weights.
3 Generalized Contrastive Alignment (GCA)
In this section, we will introduce a new framework for generalized contrastive alignment and demonstrate the connections between contrastive learning and optimal transport.
3.1 Problem Formulation
Traditional contrastive learning methods focus on bringing positive examples, such as augmentations of the same sample, closer together in representation space. In contrast, our approach reframes contrastive learning as a distributional alignment problem, allowing flexible control over how pairs are matched by imposing specific constraints on the target transport plan, Pigt-
Our objective is to learn an encoder fe that minimizes the transport cost between positive samples. By defining Ptge with specific alignment rules, such domain-specific hierarchical constraints, we can influence how samples are organized in the latent space. For instance, setting Ptgs to resemble diagonal matrix encourages each positive to align primarily with itself or its augmentations, minimizing PP) x O, where div measures the deviation from an identity matrix (e.g.. KL- divergence).
This flexibility allows us to encode more nuanced forms similarity, adapting to tasks where alignment structure varies based on domain, class, other high-level constraints. By expanding contrastive learning in this way, our method enhances separation of negatives while addressing complex relational patterns, making it suitable for wider range of learning tasks.

---

Defining the Kernel Space. Before formally stating our objective, we first need to define the concept of an augmentation kernel for our positive and negative examples.
Definition (Augmentation Kernel), Let fe denote an encoder with parameters and let (x,, x ~ A be two views drawn from the family of augmentations. A. The augmentation kernel for the encoder 0 is defined as Ko(x, xx) exp(-dist( (x),, fo /e))))) where dist(-) can be arbitrary distance metric, and fo(x)) is the normalized output of fa. and E is the regularization parameter.
Main Objective. With this definition in hand, we can now formalize our objective as follows:
min das (Pigl PgllPPo) with Pa = arg min (h(P) + {P)(
(8)
PEB
where Kg is the augmentation kernel defined in Definition 33. h(x) is a convex function (typically an indicator function). B is a closed convex constraint set (i.e. Birkhoff polytope) that defines the constraints of proximal operators, dr is Bregman divergence that is used to find the nearest points Ps on the constraint set B of Kg, das is a convex function (e.g., KL-divergence) that measures divergence between Py and the target coupling plan Pige-
Our objective bi-level optimization problem which aims to learn a representation that minimizes the divergence between the transport plan Pa with the target alignment plan Pig that encodes the matching constraints. When we consider a standard contrastive learning setup where we have pairs of positive examples the source and target distribution, then the target Pigi is the identity matrix However, we will show later that other alignment constraints can be considered. Moreover, when B is the intersection of more constraint sets like CJ n C% in Equation O nature way to get the approximation of the nearest points Pa of Kg is to run iterative projections algorithm [3], which could be extended into the intersection of several constraint sets like {n Cj}, resulting in multi-marginal problem [1
3.2 A Proximal Point Algorithm for GCA
In practice we can solve the alignment problem above by iteratively updating the two main com- ponents in bi-level objective. First, for fixed encoder parameters 0, we obtain the transport coupling Pe through our corresponding proximal operator. Second, we measure the deviation be- tween the transport plan Pe with the target igr that encodes our matching constraints, which denotes the ideal alignment plan on the intersection of the constraint sets. We provide pseudocode for this iterative approach in Algorithm [ which we refer as generalized contrastive alignment or GCA. The implementation of our methods is in
Algorithm 1 Proximal-Point Algorithm for Generalized Contrastive Alignment (GCA)
Initial encoder parameters 0, target transport plan P igg kernel function Kg, the function h(x),, divergences di and das (KL or W1). Initialize transport plan Pa based on f. 2: Compute the transport coupling Pa: Update Po using the proximal operator scaling for fixed
1: Initialization:
described in Eq. R):
Pg arg min min dr (P Ko)). PEB
3: Calculate the loss: Calculate deviation between the target and current transport plans
LoCa = (Pe,
Update networks fo (encoder) and go (projector) to minimize LOCA- 4: Repeat until convergence: Repeat steps and 3 until convergence.
2 (forward-pass) in GCA algorithms could be treated as a specific type of Dykstra's projection algorithms [5], which computes the iterative projection on the intersection of affine convex sets [3] 42 The proofs of convergence are provided in Appendix [B.]
Computing the transport coupling Pe
"With a single constraint set like Of in Equation computing the proximal point only involves a single projection. However, if there are intersecting constraint sets like C n C% solving for the proximal point requires multiple projections before approach the nearest point their intersection.
5

---

3.3 GCA-UOT Method
We can also benefit from the rich literature on optimal transport to build different relaxations of our objective 43 33
35 In particular, we choose to leverage a formulation of unbalanced optimal
transport (UOT) to further relax the marginal constraints [1] in our objective.
In this case, we can add the dual form of dr to the Equation 8 and reformulate our objective as:
min dad (g (Po) + A1h (Pb)(()))))) A (P) 1| + +Ho)
0
Here h and hgg can be different divergence measures (e.g., KL divergence) that penalize deviations from the desired marginals f and v, and λ and A2 are regularization parameters that control the trade- off between the transport cost and the divergence penalties. This relaxation leads to different types of proximal operators which we outline in Appendix B.2 The impact of the entropy regularization parameter on the coupling matrix is studied in Figure A5 along with the number of iterations and corresponding sensitivity is provided in Figure A6
3.4 Modifying the Target Transport Plan to Encode Matching Constraints
Contrastive learning objectives can be cast as minimization of the deviations between the transport plan Pe and the identity matrix, i.c., Pigs
However, our GCA formulation enables learning representations that extend beyond this one-to-one matching constraint. This flexibility allows us to incorporate additional matching constraints informed by domain-specific knowledge. For example, in domain generalization scenarios [23 28 where each batch contains samples from multiple domains, the target alignment plan can be structured as:
P J = [i..]] + a I(D; = Dj,
I.,
Where I(.) is the indicator function, which equals 1 if the condition inside is true and 0 otherwise. D; represents the domain of sample where 2 0 and 2 0. In this case, we can improve the representation by building the block constraints which encode either class information (in supervised setting) or domain information (in across domain generalization, visualized Figure
3.5 Computational Complexity
The forward-pass pass only involves the scaling operations in Equation 7 and doesn't 't affect the complex- ity of the backward- pass. Therefore, GCA methods can be thought of as form of batch normalization operations with adaptive scaling. An analysis of the complexity is provided along with experiments in Appendix [B.1] Our results show that GCA iterations only slightly increase the computational complexity when compared with their single step equivalent (GCA-INCE vs. INCE). However, we found that GCA-UOT is faster than INCE due to the improved symmetry and smoothness of the loss. Moreover, we record the floating point operations per second (Flops) of running GCA methods. We find that GCA-INCE (6.65 MFlops) has 5% more Flops than INCE (6.31 MFlops), while GCA-UOT saves 30% Flops (4.54 MFlops). These results show that our GCA-UOT method is not only superior terms of accuracy but also speed.
4 Building Connections to Different CL Objectives
In this section, we show how the modifica- tion of the different parts of our main objec- tive (dr, das, B, K2) in Equation (8) can be con- nected to different contrastive losses. See Table I for a summary of how different losses can be mapped back our formulation.
Table Compurison of different contrastive alignment objectives. Here we have Cf and C5 as constraint sets (denoted as B3) defined in Equation 4 with their cor- responding indicator function. "Iter" refers to iterative methods.
Methods
du
dr
B
Iter
INCR
KI.
KL.
C
4.1 Connection INCE
GCA-INCE
KI.
KL.
CP C
RINCE (q=l)
W1
KL.
CH
An interesting connection that we can make be- tween GCA main objective and contrastive learn- ing is that we can interpret INCE as a single step iterative GCA objective 511] theorem.
GCA-RINCE (q=1)
W1
KL
n
BYOL
KL
L2
RBxB
This connection can be further summarized through the following

---

Theorem (INCE Equivalence), Let Kg denote the augmentation kernel as in Definition with cosine similarity, dr and d equal to KL-divergence, and constraint set as Of Equation The INCE objective in Equation ( can be re-expressed as a GCA problem in Equation 8 as follows:
min KL (1) Prox (Ko)).
KL
(10)
The proof is contained in Appendix B.3 Theorem 1 shows that the INCE loss can be viewed
solving the matching problems in Equation 3 with row normalization constraints C₁. This connection between GCA and INCE allows us to derive the iterative algorithm for GCA-INCE by running Bregman projection iteratively on both row and column normalization sets
4.2 Connection to RINCE
We introduce the following result to build the connection between our framework and RINCE 12
Theorem (RINCE Equivalence). Let Kg denote the augmentation kernel as in Definition 43. Set
target plan P sgt I, dr equal to the Ki--iivegeeee, (I ( (P)) y + [AI) with
u
q. and diag (
). and constraint set C defined in Equation 44. The RINCE objective in Equation can be re- re-expressed as a GCA problem as follows:
min a (T)) with Pa = Proxee KL (Ko),
(11)
The proof is provided in Appendix B.4.1] As can see, RINCE introduces adjustable parameters q and X, with X controlling the weight of negative samples. while 9 E (0, serves to switch between KL divergence and Wasserstein discrepancy. When q = 1, we have the following theorem:
Theorem 3 (W1 Equivalence). Let Kg denote the augmentation kernel as in Definition
with
cosine similarity. Set target plan Pigs I, dr equal the KL-divergence, das equal the Wasserstein distance (W1) in Definition 2 and the constraint set as C+t defined in Equation (4.
The RINCE object in Equation (2) with q can be re-expressed as GCA problem as follows:
min W1 (Pr |Prox KL (Ko)).
(12)
See Appendix B.5 for the proof.
This connection to RINCE suggests an extended iterative formulation to calculate the coupling plan as the projection point p(x) Pr (Kg) of K,, on the constraint set Cf n C%. In this case, we can write an iterative algorithm for robust alignment called GCA-RINCE as follows:
GCA-RINCE = min - -1 (diag(P) 2t (2-11)) //l-"))/uuuu /// + q (APrgr/u(0)),
(13)
8
where À and are hyperparameters, p(1) = diag(u!) (o diag(v(")), and t is the number of iterations.
4.3 Connection to BYOL
Our framework also allows us to make connections to BYOL (22). BYOL learns by encouraging similarity between positive image pairs, without explicitly conditioning on negative examples. To build this connection, recall that BYOL has the online network parameterized by 0 and target network parameterized by <, where z'p-
fo('') and z = fx() are the normalized outputs of the online and target networks, respectively. A simplified version of the BYOL loss can be written as: LBYOL- = Too (2)) z² Il where qo (2')) is the normalized output after online network and qo the predictor In this case, we can provide the following connection between GCAand BYOL as follows.
Theorem 4 YOL Equivalence). Let (x(x x")) = apl-llãe z I)) denote the augmentation kernel. Set the target plan Pigs = I. dr equal to the L2 di-distance, dıs equal to the di-iivergence, and constraint set as RBx B The BYOL objective can be re-expressed as a GCA problem as follows:
min KL. (I][sa), with S. = Prex IIl
Rax (Sa).
(14)
See the proof in Appendix B.6
'In practice. BYOL also switches the order of views to symmetrize the loss. For ease of discussion, we consider just one pair of views but the same could be argued for the full symmetric version.
7

---

5 Theoretical Analysis
In this section, we aim show how the GCA-methods can improve alignment and uniformity in the latent space [57]. Here, alignment means that the features of the positive samples are as close as possible, while uniformity means that the features of negative samples are uniformly distributed on latent space (see Appendix [CI for formal definitions), These quantities have been studied in a number of related works (57 45). where one can show that improved alignment and uniformity can lead to different benefits in representation learning.
5.1 Improved alignment with GCA
Contrastive learning minimizes the deviation between the target alignment plan with the transport plan in Definition 3 through empirical risk minimization (ERM), Therefore, a tighter bound on the empirical risk corresponds to smaller difference between the ideal alignment with the coupling matrix. We show that this in turn leads to better alignment of the positive views.
Analysis of INCE vs GCA-INCE. GCA-INCE ensures that the final transport plan poo) is closer to the ideal identity matrix compared to the INCE, as we show in the following theorem.
Theorem 5 (Improved Alignment with INCE), Let Kg denote the augmentation kernel in Defini-
tion 3 Set da and dr to the KL-divergence, and Pig: = I. The GCA-INCE loss with converged plan (o0) is lower than the GCA-INCE loss with p(.) in Equation (6 for all t.
The full proof provided in Appendix C.I.I The above theorem tells us that solving Equation (8)
with iterative projection will converge to transport plans P (00) with lower KL divergence than the one-step solution provided by INCE. We can establish the convergence of the p(t) + p(()), based on the convergence of Bregman projection.
Analysis of RINCE GCA-RINCE. GCA also benefits from other Bregman divergences,
like the WDM in RINCE, which provides robustness against distribution shift compared to the
KL-divergence in INCE. GCA-RINCE provides lower bound on the RINCE loss in Equation (2). which allows us to develop tighter bound with p(x) obtained by several proximal steps with GCA.
Theorem 6 (Improved Alignment with RINCE). GCA-RINCE loss with P(() in Equation
13 is
lower than the loss in the Theorem as GCA-RRINCE (P!)) S L (P(2))
A===1
See Appendix C.1.1 for the full proof and an analysis of GCA methods for different choices of dM-
5.2 Improved Uniformity of Representations Through GCA
The improved alignment of GCA-methods comes from maximization of the uniformity under the constraint of intersection Cf' n C% in Equation 44. rather than the constraint set Cf in INCE (see Table Finding the projection of K2 on set of Cf¹ n C% through proximal steps equivalent to solving the dual problem of EOT, which can be summarized through the following theorem.
Theorem 7 (Improved Uniformity). Given the constraint sets Equation ( the optimal transport coupling upon convergence of Equation (6) denoted p(x)) achieves higher uniformity loss
compared to the single-step transport plan (()) obtained by INCE.
The proof is provided in the Appendix C.2 Through loss propagation, we show that the alignment
plan offered by p(x)) will guide the subsequent iterations towards more uniform representations.
5.3 Impacts of GCA on a downstream classification task
We take this one step further and examine the impact of GCA on downstream classification task. For a classification task, using a labeled dataset D = {(x,, y:)} E X x ) where y = [1,..., ......,, M) with M classes, we consider a fixed, pre-trained encoder fo E F X - S. Assume that positive and negative views of n original samples (x, (E(... C. X' are sampled from the data distribution p(x).
In this case, the uniformity loss is equivalent optimizing the downstream supervised classification tasks with cross-entropy (CE) loss when the following two assumptions are satisfied [66].
Assumption 1 (Expressivity of the Encoder). Let define H. is the RKHS associated with the kernel K defined on X, and (H Ka) defined on X with augmentation kernel K,- = (Jo()), So( )re in Definition 3 And we assume that Vg E Ho A (z(z) E H.x.
8

---

Standard Setting
Noisy Setting
Method
CIFAR-10
CIFAR-100
SVHN
ImageNet
CI
CIFAR-100 (Ex)
CIFAR-IOC
CIFAR-IOC (Ex)
INCF
92.01 0.40
70.07 + 0.42
90.60
73.01 +661
82.03 0.33
54.70 004
87.20
74,84
GCA-INCE A
92.36.0244
00.5
00.4
#2.18
54.91
7
-0.35
+0.04
-0.20
40.03
+0.15
40.21
+0.14
+1.16
RINCE
91.05
69.06
90,97.0.19
71.91
1
55.43
88 68.62
77.05 0
OCA-RINCE
92.09 0.22
69.72 4
21.45
23.44
+..54
82.76 0.49
=0.1
0.22 S8.
72.23 0.06
-1.04
+0.66
+0.48
+1.53
+0.16
+0.47
+0.14
+0.18
SimCLR
0111
69.95 a 0.14
90.24
72.20 a 0.78
81.87
0.79
86.98 a 1.59
73.79 0.32
BYOL
90.56 0.59
69.7550..5
00.5000.4
69.75 0.83
81.55 0.50
54,18 0.46
87.88 + 1.02
69.40 1.11
67.19
0.21
72.27
52.40
58.75
ICTT-nii
R1
90.89 0.57
67.03 + 0.40
90.54
0.20
72.88 0.71
80.79+0..3
53.04 0.52
69.58 1.25
59.05 1.86
92.61
71,45
74,09 0,.00
83.18
56.30 0.51
89.61
77.60
Table 2: Test accuracy (%) on a downstream classification task after pretraining. Results are provided for
CIFAR-IO (ResNet18). CIFAR-100 (ResNet18), SVHN (ResNet50), and ImageNet100 (ResNet50) under standard and extreme (Ex) augmentation conditions (averaged over 5 seeds). The top model is bold and the second-place model is underlined. For INCE and RINCE, we also provide the improvement A by adding GCA each method.
Assumption 2 (Small Intra-Class Variance). For + y', the intra-class variance d;, 65 are negligible
compared to the distance among different class centroids, Pay Hyf as 1ily - py || > |6,, - 6,ll.
Claim 1. If Assumption 7 and Assumption 2 hold, then maximizing the uniformity is equivalent to
minimizing the downstream CE loss.
The proof is provided in Appendix C.2 Optimizing the self-supervised loss under ideal conditions improves downstream CE tasks and helps to explain why maximizing uniformity aids classification.
Remark.. Maximizing uniformity can enhance downstream classification but risks "feature suppres- sion" by encouraging shortcut features that harm generalization [8]. In GCA-UOT, adding penalties modifies the transport plan from that of a pure uniformity loss, helping to avoid feature suppres- sion. We find empirical evidence that UOT provides a more robust transport plan which appears to circumvent some of these shortcut features from being learned (Figure A4 in Appendix C.3
6
Experiments
In this section, we conduct empirical evaluations to study the performance of our approach both handling noisy and corrupted views and in domain generalization tasks.
6.1 Comparison with CL Baselines
Experiment setup. To examine the robustness of our framework, we trained INCE and RINCE as baselines, and developed their GCA-based alternatives (+GCA). In addition, we also compared with our novel GCA-UOT method, two variants IOT established in [51], and other CL baselines, including BYOL and SimCLR. For experiments with SVHN (36] and ImageNet100 [15] we use the ResNet-50 encoder as the backbone and use a ResNet- 18 encoder as the backbone for 10, CIFAR- 100 [29] and corrupted version of CIFAR called CIFAR- (25).
In all of these cases, we follow the standard sulf-supervised learning evaluation protocol [ where we train the encoder on the training set in an unsupervised manner and then train a linear layer on top of the frozen representations to obtain the final accuracy on the test set. In addition to standard data augmentation policies commonly used [12), we also apply three different extreme augmentation policies to examine the robustness of GCA towards noisy views (details in Appendix [.2). Learning rates and other training details for CIFAR-10, CIFAR- 100, SVHN, and ImageNet1000 100 are provided in Appendix D. while specific training details for CIFAR- are included in Appendix D.2
Results on Standard Augmentations. First, performed experiments on CIFAR-10. 10, CIFAR- 100, SVHN and ImageNet100 using standard sets of augmentations that are applied to achieve state-of-the-art performance (Table 2 Standard Setting). We found the +GCA versions of INCE and RINCE exhibit performance gains in almost all settings except for SVHN, with bigger gains observed when adding GCA to RINCE. Additionally, we find that our unbalanced OT method, GCA-UOT, achieves the top performance across the board, on all four datasets tested. The transport plans obtained by each methods are provided in Figure A4 along with a study of the sensitivity of the methods to hyperparameters (Appendix (77).
Results on Corrupted Data and Extreme Augmentations. noisy settings. In the first set of experiments, we apply extreme augmentations to CIFAR-10 (Ex) and
Next, we tested the methods two

---

CIFAR- 100 (Ex) (see Appendix D.2) to introduce noisy views during training. In the second set of experiments, we used the CIFAR- IOC to further test the ability of our method to work in noisy settings. Our experimental results demonstrate that the GCA-based strategy effectively enhances the model's generalization ability and adaptability to aggressive data augmentations. classification accuracy, the GCA-based methods also improve the representational alignment and uniformity, as shown in Appendix E.2 This observation in line with our theoretical analysis in Section 5.2. where we show that the obtained representations provide better overall alignment of positive views and better spread terms of uniformity [57].
addition to improving
6.2 Block Diagonal Transport in Domain Generalization
(A)
In a final experiment, we aimed to demonstrate the flexi- bility and robustness of our framework by applying it to a domain generalization task, where samples originate from different domains (e.g., Photo, Cartoon, Sketch, Art). We explored the effects of introducing domain-speciffe align- ment constraints in our transport plan, hypothesizing that this could enhance the latent space organization to capture more nuanced domain similarities.
DI
D2
D4
U22
14
(B)
0.90
0.30
Our approach enables additional contextual information to be seamlessly integrated into the transport process. In this case, domain information was incorporated to distinguish the alignment of samples from the same versus different do- mains. To achieve this, we adjusted the target transport plan Pigt. selectively modifying parameters (a, B) to vary the influence of domain-based alignment constraints as shown in Figure AA) Specifically, set {a 0, 0} to prior-
0.85
0.35
-0.4
-0.2
0.2
(ct (3)
Figure 1: Incorporating different priors into learning across multiple domains. (A) Example target alignment plan Pige. where the target over all samples from the same domain set a, the diagonal values set to 1, and across-domain samples are set to B. (B) The domain classification accuracy (red) and overall class accuracy (blue) with (a () increases.
itize cross- rros--domann alignment and {a > 0, B = 0} to focus on intra-domain aligmment.
The training was conducted on the PACS dataset 31 using a ResNet-18 18 encoder with the GCA-INCE objective. After training the encoder in an unsupervised manner, we freeze the encoder and then train linear readout layer predict either the sample's s class or the domain it belonged to. This setup allowed us to isolate the effect of our transport adjustments on the latent space's capacity to encode both class and domain information.
The results, displayed in Figure B), revealed that increasing the domain alignment weight enhances the accuracy of domain classification (from 72. 1% to 15.16%) without diminishing classification performance. This outcome suggests that GCA can effectively encode both domain and class information in a single latent representation. The ability to adjust alignment constraints provides powerful tool for domain generalization tasks, enabling multiple types of similarity to be jointly encoded. This flexibility can potentially alleviate issues related to information loss from data augmentation, especially in fine-grained classification settings, by retaining essential domain-specific characteristics across transformations.
7 Conclusion
In this work, we introduced generalized contrastive alignment (GCA), a flexible framework that redefines contrastive learning a distributional alignment problem using optimal transport to control alignment. By allowing targeted control over alignment objectives, GCA demonstrates strong performance across both standard and challenging settings, such as noisy views and domain generalization tasks. This work opens up broader possibilities for learning robust representations in real-world scenarios, where data is often diverse, noisy, or comes from multiple domains.
Future work includes applications GCA to graphs and time series data, as well as multi-modal settings where our approach can integrate various forms of similarity. As alignment strategies become integral to contrastive learning, GCA offers a promising foundation for more adaptive and expressive self-supervised models.
10

---

Acknowledgements
We would like to thank Mehdi Azabou, Divyansha, Vinam Arora, Shivashriganesh Mahato, and Ian Knight for their valuable feedback on the work. This work funded through NSF IIS-2212182, NSF IIS 2039741, and the support from the Canadian Institute for Advanced Research (CIFAR). We would also like to acknowledge the use of ChatGPT for providing useful feedback and suggestions on the writing of the paper.
References
[1] Dongsheng An, Na Lei, Xiaoyin Xu, and Xianfeng Gu. Efficient optimal transport algorithm by accelerated gradient descent. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 10119-10128, 2022.
[2] Martin Arjovsky, Soumith Chintala, and Léon Bottou. Wasserstein generative adversarial networks. International conference on machine learning, pages 214-223 PMLR, 2017.
[3] Jean-David Benamou, Guillaume Carlier, Marco Cuturi, Luca Nenna, and Gabriel Peyré. Itera- tive bregman projections for regularized transportation problems. SIAM Journal on Scientific Computing, 37(2):A1111AA 11111111111,,1111111111111111111 2015.
[4] Robert Berman. The sinkhorn algorithm, parabolic optimal transport and geometric monge- ampère equations. Numerische Mathematik, 145(4) 2020.
[5] Lev M Bregman. The relaxation method of finding the common point of convex sets and its application to the solution of problems in convex programming. USSR computational mathematics and mathematical physics, 7(3):200-217, 1967.
[6] Xing-Ju Cai, Ke Guo, Fan Jiang, Kai Wang, Zhong-Ming Wu, and De-Ren Han. The devel- opments of proximal point algorithms. Journal of the Operations Research Society of China, 10(2):197-239, 2022.
[7] Liqun Chen, Dong Wang, Zhe Gan, Jingjing Liu, Ricardo Henao, and Lawrence Carin. Wasser- stein contrastive representation distillation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 16296-16305, 2021.
[8] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference machine learning, pages 1597-1607. PMLR, 2020.
[9] Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 15750-15758, 2021.
[10] Zining Chen. Weiqiu Wang, Zhicheng Zhao, Fei Su, Aidong Men, and Yuan Dong. Instance paradigm contrastive learning for domain generalization. IEEE Transactions on Circuits and Systems for Video Technology, 34(2): 2023.
(11) Lenaic Chizat, Gabriel Peyré, Bernhard Schmitzer, and François-Xavier Vialard. Scaling algo- rithms for unbalanced optimal transport problems. Mathematics of Computation, 87(314):2563- 2609, 2018
[12] Ching- Yao Chuang, R Devon Hjelm, Xin Wang, Vibhav Vineet, Neel Joshi, Antonio Torralba, Stefanie Jegelka, and Yale Song. Robust contrastive learning against noisy views. In Proceedings of the IEEE/CVF Conference on Computer Vision and Patlern Recognition, pages 16670-16681, 2022.
[13] Ching-Yao Yao Chuang. Joshua Robinson, Yen-Chen Lin, Antonio Torralba, and Stefanie Jegelka. Debiased contrastive learning. Advances neural information processing systems, 33:8765- 8775, 2020.

---

[14] Nicolas Courty, Rémi Flamary, and Devis Tuia. Domain adaptation with regularized optimal transport. In Machine Learning and Knowledge Discovery in Databases: European Conference, ECMI. PKDD 2014, Nancy, France, September 15-19, 2014. Proceedings, Part 14, pages 274-289. Springer, 2014.
[15] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large- scale hierarchical image database. In 2009 IEEE Conference on Computer Vision and Pattern Recognition, pages 248-255. Icce, 2009.
[16] Benoit Dufumier, Carlo Alberto Barbano, Robin Louiset, Edouard Duchesnay, and Pietro Gori Integrating prior knowledge in contrastive learning with kernel. In International Conference on Machine Learning, pages 8851-8878. PMLR, 2023.
[17] Marvin Eisenberger, Aysim Toker, Laura Leal Taixé, Florian Berarrd, and Daniel Cremers. A unified framework for implicit sinkhorn differentiation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 509-518, 2022.
[18] Zhenghan Fang, Sam Buchanan, and Jeremias Sulam. What's in a prior? learned proximal networks for inverse problems. arXiv preprint arXiv:2310. 14344, 2023.
[19] Promit Ghosal and Marcel Nutz. On the convergence rate of sinkhorn's algorithm. arXiv preprint arXiv:2212.0600 2022.
[20] Aritra Ghosh, Naresh Manwani, and PS Sastry. Making risk minimization tolerant to label noise. Neurocomputiing. 160:93-107, 2015.
[21] Will Grathwohl, Kuan-Chich Wang, Jörn- törn-Henrik Jacobsen, David Duvenaud, Mohammad Norouzi, and Kevin Swersky. Your classifier is secretly an energy based model and you should treat it like one. arXiv preprint arXiv: 1912.03223 2019.
[22] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural information processing systems, 13:21271-21233 2020.
[23] Ishaan Gulrajani and David Lopez-Paz. In search of lost domain generalization, arXiv preprint MrXiv:2007.01434 2020.
[24] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9729- 9738, 2020.
[25] Dan Hendrycks and Thomas G Dietterich. Benchmarking neural network robustness to common corruptions and perturbations. In International Conference on Learning Representations (ICLR), 2019.
[26] 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.
[27] Longlong Jing and Yingli Tian. Self-supervised visual feature learning with deep neural net- works: A survey, IEEE transactions on pattern analysis and machine intelligence, 43(11):4037- 4058, 2020.
[28] Dachee Kim, Youngjun Yoo, Seunghyun Park, Jinkyu Kim, and Jackoo Lee. Selfreg: Self- supervised contrastive regularization for domain generalization. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pages
2021.
[29] Alex Krizhevsky and Geoffrey Hinton. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
[30] John Lee, Max Dabagia, Eva Dyer, and Christopher Rozell. Hierarchical optimal transport for multimodal distribution alignment. Advances in neural information processing systems, 32, 2019.
12

---

[31] Da Li, Yongxin Yang, Yi-Zhe Song, and Timothy M Hospedales. Deeper, broader and artier domain generalization. In Proceedings of the IEEE international conferer on computer vision, pages -5520 2017.
[32] Ruilin Li, Xiaojing Ye, Haomin Zhou, and Hongyuan Zha. Learning to match via inverse optimal transport. Journal of machine learning research, 20(80):1 2019.
[33] Chi-Heng Lin, Mchdi Azabou, and Eva L Dyer. Making transport more robust and interpretable by moving data through small number of anchor points. Proceedings of machine learning research, 139:6631, 2021.
[34] Ran Liu, Mehdi Azabou, Max Dabagia, Chi-Heng Lin, Mohammad Gheshlaghi Azar, Keith Hengen, Michal Valko, and Eva Dyer. Drop, swap, and generate: A self-supervised approach for generating neural activity, Advances in neural information processing systems, 34: 10587-1599. 2021.
[35] Eduardo Fernandes Montesuma, Fred Ngole Mboula, and Antoine Souloumiac. Recent advances in optimal transport for machine learning. arXiv preprini arXiv:2306. /6/56, 2023.
[36] Yuval Netzer, Tao Wang. Adam Coates, Alessandro Bissacco, Bo Wu, and Andrew Y Ng. Reading digits in natural images with unsupervised feature learning. In NIPS Workshop on Deep Learning and Unsupervised Feature Learning, volume 2011,2011
[37] XuanLong Nguyen, Martin J Wainwright, and Michael Jordan. Estimating divergence func- tionals and the likelihood ratio by convex risk minimization. IEEE Transactions on Information T'heory 56(1 1):5847-5888 58611, 2010.
[38] Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv:1807.03744 2018
[39] Sherjil Ozair, Corey Lynch, Yoshua Bengio, Aaron Van den Oord, Sergey Levine, and Pierre Sermanet. Wasserstein dependency measure for representation learning. Advances in Neural Information Processing Systems, 32, 2019.
[40] Neal Parikh, Stephen Boyd, et al. Proximal algorithms. Foundations and trends@ in Optimiza- tion, (3))1277239 2014.
[41] Brendan Pass. Multi marginal optimal transport: theory and applications. ESAIM: Mathematical Modelling and Numerical Analysis, 49(6):177--1790, 2015.
[42] Gabriel Peyré. Entropic approximation of wasserstein gradient flows. SIAM Journal on Imaging Sciences, 8(4):2323-2351, 2015.
[43] Gabriel Peyré, Marco Cuturi, et al. Computational optimal transport: With applications to data science. Foundations and Trends@ in Machine Learning. 1165-6):35--67,, 2019.
[44] Khiem Pham, Khang Le, Nhat Ho, Tung Pham, and Hung Bui. On unbalanced optimal transport: An analysis of sinkhorn algorithm. In International Conference on Machine Learning, pages 7673-7682. PMLR, 2020.
[45] Shi Pu, Kaili Zhao, and Mao Zheng. uligament-uifforiity aware representation learning for zero-shot video classification. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 19968-19977, 2022.
[46] 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. pages 8748-87666 8763. PMLR, 2021.
[47] Joshua Robinson, Ching Yao Chuang, Suvrit Sra, and Stefanie Jegelka. Contrastive learning with hard negative samples. arXiv preprint arXiv:2010.04592 2020.
[48] Joshua Robinson, Li Sun, Ke Yu, Kayhan Batmanghelich, Stefanie Jegelka, and Suvrit Sra. Can contrastive learning avoid shorteut solutions? Advances neural information processing systems, 34:4974 49966 1021.
13

---

[49] Litu Rout, Alexander Korotin, and Evgeny Burnaev. Generative modeling with optimal transport maps, arXiv preprint arXiv:2110.02999, 2021.
[50] Nikunj Saunshi, Jordan Ash, Surbhi Goel, Dipendra Misra, Cyril Zhang, Sanjeev Arora, Sham Kakade, and Akshay Krishnamurthy. Understanding contrastive learning requires incorporating inductive biases. In International Conference on Machine Learning, pages 19250-19286. PMLR, 2022.
[51] Liangliang Shi, Gu Zhang, Haoyu Zhen, Jintao Fan, and Junchi Yan. Understanding and generalizing contrastive learning from the inverse optimal transport perspective. In International Conference Machine Learning, pages 31408-31421. PMLR, 2023.
[52] Saeed Shurrab and Rehab Duwairi. Self-supervised learning methods and applications in medical imaging analysis: A survey. PeerJ Computer Science, 8:e1045, 2022.
[53] Andrew M Stuart and Marie-Therese Wolfram. Inverse optimal transport. SIAM Journal on Applied Mathematics, 80(1):599-619 2020.
[54] Fariborz Taherkhani, Ali Dabouei, Sobhan Soleymani, Jeremy Dawson, and Nasser M Nasrabadi. Self-supervised wasserstein pscudo-labeling for semi-supervised image classification. In Pro- ceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 12267-12277, 2021.
[55] lya Tolstikhin, Olivier Bousquet, Sylvain Gelly, and Bernhard Schoelkopf. Wasserstein auto- encoders. arXiv preprint 17X1v:/711..05555: 2017.
[56] Ankit Vishnubhotla, Charlotte Loh, Akash Srivastava, Liam Paninski, and Cole Hurwitz. Towards robust and generalizable representations of extracellular data using contrastive learning. Advances in Neural Information Processing Systems, 36, 2024.
[57] Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through alignment and uniformity on the hypersphere. International Conference on Machine Learning, pages 9929-99999999 9929999399....199999999999999999 PMLR, 2020.
[58] Yule Wang, Chengrui Li, Weihan Li, and Anqi Wu. Exploring behavior-relevant and disen- tangled neural dynamics with generative diffusion models. arXiv preprint arXiv:2410.09614, 2024.
[59] Yule Wang. Zijing Wu, Chengrui Li, and Anqi Wu, Extraction and recovery of spatio-temporal structure in latent dynamics alignment with diffusion model. Advances in Neural Information Processing Systems, 36, 2024.
14

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: This paper is theoretical paper which discusses the generalized framework for contrastive learning, which involves to convert them into series proximal algorithms. The abstract and instruction illustrate this point properly.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results
be expected to generalize other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Contrastive learning is still used in an ad hoc manner, with augmentations often causing harmful effects. This in turn can introduce bias or hallucinations which can negative societal impacts. We will expand on these topics in future revision. Since we disscuss the limitations and will introduce future works the conclusion, the answer should be yes.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model vell-specificattin, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
15

---

3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [Yes]
Justification: Please check the theoretical part in the main text, which is mainly contained in section 3 3 and Their proofs are provided in the Appendix, where each theory has its corresponding proof.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: Please check Section and Appendix for experimental details.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
also be provided via detailed
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors welcome describe the particular way they provide for reproducibility.
16

---

In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, described in supplemental material?
Answer: [Yes]
Justification:
We will release all codes after review.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (nttps://niis.cc)
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
7 //nips. /twpps.ccppublli/// httpcc/pccccccceeeeeeeeeeeeeeeee wwtpc//guiies/e/eee http:/uuddeesooeduubiiiiiiei
for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data,
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/eetails
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Question:
Answer: [Yes]
Justification: Please check Section 6 and Appendix for details.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: Please check Section Gnd Appendix for details.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by emor bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
17

---

The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification:
We provide GPU information in Appendix Dalong with experimental details.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: http:::///eeurreee nwwrips....... ctpss.cc/pubbee cccp
Answer: [Yes]
Justification:
This is theoretical work and fullfills the Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- cration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: By providing more provable and robust methods for representation learning, this work can have impact in fairness and help to reduce uncertainty in decision making.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
18

---

If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This thcoretical paper has no risk that we are aware of.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Are the creators or original owners of assets (e.g.. code, data, models), used in the paper, properly credited and the license and terms of use explicitly mentioned and properly respected?
Question:
Answer: [Yes]
Justification: Please check references.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
19

---

If assets are released, the license, copyright information, and terms of use in the thcode. has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
package should be provided. For popular datasets, paperswi
latasets.paperssit
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information not available online, the authors are encouraged to reach out to the asset's s creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not introduce new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moode as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: This paper does not include crowdsourcing or human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country institution) were obtained?
Answer: [NA]
Justification: This paper does not include human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
20

---

We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
21


Reviews:
[{"review_id": "KdsgVsLMJS", "reviewer": "Reviewer_96Mt", "summary": "This paper proposed to view contrastive learning (CL), a popular framework for learning data representation in machine learning. Specifically, the work builds on some recent works that view CL as an alignment problem with optimal transport, showing some previous popular CL frameworks are the special form of this new framework, named generalized contrastive alignment (GCA). Leveraging the theory of unbalanced optimal transport, the authors also introduced an unbalanced CL loss to handle outliers. Empirical benchmarks on the CIFAR-10 and CIFAR-100 datasets are performed to show the effectiveness of GCA.", "questions": "See weaknesses.", "limitations": "Limitation is discussed, but an important point I raised on the weaknesses section is missing.", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 1, "strengths": "1. The methodology and theory introduced is sound and based on the existing theory of optimal transport and its entropic regularization version. \n\n2. The writing is good and easy to follow. The authors also seem to have done a thorough literature review related to their work.", "weaknesses": "1. **Major: there is a large overlap, not in the writing, but in the idea and content of this work and Shi et al (2023).** In particular, the most important one that I want to point out is Algorithm 1 in this work and Algorithm 1 in Shi et al (2023). If one replaces the Bregman div $d_\\Gamma$ and $d_M$ by KL divergence into the proximal loss of the inner optimization (eq 8), and $h(P)$ as in the indicator function, one recovers exactly the proposed loss in Shi et al (2023), Eq 8. In the empirical evaluation part, I believe the authors of this work used the aforementioned setting as well, hence I wonder what is the major difference between the main proposed method and that of Shi et al (2023). The idea of using unbalanced OT and its connection to MoCo also has been mentioned in Shi et al (2023, end of Section 3.2).\n\n2. The empirical evaluation is interesting, but I do not find it enough. I believe the authors should have included Shi et al. (2023) as a baseline. The evaluation also should have been done equally on larger dataset such as ImageNet. \n\n3. This is more philosophical point: I am not sure whether if we want the transport plan $P$ to be smooth. Therefore I wonder how did the authors tune their entropic regularization $\\varepsilon$ parameter? In computational optimal transport there is a well-known tradeoff between the sparsity of the solved transport plan (in the proximal step of the inner problem) and the stability of scaling operation: if the value $\\varepsilon$ (which makes the plan more sparse) the operation will be quite unstable. The Sinkhorn update will run better and faster with larger $\\varepsilon$, but the transport is not sparse anymore. I also suggest the authors do a visualization of the quality of solution $P$, similar to what has been done in Figure 3 of Shi et al (2023). \n\n- Small nitpick: Eq (1) missing expectation wrt to samples.\n\n\nLiangliang Shi, Gu Zhang, Haoyu Zhen, Jintao Fan, and Junchi Yan. Understanding and generalizing contrastive learning from the inverse optimal transport perspective. In International Conference on Machine Learning, pages 31408–31421. PMLR, 2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your comments and detailed feedback. We appreciate the opportunity to discuss the contributions of our work in connection to Shi et al. and also expand our comparisons based upon your suggestion.\n\n1. Discussion of overlap with Shi et al.\n\n**Reply:**  Thanks for your comment and the opportunity to discuss the differences between our work and theirs. As you point out, Shi et al. have provided initial connections between INCE and optimal transport to build a multistep algorithm for alignment using Sinkhorn iterations. In our work, in order to incorporate more general losses and alignment objectives, we have made a number of novel contributions:\n\n- **New algorithm for generalized alignment for contrastive learning:** To allow for more generalized losses, our algorithm allows the intersection of new constraint sets to be iteratively solved, while previous work mainly focuses on the solution of alignment exclusively through Sinkhorn iterations. \n\n- **Novel approach for unbalanced OT-based alignment:** We leverage a rich body of work in OT to introduce a variant of GCA that relaxes the constraints on the distribution penalty (Sec. 3.3). By converting the hard penalty (constraint sets) into the soft regularization terms, our GCA-UOT method achieves high classification accuracy (Table 2) and faster convergence than INCE (Fig. A3), linking OT literature with optimization and contrastive learning.\n\n- **Connections and a multistep variant of RINCE:**  By building a more generalized form of alignment, we demonstrate that it's possible to develop connections to the Robust INCE loss, RINCE. This equivalence enables us to develop a multistep RINCE variant that performs better with corrupted views.\n\n- **Novel results in domain generalization through block-diagonal matching constraints:** By changing the target plan to have block diagonal structure, we can absorb the domain information (Sec. 6). Adding domain-specific matching constraints can improve the pre-training model and enhance classification accuracy in cross-domain generalization tasks.\n\n- **New theory and insights:**  We provide the illustration and prove for the convergence of our more generalized algorithms, not just for the KL divergence in sinkhorn situations, but for other Bregman divergences, which is not shown in previous algorithms and develop new results to explain why running GCA could lead to better uniformity and alignment. \n\nIn summary, our GCA framework provides a foundation for addressing a wider range of potential issues in contrastive learning.\n\n2. Difference between GCA-UOT and the idea of “unbalanced matching” in Shi et al.\n\n**Reply:**   Thanks for your question. The usage is actually quite different. In our case, we are using the term “unbalanced OT” to refer to a large body of work [1,2] in the OT literature that focuses on relaxing the constraint on the distribution matching across the source and the target domain. This often involves building an unconstrained optimization objective and relaxing our constraints on distribution matching. We find that this relaxation has two main advantages: (i) improved accuracy  and (ii) improved complexity and faster speed of convergence (Fig. A3).\n\nIn contrast to our method, Shi et al. introduced the idea of “unbalanced matching” which models the fact that the two different encoders in twin network approaches like Moco consider views from two different encoders (often tied via a momentum term).  Our approach, GCA-UOT, employs unbalanced optimal transport which aims to convert hard penalty (constraint sets) into the soft penalty like L2 regularization terms. \n\n[1] Xu, M., & Gould, S. (2024). Temporally Consistent Unbalanced Optimal Transport for Unsupervised Action Segmentation. arXiv:2404.01518.\n\n[2] De Plaen, H., et al. (2023). Unbalanced optimal transport: A unified framework for object detection. CVPR, 3198-3207.\n\n3. “I believe the authors should have included Shi et al. (2023) as a baseline.”\n\n**Reply:**  Thanks. Indeed, there are some differences in the implementation of Shi et al. version of MS-INCE in comparison to ours as we use a dual form of the OT objective and Shi uses the Sinkhorn algorithm. We tried to obtain their exact implementation but couldn’t find code online so we implemented it and ran this method as another baseline. Please see the (see Table R1, general response) below for the comparison with their IOT method on ImageNet100, SVHN. We are working on adding this baseline for all of the results in the paper.\n\n4. “Evaluation also should have been done equally on a larger dataset such as ImageNet.”\n\n**Reply:**  Thanks for your suggestion. Rather than focusing on larger datasets with more classes like ImageNet, we chose to explore new settings like domain generalization and robust variants of CL where we could demonstrate the versatility of the approach. Based upon the reviewer's suggestions, we ran additional experiments on ImageNet-100 and SVHN (see Table R1 in the general response). Our results show that our alignment methods improve over standard contrastive losses for the same backbone. Consistently, we find that UOT gives even further improvements over the other methods. \n\n5. “How did authors tune their entropic regularization ε  parameter? [...] I also suggest the authors do a visualization of the quality of solution P, similar to what has been done in Figure 3 of Shi et al”\n\n**Reply:**  Thank you. We have computed the same plots (Fig. U1) and include an analysis of how the epsilon changes performance.  A smaller epsilon results in a sparser optimal transport plan but potentially leads to numerical instability. We chose a small epsilon (ε=0.2) for CIFAR-10 through cross-validation, however, we find that our alignment method is not very sensitive to this choice. We also examined different ε values for our unbalanced OT method on CIFAR-10 (Fig. U2). The results suggest that an ε range of 0.2-0.6 achieves good compactness and high accuracy.", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "Pd8mjQMa3P", "reviewer": "Reviewer_ih3d", "summary": "This paper introduces a framework called Generalized Contrastive Alignment (GCA) that connects contrastive learning to distribution alignment using optimal transport. The key contributions include:\n\n1. Establishing a novel class of losses and algorithms for representation learning through GCA, showing how different contrastive losses can be interpreted as variants of a generalized distribution alignment objective.\n\n2. Proving convergence of GCA-based methods and demonstrating theoretically that the alignment objective can improve the quality of learned representations.\n\n3. Empirically validating GCA's effectiveness in image classification and domain generalization tasks, showing it can achieve superior performance over baseline methods.\n\n4. Demonstrating how GCA allows building unbalanced losses using tools from optimal transport, which can handle noisy views and customize representations.\n\n5. Providing a unified framework that connects existing contrastive learning methods like InfoNCE, RINCE, and BYOL to optimal transport formulations.\n\n6. Showing how modifying the target alignment plan in GCA can flexibly control the amount of domain knowledge incorporated into representations for domain generalization tasks.\n\nThe paper provides both theoretical analysis and experimental results to support the benefits of the GCA framework in improving representation learning and offering more flexibility compared to standard contrastive learning approaches. The authors position this work as providing new insights into the connections between self-supervised learning models and offering tools to more easily incorporate domain knowledge into learning.", "questions": "- Can you provide a detailed analysis of the computational overhead introduced by GCA, particularly for the multi-step variants?\n- How do the training time and memory usage compare to standard contrastive learning methods?\n- Is there a point of diminishing returns regarding performance improvement vs. computational cost?\n- What guidelines can you provide for setting the new hyperparameters introduced in GCA (e.g., number of iterations, α and β in the customized transport plan)?\n- Have you observed any patterns in how these hyperparameters affect performance across different tasks or datasets?\n- Can you provide an ablation study showing how performance changes with different numbers of iterations in the forward pass?\n- Have you considered comparing GCA with other alignment-based or optimal transport-based representation learning methods? If so, what were the results?\n- How does GCA differentiate itself from or improve upon these existing alignment-based approaches?\n- Are there any scenarios in which the multi-step approach might lead to worse performance or slower convergence than single-step methods?\n- Can you elaborate on any potential limitations or failure cases of the GCA approach?\n- Can you provide more details or experiments demonstrating the specific benefits of unbalanced OT over balanced OT in different scenarios?\n- Are there particular types of tasks or data where unbalanced OT shows the most significant improvements?\n- Can you provide more concrete guidelines on how practitioners should choose between different GCA variants (GCA-INCE, GCA-RINCE, GCA-UOT) for a given task?\n- Are there specific task characteristics that make one variant more suitable than others?\n- In the domain generalization experiments, how sensitive is the performance to the choice of α and β in the customized transport plan?\n- Have you explored using different transport plans for domains or tasks within the same dataset?\n- Can you provide more details on how GCA improves robustness to noisy or extreme data augmentations compared to baseline methods?\n- Are there specific types of noise or augmentations where GCA shows the most significant improvements?", "limitations": "The authors have not explicitly addressed limitations or potential negative societal impacts in the paper as it currently stands. While they do mention some areas for future work in the conclusion, a more comprehensive discussion of limitations and societal impacts would strengthen the paper. Here are some constructive suggestions for improvement:\n\n1. Limitations:\n   - The authors could add a dedicated \"Limitations\" section discussing:\n     a) The current scope of experiments (e.g., limited to certain datasets and architectures)\n     b) Potential computational overhead of the multi-step approach\n     c) Challenges in hyperparameter tuning for the new parameters introduced\n     d) Any scenarios where GCA might not be applicable or beneficial\n\n2. Potential negative societal impacts:\n   - The authors should consider adding a brief discussion on potential societal impacts, such as:\n     a) Increased energy consumption due to potentially higher computational requirements\n     b) Possible biases in learned representations, especially when using customized transport plans\n     c) Potential misuse of the technique in privacy-sensitive applications\n\n3. Ethical considerations:\n   - A brief note on any ethical considerations in data usage or potential applications of the method would be valuable.\n\n4. Future work:\n   - Expand the current mention of future work to include specific directions for addressing identified limitations.\n\nIt's important to note that the absence of these discussions doesn't necessarily indicate oversight by the authors, but rather an opportunity to enhance the paper's comprehensiveness. Adding these elements would align well with NeurIPS guidelines and contribute to responsible research practices in the field.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper demonstrates several strengths across the dimensions of originality, quality, clarity, and significance:\n\n**Originality:**\n1. The paper presents a novel framework (GCA) that creatively bridges contrastive learning and optimal transport. This connection is not entirely new, but the comprehensive treatment and generalizations provided are original.\n2. The formulation of contrastive learning as a distribution alignment problem offers a fresh perspective on a widely-studied topic.\n3. The extension to unbalanced optimal transport and the ability to customize target transport plans are innovative additions to the contrastive learning toolkit.\n\n**Quality:**\n1. The theoretical analysis is rigorous, with clear proofs for the connections between GCA and existing methods (InfoNCE, RINCE, BYOL).\n2. The empirical evaluation is comprehensive, covering multiple datasets (CIFAR-10, CIFAR-100, CIFAR-10C, PACS) and scenarios (standard classification, extreme data augmentation, domain generalization).\n3. The ablation studies and comparisons against baseline methods are thorough and well-presented.\n\n**Clarity:**\n1. The paper is well-structured, with a clear progression from problem formulation to theoretical analysis and empirical validation.\n2. The use of diagrams (e.g., Figure 1) helps illustrate complex concepts like customized transport plans.\n3. The authors provide clear pseudocode (Algorithm 1) for the GCA method, enhancing reproducibility.\n\n**Significance:**\n1. The GCA framework provides a unifying perspective on several popular contrastive learning methods, which could facilitate further theoretical developments in the field.\n2. The improved performance on classification tasks, especially under extreme data augmentation, suggests practical benefits for real-world applications.\n3. The flexibility offered by customizable transport plans in domain generalization scenarios opens up new possibilities for tailoring representations to specific tasks or domains.", "weaknesses": "While the paper has many strengths, there are several areas where it could be improved:\n\n1. Limited scope of empirical evaluation:\n   - The experiments are primarily conducted on relatively small datasets (CIFAR-10, CIFAR-100, PACS). While these are standard benchmarks, the absence of results on larger, more complex datasets like ImageNet limits the assessment of GCA's scalability and practical impact.\n   - The authors acknowledge this limitation in their conclusion, but providing some preliminary analysis or discussion on potential challenges in scaling to larger datasets would be beneficial.\n\n2. Computational complexity:\n   - The paper lacks a detailed analysis of the computational overhead introduced by the GCA framework, particularly for the multi-step variants.\n   - While Algorithm 1 is provided, there's no discussion on how the additional forward passes impact training time or memory requirements compared to standard contrastive learning methods.\n\n3. Hyperparameter sensitivity:\n   - The paper introduces several new hyperparameters (e.g., number of iterations, α and β in the customized transport plan), but there's limited discussion on their impact or guidelines for setting them.\n   - An ablation study or sensitivity analysis for these parameters would provide valuable insights for practitioners looking to implement GCA.\n\n4. Comparison with other alignment-based methods:\n   While the paper compares GCA with standard contrastive learning baselines, it does not compare it with other alignment-based or optimal transport-based representation learning methods (e.g., [1], [2]).\n   - Such comparisons would help better contextualize the contributions of GCA within the broader landscape of alignment-based approaches.\n\n5. Theoretical limitations:\n   - The theoretical analysis, while thorough, focuses primarily on convergence and improved alignment. Discussing potential limitations or failure cases of the GCA approach would be beneficial.\n   - For instance, are there scenarios where the multi-step approach might lead to worse performance or slower convergence?\n\n6. Ablation on the number of alignment steps:\n   - While the paper mentions using 5 iterations in the forward pass, there's no analysis on how performance changes with different numbers of iterations.\n   - An ablation study showing the trade-off between computational cost and performance improvement as the number of iterations increases would be valuable.\n\n7. Limited exploration of unbalanced OT:\n   - While the paper introduces unbalanced OT as a potential extension, the empirical evaluation of this approach is limited.\n   - More extensive experiments or analysis demonstrating the specific benefits of unbalanced OT over balanced OT in different scenarios would strengthen this contribution.\n\n8. Clarity on practical implementation:\n   - While the paper provides theoretical foundations, it could benefit from more practical guidance on implementing GCA in real-world scenarios.\n   - For instance, how should practitioners choose between different variants (GCA-INCE, GCA-RINCE, GCA-UOT) for a given task?\n\n[1] W. Wang, etal, Zero-Shot Recognition via Optimal Transport   \n[2] Y. Balaji, etal Normalized wasserstein distance for mixture distributions with applications in adversarial learning and domain adaptation", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thanks so much for your detailed feedback! We will now provide a point-by-point response to your questions.\n\n1. “Scaling to larger datasets would be beneficial.”\n\n**Reply:**  Thanks for your suggestion. Please see that we have added the ImageNet100 and SVHN evaluation results to Table R1 in the general response. \n\n2. “Can you provide a detailed analysis of the computational overhead [xxx]? memory usage [...]?”\n\n**Reply:**  Our summary of the computational complexity is in Sec. 3.2, Line 188, and further analyzed in Appendix C.1. The backward pass is not significantly impacted by the number of iterations (Fig. A3); 10 iterations of GCA-INCE require only 5% more FLOPS in general response (#2). GCA-UOT even reduces FLOPS by 30% compared to INCE. Regarding memory usage, GCA-INCE (99.82 MB) and GCA-UOT (51.28 MB) are similar to INCE (51.28 MB) using the same ResNet18 model.\n\n3. “Hyperparameter sensitivity”\n\n**Reply:** Thank you. In the main text, we summarize the ablation study of the sensitivity of α and β in Sec. 6.2. Additionally, we provide the ablation study of epsilon in Fig. U2. We show the impact of the iteration number and provide the discussion for chosen λ and q in GCA-RINCE in Fig. U3. We hope our efforts alleviate your concern about the hyperparameters sensitivity.\n\n4. “In the domain generalization [...], how sensitive is the performance to the choice of α and β in the customized transport plan?”\n\n**Reply:**   As shown in Fig. 2, increasing the difference between α-β enhances classification accuracy, where α is the weight for same-domain samples and β for different domains. This is likely due to domain-informative constraints guiding the model during pretraining. Larger weights on same-domain samples benefit domain generalization.\n\n5. “Did you try comparing GCA with other alignment-based [...] methods? If so, what were the results? How does GCA differentiate itself [...]?”\n\n**Reply:**   Thanks. In Wang, et al, they consider a different setup where they use OT to align seen and unseen images to achieve zero-shot performance. In Balaji et al, they introduce a Wasserstein metric to perform adversarial learning and domain adaptation. It’s hard to compare their work with us because they are not following the same  contrastive learning framework. However, their methods reveal additional settings where contrastive learning can be combined with different adversarial or generation tasks. \n\n6. “ Can you elaborate on any potential limitations or failure cases of the GCA approach?”\n\n**Reply:** For large datasets, GCA methods may consume more resources due to numerous alignment plans. Single-step methods might perform better and converge faster. Additionally, incorrect hyperparameters like epsilon can cause GCA to fail.\n\n7. “Ablation on the number of alignment steps.”\n\n**Reply:**  We use a convergence criterion for stopping the alignment. However, we found that typically 5 steps are sufficient for convergence.  Based upon your suggestions, we conducted an ablation study on CIFAR-10 (Fig. U2) showed that accuracy and cluster compactness stabilize after 5 iterations.\n\n8. “The empirical evaluation of unbalanced OT  is limited.  [...]. More extensive experiments or analysis demonstrating the specific benefits of unbalanced OT over balanced OT in different scenarios”\n\n**Reply:**   Thanks. We ran additional long-tail classification experiments to show the performance of unbalanced OT (Table R2 in general response). These experiments show that GCA-UOT outperforms other baselines and highlight another application where unbalanced OT outperforms other balanced alignment approaches.\n\n9.  “ how should practitioners choose between different variants [...]? Can you provide more concrete guidelines [...]? Are there specific task characteristics [...]?”\n\n**Reply:** Choosing the right method depends on specific constraints. GCA-UOT consistently performs best across tasks, and GCA-RINCE generally outperforms GCA-INCE, especially with noisy data. With fewer classes and good augmentations, all methods perform comparably (see CIFAR-10 in Table 1). However, for noisy or unbalanced views, GCA-UOT offers the flexibility to add accommodating constraints.\n\n10.  “Is there a point of diminishing returns [...]”\n\n**Reply:** Yes, increasing iterations in our multistep objective shows no performance improvement beyond a certain point (Fig. U2). On CIFAR-10, this plateau occurs early, typically around 5 iterations.\n\n11. “Have you explored using different transport plans [...]?”\n\n**Reply:**  In our experiments, we use a single transport plan based on matching constraints: a diagonal target for most experiments and a block constraint for domain adaptation. Using multiple matching constraints for one dataset is an interesting idea!\n\n\n12. “Can you provide more details on how GCA improves robustness [...]?” \n\n**Reply:** Thanks for the questions. Fig. U3 shows that tuning the hyperparameter q close to 1 gives GCA-RINCE symmetry properties, making it more robust to strong augmentation. Additionally, setting the divergence in GCA as symmetry loss (GCA-RINCE) , we see through Lemma 7 (Appendix A.7) that we can enhance robustness to noisy augmentations. This is supported both theoretically and empirically.\n\n13.  “Are there specific types of noise [...]”\n\n**Reply:** Thanks. Our results in Tables A2-A4 (Appendix B.2) show that GCA significantly improves under strong crop and large erase conditions, indicating its effectiveness in mask distribution recovery.\n\n14.  “The authors have not explicitly addressed limitations [...]”\n\n**Reply:**  Thanks. We have updated our discussion of limitations by incorporating your points: 1) extra computational overhead potentially increasing energy consumption, and 2) potential misuse in privacy-sensitive applications.", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "P0TWpC43n6", "reviewer": "Reviewer_TBTc", "summary": "The paper recasts several self-supervised learning (SSL) paradigms, such as SimCLR, in an optimal transport framework. In many SSL variants each batch contains two views of the same data sample. This means that the embeddings of a batch can be viewed as the union of two sets, each one containing only one of the two views of each data sample. The key insight of the paper is that the typical SSL losses can be viewed as the discrepancy between approximations of the optimal transport plans between these two sets from the target transport plan that simply matches the two views of each sample. From this OT perspective the authors propose variants of self-supervised learning losses, improving the approximation to the optimal transport plan, generalizing to weighted batches, and relaxing the target transport plan to incorporate prior domain knowledge.", "questions": "- *Q1* There seem to be multiple ways for enforcing / penalizing constraints: In eq (8) the valid choices for $P$ are restricted both by $\\mathcal{B}$ and $h(P)$, especially if the latter is an indicator function. Why use both ways of representing constraints and not include both either in $\\mathcal{B}$ or $h$? Similarly in eq (9) there are additional divergences to achieve non-uniform marginals of the transport plan. Could this not also be subsumed by $h$ or $\\mathcal{B}$? \n\n- *Q2* The convergence proofs mentioned in line 59, line 188 and A4 only refer to the inner optimization loop of finding the optimal transport plan, right? In particular, they do not show that the parameters of the neural network will converge independent of the data. If so, rephrase the statements in lines 59 and 188 accordingly. Currently, they can be misunderstood as an overstatement.\n\n- *Q3* In the full GCA setup, where the optimal transportation plan is computed (not just one step), the backwards pass needs to unroll all the Sinkhorn steps, right? Does this not lead to high complexity and potentially exploding / vanishing gradients? In particular, I was surprised by line 186 stating that the forward-pass does not affect the computational complexity of the backward pass.\n\n- *Q4* How are the marginal distributions $\\mu$ and $\\nu$ chosen for unbalanced GCA in section 6.2 and Table 2?\n\n- *Q5* What is the value of $\\mu$ in Thm 1, 2, and 3? Is is simply the vector of all ones?\n\n- *Q6* While the statements in Theorem 5- 8 are interesting, I wonder why having higher / lower losses should help with better representations?. The authors do show this empirically, but I wonder why one would expect this from just from having higher / lower losses. Similarly, I wonder why higher uniformity loss implies improved uniformity (Thm 5) and why lower general loss implies improved alignment (Thm 6). For instance, could one not argue that since the loss is already lower, there is less learning signal and thus worse alignment? Also it is the full GCA loss that is lower than the full INCE loss in Thm 6 not just the attractive (alignment) part.\n\n- *Q7* What is the unit for the y-axis in Fig A3? Is it seconds?", "limitations": "Limitations are not explicitly discussed, but I also do not see obvious limitations that need to be stated.", "rating": 4, "confidence": 4, "soundness": 4, "presentation": 1, "contribution": 3, "strengths": "Originality\n- The connection between SSL and OT discussed in this paper seems to extend prior works on the relation of SSL and OT.\n\nQuality\n- Connecting two active areas of research, such as SSL and OT, is very useful.\n- Phrasing SSL objectives in an OT framework provides an interesting new perspective on what SSL does.\n\nClarity:\n- The authors provide a dedicated appendix elaborating the concept of proximal operators, which is helpful as this is more of a niche topic for the general ML community.\n\nSoundness:\n- Their reformulation of SSL losses is correct and most statements are precisely formulated.", "weaknesses": "I appreciate the new perspective the paper presents a lot. However, the presentation needs to be significantly improved.\n\n- *W1 Presentation*\n\t- To me the motivation for introducing all the notation and machinery in sections 2.2-3.3 was not sufficiently clear when reading the paper for the first time. Many abstract concepts are introduced in great generality (lots of choices for divergences, constraint sets etc), but the first real benefit of using this framework appears only on page 6 (Thm 1). I would suggest to drastically restructure the paper and use recasting INCE in the OT framework as the red thread. Flesh out the proof of Thm 1 by explaining how the term inside the logarithm of INCE can be seen as half a Sinkhorn step and use this to motivate the introduction of proximal operators. Then explain how one gets the full INCE loss when using the KL divergence between the approximately optimal transport and the target transport plan. This way every ingredient of the OT framework is directly motivated. Once the INCE case is discussed you can generalize the OT perspective to also incorporate other SSL losses and generalize them.\n\t- My understanding is that the terms \"half-Sinkhorn step\", \"Bregman projection\", and \"proximal operator\" are pretty much synonyms, at least in the context of this work. If so, I would recommend to stick to one of the three terms (while possibly maintaining an appendix section that explains their relations). This would make the exposition much more accessible. My preference would be \"half-Sinkhorn step\" as Sinkhorn operations are most widely known in the ML community. I would even suggest simplifying the terminology at the expense of generality (maybe the BYOL connection really needs proximal operators rather than half-Sinkhorn steps) and defer the fully general setting (explaining the BYOL connection) to the appendix if needed.\n\t- Overall, I would recommend to introduce as little jargon as necessary to present the results. The most general setting can be discussed later in the paper or in the appendix.  \n\t- I recommend phrasing the main idea more clearly early on in the paper (similar to my summary). In particular, lines 30-32 made me think that the transport plans are restricted to only match positive pairs. But instead they are only penalized if they do not. I only understand the setup correctly on page 4.\n\t- The notation used in the main paper needs to be properly introduced in the main paper. For instance, what do the subscripts $\\mathcal{F}$ and $\\mathcal{G}$ stand for in line 199? Define what the target coupling plan $P_{tgt}$ is supposed to do in line 167. What is $R^{B\\times B}$ in line 265? If $R$ stands for the real numbers, rather use $\\mathbb{R}$. \n\t\n- *W2* There is no related work section. It would be useful to at least explain in which ways the present work extends prior works on the connection between SSL and OT, such as [39]\n\t\n**Minor:**\n- *W3* Line 107:  The equation given for the cost is nearly the cosine similarity (but not quite due to the absolute value) but the sentence states that the cost often encodes the L2 distance. This is confusing as the provided formula does not encode the L2-distance.\n- *W4* Several cross-reference links are faulty. For instance, the links to Appendix A.4 and algorithm A1 in lines 188, 189 all point to algorithm 1.\n- *W5* Combine Thm 7 and 6 into one as Thm 7 is a strict generalization of 6. \n- *W6* The INCE result on CIFAR-10 is unusually low. Other sources (Damrich et al. 2023 or https://github.com/p3i0t/SimCLR-CIFAR10) report linear classifier accurarcies above 92% for SimCLR on Cifar-10, much higher than for any SSL variant in this paper.\n- *W7* Casting BYOL into the GCA framework seems pretty forced: The proximal operator is the identity, the result $S_\\theta$ is really not a transport plan, and applying the KL loss seems odd as there is no normalization constraint on $S_\\theta$. I think the ultimate reason for this less on the GCA side and has more to do with the lack of a repulsive force in BYOL (and the existence of a degenerate optimum with a constant encoder). I would perhaps recommend deferring this finding to the appendix, which might also allow to reduce the level of generality required in the main paper, see above.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We sincerely thank you for your thorough evaluation and insightful feedback on our manuscript. Based upon your suggestions, we plan to make major revisions to this paper to improve the quality of the presentation. Below we provide replies to your other questions and concerns.\n\n1. Suggestions on restructuring the paper\n\n**Reply:**  We sincerely thank you for the suggestions to rearrange our paper. We thought long and hard about this and originally started with INCE and Sinkhorn in a previous version. However, we decided to go with this organization as it sets up the problem in full generality before diving into specific examples of the idea in action. \n\n2. Terminology surrounding alignment iterations\n\n**Reply:**  Since the use of the Sinkhorn algorithm is a specific case in our framework, only applied when using INCE, we have referred to it as a \"half-Sinkhorn\" step. As we extend to more general divergences and losses, proximal operators provide a more generalized framework for this optimal transport approach. We will simplify our terminology and incorporate your ideas to clarify our motivation. We appreciate your detailed suggestions!\n\n3. “What do the subscripts $\\mathcal{F}$ and $\\mathcal{G}$ in line 199  [..] what the target coupling plan Ptgt in line 167 [...]. What is $R^{BxB}$ in line 265 [...]?”\n\n**Reply:**  Thank you. In line 199, $\\mathcal{F}$ and $\\mathcal{G}$ are constraint sets for divergences $h_F$ and $h_G$. In line 167, $P_{\\text{tgt}}$ is the target distribution, set as the identity matrix (line 171) or other relaxed forms (line 208). In line 265, $\\mathbb{R}^{B \\times B}$ is correct. We will clarify these terms in our revision.\n\n4. Discussion of related work\n\n**Reply:**   Thanks. We added some related work in the background section but agree that a more thorough discussion is needed. We plan to expand this and clarify the differences from [39] in our revision.\n\n5. “Line 107: The equation given for the cost [...]”\n\n**Reply:**  Thanks. It was a typo in line 107, and we've corrected it to cosine similarity. \n\n6. “Combine Thm 7 and 6 [...]”\n\n**Reply:**  Thanks, we agree with your suggestion. We plan to combine the two theorems as you suggested.\n\n7. “The INCE result on CIFAR-10 is unusually low.[...]”\n\n**Reply:** Thanks. Our results are based on 400 epochs with an SGD optimizer for the linear layer. After training for 1000 epochs with a LARS optimizer, we achieved 90.42% accuracy. Our approach only considers pairwise matching across B samples, not the full SimCLR implementation with 2B samples, causing a slight decrease in performance compared to reported SimCLR results.\n\n8. “Casting BYOL into the GCA framework seems pretty forced. [...] ”\n\n**Reply:**  Thanks. While our formulation may not be the most natural for expressing BYOL, we believe BYOL demonstrates the flexibility of our proximal theory framework by simply changing the kernel \\(\\bf K_\\theta\\). We will consider your suggestion for the final submission.\n\n9. “Why use both ways of representing constraints and not include both either in B or h?”\n\n**Reply:**  Thanks. The set $\\mathcal{B}$ is the set of feasible solutions, while $h(P)$ encodes the penalty. Thus, finding the solution for $\\mathbf{P}$ is influenced by both. In Eq. (9), we show that constraints in $\\mathcal{B}$ can be converted into a soft penalty via functions $h_i$ (where $i = 1, 2, \\ldots, n$) by finding their dual formulation.\n\n10. “The convergence proofs only refer to the inner optimization loop [...] rephrase the statements in lines 59 and 188 accordingly. ”\n\n**Reply:** Thanks. Yes, it only refers to the inner loop. We will rephrase the statements in lines 59 and 188 to make this clear. Our theory for this is detailed in Theorems 5-8.\n\n11. “The backwards pass needs to unroll all the Sinkhorn steps, right?” [...] “Does this not lead to potentially exploding/vanishing gradients”\n\n\n**Reply:**  No, the backward pass doesn't need to unroll all the Sinkhorn steps (or iterations for general losses). See our general response (#2) on complexity. As shown in Fig. A3, the computational resources for the backward pass aren't significantly affected by the number of iterations. Regarding the impact on gradients, incorrect parameters like epsilon could influence them. This can be mitigated by selecting appropriate regularization parameters. \n\n12.  “How are the marginal distributions $\\mu$ and $\\nu$ are chosen for unbalanced GCA”\n\n**Reply:**  In Table 2, the target transport plan for GCA-UOT is set to the identity matrix. In Sec. 6.2, we explore different constraints by varying the alpha and beta values in the target transport plan matrix, so the values of $\\mu$ and $\\nu$ change with it.\n\n13. “What is the value of \\mu” in Thm 1, 2, and 3? Is it simply the vector of all ones?\n\n**Reply:** Yes, they are vectors of ones. We will clarify this in the final version.\n\n14. “[...] why having higher / lower losses should help with better representations?. [...]”\n\n**Reply:** That's an excellent question. Proving the benefits for the lower bound is challenging due to training dynamics. In Appendix A10 and Fig. A2, we explain that INCE aligns $\\bf{P}$ in one direction with $\\bf{P}\\bf{1}_m=\\bf{a}$, while GCA aligns samples in two directions with $\\bf{P}^T\\bf{1}_n=\\bf{b}$, leading to more uniform results. In Theorem 6, we show that our approach improves alignment loss because optimizing INCE loss is akin to optimizing alignment loss. We will clarify this in our final submission.\n\n15. “What is the unit for the y-axis in Fig A3? Is it seconds?”\n\n**Reply:**  Yes, we have updated the units in Fig A3 to seconds.", "claims_and_evidence": "", "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": "", "extra_scores": {}}, {"review_id": "PRF84Rezlx", "reviewer": "Reviewer_VXuY", "summary": "The paper presents a novel perspective on contrastive learning (CL) by framing it as a distribution alignment problem using entropic optimal transport (OT). It trains an encoder network $f_\\theta$ by iteratively updating encoder parameter $\\theta$ and corresponding transport plan $P$ among encoded augmentations of samples. The authors establish connections between noise contrastive estimation losses widely used in CL and distribution alignment with OT. This novel connection allows for the development of various loss functions and multi-step variants for existing CL methods. The theoretical insights and experimental evidence provided demonstrate the benefits of this approach in improving the generalization and robustness of contrastive alignment in both clean and noisy settings.", "questions": "In your experiments, you utilize a fixed number of epochs. How do you ensure that the convergence of the algorithm is achieved within this fixed number of iterations? Would it be feasible to implement a convergence criterion to replace the fixed number of iterations?\n\nOptimal transport is known to require substantial computing resources, especially with large sample sizes. Could you provide a runtime analysis for your algorithm? Does the algorithm face significant computational challenges when dealing with large datasets?\n\nYou utilized proximal operators for T steps to obtain the transport plan P, with the differentiation of the parameter θ computed via auto-differentiation from the loss, propagating back through T steps. Could the auto-differentiation process impose a significant computational burden when a large number of steps T is required in the proximal operator? How do you mitigate this potential issue?\n\nWould you method be able to handle more than 2 types of augmentation? What would be your target transport plan if more than 2 augmentation is used in the contrastive learning?", "limitations": "N/A", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 4, "strengths": "The paper offers a fresh view on contrastive learning by linking it to optimal transport, providing a solid theoretical foundation for understanding and improving CL methods. The authors provide rigorous theoretical analysis and proofs for the convergence of their proposed methods, offering strong support for their claims. The proposed Generalized Contrastive Alignment (GCA) framework is versatile, allowing for the incorporation of domain-specific knowledge and customization of representation spaces.\n\nOriginality\n\nThe paper presents an innovative connection between contrastive learning and optimal transport, a novel perspective that has not been extensively explored before. The introduction of the Generalized Contrastive Alignment (GCA) framework offers a fresh approach to enhancing contrastive learning methods.\n\nQuality\n\nTheoretical insights are rigorously developed, providing a solid foundation for the proposed connections and methodologies. The experiments demonstrate the benefits of the GCA approach from a few different perspectives.\n\nClarity\n\nThe paper is logically structured, and key concepts and methods are clearly explained.\n\nSignificance\n\nBy bridging the gap between contrastive learning and optimal transport, the paper opens up new prospects for research and application in self-supervised learning. The GCA framework has the potential to improve the expressiveness and robustness of representations in various domain generalization settings, which is relevant for real-world applications.", "weaknesses": "While the theoretical convergence of the algorithm is established, the specific criteria for convergence are not clearly defined. The paper's use of proximal operators for T steps and auto-differentiation for parameter θ could impose significant computational burdens, especially with large T. This potential issue is not sufficiently addressed, raising concerns about the algorithm's efficiency in large-scale applications.The experiments lack a detailed runtime analysis, leaving the computational efficiency of the algorithm unexamined. The computational resources necessary for implementing the algorithm are not clearly outlined, raising concerns about its scalability and practicality for large-scale applications.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback. We appreciate your time and suggestions. Our specific responses to your questions are provided below.\n\n\n1. “ the specific criteria for convergence are not clearly defined. [...] concerns about the algorithm's efficiency in large-scale applications. The experiments lack a detailed runtime analysis,[...]..”\n\n**Reply:**  Thanks for your questions. We summarize our computational analysis in line Line 188 in our main text and give the run time analysis in Sec. 3.1. We show the cost of MSINCE in 10 iterations is only 5% more operations in Flops than INCE, and GCA-UOT is even lower than the INCE with a  30% reduction in Flops in general response (#2).\n\n\n2. “ How do you ensure that the convergence of the algorithm is achieved within this fixed number of iterations? Would it be feasible to implement a convergence criterion to replace the fixed number of iterations?”\n\n**Reply:**  Thanks for your questions. In practice, we use a simple convergence criteria for automatically terminating the multistep algorithm. For CIFAR-10, we found that we could also set the maximum iterations to 5 without any loss in performance. \nBased upon the reviewers comments, we ran an experiment to examine the impact of the number of iterations on the accuracy and compactness of the classes (Fig. U2). We found that the accuracy was not very sensitive to the exact choice for the number of fixed iterations with comparable performance for anywhere from 5-11 iterations.\n\n3. “Optimal transport is known to require substantial computing resources, especially with large sample sizes. Could you provide a runtime analysis for your algorithm? Does the algorithm face significant computational challenges when dealing with large datasets?”\n\n**Reply:**  We provide an analysis of computational complexity and empirical results in the Appendix in Sec. C.1. The computation efficiency is also discussed in our general response, where we show the small amount of overhead for MS-INCE (5% more flops) and a 30% reduction in flops and running time for GCA-UOT.\n\nIn terms of dealing with larger datasets, please see our results on ImageNet-100 and SVHN in the general response (Table R1). There we show that the model scales to larger datasets and still performs on par with baseline methods.\n\n4. “You utilized proximal operators for T steps to obtain the transport plan P, with the differentiation of the parameter θ computed via auto-differentiation from the loss, propagating back through T steps. Could the auto-differentiation process impose a significant computational burden when a large number of steps T is required in the proximal operator? How do you mitigate this potential issue?”\n\n**Reply:** Thank you for your question. We don’t backpropagate the loss at different iterations in our alignment objective. Instead, we compute a final transport plan, and then backpropagate the loss, similar to the way BatchNormalization operates. While the number of T steps could become substantial when dealing with large datasets and small mini batch, we have implemented several strategies to address this:\n- We perform the optimal transport (OT) computation in the latent space rather than the input space, which reduces the data's dimensionality involved in the OT problem, thereby significantly decreasing both computational cost and memory requirements.\n\n- We set a convergence threshold to prevent an excessively large number of T steps. This ensures that the algorithm stops as soon as an adequate solution is found, further optimizing computational efficiency.\n\n5. “Would you be able to handle more than 2 types of augmentation? What would be your target transport plan if more than 2 augmentation is used in contrastive learning?”\n\n**Reply:**   By having more than two types of augmentations, do you mean having multiple views of the same example in the same batch? If so, then yes, we could handle that case with our framework! We show how this many-to-one matching can be implemented as a block diagonal constraint on the transport plan in our domain generalization experiments. \nWe hope to include this as an example for future work. Thanks for your suggestion.", "claims_and_evidence": "", "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": "", "extra_scores": {}}]

Author rebuttal:
We thank the reviewers for their feedback. We appreciate that the reviewers thought that our work “provides an interesting new perspective on what SSL does” (Reviewer **TBTc**) and “the methodology and theory introduced is sound”  (Reviewer **96Mt**), provides “ a solid theoretical foundation”  (Reviewer **VXuY**). Additionally, the reviewers remark that approaches “flexibility [...] opens up new possibilities for tailoring representations to specific tasks or domains.” (Reviewer **ih3d**).

In our general response, we would like to: (1) highlight the key differences between our approach and prior work [reviewers **96Mt, TBTc**]; (2) address the concern of computational complexity [reviewers **TBTc, ih3d, VXuY**]; (3) provide additional larger-scale experiments [reviewers **96Mt, VXuY**]. 

**1. Clarify contributions from related work [96Mt, TBTc]:**

Previous work from Shi et al. provides connections between InfoNCE and optimal transport using the Sinkhorn method. In contrast to this work, we make a number of novel contributions:

-  **New algorithm for generalized alignment for contrastive learning:**  Our algorithm handles more generalized losses by iteratively solving new constraint set intersections, whereas previous work focused solely on alignment via Sinkhorn iterations.

- **Novel approach for unbalanced OT-based alignment:**  We leverage a rich body of work in OT to introduce a variant of GCA that relaxes the constraints on the distribution penalty (Sec. 3.3). By converting the hard penalty (constraint sets) into the soft regularization terms, our GCA-UOT method achieves high classification accuracy (Table 2) and faster convergence than INCE (Fig. A3), linking OT literature with optimization and contrastive learning.

- **Connections and a multistep variant of RINCE:**  By building a more generalized form of alignment, we demonstrate that it's possible to develop connections to the Robust INCE loss, RINCE. This equivalence enables us to develop a multistep RINCE variant that performs better with corrupted views.

- **Novel results in domain generalization through block-diagonal matching constraints:**  By changing the target plan to have block diagonal structure, we can absorb the domain information (Sec. 6). Adding domain-specific matching constraints can improve the pre-training model and enhance classification accuracy in cross-domain generalization tasks.

- **New theory and insights**:   We prove the convergence of our more generalized algorithms, not just for the KL divergence in sinkhorn situations, but for other Bregman divergences, which is not shown in previous algorithms. And we develop new theorems to explain why running GCA could lead to better uniformity and alignment.

In summary, our GCA framework provides a foundation for addressing a wider range of potential issues in contrastive learning, by demonstrating how to create a variety of different contrastive methods and bringing advanced OT approaches to bear in representation learning. 

**2. Computational complexity [TBTc, VXuY, ih3d]:**  

About the additional complexity of our method, we provide analysis and empirical evidence to show the complexity (Sec. C1) and compare the running time of different methods (Fig. A3) in the appendix. Our results show that GCA iterations only slightly increase the computational complexity, while GCA-UOT is faster than INCE due to the improved symmetry and smoothness of the loss. To add to this, we record the floating point operations per second (Flops) of running GCA methods. We find that GCA-INCE (6.65 MFlops) has 5% more Flops than INCE (6.31 MFlops), while GCA-UOT saves 30% Flops (4.54 MFlops).  These results prove that our GCA-UOT method is not only superior in terms of accuracy but also in speed.

**3. Additional Experiments**

- **Experiments on larger datasets:**  Following the reviewer's suggestions, we ran additional experiments on ImageNet-100 (Table R1) and SVHN (Table R1), and compared with both INCE and Shi et al. (IOT) methods as a baseline. There is no code available from Shi et al. so we implemented their method. We found that our GCA-UOT achieves the highest accuracy on both datasets.

- **Sensitivity analysis:**  Additionally, we provide ablation studies to hyperparameters including the epsilon effect on transport plan (Fig.U1) **[96Mt]**, the iteration number (Fig. U2) **[ih3d, VXuY]**, the effect of q and λ on strong noisy augmentations Fig. U3 **[VXuY]**. 

- **Additional examples of UOT:**  To provide more scenarios of GCA-UOT over balanced OT **[VXuY]**, we test our method on the problem of longtail image recognition with CIFAR100-LT (Table R2). 

We hope our additional experiments would address your concerns. Please let us know if there are additional results that you would like to see and we will do our best to provide them.

**Results**

**Table R1. Comparison of GCA methods with other baselines on SVHN and ImageNet100 datasets.** For SVHN, we pre-trained ResNet50 models using the Adam optimizer with a learning rate of 3e-4 for both 200 and 500 epochs. For ImageNet100, we pre-trained a ResNet50 model for 100 epochs using the same optimizer and learning rate. The classification accuracy was then evaluated using a linear layer readout trained 100 epochs.

|  | **SVHN**  |    | **ImageNet100**  |
| -- | -- | -- | -- |
|**Epochs**| **200**|**500**| **100**|
|**UOT**| **89.98**|**91.85**|**68.63** |
|**MSINCE**|86.00|89.86|67.60|
|**INCE**|86.60|89.79|67.93|
|**IOT(Shi​​​​)**|85.19|90.01|66.07|

**Table R2. Comparison of GCA methods with other baselines on longtail recognition (CIFAR100-LT).** We pre-trained models for 500 epochs using the Adam optimizer with an imbalance factor of 0.1 on the CIFAR100 long-tail dataset. The classification accuracy was then evaluated using a linear layer readout trained 200 epochs.
|**Method**|**CIFAR100LT**|      
| -- | -- |
|**UOT**|**35.96**|
|**MSINCE**|33.20|
|**INCE**|33.72|
|**ShiINCE**|32.55|

Meta-review: This paper focuses on the theoretical foundations of contrastive learning (CL) for learning data representations. The authors propose a framework called Generalized Contrastive Alignment (GCA), which connects the noise contrastive estimation loss to distribution alignment using entropic optimal transport. Building on this, the authors introduce variants of self-supervised learning losses to handle outliers. The empirical evaluation on CIFAR-10, CIFAR-100, ImageNet-100, and SVHN (added during the rebuttal) demonstrates the effectiveness of the proposed method.

The reviewers raise concerns about the overlap with existing work and the absence of a discussion on related work (Reviewers 96Mt, TBTc), the structure of the paper (Reviewer TBTc), computational complexity (Reviewers TBTc, VXuY, ih3d), and the practical applicability of the proposed algorithm (Reviewer VXuY), among others. During the rebuttal, these concerns were thoroughly discussed, and the authors effectively addressed many of them. Given the overall positive reception of the proposed framework, we recommend accepting the paper. However, if accepted, the authors should include the requested revisions from the reviewers, such as the discussion of overlap with related work and improvements to the paper's structure.
Decision: Accept (poster)