Paper: Structuring Representation Geometry with Rotationally Equivariant Contrastive Learning (ICLR-2024)

Reviewer comment (Experiments):
The experiments that use Wahba’s Problem and the Relative Rotation Equivariance should contain some stronger baselines. It is expected that CARE achieves better results in terms of equivariance compared to SimCLR (that does not try to achieve orthogonal equivariance). While SimCLR could be a lower bound, it is difficult to say if the results shown by CARE are "good" as no other strong comparison is provided. A stronger result would be to show that CARE is comparable or superior to other equivariant-inducing algorithms.

Author response:
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. 

Algorithm | CIFAR10 | CIFAR100 | STL10
|----------|----------|----------|----------|
Barlow Twins | 84.54 $\pm$ 0.02 | 55.54 $\pm$ 0.05 | 90.62 $\pm$ 0.02
$\text{CARE}_{\text{Barlow Twins}}$ | **85.65 $\pm$ 0.05** | **56.76 $\pm$ 0.02** | **90.92 $\pm$ 0.01** |

Reviewer comment (Experiments):
The results on linear probing seem to report results on an uncommon setting (training on CIFAR-10, 100, STL, ImageNet100, as opposed to training on ImageNet or ImageNetTiny and testing a linear probe on those datasets). Previous works such as BYOL, SimCLR, and Duet report much stronger results in the transfer learning setting. These transfer learning results are crucial for this work because they could support the claim that using equivariance allows learning more generic and transferable features. Training and testing on the same dataset is less interesting.

Author response:
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).

Reviewer comment (Experiments):
For results that can be found on training SimCLR on CIFAR-10, the results reported in the paper seem below previously reported numbers. The original SimCLR paper reports 94% accuracy on CIFAR-10 (appendix), compared to 90.98% reported here. In contrastive frameworks, once the pre-training pipeline is optimized and training scales up (in terms of datasets, batch size, and epochs), initial advantages from model tweaks reduce or disappear. It would be valuable to show that this does not happen with equivariance.

Author response:
We concur with the reviewer that testing at large scale is indeed important. However, it's important to note two key points:

- We cannot definitely predict if the gains because of CARE won’t persist in the large scale regime. In fact, we emphasis that our newly added baseline, Barlow Twins and $\text{CARE}_{\text{Barlow Twins}}$ were trained for 800 pretraining epochs. Notably, CARE consistently outperformed its invariant baseline, Barlow Twins across the three benchmark datasets, CIFAR10, CIFAR100 and STL10. 

- While the performance of SimCLR is known to be highly sensitive to batch size and pretraining epochs, other algorithms like MoCo-v2 [1] and Barlow Twins [2] have shown robustness to changes in batch size, performing well even with a batch size of 256.

Following the reviewer’s suggestion and considering the constraints of available compute, we would be happy to include the baselines for 800 pretraining epochs in the camera-ready version (if requested). However, we emphasize that SimCLR on CIFAR10 achieves a top-1 accuracy of 91.1\% for a batch size of 512 and 800 pre-training epochs [3],  which is only 0.12% higher than our reported numbers.

[1] Chen, Xinlei, et al. "Improved baselines with momentum contrastive learning." arXiv preprint arXiv:2003.04297 (2020).

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

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

Reviewer comment (Experiments):
In the protein point clouds experiments, it is unclear what the desired (correct) trajectory is. The manuscript could explain more clearly the expected result (Figure 3). Given that SimCLR requires large batches and datasets, it is questionable whether this dataset’s size makes it suitable for pre-training with SimCLR. A better test might be to pre-train on a larger dataset and fine-tune on this one.

Author response:
- **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.

Reviewer comment (Evaluation):
The qualitative results on the image retrieval tasks are difficult to assess. The results from SimCLR appeared better. This task should be explained more clearly, and it would be better to employ quantitative evaluation if possible. The role of the “input” compared to the “query” is also unclear.

Author response:
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.

