Structuring Representation Geometry with Rotationally Equivariant Contrastive Learning

Number reported on some datasets such as CIFAR10 are below the previously reported numbers

The 94% top-1 accuracy reported in the appendix of the original SimCLR paper [1] corresponds to training the model with a batch size of 1024 and over 1000 epochs. Figure B.7 in this paper reports a performance of approximately 92% using a batch size of 256 and 400 as the number of pre training epochs. A popularly referenced paper for CIFAR10 baseline numbers [2], reports a top-1 accuracy of 91.1% for a batch size of 512 and 800 pre-training epochs. EquiMOD [3], another paper in the literature of equivariance in contrastive learning reports an accuracy of 90.98% with a batch size of 512 and 800 pre-training epochs.

As rightly highlighted in your review, the quality of representation learned by SimCLR and its performance on the downstream tasks are significantly influenced by the training batch size, choice of architecture, number of training epochs etc. Due to computational constraints, we experimented with a fixed batch size of 256 and 400 training epochs across experiments. Given our hyper-parameters, we believe our reported number of 90.98% are consistent with the reported numbers in the above works.

[1] Chen, Ting, et al. "A simple framework for contrastive learning of visual representations." International conference on machine learning. PMLR, 2020.

[2] Chen, Xinlei, and Kaiming He. "Exploring simple siamese representation learning." Proceedings of the IEEE/CVF conference on computer vision and pattern recognition. 2021.

[3] Devillers, Alexandre, and Mathieu Lefort. "Equimod: An equivariance module to improve self-supervised learning." arXiv preprint arXiv:2211.01244 (2022).

Qualitative results on the image retrieval tasks are difficult to assess. What is the role of “input” as opposed to “query” ​​

In Figure 9, the “input” ($x$) corresponds to the original batch of images from the dataset, while “query” ($\tilde{x}$) refers to the transformed batch obtained using a randomly sampled augmentation $a$. The trained model receives $\tilde{x} = a(x)$ as input, and the top 5 nearest neighbors of its embedding i.e. $f(a(x))$ are extracted based on the Euclidean distance. The images corresponding to these top 5 nearest neighbor embeddings are displayed for both CARE and SimCLR.

Ideally, a representation equivariant to color variations in the input space must change the retrieved results (top 5 neighbors) in response to a change in query color. From Figure 9, CARE qualitatively performs better than SimCLR in the visual retrieval task, as its retrieved neighbors have colors similar to that of the query. In contrast, the retrievals for SimCLR remain relatively static.

Use of equivariant loss before the projection head and invariance loss after the projection head

Indeed, this is an excellent point. However, we tried this variant early in our testing, but found that it did not improve downstream performance. For instance we found that with ResNet18 on CIFAR10, we achieve the following kNN accuracies:

On the consistency under compositions exhibited by CARE

Thank you for raising this excellent question. We have revised the manuscript to include an approximate equivariance bound (Section A.1 and B.1) when the equivariant loss is small but non zero. We briefly describe the results below.

• To generalize Corollary 3 in our submission, assume that the loss is bounded (not exactly zero), so $||AA^\top - BB^\top|| < \epsilon^2$ in the 2-norm for some error $\epsilon^2 > 0$. Then, using Theorem 1 in [1], we can demonstrate that $\min_{R \in O(d)} ||A - BR|| < o(\epsilon)$, where $o(\epsilon) \to 0$ as $\epsilon \to 0$. As $\epsilon$ (and consequently the loss) approaches zero, we recover our exact equivariance result. For $\epsilon > 0$, we achieve an approximate equivariance up to $o(\epsilon)$.

• We further show that CARE enjoys consistency under composition of transformations even under low equivariant loss. Mathematically, $\rho : \mathcal A \rightarrow O(d)$ given by $\rho(a) = R_a$ satisfies $||\rho(a' \circ a) - \rho(a')\rho(a)|| \leq o(\epsilon^2)$ for almost all $a,a' \in \mathcal{A}$, where $o(\epsilon) \to 0$ as the error $\epsilon \to 0$.

Furthermore, we would like to clarify that CARE is tested within the standard contrastive learning framework, where each input transformation is already a composition of random crop, random horizontal flip, color jitter, random grayscaling, and Gaussian blurring.

[1] Arias-Castro, Ery, Adel Javanmard, and Bruno Pelletier. "Perturbation bounds for procrustes, classical scaling, and trilateration, with applications to manifold learning." Journal of machine learning research 21 (2020).

We appreciate your critical and thorough feedback, which has been valuable in enhancing our work. We are glad you appreciated the technical novelty and analysis, and believe this could form the basis for a positive review. Below, we aim to answer their questions as concretely as possible. The changes have also been incorporated into the revised manuscript and are highlighted in magenta.

Comparison of CARE with another equivariant baseline for the Wahba’s problem

