Understanding Contrastive Learning via Gaussian Mixture Models

Remark 1 : How do the results (and Fisher optimality) extend to real encoder architectures with a more non-linear setup

We show that for linear setting, InfoNCE loss learns the Fisher subspace, which we motivated as the optimal subspace for clustering. We do not know of an equivalent definition of optimality for representations under non-linear mappings, which also seems prohibitively difficult to come up with without assumptions on the mappings.

Remark 2 : Lack of demonstrations for large scale vision / vision-language benchmarks.

We realised the toy nature of our setup and hence have extended our experimental setup to the ImageNet dataset. This setup helps us argue about distributions more general than shared-covariance GMM. Specifically, we take images corresponding to 20 random classes from the ImageNet and use a pretrained ResNet-50 model to get non-linear mappings for the images. We then learn a linear map from the ResNet embedding space (2048 dimensional) to our target space. We present the results below (presented as ARI/AMI).

Remark 3 : All theoretical guarantees assume access to the true data distribution (infinite samples). It remains unclear how many real examples or negatives are needed

We acknowledge that our analysis is designed for the infinite sample setting and there does not seem to be a trivial extension to finite samples. The main technical difficulty lies with the LogSumExp term in the loss function of InfoNCE. We believe that our reviewer would agree that such an analysis is still non-trivial in the linear setup we propose. That said, we are exploring possible directions which would enable us to make meaningful comments about the establishing bounds on samples for learning subspaces that are comparable fo Fisher subspace. Our goal in this paper is to quantify the absolute limits of contrastive learning in the limit.

Finite-sample analysis with sample complexities and convergence rate analysis would also require a careful optimization algorithm design to solve the InfoNCE or other losses, which is not trivial and beyond the scope of our infinite sample analysis. We do leave the finite-sample analysis to future work, which we believe is an independent interest to the community.

Remark 4 : Practical benefits following from our results.

Our analysis seems to show that while both InfoNCE and CLIP loss are effective at filtering out the noise directions, InfoNCE loss (when considering perfect augmentations) can provably recover all discriminative directions (i.e. learns the complete Fisher subspace). We prove that CLIP loss (at least in the linear setting) is not able to capture the complete subspace and hence the representations are possibly less suited for image-only tasks like clustering.

Remark 5 : How sensitive is the model to the augmentation noise parameter in practice?

We believe the model is robust to the choice of the noise parameter and we empirically verify it with the synthetic shared-covariance GMM experiments. In Figure 2 (a) in the paper, we plot the the clustering performance (AMI/ARI) w.r.t. different augmentation noise $\delta$. We can see that InfoNCE achieves almost perfect clustering even at relatively high noise levels ($\delta = 0.4$). Please let us know if you would be interested in further experiments on this for the discussion period.

Remark 5 : For multi-modal clip what impact does not recovering the Fisher subspace have for downstream performance? How can we identify or recover the missing directions?

Our results seem to suggest that for a given model capacity, not recovering the full Fisher subspace can possibly lead to collapsing. Since Fisher subspace ensures the maximal separability (non-necessarily perfect separability), there might exist a subspace where certain components are almost indistinguishable, leading to poor performance.

One possible solution for recovering the missing directions might be to first train image and text encoders separately, and then fine-tune them for CLIP loss. Then, training the image encoder with the InfoNCE loss would help with learning discriminative representations.

Remark 6 : How the CLIP-GMM theory holds up when applied to actual CLIP embeddings on image-text retrieval datasets

Due to the size of the CLIP dataset, we were not able to complete experiments on it. We are working on a subset of the dataset and hopefully, we will try to make the results ready for the discussion period. Moreover, our theory assumes a linear map and the results could be a result of the limited mapping capacity of these linear maps. It is unclear whether our findings would transfer to representations learnt on highly non-linear networks.

Remark 7 : Discrepancy between the ARI and AMI scores for LDA and InfoNCE

A core difference between ARI and AMI is that ARI emphasizes pairwise clustering performance; it looks into how many pairs of points originally belonging to the same cluster ended up together. ARI could be smaller than AMI if the original cluster sizes are not balanced or the final clusters have size imbalance.

Remark 8 : Minor comments

