Your contrastive learning problem is secretly a distribution alignment problem

Thank you for your feedback. We appreciate your time and suggestions. Our specific responses to your questions are provided below.

1. “ 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,[...]..”

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).

2. “ 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?”

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. Based 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.

3. “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?”

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.

In 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.

4. “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?”

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:

• 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.

• 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.

5. “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?”

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. We hope to include this as an example for future work. Thanks for your suggestion.

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.

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.

1. Suggestions on restructuring the paper

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.

2. Terminology surrounding alignment iterations

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!

3. “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 [...]?”

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.

4. Discussion of related work

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.

5. “Line 107: The equation given for the cost [...]”

Reply: Thanks. It was a typo in line 107, and we've corrected it to cosine similarity.

6. “Combine Thm 7 and 6 [...]”

Reply: Thanks, we agree with your suggestion. We plan to combine the two theorems as you suggested.

7. “The INCE result on CIFAR-10 is unusually low.[...]”

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.

8. “Casting BYOL into the GCA framework seems pretty forced. [...] ”

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.

9. “Why use both ways of representing constraints and not include both either in B or h?”

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.

10. “The convergence proofs only refer to the inner optimization loop [...] rephrase the statements in lines 59 and 188 accordingly. ”

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.

11. “The backwards pass needs to unroll all the Sinkhorn steps, right?” [...] “Does this not lead to potentially exploding/vanishing gradients”

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.

12. “How are the marginal distributions $\mu$ and $\nu$ are chosen for unbalanced GCA”

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.

13. “What is the value of \mu” in Thm 1, 2, and 3? Is it simply the vector of all ones?

Reply: Yes, they are vectors of ones. We will clarify this in the final version.

14. “[...] why having higher / lower losses should help with better representations?. [...]”

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.

15. “What is the unit for the y-axis in Fig A3? Is it seconds?”

Reply: Yes, we have updated the units in Fig A3 to seconds.

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.

1. Discussion of overlap with Shi et al.

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:

• 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.

• 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 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.

In summary, our GCA framework provides a foundation for addressing a wider range of potential issues in contrastive learning.

2. Difference between GCA-UOT and the idea of “unbalanced matching” in Shi et al.

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).

In 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.

[1] Xu, M., & Gould, S. (2024). Temporally Consistent Unbalanced Optimal Transport for Unsupervised Action Segmentation. arXiv:2404.01518.

[2] De Plaen, H., et al. (2023). Unbalanced optimal transport: A unified framework for object detection. CVPR, 3198-3207.

3. “I believe the authors should have included Shi et al. (2023) as a baseline.”

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.

4. “Evaluation also should have been done equally on a larger dataset such as ImageNet.”

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.

5. “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”

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.

Thanks so much for your detailed feedback! We will now provide a point-by-point response to your questions.

1. “Scaling to larger datasets would be beneficial.”

Reply: Thanks for your suggestion. Please see that we have added the ImageNet100 and SVHN evaluation results to Table R1 in the general response.

2. “Can you provide a detailed analysis of the computational overhead [xxx]? memory usage [...]?”

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.

3. “Hyperparameter sensitivity”

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.

4. “In the domain generalization [...], how sensitive is the performance to the choice of α and β in the customized transport plan?”

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.

5. “Did you try comparing GCA with other alignment-based [...] methods? If so, what were the results? How does GCA differentiate itself [...]?”

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.

6. “ Can you elaborate on any potential limitations or failure cases of the GCA approach?”

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.

7. “Ablation on the number of alignment steps.”

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.

8. “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”

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.

9. “ how should practitioners choose between different variants [...]? Can you provide more concrete guidelines [...]? Are there specific task characteristics [...]?”

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.

10. “Is there a point of diminishing returns [...]”

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.

11. “Have you explored using different transport plans [...]?”

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!

12. “Can you provide more details on how GCA improves robustness [...]?”

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.

13. “Are there specific types of noise [...]”

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.

14. “The authors have not explicitly addressed limitations [...]”

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.