Thank you for this valuable suggestion. We have included the following additional experiments in our revised manuscript (highlighted in magenta) to address your comments about our experimental protocols as concretely as possible.

• Wahba error for a model trained using EquiMOD on CIFAR10 in Figure 5. As shown in the figure, CARE achieves the lowest error on Wahba's problem, highlighting its ability to learn an orthogonally equivariant representation.

• EquiMod, an equivariant contrastive learning baseline, in our protein point cloud experiment (Figures 3 and 4, as shown in the updated manuscript). Figure 3 provides a qualitative comparison of trajectories, demonstrating that EquiMod does not exhibit a rotational structure and is qualitatively closer to SimCLR than CARE. Figure 4 presents quantitative results, showing that CARE outperforms EquiMod on the principal component prediction task.

• Qualitative assessment of the representation learned by EquiMOD in Figure 21 on Flowers102. Both CARE and EquiMOD, being equivalent baselines, show sensitivity to color. However, EquiMOD's representation exhibits nearest neighbors with significantly different shades (e.g., red and orange) compared to those learned by CARE, which are closer in color to the query images. Note that this experiment assesses CARE's ability to learn equivariance and not orthogonal equivariance. Thus, any equivariant baseline would exhibit this sensitivity to input transformations (color variations in this case).

• Additionally, we examine the quality of features learned by training the Barlow Twins [1] invariant baseline with our equivariant loss $\mathcal{L}\_{\text{equiv}}$. As shown below and in Table 1 of the revised manuscript, $\text{CARE}\_{\text{Barlow Twins}}$ outperforms its invariant counterpart on all three datasets, CIFAR10, STL10 and CIFAR100.

Transfer learning experiments

We acknowledge the reviewer's point regarding the common practice of evaluating linear probe performance for a model trained with ImageNet on other datasets like CIFAR10, CIFAR100, and STL10 when proposing invariant baselines. However, it's worth noting that existing equivariant research [1, 2] often conducts experiments involving pretraining and testing on the same datasets, with CIFAR10 and ImageNet being popular choices. In line with this convention, we test our approach on smaller benchmarks (CIFAR10, CIFAR100, STL10) and a larger dataset (ImageNet100) as a computationally efficient alternative to conducting experiments directly on ImageNet, which is beyond our resource limits.

[1] Dangovski, Rumen, et al. "Equivariant contrastive learning." arXiv preprint arXiv:2111.00899 (2021).

[2] Devillers, Alexandre, and Mathieu Lefort. "Equimod: An equivariance module to improve self-supervised learning." arXiv preprint arXiv:2211.01244 (2022).

Desired correct trajectory for the protein experiment. Could the size of the dataset or batch be responsible for lower performance in this experiment for SimCLR.

• About the dataset and desired trajectory. In this experiment, we use the Protein Data Bank (PDB)—the single global repository of experimentally determined 3D structures of biological macromolecules and their complexes, containing approximately 130,000 samples. The core objective is to assess whether our method can effectively encode the SO(3) manifold, a highly challenging and critical aspect in drug discovery. Each input point cloud is transformed through action of the SO(3) group. Consequently, the desirect 2D trajectory is a circle corresponding to rotation along each of three orthogonal axes. This is consistent with the structure of the SO(3) manifold. FIgures 3, 17, 18, 19 and 20 illustrate additional trajectories observed through the embedding space, showing that CARE is adept at learning the SO(3) manifold. We had added additional clarification in the revised manuscript

• Failure of SimCLR. The reason for low performance for SimCLR is its inability to encode any information about the input transformation i.e. rotation in 3D. As a consequence, it is unable to learn the orientation of a given protein. Note that the batch sizes and architectural choices were fixed for both CARE and SimCLR.

We appreciate the reviewer's thorough and insightful review, along with their positive feedback regarding the strengths of our work, including our theoretical analysis and discussions on extending the approach to other groups and representation space geometries. In response to the review, we have highlighted the changes in magenta in the revised manuscript and address the questions raised below.

Why not optimize the Wahba’s Problem as an alternate loss to our equivariance loss.

This is indeed an excellent question. While Wahba's error, denoted as $\mathcal{W}\_f$ = $\min_{R \in SO(d)} || RF - F_a ||_{\text{Fro}}$, seems like a plausible loss function to enforce orthogonal equivariance, it faces practical challenges in real-world settings for two important reasons:

• Computational Inefficiency: Solving the Wahba's problem involves obtaining a closed-form solution through the singular value decomposition of one of the two point clouds. Minimizing the Wahba's error in each iteration is computationally intensive, with a complexity of $O(m^2n, mn^2, n^3)$, where $m$ and $n$ are the dimensions of the involved matrices. In our context, the computational complexity is dominated by $O(bs * dim^2)$, where $bs$ is the batch size, and $dim$ is the dimensionality of the representation (e.g., 2048 for ResNet50). In contrast, our proposed loss achieves the same objective by operating on pairs, with the same complexity as standard invariant contrastive baselines.