We will work on them as pointed out by the reviewer. Regarding the comment on Figure 1, in this figure we wanted to show how the separability of representations after projection onto their SVD subspace vs onto their fisher subspace. The GMMs are well separated in their ambient space in both Fig. 1 (a) and 1 (b). Though for Fig1 (b), projection onto the SVD subspace (y-axis) leads to a mode collapse, while projection onto the Fisher subspace (x-axis) ensures that the underlying components still remain well separated.

Remark 1 : Restriction of analysis to Linear Mapping Assumption

We show that for linear setting, InfoNCE loss learns the Fisher subspace, which we motivated as the optimal subspace for clustering. We do not know of an equivalent definition of optimality for representations under non-linear mappings, which also seems prohibitively difficult to come up with without assumptions on the mappings.

Remark 2 : Simplified SimSiam analysis

The reviewer raises a valid point. The analysis including the stop-grad and the prediction head introduces new challenges we cannot tackle without loosening the analysis. Our paper deals with the fixed point analysis of these loss functions. Such fixed point analyses are not affected by the stop-grad and the prediction head component. Though we agree these are important components essential to the success of SimSiam, for our results (i.e., learnt representations exist in the Fisher subspace) the presence of these components does not affect the results. To study the contribution of the stop-grad and the prediction head, one would have to take a dynamical systems approach to parameter learning, which is more complex than our setting. It would also be unclear whether there would be a well-defined notion of optimality for this setting.

Remark 3 : Proof for noise augmentations is only conjectured.

In our analysis, we try to show that for any direction in the Fisher subspace not included in the optimal solution, the gradient of the loss function is negative. This leads us to upper bound the gradient of the loss by $- \delta \cdot \sum_k w_k a_k^2 + \sum_k w_k a_k^2$. This upper bound comes from the Perron-Frobenius theorem (at the end of Supplementary B). The gradient is indeed non-positive if $\delta = 1$. When $\delta <1$, our analysis cannot guarantee this, but we believe that the same is true even for relatively large values of $\delta$. However, this would require a more fine-grained analysis that takes into account the properties of mean subspace and the covariance matrix.

Additionally, we do study the effect of noise empirically as shown by our synthetic experiments. We believe that introducing noise doesn’t change the optimal subspace in the infinite sample limit, though we expect to have adverse effects on the number of samples required for it. This ties to the finite sample analysis question, which is on our future works list.

Remark 4 : Non-linear extensions

We show that for linear setting, InfoNCE loss learns the Fisher subspace, which we motivated as the optimal subspace for clustering. We do not know of an equivalent definition of optimality for representations under non-linear mappings, which also seems prohibitively difficult to come up with without assumptions on the mappings.

Remark 5 : Finite-sample bounds

We acknowledge that our analysis is designed for the infinite sample setting and there does not seem to be a trivial extension to finite samples. The main technical difficulty lies with the LogSumExp term in the loss function of InfoNCE. We believe that our reviewer would agree that such an analysis is still non-trivial in the linear setup we propose. That said, we are exploring possible directions which would enable us to make meaningful comments about the establishing bounds on samples for learning subspaces that are comparable of Fisher subspace. Our goal in this paper is to quantify the absolute limits of contrastive learning in the limit and we defer the finite sample study for our future work.

Finite-sample analysis with sample complexities and convergence rate analysis would also require a careful optimization algorithm design to solve the InfoNCE or other losses, which is not trivial and beyond the scope of our infinite sample analysis. We do leave the finite-sample analysis to future work, which we believe is an independent interest to the community.

Remark 6 : Role of Temperature and Batch Size

This is indeed a great questions, and it is relevant to the finite-sample analysis and the convergence rate of the optimization process. Clearly, batch size is not applicable to our analysis framework as we focus on the limiting behavior. Although temperature is more relevant, it should not effect our conclusion, beyond constants in the calculations, that InfoNCE would learn the Fisher subspace under perfect augmentations.

Dear Reviewer 7TSE,

With the Author-Reviewer discussion ending in less than a day, could you please let the authors know if you are satisfied with their response?

Sincerely,

AC

Remark 1 : Theorem 4.1 is restricted to the ideal case of perfect augmentation. Analysis doesn’t account for more realistic scenarios.

In our analysis, we try to show that for any direction in the Fisher subspace not included in the optimal solution, the gradient of the loss function is negative. This leads us to upper bound the gradient of the loss by $- \delta \cdot \sum_k w_k a_k^2 + \sum_k w_k a_k^2$. This upper bound comes from the Perron-Frobenius theorem (at the end of Supplementary B). The gradient is indeed non-positive if $\delta = 1$. When $\delta <1$, our analysis cannot guarantee this, but we believe that the same is true even for relatively large values of $\delta$. However, this would require a more fine-grained analysis that takes into account the properties of mean subspace and the covariance matrix.

