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 {\em 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.
Paper
Similar papers
Peer review
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.
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. Originality The 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. Quality Theoretical 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. Clarity The paper is logically structured, and key concepts and methods are clearly explained. Significance By 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.
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? 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? 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? Would 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?
Rating
6
Confidence
4
Soundness
4
Presentation
3
Contribution
4
Limitations
N/A
Dear Reviewer VXuY, We greatly appreciate your constructive comments and are glad that the additional analysis addressed the issues you raised. We kindly hope that you might consider increasing your score based upon the discussion.
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.
Strengths
Originality - The connection between SSL and OT discussed in this paper seems to extend prior works on the relation of SSL and OT. Quality - Connecting two active areas of research, such as SSL and OT, is very useful. - Phrasing SSL objectives in an OT framework provides an interesting new perspective on what SSL does. Clarity: - 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. Soundness: - 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. - *W1 Presentation* - 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. - 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. - 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. - 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. - 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}$. - *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] **Minor:** - *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. - *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. - *W5* Combine Thm 7 and 6 into one as Thm 7 is a strict generalization of 6. - *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. - *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.
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}$? - *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. - *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. - *Q4* How are the marginal distributions $\mu$ and $\nu$ chosen for unbalanced GCA in section 6.2 and Table 2? - *Q5* What is the value of $\mu$ in Thm 1, 2, and 3? Is is simply the vector of all ones? - *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. - *Q7* What is the unit for the y-axis in Fig A3? Is it seconds?
Rating
4
Confidence
4
Soundness
4
Presentation
1
Contribution
3
Limitations
Limitations are not explicitly discussed, but I also do not see obvious limitations that need to be stated.
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: 1. 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. 2. Proving convergence of GCA-based methods and demonstrating theoretically that the alignment objective can improve the quality of learned representations. 3. Empirically validating GCA's effectiveness in image classification and domain generalization tasks, showing it can achieve superior performance over baseline methods. 4. Demonstrating how GCA allows building unbalanced losses using tools from optimal transport, which can handle noisy views and customize representations. 5. Providing a unified framework that connects existing contrastive learning methods like InfoNCE, RINCE, and BYOL to optimal transport formulations. 6. Showing how modifying the target alignment plan in GCA can flexibly control the amount of domain knowledge incorporated into representations for domain generalization tasks. The 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.
Strengths
This paper demonstrates several strengths across the dimensions of originality, quality, clarity, and significance: **Originality:** 1. 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. 2. The formulation of contrastive learning as a distribution alignment problem offers a fresh perspective on a widely-studied topic. 3. The extension to unbalanced optimal transport and the ability to customize target transport plans are innovative additions to the contrastive learning toolkit. **Quality:** 1. The theoretical analysis is rigorous, with clear proofs for the connections between GCA and existing methods (InfoNCE, RINCE, BYOL). 2. The empirical evaluation is comprehensive, covering multiple datasets (CIFAR-10, CIFAR-100, CIFAR-10C, PACS) and scenarios (standard classification, extreme data augmentation, domain generalization). 3. The ablation studies and comparisons against baseline methods are thorough and well-presented. **Clarity:** 1. The paper is well-structured, with a clear progression from problem formulation to theoretical analysis and empirical validation. 2. The use of diagrams (e.g., Figure 1) helps illustrate complex concepts like customized transport plans. 3. The authors provide clear pseudocode (Algorithm 1) for the GCA method, enhancing reproducibility. **Significance:** 1. The GCA framework provides a unifying perspective on several popular contrastive learning methods, which could facilitate further theoretical developments in the field. 2. The improved performance on classification tasks, especially under extreme data augmentation, suggests practical benefits for real-world applications. 3. 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: 1. Limited scope of empirical evaluation: - 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. - 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. 2. Computational complexity: - The paper lacks a detailed analysis of the computational overhead introduced by the GCA framework, particularly for the multi-step variants. - 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. 3. Hyperparameter sensitivity: - 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. - An ablation study or sensitivity analysis for these parameters would provide valuable insights for practitioners looking to implement GCA. 4. Comparison with other alignment-based methods: 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]). - Such comparisons would help better contextualize the contributions of GCA within the broader landscape of alignment-based approaches. 5. Theoretical limitations: - 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. - For instance, are there scenarios where the multi-step approach might lead to worse performance or slower convergence? 6. Ablation on the number of alignment steps: - While the paper mentions using 5 iterations in the forward pass, there's no analysis on how performance changes with different numbers of iterations. - An ablation study showing the trade-off between computational cost and performance improvement as the number of iterations increases would be valuable. 7. Limited exploration of unbalanced OT: - While the paper introduces unbalanced OT as a potential extension, the empirical evaluation of this approach is limited. - More extensive experiments or analysis demonstrating the specific benefits of unbalanced OT over balanced OT in different scenarios would strengthen this contribution. 8. Clarity on practical implementation: - While the paper provides theoretical foundations, it could benefit from more practical guidance on implementing GCA in real-world scenarios. - For instance, how should practitioners choose between different variants (GCA-INCE, GCA-RINCE, GCA-UOT) for a given task? [1] W. Wang, etal, Zero-Shot Recognition via Optimal Transport [2] Y. Balaji, etal Normalized wasserstein distance for mixture distributions with applications in adversarial learning and domain adaptation
Questions
- Can you provide a detailed analysis of the computational overhead introduced by GCA, particularly for the multi-step variants? - How do the training time and memory usage compare to standard contrastive learning methods? - Is there a point of diminishing returns regarding performance improvement vs. computational cost? - What guidelines can you provide for setting the new hyperparameters introduced in GCA (e.g., number of iterations, α and β in the customized transport plan)? - Have you observed any patterns in how these hyperparameters affect performance across different tasks or datasets? - Can you provide an ablation study showing how performance changes with different numbers of iterations in the forward pass? - Have you considered comparing GCA with other alignment-based or optimal transport-based representation learning methods? If so, what were the results? - How does GCA differentiate itself from or improve upon these existing alignment-based approaches? - Are there any scenarios in which the multi-step approach might lead to worse performance or slower convergence than single-step methods? - Can you elaborate on any potential limitations or failure cases of the GCA approach? - Can you provide more details or experiments demonstrating the specific benefits of unbalanced OT over balanced OT in different scenarios? - Are there particular types of tasks or data where unbalanced OT shows the most significant improvements? - 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? - Are there specific task characteristics that make one variant more suitable than others? - In the domain generalization experiments, how sensitive is the performance to the choice of α and β in the customized transport plan? - Have you explored using different transport plans for domains or tasks within the same dataset? - Can you provide more details on how GCA improves robustness to noisy or extreme data augmentations compared to baseline methods? - Are there specific types of noise or augmentations where GCA shows the most significant improvements?
Rating
6
Confidence
3
Soundness
2
Presentation
2
Contribution
2
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: 1. Limitations: - The authors could add a dedicated "Limitations" section discussing: a) The current scope of experiments (e.g., limited to certain datasets and architectures) b) Potential computational overhead of the multi-step approach c) Challenges in hyperparameter tuning for the new parameters introduced d) Any scenarios where GCA might not be applicable or beneficial 2. Potential negative societal impacts: - The authors should consider adding a brief discussion on potential societal impacts, such as: a) Increased energy consumption due to potentially higher computational requirements b) Possible biases in learned representations, especially when using customized transport plans c) Potential misuse of the technique in privacy-sensitive applications 3. Ethical considerations: - A brief note on any ethical considerations in data usage or potential applications of the method would be valuable. 4. Future work: - Expand the current mention of future work to include specific directions for addressing identified limitations. It'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.
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.
Strengths
1. The methodology and theory introduced is sound and based on the existing theory of optimal transport and its entropic regularization version. 2. 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). 2. 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. 3. 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). - Small nitpick: Eq (1) missing expectation wrt to samples. 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 on Machine Learning, pages 31408–31421. PMLR, 2023.
Questions
See weaknesses.
Rating
4
Confidence
4
Soundness
2
Presentation
2
Contribution
1
Limitations
Limitation is discussed, but an important point I raised on the weaknesses section is missing.
Thank you for the detailed reply and the intention of addressing several of my points in the revision. Please find some follow-up questions / comments below. **1. Structure:** Restructuring the paper is clearly a lot of work. Nevertheless, I still think that you original plan is more friendly for a wider audience of readers. The presentation is the main weakness in my mind. **9. Use of $\mathcal{B}$ and $h$:** But you write in line 163 that $h$ is typically an indicator function. In line 896 you directly translate between $\mathcal{B}$ and $h$ with infinite penalty outside $\mathcal{B}$. Does *indicator function* imply infinite penalty in your setting? If so, I really do not see the difference between constraining via $\mathcal{B}$ and $h$. If you usually use other functions than indicators for $h$, perhaps rephrase line 163. **11. Complexity:** Yes, I see that empirically your method is similarly fast as other SSL methods. However, I do not understand why you do not need to unroll Sinkhorn operations to compute the gradients. How are gradients computed when the forward pass includes an inner optimization line in step 2 of Algorithm 1. **12. Value of $\mu$ and $\nu$:** Is it correct that $\mu$ and $\nu$ are the marginals of $P_{tgt}$, which in turn depends on $\alpha$ and $\beta$? **14. Higher alignment loss:** To be honest, I still do not understand the argument. On a high-level it sounds as if the INCE loss is less informative than the GCA loss, because the former has one and the latter two requirements on the marginals. But I cannot connect this with either Figure A2 or Appendix A10. I also do not understand what this has to do with the relative size of the loss values. Here are some specific questions regarding Figure A2: Are the blue / red points in the first panel the current batch or some random selection of positive pairs? Should the arrows in the second and third panel point to the blue / red points (they do not: for instance the orange arrows in the middle and the right panels are not the same). What are the arrows supposed to mean? Are there any projections to tangent spaces involved? Are the fat arrows in the middle plot contributions to a gradient? If so, of which point is this supposed to be the gradient? I do not get at all why there is a line for INCE but a plane for GCA INCE. How is this connected to the constraints on the marginals? You write in line 1315 that you perform supervised training. But INCE and GCA-INCE are self-supervised. You write in line 1316 that you perform PCA. Was this of the unnormalized points in high-dimensional space or of the normalized points? In either case the result are non-normalized. Why should PCA + renormalization to the sphere encode information about the structure in high-dimensional space? Why not retrain a resnet-18 with a 3D output space, so that you can directly visualize it.
Thanks for your questions and rigorous review. 9. Use of $\mathcal{B}$ and $h$: But you write in line 163 that $h$ is typically an indicator function. In line 896 you directly translate between $\mathcal{B}$ and $h$ with an infinite penalty outside $\mathcal{B}$. Does indicator function imply infinite penalty in your setting? If so, I really do not see the difference between constraining via $\mathcal{B}$ and $h$. If you usually use other functions than indicators for $h$, perhaps rephrase line 163. Reply: Thank you for your question. The penalty function h can provide a hard constraint (using an indicator function) or it can be a different function that measures the deviation from the constraint. In our implementation of GCA-UOT, we use a KL-divergence for h to relax our constraints in balanced OT from a hard to a soft penalty. We will rephrase line 163 to make this clear. 11. Complexity: Yes, I see that empirically your method is similarly fast as other SSL methods. However, I do not understand why you do not need to unroll Sinkhorn operations to compute the gradients. Reply: Our optimization process is decoupled and we can solve it through two different optimization procedures, one to compute the representations and the other to find the optimal transport plan. You can think of the Sinkhorn algorithm as a process that provides us with a static transport plan between two sets of points in the latent space, and what we are actually differentiating over is the cost matrix. 12. Value of $\mu$ and $\nu$: Is it correct that $\mu$ and $\nu$ are the marginals of $P_{tgt}$, which in turn depends on $\alpha$ and $\beta$? Reply: Yes, in the domain adaptation setting, the marginals depend on alpha and beta. 14a. Higher alignment loss: To be honest, I still do not understand the argument. On a high-level it sounds as if the INCE loss is less informative than the GCA loss, because the former has one and the latter two requirements on the marginals. But I cannot connect this with either Figure A2 or Appendix A10. I also do not understand what this has to do with the relative size of the loss values. Reply: The CL loss in Eq.(1) can be decomposed into two terms, the alignment loss and uniformity loss [1] which corresponds to the entropy $-\epsilon H(P)$ and $<P,C>$. Under the perfectly aligned condition where the alignment loss will be zero, we show that the uniformity loss in GCA will be lower than the original INCE objective. We can use this, along with recent work from [2] where they show that a tighter bound on the uniformity can benefit the downstream task like classification, to reason that lower loss can lead to provable benefits in learning. [1] Wang T, Isola P. Understanding contrastive representation learning through alignment and uniformity on the hypersphere, International conference on machine learning. PMLR, 2020: 9929-9939. [2] Dufumier B, Barbano C A, Louiset R, et al. Integrating prior knowledge in contrastive learning with kernel, International Conference on Machine Learning. PMLR, 2023: 8851-8878. 14b. Here are some specific questions regarding Figure A2: Are the blue / red points [...]Are the fat arrows in the middle plot contributions to a gradient? If so, at which point is this supposed to be the gradient? Reply: We are sorry for the confusion in Figure A2. Our goal was to show that by imposing two constraints on the row and column spaces with GCA, we would get more uniformly distributed latents. The subspace shows the two sets of constraints. 14d. You write in line 1315 that you perform supervised training. But INCE and GCA-INCE are self-supervised. Reply: That’s a typo. Sorry, we want to say we train a linear layer for evaluating the representations. 14e. You write in line 1316 that you perform PCA. Was this of the unnormalized points in high-dimensional space or of the normalized points? In either case the results are non-normalized. Reply: The PCA in line 1316 is for the normalized points that are mapped to the unit sphere in 3D. It may look like the points are un-normalized because you can see points on the other side of the sphere. We are sorry this figure is confusing. We will remove it from the final version of the paper as we believe it adds more confusion.
Thanks for the additional explanations! **Complexity:** Unfortunately, I still do not understand your argument. By Alg 1, I thought the forward pass is as follows $\theta \to f_\theta, g_\theta \to K_\theta \to P_\theta \to \mathcal{L}_{GCA}$ where the step $K_\theta \to P_\theta$ is an inner optimization, e.g., computed via Sinkhorn. To compute the gradients of $\mathcal{L}$ with respect to $\theta$, you need to backprop through the entire forward pass, including the step from $K_\theta$ to $P_\theta$. So yes, you ultimately want to differentiate the cost matrix $K_\theta$, but your loss depends on it only through the transport plan $P_\theta$. What am I missing? **Lower uniformity loss:** Thank you for providing the reference to Dufumier et al. which I recommend adding to section 5. Please note that the theorem they proved on the usefulness of lower uniformity loss is for their specific form of the uniformity loss. It might be interesting future work to derive a similar statement for your version of the uniformity loss. **Figure A2:** Yes, I recommend either to omit figure A2 or to overhaul it and the writing in appendix A10 on page 38 significantly.
Dear Reviewer TBTc, Thank you for your follow-up question. In the Sinkhorn algorithm, the transport plan $P_\theta $ is computed as: $P_\theta = \exp(f + g - C_\theta)/\epsilon$ Here, $f$ and $g$ are the dual variables that are iteratively updated during the Sinkhorn algorithm, but crucially, they do not involve gradients with respect to $\theta$. The optimization process in Sinkhorn essentially involves scaling the rows and columns of $P$ to satisfy the marginal constraints, which can be viewed as element-wise operations (scaling and exponentiation) on the cost matrix $C_\theta$. Because $P_\theta$ is computed using the fixed-point iteration of $f$ and $g$ that depend only on the current values of $ C_\theta $, the gradient backpropagation process is simplified. Specifically, the gradient of the loss with respect to the cost matrix $ C_\theta $ is the key part that needs to be differentiated, not the intermediate steps involving $ f $ and $g$. A typical workflow of these algorithms was shown in Figure 2 of [1], the gradient flow primarily involves differentiating through $ C_\theta $, which is done only once, and not through each step of the Sinkhorn iterations. This approach reduces computational complexity and avoids the need for backpropagation through every iterative update within the Sinkhorn algorithm, which might otherwise be computationally expensive. We hope this clarifies the process. Please let us know if you have any further questions! Your insights have been instrumental in improving our work, and we kindly hope you might reconsider your evaluation in light of these discussions and forthcoming changes. [1] Eisenberger M, Toker A, Leal-Taixé L, et al. A unified framework for implicit sinkhorn differentiation. Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2022: 509-518.
Thank you for the explanation! So you are performing implicit differentiation as in Eisenberger et al.? This would resolve my issue. Regarding your explanation, I still have the following question: During the fixed point iterations, you update $f$ and $g$ based on their previous values and $C_\theta$. This makes each iteration of $f$ and $g$ depend on $C_\theta$. Normal backprop would have to trace these dependencies on $\theta$ (this is what I meant by "unrolling the Sinkhorn steps"). I appreciate that it is close to the discussion deadline, so I understand if you cannot reply in time. I just wanted to make this comment to clarify my original concern. The rebuttal resolved several of my smaller points. My main issue, the structure of the paper, remains unaddressed and I still believe that restructuring would make the paper significantly more accessible. Therefore, I will keep my score. The other reviewers seem to not mind the presentation, however, and I would be fine with the paper being accepted.
Dear Reviewer TBTc, Thanks for your agreement for our paper to be published! To clarify further, while the gradient backpropagation does trace these dependencies, the fact that $\theta$ remains unchanged during the Sinkhorn iterations ensures that the computational complexity for the model parameters does not increase. We greatly appreciate your suggestions on improving the presentation. We plan to incorporate your feedback into the final submission to enhance clarity. Given the current presentation and the improvements we plan to make, we kindly encourage you to reconsider your evaluation of the paper.
Thanks for the author's reply and the detailed runtime analysis provided. This has alleviated my concerns about the practical applicability of the proposed algorithm.
Thank you for your rebuttal. I increased my evaluation to 4 to reflect the newly added materials; the evaluation still leaning Reject for the following reason: 1. The result in Table R1 is actually quite mixed: one could see that given enough pretraining iterations, the performance of all 4 methods are very close (SHVN on 500 epoch, third column -- note I have no idea why the authors skip the CIFAR100 dataset). For ImageNet I can understand it takes quite a bit of time to train the network, but I suspect if taking ResNet50 at 500 epochs results would be the same. I also do not understand why IOT of Shi et al. (2023) could perform so badly compared to other 3 on ImageNet. 2. I still hold my opinion that the novelty here is somewhat limited, in the sense that while extending the framework of Shi et al. (2023) is non trivial, it is as hard as extending from classical OT to unbalanced OT setting. I also observed other reviewers have some interesting questions and comments, especially Reviewer TBTc. I therefore look forward to see the answers to those comments as well.
Dear Reviewer 96Mt, Thank you for your continued engagement with our work and for revisiting your evaluation based on our rebuttal. We appreciate the constructive feedback and would like to address the specific concerns you have raised. 1a. “ given enough pretraining iterations, the performance of all 4 methods are very close (SHVN on 500 epoch, third column -- note I have no idea why the authors skip the CIFAR100 dataset). For ImageNet I can understand it takes quite a bit of time to train the network, but I suspect if taking ResNet50 at 500 epochs results would be the same. While it is true that the INCE-based methods provide similar results, we find that our unbalanced OT approach, GCA-UOT, consistently provides improvements over these other methods. We note that this is out-of-the-box without tuning of the hyperparameters for the alignment method. Thus, this demonstrates that our approach generalizes well and can be used with little to no tuning in new datasets. Please note that we did not add CIFAR100 to the response because we provide these results in Table 2 in the main text. 1b. “I also do not understand why IOT of Shi et al. (2023) could perform so badly compared to other 3 on ImageNet.” We applied both our method and the Shi method out of the box, using the same parameters as described in their paper, without tuning further. Perhaps with some small tuning their implementation of MS-INCE and our dual form might be more comparable. 2. Novelty of the Contribution By providing a more general framework for contrastive alignment, our work allows for numerous extensions, including UOT and robust INCE-based losses, as well as block wise constraints in domain generalization. We think this general formulation is key for unlocking additional applications in the future as well. Shi’s work has only considered Sinkhorn-based INCE alignment and thus we see the possibilities with our framework as providing a major advance over the existing work.
> Please note that we did not add CIFAR100 to the response because we provide these results in Table 2 in the main text. Just a comment: Table R1 in the rebuttal is the result for Resnet-50; Table 2 is the result for Resnet-18, and there is no IOT of Shi et al. 2023 for comparison in that table in the main text. There is a strong discrepancy in prediction accuracy between a pretrained Resnet-18 and Resnet-50 -- this is now a well-known fact in the deep learning community.
Thank you for your comment. We will include the CIFAR100 experiments and comparisons with IOT of Shi et al. (2023) in the main text in our revised submission.
Decision
Accept (poster)