• Noise and instability: The closed-form solution to the Wahba's problem, obtained through singular value decomposition, can introduce instability and noise in calculations. These issues can accumulate over iterations, particularly when dealing with noisy data. In some cases, even minor data perturbations can lead to substantial changes in the decomposed components, especially in the smallest singular values and corresponding vectors. This noise becomes more pronounced when working with ill-conditioned matrices.

Both qualitative and quantitative comparisons to equivariant methods

Thank you for this valuable suggestion. We have included the following additional experiments in our revised manuscript (in magenta)

• EquiMod, an equivariant contrastive learning baseline, in our protein point cloud experiment (Figures 3 and 4, as shown in the updated manuscript). Figure 3 provides a qualitative comparison of trajectories, demonstrating that EquiMod does not exhibit a rotational structure and is qualitatively closer to SimCLR than CARE. Figure 4 presents quantitative results, showing that CARE outperforms EquiMod on the principal component prediction task.

• Wahba error for a model trained using EquiMOD on CIFAR10 in Figure 5. As shown in the figure, CARE achieves the lowest error on Wahba's problem, highlighting its ability to learn an orthogonally equivariant representation.

• Qualitative assessment of the representation learned by EquiMOD in Figure 21 on Flowers102. Both CARE and EquiMOD, being equivalent baselines, show sensitivity to color. However, EquiMOD's representation exhibits nearest neighbors with significantly different shades (e.g., red and orange) compared to those learned by CARE, which are closer in color to the query images. Note that this experiment assesses CARE's ability to learn equivariance and not orthogonal equivariance. Thus, any equivariant baseline would exhibit this sensitivity to input transformations (color variations in this case).

• Additionally, we examine the quality of features learned by training Barlow Twins [1], an invariant baseline with our equivariant loss $\mathcal{L}\_{\text{equiv}}$. As shown below and in Table 1 of the revised manuscript, $\text{CARE}_{\text{Barlow Twins}}$ outperforms its invariant counterpart on all three datasets, CIFAR10, STL10 and CIFAR100.

[1] Zbontar, Jure, et al. "Barlow twins: Self-supervised learning via redundancy reduction." International Conference on Machine Learning. PMLR, 2021.

Comparison with explicit parametrization of $R_a$.

We agree that comparing with the direct parameterization of augmentations $a$ would be an interesting avenue of research. However, one of the primary objectives of our work was to design an equivariant approach that did not rely on parameterizing augmentations $a$ and instead learns to be equivariant to them directly from data pairs. We view this as a possible disadvantage since it relies on having convenient features describing a particular augmentation sample. Although not for rotational equivariance, this approach has been explored in related works such as EquiMod, which we compare to in Table 1.

Figure 8 and the paragraph discussing relative rotational equivariance

We apologize for the confusion regarding Figure 8, which has been updated to address the initial duplication issue. It illustrates three key metrics from left to right: the relative equivariance metric $\gamma_f$, the rotational equivariance measure, and the invariance measure (as detailed in Section 4).

Here are the two key observations from this plot:

• The invariance measure, $\mathcal{L}_{\text{inv}}$, is comparable for both SimCLR and $\text{CARE}\_{\text{SimCLR}}$ and is non-zero, implying approximate invariance.
• The rotational equivariance measure is significantly lower for $\text{CARE}\_{\text{SimCLR}}$ compared to SimCLR, indicating orthogonal equivariance in $\text{CARE}\_{\text{SimCLR}}$.

These observations suggest that $\text{CARE}_{\text{SimCLR}}$ is not merely achieving lower equivariance error by collapsing into invariance, which would be a trivial solution of equivariance. We have included this clarification in the revised manuscript for better understanding.

We are grateful to the reviewer for the time they put in to review our work. We are glad to see that they recognize several strengths in our work, including theoretical guarantees, comprehensive empirical evaluation, and qualitative measures for assessing the structure of learned representation. Below, we share our thoughts on the questions asked.

Can the structures learned be effectively transferred to other classes or varied transformation parameter ranges?

Thank you for your valuable suggestion. We would like to emphasize that the key aim of this paper is to learn a truly equivariant representation which maps complex input transformations to isometries in the embedding space. Assessing the generalization capabilities of the learned representation across a wide range of transformation parameters is indeed a very intriguing future direction and warrants thorough investigation.

Reporting variance along with the mean due to small performance differences on CIFAR10 and STL10.

This is indeed a crucial point, and we have updated the manuscript to include both the average performance and standard deviations for clarity. As shown in the results summarized below, CARE outperforms its invariance counterparts by more than the corresponding standard deviation.