Additionally, we do study the effect of noise empirically as shown by our synthetic experiments. We believe that introducing noise doesn’t change the optimal subspace in the infinite sample limit, though we expect to have adverse effects on the number of samples required for it. This ties to the finite sample analysis question, which is on our future works list.

Remark 2 : Reliance on shared covariance assumptions

We wanted to expand the typical spherical (identity) covariance assumption and consider a general covariance matrix shared among components. The assumption though is necessary for the theoretical analysis. One of the steps in our proof requires an affine transformation to make each component’s covariances isotropic, and the shared (non-spherical) covariance assumption is the most general we could come up with (please see Supplementary C). Moreover, the Bayes optimality of the Fisher subspace is only true for shared covariance GMMs and we do not have a reliable measure of optimality beyond that.

Remark 3 : Lack of experimental validation for multi-modal scenario

Owing to lack of a good small scale multi-modal dataset and a time crunch, we unfortunately could not make it to the rebuttal period. A possible experimental setup would be to take ImageNet images and generate synthetic text captions for them using an LLM. We believe this is a toy scenario and we were not sure about its value before the reviewer’s eye. If you would be willing to see such results, we will conduct a small scale experiment in the above-described manner and make them ready for the discussion period. Alternatively, we could provide an analysis of (a subset of) real CLIP embeddings, which we were not able to complete due to the limited period allocated for the rebuttal preparation. We will be working on this in the meantime.

Remark 4 : Analysis doesn’t account for more realistic scenarios like a limited number of positives and negatives. No discussion on number of positives and negatives required for effective contrastive learning

We acknowledge that our analysis is designed for the infinite sample setting and there does not seem to be a trivial extension to finite samples. The main technical difficulty lies with the LogSumExp term in the loss function of InfoNCE. We believe that our reviewer would agree that such an analysis is still non-trivial in the linear setup we propose. That said, we are exploring possible directions which would enable us to make meaningful comments about the establishing bounds on samples for learning subspaces that are comparable of Fisher subspace. Our goal in this paper is to quantify the absolute limits of contrastive learning in the limit and we defer the finite sample study for our future work.

Finite-sample analysis with sample complexities and convergence rate analysis would also require a careful optimization algorithm design to solve the InfoNCE or other losses, which is not trivial and beyond the scope of our infinite sample analysis. We do leave the finite-sample analysis to future work, which we believe is an independent interest to the community.

Remark 5 : Limited experimental analysis

We agree with your comment on our CIFAR-100 experiments. Following your suggestion, we conducted another set of experiments using the ImageNet dataset. Due to character limits, please refer to Remark 2 for Reviewer nsJs for the new result.

Remark 6 : Experimental details are not quite clear

We apologize for the lack of clarity and detail, and thank you for your constructive comment. We will resolve the ambiguities and also fill out any incomplete details in the supplementary sections as pointed out. To summarize the details; for the synthetic experiments we consider a 10-component shared-covariance GMM, with an ambient dimension of 100. The covariance matrix is a diagonal matrix. To simulate a parallel pancake-like setup, we increase the variance in directions orthogonal to the mean subspace which leads to poor clustering performance with methods like PCA. We particularly choose this setup, as unsupervised methods struggle with this kind of data.

Remark 7 : Would it not be more realistic to model the augmentation process as a conditional distribution. How would this change the theoretical conclusions, especially under statistical dependencies introduced by type-based augmentations?

Conditionally generated samples, as studied in [HaoChen et al., 2021], serve as an alternative framework for analyzing self-supervised learning. Prior work exploring these directions often deal with theory-friendly approximations of InfoNCE loss, i.e., replacing LogSumExp by a squared term and are not relatable to some Bayes-optimal notion of representation that our work deals with. Using probabilistic modelling is a more realistic approach, however, the conclusions will not be enough to explain possible advantages of contrastive learning against, for instance, fully supervised learning.