Reviewer comment (Theory):
In the proposed algorithm, all three terms act on the same representation space (after the projection head). This introduces ambiguity—whether the objective is to induce invariance or equivariance. While invariance satisfies equivariance, achieving equivariance through invariance results in information loss typical of invariance-inducing algorithms. Common contrastive learning algorithms based on invariance impose the constraint after the projection head, using the representation before it (where invariance is hopefully not strongly achieved). In the current setup, the equivariance term seems to act only as a regularization effect (especially since the λ value is small). A potentially more interesting use of the contrastive framework would be to impose equivariance on the embedding before the projection head (used for downstream tasks) and invariance after the projection head, potentially leveraging the benefits of both properties.

Author response:
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:

Algorithm | Top-1 kNN accuracy| 
|----------|----------|
Loss before projection | 86.0% | 
Loss after projection | **86.7%** |

Reviewer comment (Evaluation):
While observing the performance metrics, one notices that the performance gap between CARE and SimCLR on CIFAR10 and STL10 is <1%. It raises the question of the actual significance and practical implications of this difference. For a more comprehensive understanding, it would be great if Table 1 could also report the variance alongside the mean.

Author response:
> 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.  

Algorithm | CIFAR10 | CIFAR100 |STL10 | ImageNet100 |
|----------|----------|----------|----------|----------|
SimCLR | 90.98 $\pm$ 0.10| 66.77 $\pm$ 0.34 | 84.19 $\pm$ 0.13| 72.79 $\pm$ 0.08
$\text{CARE}_{\text{SimCLR}}$ | **91.92 $\pm$ 0.12** | **68.05 $\pm$ 0.28** | **84.64 $\pm$ 0.29** |**76.69 $\pm$ 0.08**
| | | | |
MoCo-v2 | 91.95 $\pm$ 0.05| 69.88 $\pm$ 0.23| - | 73.50 $\pm$ 0.19
$\text{CARE}_{\text{MoCo-v2}}$| **92.19 $\pm$ 0.01** | **70.56 $\pm$ 0.15** | 88.97 $\pm$ 0.48 | **74.30 $\pm$ 0.07**

Reviewer comment (Experiments):
The analysis is limited to ResNet networks. How well do the findings generalize to other architectures? How does the architecture affect the properties of the learned representations? How well does the proposed method generalize to other domains beyond the standard computer vision datasets?

Author response:
> 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

Algorithm | Linear Probe Accuracy
|----------|----------|
SimCLR | 82.45
$\text{CARE}_{\text{SimCLR}}$ | **83.56**

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.

Reviewer comment (Presentation):
Fig. 8 is confusing: no legend, all three plots are the same (?).

Author response:
> Fig. 8 is confusing

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

Reviewer comment (Theory):
The use of both an invariant and equivariant loss in the overall objective function seems conceptually strange (although the ablations show it leads to superior performance). I would like to understand what might be being learned with this combination of losses -- my reading of the explanation in section 3 is that, among embeddings that *minimize* the equivariant regularizer (following proposition 1), those that achieve a small invariant loss will prefer small, rather than large, orthogonal transformations. However, it is not clear to me why this setting should arise in experiments -- why not a situation where the invariance loss is minimized, and the equivariant regularizer is only small? I am also unsure how this connects with the discussion in the "Relative rotational equivariance." paragraph later.

Author response:
> 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.

Reviewer comment (Theory):
It would be ideal if the authors could assert theoretically some degree of approximate invariance given approximate minimization of their regularizer (see a comment to this effect below). It seems important to precisely understand what happens in this setting, given the fact that a mixture of equivariant and invariant losses are required for strong practical performance.

Author response:
>  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).

Reviewer comment (Writing):
After equation (4): better to not reference a figure in the appendix without adding something like "In the appendix, we show ..."

Author response:
- *Referencing of figures in appendix*: Thank you for pointing this out. We have fixed this in the revised manuscript.

Reviewer comment (Writing):
Bottom of page 5: principle -> principal

Author response:
- *Typographical error*: We are changed ‘principle’ to ‘principal’ at the bottom of page 5

Reviewer comment (Theory):
Is there an "approximate" version of the results from invariant theory that the authors use to establish their theoretical results? If one digs into the proof of the result from invariant theory (say, specialized to $\mathrm{SO}(n)$), what obstructions are there to having an "approximate" version of the result?

Author response:
>  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).