On the generalization to other architectures and domains beyond computer vision.

ResNets are a popular and common choice of architecture for contrastive learning and are also used in our experiments. To address the concern of the reviewer, we run additional baselines on the CIFAR10 dataset using both SimCLR and $\text{CARE}_{\text{SimCLR}}$ with DenseNet121. CARE outperforms the baseline on top-1 linear probe accuracy as shown below

Further, in the original manuscript, we also demonstrate the effectiveness of our approach in learning the SO(3) manifold on 3D protein structures using a DeepSet architecture, which is significantly different from 2D vision models. Qualitative analysis to assess the learned structure in the representation is also conducted for the Flowers 102 dataset apart from the standard linear probe experiment on the vision datasets.

In many practical applications input transformations may not preserve distances and angles. What does CARE learn in this case?

In our setup, we train a model to map complex input transformations like random cropping, color jitter, and random flipping to orthogonal transformations in the embedding space. Note that these input transformations typically do not preserve distances or angles for a general encoder. Even so, the underlying mechanism of CARE - to transform data into an embedding space where the complex input transformations become orthogonal transformations. This is provably possible for any transformation that belongs to a compact Lie group, so long as the embedding space dimension is large enough (this is also discussed after Corollary 1).

Fig. 8 is confusing

We apologize for the confusion regarding Figure 8, which has been updated to address the initial duplication issue.

We thank the reviewer for their thought-provoking review and encouraging words of appreciation for our work. We have endeavored to address your concerns as concretely as possible, with the changes in the revised manuscript highlighted in magenta.

The use of both an invariant and equivariant loss in the overall objective function seems conceptually strange

• As highlighted in our original submission, the core idea behind this conceptualization is based on our hypothesis that small perceptual changes in data, resulting from data augmentations, should correspond to small perturbations in embeddings. However, solely minimizing the equivariance and uniformity loss $\mathcal{L}{\text{equiv}} + \mathcal{L}{\text{unif}}$ does not guarantee this. To address this, we introduce an additional constraint by enforcing localized transformations in the representation space. A natural choice in contrastive learning is to minimize the invariance loss $\mathcal{L}_{\text{inv}}$ to push the embeddings of pairs of data closeby. In essence, our model aims to learn approximate equivariance, where augmentations in input space translate to small and orthogonal transformations in the representation space.

• Relative rotational equivariance $\gamma_f$ is the ratio of an orthogonally equivariant measure to the invariance loss. Specifically, since invariance $f(a(x)) = f(x)$ is a trivial solution of our equivariant loss, through $\gamma_f$, we measure the degree to which $f(a(x)) = R_af(x)$ where $R_a \neq I_d$. A lower value of $\gamma_f$ indicates a higher degree of orthogonal equivariance, signifying non-trivial equivariance.

Theoretical guarantees of CARE given approximate (and not exact) minimization of the regularize

We thank the reviewer for raising this excellent question. Indeed, we can provide an approximate equivariance bound when the equivariant loss is non zero. The changes are highlighted in Section A.1 and B.1 of the Appendix (in magenta) in the revised manuscript and are briefly described below.

• To generalize Corollary 3 in our submission, assume that the loss is bounded (not exactly zero), so $||AA^\top - BB^\top|| < \epsilon^2$ in the 2-norm for some error $\epsilon^2 > 0$. Then, using Theorem 1 in [1], we can demonstrate that $\min_{R \in O(d)} ||A - BR|| < o(\epsilon)$, where $o(\epsilon) \to 0$ as $\epsilon \to 0$. As $\epsilon$ (and consequently the loss) approaches zero, we recover our exact equivariance result. For $\epsilon > 0$, we achieve an approximate equivariance up to $o(\epsilon)$.

• We further show that CARE enjoys consistency under composition of transformations even under low equivariant loss. Mathematically, $\rho : \mathcal A \rightarrow O(d)$ given by $\rho(a) = R_a$ satisfies $||\rho(a' \circ a) - \rho(a')\rho(a)|| \leq o(\epsilon^2)$ for almost all $a,a' \in \mathcal{A}$, where $o(\epsilon) \to 0$ as the error $\epsilon \to 0$.

[1] Arias-Castro, Ery, Adel Javanmard, and Bruno Pelletier. "Perturbation bounds for procrustes, classical scaling, and trilateration, with applications to manifold learning." Journal of Machine Learning Research 21 (2020).

Minor issues and typographical errors

We thank you for your attention to detail. We have made the following corrections in the revised manuscript.

• Referencing of figures in appendix: Thank you for pointing this out. We have fixed this in the revised manuscript.
• Regarding low CARE loss guarantees: We concur with the reviewer and have made the correction in the revised manuscript.
• Typographical error: We are changed ‘principle’ to ‘principal’ at the bottom of page 5