Regarding the effect of considering augmentations as conditional distributions for our setting: we also tried assuming a naive conditional distribution $\hat{x} \sim \mathcal{N}(x, \sigma)$, i.e., the augmented sample is a “noisy” version of the point $x$. Assuming this distribution, we were not able to recover the Fisher subspace with the InfoNCE loss. We posit that an ideal conditional distribution should be more informative than the original sample itself, and hence should provide some information about the cluster it belongs to . A stronger conditional distribution (our AeD) which assumes the augmented sample to be biased towards the underlying component of the original sample, is appropriately informative enough to learn the Fisher subspace. Identifying a less restrictive (when compared to AeD) conditional distribution, which can learn the Fisher subspace is indeed an interesting question.

We admittedly study a much simpler setup, but this helps us prove stronger results, i.e. showing equivalence to fully supervised methods, and our findings are partially recovered by the synthetic experiments and surprisingly validated by the real-data examples (both the old CIFAR-100 experiments and the new, non-linear ImageNet results).

HaoChen et.al., 2021. Provable Guarantees for Self-Supervised Deep Learning with Spectral Contrastive Loss, 2021

Remark 8 : Projection dimension r>K.

We tried to study a more difficult problem with fewer embedding dimensions to possibly make a stronger conclusion empirically, and we unfortunately did not have experiments with $r > K$. Following this remark, we run additional experiments with r=30 when for the 20-class dataset. We refer our reviewer to the last column of our ImageNet results in Remark 5.

Remark 9 : What does scaling within the subspace mean

Thanks for pointing this out. Our terminology was inaccurate and admittedly loose; let us be more precise. By scaling within subspace, we mean that methods such as LDA learn just a projection matrix (i.e., projection onto a subspace). While methods like InfoNCE learn an affine transformation (projection onto a subspace + affine transformation within the subspace) which leads to more suitable representations for clustering. We interpret that this is possibly the reason InfoNCE does better than LDA (Optimal in synthetic experiments) which only projects onto optimal subspace. We will make this clear in the camera ready version.

Remark 10 : Definition of a good projection and its sufficiency

We believe our reviewers refers to our definition of “favorable” mapping that maximizes Fisher discriminant. Our definition is indeed dependent on learning discriminative representations for each cluster. When each cluster corresponds to a broad class like “dogs”, these representations might be sub-optimal for certain transfer learning tasks (like distinguishing between different dog breeds). But when the clusters are finely defined (via augmentations), we get “good” representations that are discriminative enough for the sub-structures in the data. Naturally, the quality of the representations are tied to how fine-grained and/or noisy the augmentations are (which is the only information through which we could discover “classes”). Under this scenario, we believe our definition is reasonable, however, we would be delighted to hear more on this if our reviewer has better suggestions.

Remark 1 : Including an overview figure for the paper

Thanks for the suggestion. We will try to include one in the camera-ready version of the paper. We plan to add a figure capturing the data model for shared-covariance GMMs and display the properties of representations learnt by different self-supervised loss functions. We will also include a table that explains the main theoretical results for contrastive losses and the properties of the subspaces learned in comparison to optimal Fisher subspace.

Remark 2 : Comment on the shared covariance assumption

You are correct with your assessment; shared covariance means all components have the same covariance, and the data samples per component are distributed over the ambient space in the “same shape”. We wanted to expand the typical spherical (identity) covariance assumption and consider a general covariance matrix shared among components. The assumption though is necessary for the theoretical analysis. One of the steps in our proof requires an affine transformation to make each component’s covariances isotropic, and the shared (non-spherical) covariance assumption is the most general we could come up with (please see Supplementary C). Moreover, the Bayes optimality of the Fisher subspace is only true for shared covariance GMMs and we do not have a reliable measure of optimality beyond that. That said, we are exploring more general extensions of our analysis with appropriate optimality metrics.

In practice, however, we still get good numbers beyond shared covariances. For general Gaussians, we empirically show that InfoNCE learns representations that match (or even outperform) LDA on clustering metrics. We also include a new set of results on ImageNet. Specifically, we take images corresponding to 20 random classes from ImageNet-1K and use a pretrained ResNet-50 model to get the non-linear mappings. We then learn a linear map from the ResNet embedding space (2048 dimensional) to our target space. We present the results below (presented as ARI/AMI).

Remark 3 : Relevant work “I-Con: A Unifying Framework for Representation Learning”

We appreciate the reviewer pointing out other relevant work. Our method occupies the same block as InfoNCE i.e., uniform over positive pairs supervisory signal and gaussian learnt representations. Though for us the uniform pairs are defined as points belonging to the same cluster instead of a set of given positive pairs.