Structuring Representation Geometry with Rotationally Equivariant Contrastive Learning
Authors: Sharut Gupta
Joshua Robinson
Derek Lim
Soledad Villar
Stefanie Jegelka
Venue: ICLR 2024 poster (2024)
Keywords: Self supervised learning
Equivariance
Contrastive learning
Primary area: unsupervised, self-supervised, semi-supervised, and supervised representation learning
TL;DR: We add additional geometric structure to the embedding space of contrastive learning by enforcing transformations of input space to correspond to simple (i.e., linear) transformations of embedding space.

Abstract: Self-supervised learning converts raw perceptual data such as images to a compact space where simple Euclidean distances measure meaningful variations in data. In this paper, we extend this formulation by adding additional geometric structure to the embedding space by enforcing transformations of input space to correspond to simple (i.e., linear) transformations of embedding space. Specifically, in the contrastive learning setting, we introduce an equivariance objective and theoretically prove that its minima force augmentations on input space to correspond to rotations on the spherical embedding space. We show that merely combining our equivariant loss with a non-collapse term results in non-trivial representations, without requiring invariance to data augmentations. Optimal performance is achieved by also encouraging approximate invariance, where input augmentations correspond to small rotations. Our method, CARE: Contrastive Augmentation-induced Rotational Equivariance, leads to improved performance on downstream tasks and ensures sensitivity in embedding space to important variations in data (e.g., color) that standard contrastive methods do not achieve. Code is available at https://github.com/Sharut/CARE

Published as a conference paper at ICLR 2024
STRUCTURING REPRESENTATION GEOMETRY WITH Ro- TATIONALLY EQUIVARIANT CONTRASTIVE LEARNING
Sharut Gupta"
Joshua Robinson"
Derek Lim MIT CSAIL
MIT CSAIL
MIT CSAIL
sharut
joshrobemit.edu
derekl www.iiiiitete..mmmmmmmm..m edu
Soledad Villar
Stefanie jegelka MIT CSAIL stefjeemit edu
Johns Hopkins University svillar3e jhu. edu
ABSTRACT
Self-supervised learning converts raw perceptual data such as images to compact space where simple Euclidean distances measure meaningful variations in data. this paper, we extend this formulation by adding additional geometric structure the embedding space by enforcing transformations of input space to correspond to simple (i.e., linear) transformations of embedding space. Specifically, in the contrastive learning setting, we introduce an equivariance objective and theoret- ically prove that its minima forces augmentations input space to correspond to rotations on the spherical embedding space. We show that merely combining our equivariant loss with non-collapse term results in non-trivial representa- tions, without requiring invariance to data augmentations, Optimal performance is achieved by also encouraging approximate invariance, where input augmentations correspond to small rotations. Our method, CARF: Contrastive Augmentation- induced Rotational Equivariance, leads to improved performance on downstream tasks, and ensures sensitivity in embedding space to important variations in data (e.g., color) that standard contrastive methods do not achieve. Code available at com/ Sharut /CAR.
https:
INTRODUCTION
is only partially understood what structure neural network representation spaces should possess in order to enable intelligent behavior to emerge efficiently (Ma al., 2022). One known key ingredient to learn low-dimensional spaces in which simple Euclidean distances effectively measure the similarity between data, as demonstrated by powerful self-supervised methods for web-scale learning (Chen et al., 2020; Schneider et al., 2021; Radford et al., 2021). However, many use cases require the use of richer structural relationships that similarities between data cannot capture. One example that has enjoyed considerable success is the encoding of relations between objects (X is a parent of Y, A a treatment for B) as simple transformations of embeddings (e.g., translations), which has driven learning with knowledge graphs (Bordes et al., 2013; Sun et al., 2019; Yasunaga et al., 2022). But similar capabilities have been notably absent from existing self-supervised learning recipes.
Recent contrastive self-supervised learning approaches have explored ways to close this gap by ensuring representation spaces are sensitive to certain transformations of input data, such as variations in color (Dangovski et al., 2022; Devillers & Lefort, 2023; Garrido et al., 2023; Bhardwaj et al., 2023) that earlier contrastive methods lacked. Encouraging sensitivity is especially important in contrastive learning, as it is known to learn shortcuts that forget features that are not needed to solve the pretraining task (Robinson et al., 2021). This line work formalizes sensitivity in terms of equivariance: transformations of input data correspond predictable transformations in representation space. Equivariance requires specifying a family of transformations E A in the input space, corresponding transformation T, in representation space, and training f so that
Equal contribution.

---

Published as a conference paper at ICLR 2024
RyD R Biy
A
sui R.
Flxtt)
Negative Sample
Positive Pairs
Positive Pairs
Positive Pairs
Orthogonally Equivariant Contrastive Learning (CARE)
Invariant Contrastive Learning (SimC1.R)
Figure I:
CARE is equivariant contrastive learning approach that trains augmentations (cropping, blurring, etc.) of input data to correspond to orthogonal transformations of embedding space.
f(a(z)) x (z) Prior works have considered a learnable feed-forward network for Ta (Devillers Lefort, 2023; Garrido et al., 2023). However, we find that this approach suffers from geometric pathologies, such inconsistency under compositions: Tagoaz f (xx Taz Ta, (x). Furthermore, encoding the relation between the embeddings of x and a(x) in a non-linear way conflicts with the aim of learning representation space where linear transformations relate different instances.
To address these concerns we propose CARE, an equivariant contrastive learning framework that learns translate augmentations in the input space (such cropping, blurring, and jittering) into simple linear transformations in feature space. Here, we use the sphere as our feature space (the standard contrastive learning space) and hence consider the isometries of the sphere: rotations and reflections, i.c., orthogonal transformations. As orthogonal transformations are less expressive (by design), our learning problem is more constrained, meaning that prior approaches for learning non-linear transforms cannot be used (see Section 3). Instead, CARE trains f to preserve angles, i.e., [((())) (a(x')) f(x) [a²), a property that must hold if f is orthogonally equivariant. We show that achieving low error this seemingly weaker property also implies approximate equivariance and enjoys consistency under compositions. Critically, we can easily integrate CARE into contrastive learning workflows since both operate by comparing pairs of data.
The key contributions of this work include:
Introducing CARE, novel equivariant contrastive learning framework that trains transfor- mations (cropping, jittering, blurring, etc.) in input space to approximately correspond to local orthogonal transformations in representation space.
Theoretically proving and empirically demonstrating that CARE places an orthogonally equivariant structure on the embedding space.
Showing that CARE increases sensitivity to features (e.g., color) compared to invariance- based contrastive methods, and also improves performance on image recognition tasks.
RETHINKING HOW AUGMENTATIONS ARE USED IN SELF-SUPERVISED LEARNING
Given access only to samples from a marginal distribution p(x) on some input space X' such as images, the goal of representation learning is commonly to train a feature extracting model : X - 80-1 mapping to the unit sphere $d-1 {2 E R ||=||| = 1}. A common strategy to automatically generate supervision from the data to additionally introduce space of augmentations A. containing maps a : X + X' which slightly perturb inputs (blurring, cropping, jittering, etc.). Siamese self- supervised methods learn representation spaces that reflect the relationship between the embeddings of x = a(z) and x a (ø), commonly by training f to be invariant or equivariant to the augmentations in the input space (Chen & He, 2021).
Invariance to augmentation. One approach is to train f to embed x and t nearby-i.s so that f(x f(x+) is invariant augmentations. The InfoNCE loss (van den Oord et al., 201 8; Gutmann
2

---

Published as a conference paper at ICLR 2024
& Hyvärinen, 2010) used in contrastive learning achieves precisely this:
e f(=T\/T
LIaIoNCE = E at (x7)
log ef(n)) /(xt)/T + e
(1)
where T > 0 temperature hyperparameter, and T; p are negative samples from the marginal distribution on X. As noted by Wang & Isola (2020), the contrastive training mechanism balances invariance augmentations with a competing objective: uniformly distributing embeddings over the sphere, which rules out trivial solutions such as constant functions.
Whilst contrastive learning has produced considerable advances in large-scale learning (Radford et al., 2021), several lines of work have begun to prohe the fundamental role of invariance in contrastive learning. Two key conclusions of recent investigations include: 1) invariance limits the expressive power of features learned by f. as it removes information about features transformations that may be relevant in fine- greiggrinedd tasks (Lee et al., 2021; Xie et al., 2022), and 2) contrastive learning actually benefits from not having exact invariance. For instance, critical role of the projection head to expand the feature space so that f is not fully invariant (Jing et al., 2022), suggesting that is preferable for the embeddings of and xt to be close, but not identical.
Equivariance to augmentation. To address the limitations of invariance, recent work has addition- ally proposed to control equivariance (i.e., sensitivity) of data transformations (Dangovski et al., 2022; Devillers & Lefort, 2023; Garrido et al., 2023). Prior works can broadly be viewed as training set of features f (sometimes alongside the usual invariant features) so that [a((=) (x)(x) for samples ~ p from the data distribution where Ta is some transformation of the embedding space. A common choice to take Taf(() = (f)(f(x).. ) learnable feed-forward network, and optimize a loss [MLPP([(z) about a encoded into the embedding of a(x), it permits complex non-linear relations between embeddings and hence does not necessarily encode relations in a linearly separable way. Furthermore, it does not enjoy the beneficial properties of equivariance the formal group-theoretic sense, such as consistency under compositions in general: Tagoo fz) + TooTaa (zz).
[((xx) |2- Whilst learnable MIP ensures that information
Instead, this work introduces CARE, an equivariant contrastive learning approach respecting two key design principles:
Principle 1. The map Ta satisfying [a(x)) = (x) should be linear.
Principle 2. Equivariance should be learned from pairs of data, in invariant contrastive learning.
The first principle asks that f converts complex perturbations a of input data into much simpler (i.e.. linear) transformations in embedding space. Specifically, we constrain the complexity of T. by considering isometrics of the sphere, (dd) = {Q E ddd'd" ; QQ" = Q"Q = I). containing all rotations and reflections. Throughout this paper define f(((1)
(x= for Ta O(d) to be orthogonal equivariance. This approach draws heavily from ideas in linear representation theory (Curtis& & Reiner, 1966; Serre et al., 1977), which studies how convert abstract group structures into matrix spaces equipped with standard matrix multiplication as the group operation.
The second principle stipulates how want learn orthogonal equivariance. Naively following previous non-linear approaches is challenging as our learning problem is more constrained, requiring learning a mapping + R. to orthogonal matrices. Furthermore, for a single (a, x) pair, the orthogonal matrix R2 such that [a(x)) = [z) is not unique, making it hard to directly learn R.- We sidestep these challenges by, instead of explicitly learning Ra, training so that an augmentation applied to two different inputs z, x produces the same change in embedding space.
Our method, CARE, encodes data augmentations (cropping. blurring, jittering etc.) as O(d) transfor- objective function. CARE can be viewed as
mations of embeddings using an equivariance-proneneee an instance of symmetry regularization, a term introduced by Shakerinava al. (2022).
CARE: CONTRASTIVE AUGMENTATION ROTATIONAL EQUIVARIANCE
This section introduces a simple and practical approach for training a model f: X s so that is orthogonally equivariant: i.e., data augmentation a A (cropping, blurring, jittering, etc.)
3

---

Published as a conference paper at ICLR 2024
applied to any input E X' causes the embedding f() to be transformed by the same Ra E ((d)
for all x E X: f(a(z)) = (z).
To achieve this, we consider the following loss:
Lequil (f) BanABre
- f(x) s()))
Laqui = Ba [f(a(z')) f(a())) - f(a (((di)))
(3)
Since inner products describe angles on the sphere, this objective enforces the angles between the embeddings of independent samples and x' to be the same those between their transformed counterparts a(x) and a(z'). This is necessarily true if f is orthogonally equivariant or, more generally. R. E ((d) exists. But the converse- that Lequi = 0 implies orthogonal equivariance is non non-obvious. In Section 3.1 we theoretically analyze Lequi demonstrating that it does indeed enforce mapping input augmentations to orthogonal transformations of embeddings. In practice. we replace the f (a:) f((r') term with f(a((x)) (aa(zz') for a freshly sampled a' ~ A, noting that minimizing this variant also minimizes Lequi if we assume a' can be the identity function with non-zero probability. A trivial but undesirable solution that minimizes Lequi is to collapse the embeddings of all points to be the same (see Figure 10). One natural approach to avoiding trivial solutions is combine the equivariance loss with non-collapse term such as the uniformity Lumit = log Ex,x'>xx exp (f(z) f(x')) (Wang & Isola, 2020) whose optima f distribute points uniformly over the sphere:
L())- = Lequii (J) + (f)
This is directly comparable to the InfoNCE loss, which can similarly be decomposed into two terms:
LIstONCE(f) = (f) + Lumi(f))
(5)
where (f) ((l((=()(((( f(a((2))) is minimized when f is invariant i-..., f(a(z)) = f(x). Figure in the Appendix shows that training using Lequi + Luiit yields non- trivial representations. However, the performance is below that of invariance- contrastive learning approaches. We hypothesize that this is because data augmentations which make small perceptual changes data-should correspond small perturbations of embeddings, which Lequi does not enforce.
To rule out this possibility, we introduce CARE: Contrastive ciimmenttttio-inducee Rotational Equivariance. CARE additionally enforces the orthogonal transformations in embedding space to be localized by reintroducing an invariance loss term Liw to encourage f to be approximately invariant. Doing so breaks the indifference of Lequi between large and small rotations, biasing towards small. Specifically, we propose the following objective that combines our equivariant loss with InfoNCE:
CCARE (f) Cuny (f) + (f (feuu())
(6)
where À weights the equivariant loss. We note that many variations of this approach are possible. For instance, the equivariant loss and InfoNCE loss could use different augmentations, resulting in invariance to specific transformations while maintaining rotational equivariance to others, similar to Dangovski et al. (2022). The InfoNCE loss can also be replaced by other Siamese self-supervise losses. We leave further exploration of these possibilities to future work. In all, CARE consists of three components: (i) term to induce orthogonal equivariance; (ii) non-collapse term; and (iii) an invariance term to enforce localized transformations on the embedding space.
3.1 THEORETICAL PROPERTIES OF ORTHOGONALLY EQUIVARIANT LOSS
In this section, we establish that matching angles via Lequi leads to a seemingly stronger property. Specifically, Loqui 0 implies the existence of an orthogonal matrix R., Odd) for any augmentation a, such that f (a(x)) ) = R₂ (ar) holds for all x. The converse also holds and is easy to see. Indeed, suppose such an R2 (dd) exists. Then, faa(x')) [((xx) s(x)) R.! Rxx) fx)) S(x),
T which implies Lequi (f) = O. We formulate the first direction as a proposition.
Proposition Suppose Cague = O. Then for almost every a E A, there is an orthogonal matrix R. (dd such that fa((x) = Raf (x) for almost all x E X.
4

---

Published as a conference paper at ICLR 2024
Figure 1 illustrates this result. Crucially R. is independent of z, without which the Proposition 1 would be trivial. That is, a single orthogonal transformation Rg captures the impact of applying across the entire imput space X. Consequently, exact minimization of Lequi loss converts "unstruc- tured" augmentations in input space to have structured geometric interpretation as rotations in the embedding space. In Appendix 1.1, extend this result to the case where the loss is low but not exactly minimized, in which case we have a guarantee of approximate equivariance.
This result can be expressed as the existence of a mapping p A - (d) that encodes the space of augmentations within (d).. This raises a natural question: how much of the structure of A does this encoding preserve? For instance, assuming A is a cemi-group (i.e., closed under compositions a' oa E A), does this transformation respect compositions: {(d((((z) } = Ra Raf(z) This property does not hold for non non-linear actions (Devillers & Lefort, 2023), but does for orthogonal equivariance: Corollary 1. If Lequ (f) = 0 and (f(x) I E X} spans Rd, then A - ((d) given by p(a) = Ra satisfies p(a' o a) = p(a') p(a)) for almost all a, d' That is, defines a group action on gd-1 up to a of measure zero.
Formally, this result states that if A is semi-group, then P A fines a group homomorphism (or linear representation of A in the sense of rep- resentation theory (Curtis & Reiner, 1966; Serre et al., 1977). a branch of math- ematics that studies the encoding of abstract groups as spaces of linear maps). To exactly attain (fq = O, the space of augmenta- tions A needs to have certain structure, but this becomes less restrictive if is large. Assuming for simplicity that A is a group, the first isomorphism theorem for groups states that p((4) = A/ ker(p). For instance, ker(p) is trivial, the equivariant loss can be exactly zero when the group of augmentations is subgroup of the orthogonal group. Examples include orthogonal transformations rotations that fix a subspace-i.e. (dd') or (Odd") with d' S d-or subgroups of the permutation group on d elements. Furthermore, the Peter-Weyl theorem implies that any compact Lie group can be realized as a closed subgroup of (d) for some d (Peter & Weyl, 1927). As with Proposition 1, we also extend this result to the case of low loss that is not exactly zero, where we show that compositionality is approximately preserved.
O de-
100,
a
R. slet)
f(x5)
Embedding space
Figure When Lequi compositions of augmentations correspond to compo- sitions of rotations.
3.2 EXTENSIONS TO OTHER GROUPS
Proposition 1 states that perfectly optimizing Lequi = 0 produces an f that is equivariant, encoding augmentations in the input space as orthogonal transformation in the embedding space. Notably, since the computation of Lequi solely relies on pairwise data instances x' E X. it naturally aligns with the contrastive learning paradigm that already works with pairs of data.
In fact, it is possible to extend the idea of CARE and its benefits to Mathematically, invariants of the action of (dd) on n points-seen in (Rdd)" as Q (#1.... σn) = (Q 21...... ........ Q xnn) (onn-can can be expressed as function of pairs of objects (x. 2z))..==...n.- This is because the orthogonal group is defined as the stabilizer a bilinear form. In other words, letting B(x, x') x' denote the standard inner product, we have
other group actions.
(7) This argument applies more generally other groups that are defined as stabilizers of bilinear forms. the context of special relativity, can be
((d) (A e GL(d): B(Ax, Ax') = B(x, x') for all x, E Rd).
For instance, the Lorentz group. which has applications defined as the stabilizer of the Minkowski inner product. Additionally, the symplectic group, which is used to characterize Hamiltonian dynamical systems, can be defined in a similar manner.
Such extensions other groups allow to use CARE for different embedding space geometries. For instance, several recent works have used a hyperbolic space as an embedding space for sulf-supervised learners (Ge et al., 2022; Yue et al., 2023; Desai et al., 2023). If we constrain our embedding to hyperboloid model of hyperbolic space, then linear isometries this space are precisely the Lorentz group. Further discussions on extensions to other groups are given in Appendix D.
5

---

Published as a conference paper at ICLR 2024
CARE
yoast
solds
21 yoock
pvecis
V-a-asis
3 yoark
podde
SimCLR
Angle
2:yraua
3:s-anis
Unitorm
Angle
yoass
Traeas
Pracatie
21yvecx
3rawnti
y
Invariant
2.yooth
3.y
EquiiMod
Figure 3: Trajectories through embedding space three randomly sampled protein point clouds,
rotated from O to 2T in three orthogonal axes. Rows correspond to different training methods.
MEASURING ORTHOGONAL ACTION ON EMBEDDING SPACE
To prohe the geometric properties of CARF, we consider two efficiently computable metrics for empirically measuring the orthogonal equivariance in the embedding space.
Wahba's problem. A natural way to assess the equivariance
is to sample batch of data { and an augmentation a and test to what extent applying transforms the embeddings of each x; the same way. To measure this we compute a single rotation that approximates the map from f(;;) to fa((z,) for all i. Let F and F. E d have ith columns f(z;) and f(((z;)) respectively, then we compute the error Wy = Re || RF - Fa ||o. where || || denotes the Frohenius norm. R fa;;) for all This well-studied problem known as Wahba's
W, 0, then f(a(x;)) problem. The analytic solution to Wahba's 's problem computed easily. is nearly Ie UVT where UEVT T is singular value decomposition of FT However, slight modification is required as this R could have determinant =1, and therefore may not belong to SO(d). The only modification needed is to re-scale so that the determinant is one: R = diag ( (n-1). de)(()))t(()) V where 1, denotes the vector in R" of all ones.
Relative rotational equivariance. Optimizing for the CARE objective may potentially result in learning invariance since for input image [((xx) f'x) (x) for E A is a trivial optimal solution of arg min, Cequi (f). To check that our model learning non-trivial equivariance, we consider a metric similar one proposed by Bhardwaj et al. (2023)
([)(((z')) - S((x))) (ode 2())) ([[(a(=')) -
7f = BaNABa
[I[((x))) [())
Here, the denominator measures the invariance of the representation, with smaller values corre- sponding to greater invariance to the augmentations. The numerator, on the other hand, measures equivariance and can be simplified to [f(a(=')) f(a(x)) - f(x) f(=TT(a²) f(x') (i.e., Lequi (f) up constant, because f maps to the unit sphere. The ratio 7 of these two terms measures the non-trivin trivial equivariance, with lower value implying greater non-trivial orthogonal equivariance.
EXPERIMENTS
We examine the representations learned by CARE, as well as those obtained from purely invariance- based contrastive approaches. We describe experiment configurations in Appendix F.

---

Published as a conference paper at ICLR 2024
5.1 LEARNING REPRESENTATIONS OF PROTHIN POINT CLOUDS
We consider the problem of learning representations of proteins from the Protein Data Bank (Burley al.. 2021). Each protein described by a point cloud X E R To respect the permutation invariance of each point cloud, we take f to be a DeepSet (Zaheer et al., 2017), producing embeddings in 1R:46 and train CARE using random rotations of 3D space as the augmentations-ic. are positive pair for R E 5O(3). We evaluate our models on the task of predicting the first principal component of the point cloud, an important structural property of the input. This task is rotation equivariant,
X and X R
expect that CARE should outperform the invariance-based methods such as SimCLR, as verified in Figure 4.
We test equivariance to rotations by randomly sampling tein X, and a sequence of rotations { [Rg ] along each of the three orthogonal axes, evenly spaced, tracing full 360 rotation of the point cloud. We then compute 2 = f(X R,) for each i, and project them into 2D space. Each row of Figure shows the trajectory of 3 different proteins, and three rotation trajectories, for a given training method. We find that CARE exhibits a much more regular geome- try than models trained with SimCLR, Cumts or Linv Learning the SO(3) manifold is challenging, and previous works assume access to the corresponding group action (Quessard et al., 2020; Park et al., 2021) However, CARE learns it by merely using x and a(x)), without relying on the group action a.
pro-
028
024
Teat MSE
922
928
0.19
0/6
Sillerky Surda Rrulas Beckel
Figure CARE achieves the lowest error on the task of pre- dicting the first principal com- ponent of a protein
5.2 QUANTITATIVE MEASURES FOR ORTHOGONAL EQUIVARIANCE
Wahba's Problem We compare ResNet-18 models pretrained with CARE and with SimCLR on CIFAR10. For each model, we compute the optimal value Wj of Wahba's problem, introduced in Section 4, over repeated trials. In each trial, we sample single augmentation a ~ .A at random and compute Wj for plot the results in Figure 5, where the colors of dots indicate the sampled augmentation. Results show that CARE has a lower average error and worst-case error. Further, comparing point-wise for each augmentation, CARE achieves lower error in nearly all cases.
fCARE and f fSimCLR over the test data. We repeat this process 20 times and
Wahtn''s Rotation Error M
10
30
CARBECOF (0tE - 0,005)
ALequi + (1 A) SimCLL
SinCT.R (0* 0.111)
20
Donity
Accuracy
10
Accuracy SimCLR
1.0
0.5
0.0
0.5
1.0
CARE (ours) SinCLR EquiMDD
Cosine of angle belween positive puirs
X
Figure 5: Measuring equivari- Figure 6: ance using Wahba's problem. sine angles between data pairs. Figure 7: Linear readout error as Lower score is more equivariant. CARE has much lower variance. the loss weightings vary.
Histogram
co-
Analyzing structure on 2D manifold. To further study Lequi, we train encoder f that projects the input onto a unit circle $1, where orthogonal transformations are defined by angles. We measure the cosine of the angle between pairs f(x) and (a(x)) for all x in the test set, for 20 distinct sampled augmentations
A. As shown in Figure 6, Both CARE and SimCLR exhibit high density close to demonstrating approximate invariance. However, unlike CARE, SimCLR exhibits non-zero density in the region -0.5 to 1.0., indicating that the application of augmentations significantly displaces the embeddings. Further, CARE consistently exhibits lower variance of the cosine between f(x) and fa(x)) for a fixed augmentation, showing that it transforms all embeddings in the same way.
Ablation of loss terms. The CARE loss LCARE is a weighted sum of the InfoNCE loss CIstANC and the orthogonal equivariance loss Lequi- Figure 7 evaluates the performance of ResNet-50 models trained
7

---

Published as a conference paper at ICLR 2024
Relatie Vr
Rotatinnal Equivariacee
CARE (ours)
1.5
SimCLR
10-
1.0
0.
100
300
400
150
200
400
102
Epochs
Epachs
Epocht
Figure 8: Relative rotational equivariance (Lower is more equivariant). Both CARE and invariance- based contrastive methods (e.g., SimCIR) produce approximately invariant embeddings. However, CARE learns considerably more rotationally structured embedding space.
on CIFARIO using LIMONCE + Eequii for varying X, finding optimal in the range 0.01 I S 0.1. We additionally split the InfoNCE loss into constituent parts Lime, Lunif, and test different combinations of the three losses, including Leqii We find that using all three jointly is optimal. See Figure in Appendix G.1 for detailed results.
Relative rotational equivariance. We measure the relative rotational equivariance for both CARE and SimCLR over the course of pretraining by following the approach outlined in Section 4. Specifically, we compare ResNet-18 18 models trained using CARE and SimCLR on CIFAR1O. From Figure we observe that both the models produce embeddings with comparable non-zero invariance loss Linv. indicating approximate invariance. However, they differ their sensitivity to augmentations, with CARE attaining much lower relative equivariance error Ty. Importantly, this shows that CARE achieving lower equivariance error Lequi by collapsing to invariance, a trivial form of equivariance.
5.3 QUALITATIVE ASSESSMENT OF EQUIVARIANCE
key property promised by equivariant contrastive models is sensitivity to specific augmentations. To qualitatively evaluate the sensitivity, equivariance, of our models, we consider image retrieval task on the Flowers- 102 dataset (Nilsback Zisserman, 2008), as considered by Bhardwaj et al. (2023). Specifically, when presented with an input image ær we extract the top 5 nearest neighbors based on the Euclidean distance of f(z) and fa((x)) where a E A. Figure shows that retrieved results for the CARE model exhibit greater variability in response to change in query color compared to the SimCLR model, which remains largely invariant. Additionally, Figure 21 compares CARE with EquiMOD (Devillers & Lefort, 2023), an equivariant baseline.
Input
Query
CARE (ours)
SimCLR
Figure CARE exhibits sensitivity to features that invariance- based contrastive methods (e.g..
SimCLR) do not. For each query input, we retrieve top nearest neighbors in the embedding space.
8

---

Published as a conference paper at ICLR 2024
5.4 LINEAR PROBE FOR IMAGE CLASSIFICATION
Next, we examine the quality of features learned by CARE for solving image classification tasks. We train ResNet-50 models on four datasets: CIFAR10, CIFARIOO, STLI0, and ImageNer100 using SimCLR and MoCo-v2 (see Appendix F for details). We refer to the model trained using CARE with SimCLR or MoCo-v2 backbone as CARESICCLR and CAREMoCo-22 respectively. For each method and dataset, we evaluate the quality of the learned features by training linear classifier (i.c., probe (Alain & Bengio, 2017)) on the frozen features of f and report the performance on the test set in Table 1. We find consistent improvements in performance using CARE, showing the benefits of our structured embedding approach for image recognition tasks.
Table I: Top-1 and Top-5 linear probe accuracy (%) on CIFARIO, CIFARIO0, STL1O and Ima- geNet100 datasets. We report the mean performance from 3 different random initializations for the linear classifer. & denote numbers from Devillers & Lefort (2023), and ** from Zhuo et al. (2023)
Method
CIFARIO
CIFARIO0
STL10
ImageNet100
Invariant prediction approackes
SimCLR
90.9810.10
66.7710.4
84.19
72.7910.00
MoCo-v2
91. 11.5520.05
69.88+0.23
g3.50
BYOL
90.44
67.41**
Barlow Twins
84.5420.02
55.54_0005
90.6220.02
Equivariant prediction approaches
EquiModsincLL
91.28
67.59
83.67
EquiModgyor.
91.57*
CARRSimCUR
91.9210.12 (f 0.94)
68.05 10.28 (f 1.28)
84.6410.29 (1 0.45)
76.69 10.0s 3.90)
CAREMsCo-y2
92.1910.01 (f 0.24)
20.5 (f 0.68)
88.97 10.49
24.30.0.0r (1 0.80)
CAREBarow Twins
85.65.0005 (f 1.11)
56.76 ( 1.22)
90.9220.0r (1 0.30)
DISCUSSION
Converting transformations that are complex in input space into simple transformations in embedding space has many potential uses. For instance, modifying data (e.g.. in order to reason about counterfac- tuals) can be viewed as transforming one embedding to another. If the sought after transformation was simple and predictable, it may be easier to find. Similarly, generalizing out-of distiiiiibtttionn is casier when extrapolating linearly (Xu et al., 2021), suggesting that linear transformations of embedding space may facilitate more reliable generalization. This work considers several design principles that may be broadly relevant: 1) learned equivariance preserves the expressivity of backbone architectures, and in some cases may be easier for model design than hard-coded equivariance, 2) linear group actions are desirable, but require carefully designed objectives (similar in spirit to the principle of parsimony (Ma et al., 2022), also advocated for by Shakerinava et al. (2022)), and 3) orthogonal (and related) symmetries are a promising structure for Siamese network training as they can be efficiently learned using pair-wise data comparisons.
7
ACKNOWLEDGEMENTS
This research was supported by NSF award CCF-21 12665. Derek Lim is supported by National Science Foundation Graduate Research Fellowship. Soledad Villar is partially funded by the NSF-Simons Research Collaboration on the Mathematical and Scientific Foundations of Deep Learning (MoDL) (NSF DMS 203 198), NSF CISE 2212457, ONR N00014-22-1-21266 and an Amazon A12AI Faculty Research Award.
We acknowledge MIT SuperCloud and Lincoln Laboratory Supercomputing Center for providing HPC resources that have contributed this work. We wish to thank Michael Murphy for insightful discussions on extensions of our method to biology.

---

Published as a conference paper at ICLR 2024
REFERENCES
Guillaume Alain and Yoshua Bengio. Understanding intermediate layers using linear classifier probes. In International Conference on Learning Representations (ICLR), 2017.
Ery Arias-Castro, Adel Javanmard, and Bruno Pelletier. Perturbation bounds for procrustes, classical scaling, and trilateration, with applications manifold learning. Journal of machine learning research, 21, 2020.
Sangnie Bhardwaj, Willie McClinton, Tongzhou Wang, Guillaume Lajoie, Chen Sun, Phillip Isola, and Dilip Krishnan. Steerable equivariant representation learning. preprint arXiv:2302. /3349, 2023.
Ben Blum- Smith and Soledad Villar. arXiv:2209. 14991, 2022.
Equivariant maps from invariant functions.
preprint
Antoine Bordes, Nicolas Usunier, Alberto Garcia-Duran, Jason Weston, and Oksana Yakhnenko. Translating embeddings for modeling multi multi-relational data. In Advances in Neural Information Processing Systems (NeurIPS), volume 26, 2013.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric deep learning: Grids, groups, graphs, geodesics. and gauges. preprint vrXiv:2104.13478, 2021.
Stephen K Burley, Charmi Bhikadiya, Chunxiao Bi, Sebastian Bittrich, Li Chen, Gregg V Crichlow, Cole H Christie, Kenneth Dalenberg, Luigi Di Costanzo, Jose M Duarte, et al. Resh protein data bank: powerful new tools for exploring 3d structures of biological macromolecules for basic and applied research and education in fundamental biology, biomedicine, biotechnology. bioengineering and energy sciences. Nucleic acids research, 49(D)):447-D451, 2021.
Ting Chen, Simon Korbblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International Conference on Machine Learning (ICML), pp. 1597-1607, PMLR, 2020.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In JEEE Conference Computer Vision and Pattern Recognition (CVPR), pp. 15750-15758, 2021.
Bilal Chughtai, Lawrence Chan, and Neel Nanda. toy model of universality: Reverse engineering how networks learn group operations. In ICLR Workshop on Physics for Machine Learning, 2023.
Taco Cohen and Max Welling. Group equivariant convolutional networks. In International Conference on Machine Learning (ICML), pp. 2990-2999. PMLR, 2016.
Charles W Curtis and Irving Reiner. Representation theory of finite groups and associative algebras, volume 356. American Mathematical Soc., 1966.
Rumen Dangovski, Li Jing, Charlotte Loh, Seungwook Han, Akash Srivastava, Brian Cheung, Pulkit Agrawal, and Marin Soljačić. Equivariant contrastive learning. In International Conference on Learning Representations (ICLR), 2022.
Karan Desai, Maximilian Nickel, Tanmay Rajpurohit, Justin Johnson, and Ramakrishna Vedantam. Hyperbolic image-text representations. ICLR Workshop on Multimodal Representation Learning, 2023.
Alexandre Devillers and Mathieu Lefort. Equimod: An equivariance module to improve self- supervised learning. In International Conference on Learning Representations (ICLR), 2023.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthis Minderer, Georg Heigold, Sylvain Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. In International Conference on Learning Representations (ICLR), 2021.
Quentin Garrido, Laurent Najman, and Yann Lecun. Self-supervised learning of split invariant equivariant representations. preprint ar0iv:2302. 10283, 2023.
10

---

Published as a conference paper at ICLR 2024
Songwei Ge, Shlok Mishra, Simon Kornblith, Chun-Liang Li, and David Jacobs. Hyperbolic contrastive learning for visual representations beyond objects. In JEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2022.
Michael Gutmann and Aapo Hyvärinen. Noise-contrastive estimation: A estimation principle for unnormalized statistical models. In Proceedings International Conference on Artificial Intelligence and Statistics (AISTATS), pp. 297-304, 2010.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recognition. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 770-778, 2016.
Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. In International Conference on Learning Representations (ICLR), 2022.
Eskou-Oumar Kaba, Arna Kumar Mondal, Yan Zhang, Yoshua Bengio, and Siamak Ravanbakhsh. Equivariance with learned canonicalization functions. In Advances in Neural Information Process- ing Systems (NeurIPS), 2022.
Hankook Lee, Kibok Lee, Kimin Lee, Honglak Lee, and Jinwoo Shin. Improving transferability of representations via augmentation-awant self-suuervision. In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pp. 17710-17722, 2021.
Yi Ma, Doris Tsao, and the emergence of intelligence. Frontiers of Information Technology & Electronic Engineering, 23 (9):1298-1323, 2022.
Shum. On the principles of parsimony and self-consistency for
Maria-Elena Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. 2008 Sixth Indian Conference on Computer Vision, Graphics & Image Processing. pp. 722-729. IEEE, 2008
Jung Yeon Park, Ondrej Biza, Linfeng Zhao, Jan- Willem van de Meent, and Robin Walters. Learning symmetric representations for equivariant world models. 2021.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. Pytorch: An imperative style, pigh--eeformance deep learning library. In Advances in Neural Information Processing Systems (NeurIPS), volume 32, 2019.
Fritz Peter and Hermann Weyl. Die Vollständigkeit der primitiven Darstellungen einer geschlossenen kontinuierlichen Gruppe. Mathematische Annalen, 97(1):337-755, 1927.
Robin Quessard, Thomas D Barrett, and William R Clements. Learning group structure and disentan- gled representations of dynamical environments. arXiv preprint arXiv:2002.06991, 2020.
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, et al. Learning transferable visual models from natural language supervision. In International Conference on Machine Learning (ICML), pp. 8748-8766 8763. PMLR, 2021.
Joshua Robinson, Li Sun, Ke Yu, Kayhan Batmanghelich, Stefanie Jegelka, and Suvrit Sra. Can contrastive learning avoid shorteut solutions? In Advances in Neural Information Processing Systems (NeurIPS), volume 34, pp. 4974-4986 2021.
Barbara Schmid. Finite groups and invariant theory. In Topics in Invariant Theory: Séminaire d' d'Algèbre P. Dubreil et M. -P. Malliavin 1989-1990 (40ème Année), pp. 35-66. Springer, 2006.
Alexci Schneider, Alexci Baevski, Shu-wen Chen, Sanjcev Khudanpur, and Andrew Davis. Wav2vec 2.0: A framework for self-supervised learning of speech representations. In Proceedings of the 2 Conference of the International Speech Communication Association (INTERSPEECH), pp. 1657-1661, 2021.
Slan-Pierre Serre al. Linear representations of finite groups, volume 42. Springer, 1977.
11

---

Published as a conference paper at ICLR 2024
Mchran Shakerinava, Arna Kumar Mondal, and Siamak Ravanbakhsh. Structuring representations using group invariants. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Xavier Suau, Federico Danieli, T Anderson Keller, Arno Blaas, Chen Huang, Jason Ramapuram, Dan Busbridge, and Luca Zappella. Duet: 2d structured and approximately equivariant representations. 2023.
Zhiqing Sun, Zhi-Hong Deng, Jian- Yun Nie, and Jian Tang. RotatE: Knowledge graph embedding by relational rotation in complex space. In International Conference on Learning Representations (ICLR), 2019.
Aäron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. preprint arXiv: /807.0374, 2018.
Soledad Villar. David W Hogg, Kate Storey-Fisher, Weichi Yao, and Ben Blum- -mith Scalars are universal: Equivariant machine learning. structured like classical physics. pp. 28848-28863, 2021.
Tongzhou Wang and Phillip Isola. Understanding contrastive representation learning through align- ment and uniformity on the hypersphere. In International Conference on Machine Learning (ICML), pp. 9929-99999999 99399. PMIR, 2020.
Hermann Weyl. The classical groups: their invariants and representations. Princeton university press, 1946.
Tete Xiao, Xiaolong Wang, Alexei A Efros, and Trevor Darrell. What should not be contrastive in contrastive learning. In International Conference on Learning Representations (ICLR), 2021.
Yuyang Xie, Jianhong Wen, Kin Wai Lau, Yasar Abbas Ur Rehman, and Jiajun Shen. What should be equivariant in self-supervised learning. IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 4111-411111 120, 2022.
Keyulu Xu, Mozhi Zhang, Jingling Li, Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. How neural networks extrapolate: From feedforward to graph neural networks. In International Conference Learning Representations (ICLR), 2021.
Michihiro Yasunaga, Antoine Bosselut, Hongyu Ren. Xikun Zhang, Christopher D Manning, Percy Liang, and Jure Leskovec. Deep bidirectional Banguage-knowledge graph pretraining, In Advances in Neural Information Processing Systems (NeurIPS), volume 35, pp. 37309-37323, 2022.
Yun Yue, Fangzhou Lin, Kazunori D Yamada, and Ziming Zhang. Hyperbolic contrastive learning. preprint XtXX::23002.1409,, 2023.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander J Smola. Deep sets. In Advances in Neural Information Processing Systems (NeurIPS), volume 30, 2017.
Zhijian Zhuo Yifei Wang, Jinwen Ma, and Yisen Wang. Towards unified theoretical understanding of non con-contrastive learning via rank differential mechanism. In International Conference on Learning Representations (ICLR), 2023.
12

---

Published as a conference paper at ICLR 2024
A RELATED WORK
Geometry of representation space. Equivariance is key tool for encoding geometric structure- e.g., symmetries into neural network representations (Cohen & Welling, 2016; Bronstein et al., 2021). Whilst hard-coding equivariance into model architectures is very successful, approximate learned equivariance (Kaba et al., 2022; Shakerinava et al., 2022), has certain advantages: 1) it applies when the symmetry is provided only by data, with no closed form expression, 2) it can still be used when it is unclear how hard code equivariance into the architecture, and 3) it can exploit standard high capacity architectures (He et al., 2016; Dosovitskiy et al., 2021), benefiting from considerable engineering efforts to optimize their performance. Shakerinava al. (2022) also consider learning orthogonal equivariance but consider problems where both input and embedding space are acted on by (dd. Our setting differs from this in two key ways: 1) we consider a very different set of transforms of input mance-jiter. crops, etc.-and 2) CARE can be naturally integrated into contrastive learning. and 3) we theoretically study the minima of the angle-preserving loss. A related line of trtueture-posiiiee including group
work, mechanistic interpretability, hypothesizes that algorithmic symmetries- emerges naturally within network connections during training (Chughtai et al., 2023). Our approach is very different from this we directly train models to have the desired structure without relying on implicit processes.
Self-supervised learning. Prior equivariant contrastive learning approaches extend the usual setup of learning invariance by learning sensitivity to certain features known to be important for downstream tasks. For instance, Dangovski et al. (2022) learns predict the augmentation applied but only considers a discrete group of 4-fold rotations. Lee et al. (2021) learns the difference of augmentation parameters and Xiao et al. (2021) constructs separate embedding sub-spaces that capture invariances to all but one augmentation. However, these approaches do not offer a meaningful structure to the em- bedding space. Others attempt to control how this sensitivity occurs. Specifically, Devillers Lefort (2023): Garrido et al. (2023); Bhardwaj et al. (2023) learn mapping from one latent representation to another, predicting how data augmentation affects the embedding. But this does not constrain the group action on embeddings, resulting in complex non-linear augmentation maps. Finally, the recent work Suau et al. (2023) implements approximate equivariance using 2D representations.
B PROOFS OF THEORETICAL RESULTS
The aim of this section is to detail the proofs of the theoretical results presented in the main manuscript. The key theoretical tools driving our analysis are prepared separately in Section C.
Throughout our analysis, we assume that all spaces (e.g.. A and 2') are subspaces of Euclidean space and therefore admit a Lebesgue measure. We also assume that all distributions (c.g., a ~ A and X') admit density with respect to the Lebesgue measure. With these conditions in mind, recall the loss function that is the main object of study:
Lequi (f) A [f(('))) f(a(x)) (()))((())'
Next, we re-state and prove Proposition 1, our first key result. Proposition 1. Suppose Luqui (f) = O. Then for almost every E A, there is an orthogonal matrix Ra, E O(d) such that [a(x)) = R.. (x) for almost all x EX
Proof. Suppose that Lf = O. This means that f(a(x')) f(a(x))= f(x') for almost all E G. and x, x' X. Setting ga (x) f(a())). we have that (a())) (x()) f(x) f(x'). The continuous version of the First Fundamental Theorem of invariant theory for the orthogonal group
(see Proposition 5) implies that there is an R. ((d) such that f(((x)) = ga (x) (xx.
As discussed in greater detail in the main manuscript, these results show that minimizing Cequi produces a model where an augmentation a corresponds to a single orthogonal transformation of embeddings R., independent of the input. This result is continuous in flavor it studies the loss over the full data distribution p(r)). There exists a corresponding result for the finite sample loss
Lequi. (f) Ea~A
[f(a(z))) f(a(x:)) (((xE)) fxx)) {(,)))
1,1=1
13

---

Published as a conference paper at ICLR 2024
Proposition 2. Suppose Lequi, 0. Then for almost every E A, there is an orthogonal matrix Ra O(d) such that f(a(x;)) = (zf () for all = =.....,? n.
As for the population counterpart, the proof of this result directly follows from the application the First Fundamental Theorem of invariant theory for the orthogonal group.
Proof of Proposition 2. Suppose that Loqui (f) = 0. This means that for almost every a E G. and every i = 1,...,77 nu we have [a(xr)) [a((z))) - f(x))" [(xy). In other words AA = BB" where A, B E Randd are matrices whose ith rows are A, = f((()))) and B, = f(;;) respectively. This implies, by the First Fundamental Theorem of invariant theory for the orthogonal group (see Corollary 3), that there is an Rg E O(d) such that A = BRg. Considering only the ith rows of A and leads us conclude that f(a(x:)) (x.).
A corollary of Proposition is that compositions of augmentations correspond to compositions of rotations.
Corollary If [eq = and (f(x) x } spans R4, then p : - ((d) given by p(a) Rs satisfies p(a' a) = p(a)) for almost all a, d'. That is, defines group action on gd-1 up to a set of measure zero.
Proof. Applying Proposition 1 on a' o a as the sampled augmentation, we have that f(a' o a(x;)) = Ra I (I;) p(a' o (x;)) However, taking know that f(a' o a(x;)) f()((z)) = Ref(Ø) = R. fa(xs)) = R fx)(x) ) = ( That is, p(a' fz, f(a' a(z.)) = p(aa))((p))))) (zz)) Since this holds for all and the set of f(x)) spans Rdd, we have that pla' o aa) p(a")p(a)
a(x;) and applying Proposition twice also
This corollary requires us to assume that .A is a semi-group. That is, A is closed under compositions, but group elements do not necessarily have inverses and it does need to include identity element.
B.1 APPROXIMATE EQUIVARIANCE FOR NONZERO LOSS
For Proposition 1, we show that when our equivariance loss is zero, the learned function is exactly equivariant. In practice, the equivariance loss may not exactly minimized. Here, we give an approximate equivariance guarantee when the loss is small. The proof is based on Theorem I of Arias-Castro et al. (2020) (see Section C.I). For that, define
Lequi (f. = =F f(((z)) f(x)((()))))))) -fx)) f(x')}
(10)
Proposition 3. Suppose we have finite training set X {x1, , In. For augmentation a E A, suppose that
n n
Legui, (f
ΣΣ [f(a(z))) (((ea)))((((())))) f(x,) f(x;)) f(x)) Se,
(11)
and that the matrix [x f(xx) satisfies the conditions of Proposition Let A denote the pseudoinverse of A. Then there is an R. E ((d) such that
[ RoS(2:)) (at (++
1
(12)
and hence for each Vis
(((((())) Rof(2:)) (1+ v2) [
(13)
14

---

Published as a conference paper at ICLR 2024
Proof. The first incquality is a direct application of Proposition 6, where A = 1 ..... fIx)) ] and B (f((((1))....... f(((x))) The second inequality follows from the fact that || I ||ا|||
for vectors x E R".
Corollary 2. If (f S e for all a E A and A = , (xn)] satisfies the conditions in Proposition 6 then p A - ((d) given by p(a) = Ra satisfies Il o a) p(a') p(a)) s o(e²) for any a, a' A
Proof. Applying Proposition on a' o and using the triangle inequality, we have
[ - Re Raf (x) 1 Ra (x) (x)
f(a' a(x)) + f(a' oo(x)) Ro Raf(ill V || [ (x)) s(a' o
o a(x)) - Re Baf(all
S [[ o a(xx)) RoJ (1+ v2) || All1
Now, using triangle incquality again, we have
||aoo - f(a' axxx))] = If(a' o a(x)) - Re f R. - R.f
< [I o a(x)) - R. R..I((()))) || [ R. R., (x)
Taking 3 = a(x), using the submultiplicatiiity property of matrix norms and applying Proposition twice we have,
la o aa(x)) R.. [((xx))
+ [|a ((a[x)) - Ro Ra(()))
= [[(((t))- - +
R.,(
$ (f((d(z)) - + [[f(a(x)) - (=f()))
(1+ v2) | + e
Hence,
[ RR Rof (x)]l S (1+ v2) | e (2+2V2) ||
where 1 = f(...) and Ā [(a,
By our assumption, we can choose basis f(z,, ), f (f. of Rdd, where in, i2, .n [n].
We can then write any unit norm & in terms of this basis:
5=1
Letting B [f(x;;) ..., (T we have B-1 c. Thus, we can bound the 1-norm of as
||]]] s vallcl12 s V|
(14)
Then using the triangle inequality, we have
7
Ra Ro Ra
Cj (oooaf(F)))) - Ray
Ra Ru Ruf(
lggl Ratoal(Xi)) - Ray Rad (
j=1
< |cll (1 + v2) || 2
vd(1+ v2) || ||
Taking the supremum over all unit norm finishes the proof.
15

---

linear combination of the (z;). This observation allows to apply Published as a conference paper at ICLR 2024
C BACKGROUND ON INVARIANCE THEORY FOR THE ORTHOGONAL GROUP
This section recalls some classical theory on orthogonal groups and an extension that we use for proving results over continuous data distributions.
A function f ( + IR is said to be (dd) ((d)-invariant if f(Ro...... Ru,) R E O(d). Throughout this section, we are especially interested in determining easily computed statistics that characterize an O(d) invariant function f. In other words, we would like to write ʃ as function of these statistics. The following theorem was first proved by Hermann Weyl using Capelli's identity (Weyl, 1946) and shows that the inner products v, U/ suffice.
f(n....., Un) for all
Theorem 4 (First fundamental theorem of invariant theory for the orthogonal group). Suppose that f: (Idd)" + R ((d) -ivvaiiant Then there exists a finction g Tgxx -> R for which
f(11...., Un) = 2(v] (gRis=i)-
In other words, to compute f at a given input, it is not necessary to know all of ui..... Un. Computing the value of f at point can be done using only the inner products v. Uy which are invariant O(d). Letting V be the n X d matrix whose ith row is v," we may also write f(u,,..., R.,) vVV The map V VVT is known as the orthogonal projection of V.
A corollary of this result has recently been used to develop (dd) equivariant architectures in machine learning (Villar et al., 2021).
Corollary 3. Suppose that A, B are d matrices and AAT -BT. Then A BR for some R E Od)..
Villar et al. (2021) use this characterization of orthogonally equivariant functions to parameterize function classes of neural networks that have the same equivariance. This result is also useful in our context; However, we put it to use for a very different purpose: studying Lequi-
Intuitively this result says the following: given two point clouds A, B of unit length vectors with some fixed correspondence (bijection) between each point in A and a point in B, if the angles between the ith and jth points in cloud A always equal the angle between the ith and jth point in cloud B, then A and B are the same up to an orthogonal transformation.
This is the main tool
prove the finite sample version of the main result for our equivariant loss (Proposition 2). However, to analyze the population sample loss Lequi (Proposition 1), we require an extended version of this result to the continuous limit as n - 30. To this end, we develop simple but novel extension to Theorem 4 to the case of continuous data distributions. This result may be useful in other contexts independent of our setting.
Proposition S. Let X be any sel and J, h X + Rdd be functions on X. If { h(=) for all x, y E X, then there exists R E (dd) such that Rf(z) = h(x) for all x E X.
The proof of this result directly builds on the finite sample version. The key idea of the proof is that since the embedding space Ra is finit--diennsin span has maximal rank in the linear space spanned by the outputs of f. This means that any arbitrary point (z) can be written
we may select set of points (x:)); whose
the finite sample result on each f(z;) term in the sum to conclude that f(z) is also a rotation of sum of h(x;) terms. Next, we give the formal proof.
Proof of Proposition 5. Choose 2I...., Fn E X such that F= (f(z)) (.. E a and h - [h(x)) ... h(xx))) gaxd have maximal rank. Note use "I" to denote the column-wise concatenation of vectors. Note that such H, can always be chosen. Since we have FFT HH we know by Corollary that F H for some R E O(d).
[
Now consider an arbitrary x E X' and define F = [F f(x)] and H [H | h(x)] both of which belong to R(n+1)xdd Note that again we have PPT HAT so also know that F A R for some R E Od). Since x; were chosen so that F and H are of maximal rank, we know that h(x) = - c,,( for some coefficients E R. since if this were not the case then we would have rank (f) rank(H) 1.
R
16

---

Published as a conference paper at ICLR 2024
From this, know that
R h(x))) =
Σ aRR (Ra)(
1=1
Σ caft c((l)((x)))))))))))())))))(()(
RT Σ ca((a))
(RTh(z)
= fx)).
So we have that Rf(x) -RR RRTh(x) = h(x) for all E X.
C.1 APPROXIMATE EQUIVARIANCE RESULTS
Here, we consider the setting when our equivariance loss is not exactly minimized. This corresponds to when the pairwise dot products between representations not exactly match. We use generalization of Corollary 3 to this case (Arias- Cassto et al., 2020) Instead of the dot products AA and BB exactly matching, they match only up to some error. As a result, we can guarantee is close to B R for some orthogonal matrix R E (d) up to some error.
Proposition 6, (Arias- Cassso et al., 2020) Let A, E Rxdd with n 2 d, where A is full rank. Then there exists an orthogonal R E O(d) such V2
Suppose AA BB" 1|p =BT S e² and ||AF|| E S
that
||A- BRI (1 + v &|
(15)
where A is the pseudoinverse of A, ||-1pp is the Frobenius norm, and ||-|| is the maximum singular value norm.
In other words, for fixed and full rank matrix A, we have that min REO ||4 - BR] o(c). where o(c) + 0 as the eπrоr + 0.
D EXTENSIONS TO OTHER GROUPS: FURTHER DISCUSSION
In Section 3.2, we explore the possibility of formulating an equivariant loss Lequi for pairs of points that fully captures equivariance by requiring the group to be the stabilizer of a bilinear form. In this context, the invariants are generated by polynomials of degree two in two variables, and the cquivariant functions can be obtained by computing gradients of these invariants (Blum-Smith & Villar, 2022). Section 3.2 notes that this holds true not only for the orthogonal group, which the primary focus of our research but also for the Lorentz group and the symplectic group, suggesting natural extensions of our approach.
It is worth noting that the group of rotations 5O(d) does not fall into this framework. It can be defined as the set of transformations that preserve both inner products (a 2-form) and determinants (a d-form). Consequently, some of its generators have degree while others have degree d (see (Weyl, 1946), Section IIAA.9).
Weyl's theorem states that if group acts copies of vector space (in our case, (Rd) for consistency with the rest of the paper), its action can be characterized by examining how it acts (refer to
on k copies (i.e., () when the maximum degree of its irreducible components Section of (Schmid, 2006) for a precise statement of the theorem). Since our interest lies in understanding equivariance in terms of pairs of objects, desire invariants that act pairs of
17

---

Published as a conference paper at ICLR 2024
points. One way to guarantee this is to restrict ourselves groups that act through representations where the irreducible components have degrees of at most two (though this is not necessary in all cases, such as the orthogonal group ((d) that consider in the main paper). An example of such groups is the product of finite subgroups of the unitary group (2), which holds relevance in particle physics. According to Weyl's 's theorem, the corresponding invariants of degree-2 polynomials on two variables, Polarizations represent an algebraic construction that enables the expression of homogeneous polynomials in multiple variables by introducing additional variables to polynomials with fewer variables. In our case, the base polynomials consist of degree-2 polynomials in two variables, while the polarizations incorporate additional variables. Notably, an interesting open problem lies in leveraging this formulation for contrastive learning.
be expressed polarizations
E IMPLEMENTATION DETAILS
Algorithm presents pytorch-based pseudocode for implementing CARE. This implementation introduces the idea of using smaller batch size for the equivariance loss compared to the InfoNCE loss. Specifically, by definition, the equivariance loss is defined as a double expectation, one over data pairs and the other over augmentations, Empirical observations reveal that sampling one augmentation per batch leads to unstable yet superior performance when compared to standard invariant- baselines such SimCLR. Since these invariant-bascd contrastive benchmarks generally perform well with large batch sizes, we adopt the approach of splitting a batch into multiple chunks to efficiently sample multiple augmentations per batch for the equivariance loss. Each chunk of the batch is associated with a new pair of augmentations, ensuring a large batch size for the InfoNCE loss and smaller batch size for the equivariance loss.
Algorithm 1 Py/Torhh based pseudocode for CARR
I: Notations: represents the backbone encoder network,
weight on CARE loss, apply same aug function applies the same augmentation to all samples in the input batch
2: for minibatch in clat a Laaer do
3: draw butches of augmentation functions a1, d2 A
/" Functions 01. 02 apply different augmentation to each sample in batch /
5: 21 _im = {a (x)), f (a)())
6: divide x into n_split chunks to form Tchunks
7: /" Module for calculating orthogonal equivariance loss
8:
for o2 in Tehunks in parallel do
9:
draw two augmentation functions ã1, ãa E
10:
/a Functions à,, å2 apply same augmentation to each sample in batch //
11:
Zu. 2:2 = ('app-y_smme_e_ (C,, ā1 }), f(applyy_saaaeeeeeeepp same _ame_aug(ee____ aug(e;, ā2))
12:
1 Concatenate embedding vectors corresonndngg all chunks
13:
merge ža, 4,2 into zi .equiv respectively
14:
1a Luss computation
15:
LIMONCE (f) = lactencc__lss(ee ,T)
16:
Ceqpi ( ort hogonall equivariance Lossan "oss(E""" 122 n split)
epiiv
17:
CCARE (f) = (/) + X Lagiv (f)
18:
1" Optimizatiin step
19:
Leane baackardd))
20:
optimizer.step))
F SUPPLEMENTARY EXPERIMENTAL DETAILS AND ASSETS DISCLOSURE
F.1
ASSETS
We do not introduce new data in the course of this work. Instead, we use publicly available widely used image datasets for the purposes of benchmarking and comparison.
F.2 HARDWARE AND SETUP
All experiments were performed on an HPC computing cluster using 4 NVIDIA Tesla V100 GPUs with 32GB accelerator RAM for single training run. The CPUs used were Intel Xeon Gold 6248
18

---

Published as a conference paper at ICLR 2024
processors with 40 cores and 384GB RAM. All experiments the PyTorch deep learning framework (Paszke et al., 2019).
F.3 EXPERIMENTAI. PROTOCOLS
We first outline the training protocol adopted for training proposed approach variety of datasets, namely CIFARIO, CIFARI00, STLI0, and ImageNet100.
CIFARIO, CIFARI00 and STL10 All encoders have ResNet-50 backbones and are trained for 400 epochs with temperature т = 0.5 for SimCLR and r = 0.1 for MoCo-v2 The encoded features have dimension of 2048 and are further processed by a two-layer MLP projection head, producing an output dimension of 128. A batch size of 256 was used for all datasets. For CIFARIO and CIFAR100, and weight decay of le-6. For STL10, we employed the SGD optimizer with a learning rate of 0.06, utilizing cosine annealing and weight decay of 5e-4, with 10 warmup steps. We use the same set of augmentations as in SimCLR (Chen et al.. 2020). To train the encoder using ECARE-SinCLR, we use the same hyper- paper-parameters for InfoNCE loss. Additionally, we use 4, 8 and 16 batch splits for CIFAR 100, STL10 and CIFARIO, respectively. This allows us to sample multiple augmentations per batch, effectively reducing the batch size of equivariance loss whilst retaining the same for InfoNCE loss. Furthermore, for the equívariant term, we find it op- timal to use a weight of 0.01, 0.001, and 0.01 for CIFARIO, CIFAR100, and STL10, respectively.
employed the Adam optimizer with learning rate of 1e-3
ImageNet100 We use ResNet-50 as the encoder architecture and pretrain the model for 200 epochs. A base learning rate of 0.8 is used in combination with cosine annealing scheduling and a batch size of 512. For MoCo-v2, we use 0.99 as the momentum and 7 = 0.2 as the temperature. All remaining hyperparameters
maintained at their respective official defaults in the official MoCo-v2 code. While training with LCART.SinCIR & and LCCos we find it optimal to use splits of and 8 and weight of 0.005 and 0.01 respectively the equivariant term.
Linear evaluation We train a linear classifier on frozen features for 100 epochs with a batch size of 512 for CIFARIO, CIFARI00, and STLI0 datasets. To optimize the classifier, we employ the Adam optimizer with a learning rate of le and a weight decay of 1e In the case of ImageNet 100, we train the linear classifier for 60 epochs using batch size of 128. We initialize the learning rate to 30.0 and apply step scheduler with an annealing rate of 0.1 at epochs 30, 40, and 50. The remaining hyper-parametess are retained from the official code.
G ADDITIONAL EXPERIMENTS
G.1 ABLATING LOSS TERMS
G.2 HISTOGRAM FOR LOSS ABLATION.
To accompany Figure 10, this section plots the cosine similarity between positive pairs. We provide two plots for each experiment: the first plots the histogram of similarities of positive pairs drawn from the test set; the second plots the averuge positive cosine similarity throughout training. The results are reported in Figures 11, 12, 13, 14, 15, 16.
G.3 ADDITIONAL PROTEIN TRAJECTORIES
In the protein cloud experiment in Section 5.1, we use the Protein Data Bank 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 17, 18, 19 and 20 illustrate
single
hhtpps://so thub. com/
19

---

Published as a conference paper at ICLR 2024
90.98
91.93
80
Top-1 Accuracy
60
46.88
40
20:46
22
10.0
10.0
Invariance
Equix: (E)
Uniform (U)
Random
Unitorm
SimCLR (I+U)
CARE
Equiv.
((+E+U)
Figure 10: Ablating different loss terms. Combining Lequi with uniformity promoting non-collapse term suffices to learn non-trivial features. However, optimal performance is achieved when encourag- ing smaller rotations, as in CARE. ResNet-50 models pretrained CIFARIO and evaluated with linear probes.
200
1,000
Posttive collee similatiy
0.999
150
Frequency
0.998
100
0.997
0.996
Invariance Less (D
1999
119998
1.0000
50
100
1.50
200
Positive cosine similarity
epochs
Figure 11: (left) Histogram of positive cosine similarity values the end of pre-training using the invariance loss; (right) Evolution of positive cosine similarity values over pre-training epochs using the invariance loss
additional trajectories observed through the embedding space of DeepSet trained with the CARE, SimCLR, funit and Linv loss respectively.
G.4 ADDITIONAL QUALITATIVE ASSESSMENT OF EQUIVARIANCE
Figure 21 qualitatively assesses sensitivity of the representation to changes in color in the Flowers- 102 dataset. We compare CARE to popular equivariant baseline, EquiMOD (Devillers & Lefort, 2023). As anticipated, both CARE and EquiMOD demonstrate sensitivity to changes in query color and hence are equivariant to color. However, as depicted in Figure 21, EquiMOD's representation exhibits
1.0
100
Postive cosine sinilacity
0.8
Frequency
0.6
Ortiogonal Equivenance Loss
0.4
0.00
0.25
0.50
0.75
1.00
50
100
150
200
Positive similarity
Pretraining epochs
Figure 12: (left) Histogram of positive cosine similarity values at the end of pre- pre-training using the orthogonal equivariance loss; (right) Evolution of positive cosine similarity values over pre-training epochs using the orthogonal equivariance loss
20

---

Published as a conference paper at ICLR 2024
K0
Unidormity Loss (U)
Posiive castbe similarity
60
0.06
Feequeney
0.04
0.02
0.04 Positive cosine similarity
50
100
150
200
Pretraining opochs
Figure 13: (left) Histogram of positive cosine similarity values at the end of pre- pre-training using the uniforrmity loss; (right) Evolution of positive cosine similarity values over pre- vre-training epochs using the uniformity loss
30
Postive codine similurty
0.16
Frequency
10
0.10
Uniformity Equivariance
0.08
Loss
0.075
-
0.125
0.150
0.175
100
200
Positive cosime similarity
Protraining epochs
Figure 14: (left) Histogram of positive cosine similarity values the end of pre-training using the Uniformity + Equivariance loss; (right) Evolution of positive cosine similarity values over pre-training epochs using the Uniformity + Equivariance loss
0.85
Positiv casine similarity
0.80
150
Frequency
0.75
100
0.70
0.65
0.60
InfaN CE Loss (1+ U)
0.6
0.7
50
100
150
200
Positive cosine similarity
Pretraining epochs
Figure 15: (left) Histogram of positive cosine similarity values the end of pre-training using the InfoNCE (invariance uniformity) loss; (right) Evolution of positive cosine similarity values over pre-training epochs using the InfoNCE loss
200
0.85
Poritire onsine dinilarity
0.80
150
Frequency
0.75
100
0.65
50
0.60
CARE Loss (1+ -8)
0.6
0.7
01X
30
100
150
200
Positive similarity
Pretrnining epochs
Figure 16: (left) Histogram of positive cosine similarity values at the end of pre- pre-tinining using the CARE (InfoNCE orthogonal equivariance) loss; (right) Evolution of positive cosine similarity values over pre-training epochs using the CARE loss
21

---

Published as a conference paper at ICLR 2024
CARE
Angle
CARE
Angle
CARE
O
Angle
CARE
Angle
Figure 17: Additional trajectories through the embedding space of a DeepSet trained with CARE.
SimCLR
Angle
SimCLR SimCLR SimCLR
Angle
Angle
Angle
Figure 18: Additional trajectories through the embedding space of DeepSet trained with SimCLR.
Uniform Uniform Uniform Uniform
Angle
Angle
Angle
Angle
Figure 19: Additional trajectories through the embedding space a DeepSet trained with the uniformity loss Lunit-
nearest neighbors with significantly different shades (e.g., red and orange) compared to those learned by CARE, which closer in color to the query images.
22

---

Published as a conference paper at ICLR 2024
Invariant Invariant Invariant Invariant
Angle
Angle
Angle
Angle
Figure 20: Additional trajectories through the embedding space a DeepSet trained with the
invariance loss Liev.
Imput
Query
(ours)
EquiMOD (equivariant)
Figure 21: CARE exhibits sensitivity to features as with other equivariance-barr based contrastive methods (e.g., EquiMOD), For each query input, we retrieve top 5 nearest neighbors in the embedding space.
H ADDITIONAL DISCUSSION
Limitations. While our method, CARE, learns embedding spaces with many advantages over prior contrastive learning embedding spaces, there certain limitations that acknowledge here. First, we do not provide a means to directly identify the rotation corresponding to specific transformation. Instead, approach allows the recovery of the rotation by solving Wahba's problem. However, this requires solving an instance of Wahba's for each augmentation interest. Future improvements that develop techniques for quickly and easily (i.e., without needing solve an optimization problem) identifying specific rotations would be a valuable improvement, enhancing the steerability of our models. Second, it is worth noting that equivariant contrastive methods, including CARE, only achieve approximate equivariance. This is fundamental challenge shared by all such methods, as unclear how precisely encode exact equivariance. The question remains open as a) whether this approximate equivariance should be considered damaging in the first place, and if so, b) whether scaling techniques can sufficiently produce reliable approximate equivariance to enable the diverse applications that equivariance promises. Addressing this challenge is crucial area for future research and exploration the field. Each these limitations points to valuable directions for future work.
Broader impact. Through our self-supervised learning method CARE we explore foundational questions regarding the structure and nature neural network representation spaces. Currently, approaches exploratory and not ready for integration into deployed systems. However, this line work studies supervuuee learning and therefore has the potential to scale and eventually contribute systems that do interact with humans. In such cases, it is crucial to consider the usual
23

---

Published as a conference paper at ICLR 2024
safety and alignment considerations. However, beyond this, CARE, offers insights into algorithmic approaches for controlling and moderating model behavior. Specifically, CARE identifies a simple rotation of embedding space that corrssoonds to change in the attribute of the data. In principle, this transformation could be used to "canonicalize" data, preventing the model from relying on certain attributes in ppcision-making Additionally, controlled transformations of embeddings could be used to debias model responses and achieve desired variations in output. It is important to note that while our focus on the core methodology, we do explore these possibilities in this particular work.
24


Reviews:
[{"review_id": "bqfUkL0htu", "reviewer": "Reviewer_Dtcq", "summary": "The authors consider learning an embedding $f$ for high-dimensional data into a\nstructured latent space using (unsupervised) contrastive-learning-type methods.\nThey propose a new loss function (in implementation, effectively a new\nregularizer) for learning this embedding, which builds on prior work on\nequivariant regularizers. The regularizer asks not for the embedding $f$ to be\ninvariant to augmentations $a(x)$ of the input $x$, as in previous contrastive\nlearning methods such as SimCLR, but that it maps those augmentations to simple\nlinear transformations of the input, say $f(a(x)) = T_a f(x)$; it does this\nindirectly via an equivariance regularizer that the authors prove yields the\naforementioned structure on $f$ when exactly minimized. The precise\nimplementation becomes using this regularizer on top of existing contrastive\nlearning losses, such as SimCLR or InfoNCE, to prevent trivial embeddings and to\nencourage (in some sense) $a \\mapsto T_a$ to be \"continuous\". Experiments\ndemonstrate on simple datasets that the method improves the embeddings with\nrespect to two equivariance metrics; it improves linear probe performance on\nImageNet-100 scale image classification over SimCLR/MoCo-v2 (in the best case,\nby a significant margin); and that its design choices are necessary for these\nimprovements, via ablations.", "questions": "- Is there an \"approximate\" version of the results from invariant theory that\n  the authors use to establish their theoretical results? If one digs into the\n  proof of the result from invariant theory (say, specialized to\n  $\\mathrm{SO}(n)$), what obstructions are there to having an \"approximate\"\n  version of the result?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "- The paper is very well written. Conceptual explanations are clear, theoretical\n  results are precisely phrased and explained (mercifully, the representation\n  theory is kept to an absolute minimum, which seems uncommon in this area), and\n  the general writing is engaging and compelling.\n\n- The experimental evaluation is solid: it proposes reasonable metrics to assess\n  learned equivariance, shows that CARE improves over baselines on these\n  metrics, demonstrates improved linear probe performance at reasonable scale,\n  and gives some useful ablations in the appendix.\n\n- The inclusion of a theoretical (mostly conceptual-type) basis for the method\n  is valuable, and completes a well-rounded presentation of the method and its\n  motivations.", "weaknesses": "- The use of both an invariant and equivariant loss in the overall objective\n  function seems conceptually strange (although the ablations show it leads to\n  superior performance). I would like to understand what might be being learned\n  with this combination of losses -- my reading of the explanation in section 3\n  is that, among embeddings that *minimize* the equivariant regularizer (following\n  proposition 1), those that achieve a small invariant loss will prefer small,\n  rather than large, orthogonal transformations. However, it is not clear to me\n  why this setting should arise in experiments -- why not a situation where the\n  invariance loss is minimized, and the equivariant regularizer is only small?\n  I am also unsure how this connects with the discussion in the \"Relative\n  rotational equivariance.\" paragraph later.\n\n\n- It would be ideal if the authors could assert theoretically some degree of\n  approximate invariance given approximate minimization of their regularizer\n  (see a comment to this effect below). It seems important to precisely\n  understand what happens in this setting, given the fact that a mixture of\n  equivariant and invariant losses are required for strong practical\n  performance.\n\n\n### Minor issues\n\n- After equation (4): better to not reference a figure in the appendix without\n  adding something like \"In the appendix, we show ...\"\n\n- The claim after Proposition 1 that \"[c]onsequently, low [CARE] loss converts\n  'unstructured' augmentations in input space to have a structured geometric\n  interpretation as rotations in the embedding space\" does not seem to follow\n  from Proposition 1 or the preceding discussion (since the proposition requires\n  exact minimization of the loss, rather than just a low value). It is not clear\n  to me that the proof generalizes, because it relies on an external result from\n  invariant theory (which might be algebraic).\n\n- Bottom of page 5: principle -> principal", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M4fcU8BKiP", "reviewer": "Reviewer_4N3v", "summary": "The paper introduces a novel approach to self-supervised learning, adding geometric structure to the embedding space such that input transformations correspond to linear transformations in the embedding space. In the context of contrastive learning, the study presents an equivariance objective, which theoretically ensures that data augmentations in the input space align with rotations in the spherical embedding space. This method, named CARE, not only enhances performance in subsequent tasks but also captures essential data variations, like color, which standard methods overlook.", "questions": "* In many practical applications input transformations may not preserve distances and angles. What does CARE learn in this case?\n* Fig. 8 is confusing: no legend, all three plots are the same (?).", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* The proposed equivariant contrastive learning method that maps transformations of the input to local orthogonal transformations in the embedding space is new. The authors provide theoretical arguments and show empirical evidence for the desired structure of the embedding space. Both the method and its justification are novel and valuable.\n* The analysis of the structure of the learned representation space is solid and interesting. Fig. 9 is a particularly insightful and sheds light on the merits of the proposed method.", "weaknesses": "* It would be beneficial to delve deeper into the influence of the choice of A—the space of transformations experienced during training—on the structure and caliber of the representations derived. An intriguing question to address is the method's capacity to generalize: Can the structures learned be effectively transferred to other classes or varied transformation parameter ranges? Exploring these nuances could further solidify the robustness and versatility of the method.\n* 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.\n* 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?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "UomgAA86ZQ", "reviewer": "Reviewer_MvBZ", "summary": "The paper introduces a self-supervised loss that demonstrably enforces a connection between latent embeddings of related complex augmentations through \"simple\" linear (rotational) transformations.\nThe loss is theoretically analyzed, showing that rotational equivariance emerges when scalar product in the embedding space are preserved under augmentations. Furthermore, the generalization of this fact to other bilinear forms and geometries is discussed.\nThe approach is assessed and demonstrated to enhance linear probing performance on image datasets. Additionally, it is qualitatively evaluated and shown to capture equivariance as intended by the suggested loss.", "questions": "No specific questions. I would appreciate a response with respect to the weakness stated above.", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Overall the paper is well written and easy to follow.\n\nThe suggested loss simplifies existing equivariant techniques as it alleviates the need to learn the equivariance transformation, making it emerge solely as a result of an optimized loss term.\n\nThe method is theoretically analyzed. The analysis seems to be solid.  I appreciate the discussion about the possible generalizations to different geometries.", "weaknesses": "It seems that the main weakness of the paper is in the evaluation section.\nFirst, I would expect to see both qualitative (figure 9)   and quantitative comparisons to methods that learn the equivariant transformation (such as Garrido et al.). \nSecondly, the evaluation metric suggested in the Wahba’s Problem, seems to be another possible alternative loss to the suggested equivariance loss. Why shouldn’t it be used to optimize directly?\nLastly, it would also be interesting to compare to explicit parametrization of R_a.\n\nAdditionally, I found Figure 8 and the paragraph discussing relative rotational equivariance to be somewhat unclear. It appears that they could benefit from a revision to enhance clarity.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "NHP6ag6O5l", "reviewer": "Reviewer_avuo", "summary": "The paper introduces a contrastive learning framework called CARE, which stands for Contrastive Augmentation-induced Rotational Equivariance. CARE extends the InfoNCE loss by including a term that enforces an equivariant constraint with respect to the image transformations employed during training, as opposed to the traditional contrastive learning that attempts to induce invariance.\n\nThe central concept behind CARE is to ensure that transformations in the input space correspond to local orthogonal transformations in the representation space, a property known as orthogonal equivariance. This is achieved by minimizing a loss term (L_equi) that encourages the angles between pairs of images and their transformed counterparts to remain ideally the same, thus promoting orthogonal equivariance. The authors show theoretically that satisfying their proposed loss (i.e., L_equi = 0) implies orthogonal equivariance. \n\nTo avoid a common issue where all data points collapse to the same location in the representation space, the authors introduce an additional term (L_uni) inspired by the work of Isola and Wang. L_uni encourages representations within a batch to be evenly distributed. However, the combination of L_equi and L_uni alone underperforms compared to traditional contrastive learning frameworks like SimCLR. To address this, the authors introduce a term (L_inv) traditionally used to induce invariance with respect to the applied transformations. The final loss function comprises these three terms: \\lambda * L_equi (minimize angle differences) + L_inv (minimize differences among positive pairs) + L_uni (maximize uniformity).\n\n The authors assess CARE on a protein point cloud dataset (Protein Data Bank), they also demonstrate that CARE can better induce orthogonal equivariance compared to SimCLR using metrics like Whaba's problem and Relative Rotational Equivariance applied to the embedding of a Resnet-18 architecture trained on CIFAR-10. They qualitatively compare CARE and SimCLR in an image retrieval task, and finally evaluate CARE against SimCLR, MoCo-V2, and BYOL in the linear probing image classification task on various datasets: CIFAR10, CIFAR100, STL10, and ImageNet100.", "questions": "My question mostly revolve around the weaknesses listed above. Is there any misunderstanding in the list I provided above?\n\nIn addition\n- Would you agree that it make sense to compare Wahba’s Problem and the Relative Rotation Equivariance with another equivariance inducing method (in addition to SimCLR which is not inducing equivariance)? And that without this comparison it's difficult to assess \"how good\" the results shown are.\n- Would you be able to provide transfer learning results? Would you agree that those are ideal to show that CARE is learning more generic features?\n- Could you resolve the issue on the CIFAR-10 classification results? Why SOTA report 94% accuracy while your version is at 90%?\n- Could you provide more clarity on the protein folding experiment? \n- Could you provide more clarity (and maybe quantitative metrics) on the image retrieval task?\n- One of the claim is that CARE can handle composition of transformations but it is unclear if there is any experiment that support this claim. I understand that this can be proven theoretically, however, L_equi is never zero and its contribution if further reduced by the \\lambda factor, hence in practice it remains to be seen if CARE can handle equivariance of composition of transformations. Or is it the case that the augmentations are random resize crop + rotation, hence the equivariance is achieved for the combination of these two?\n- Lastly, not a question but I hope you take on the suggestion of trying to impose the equivariance constraint on the embeddings used for downstream task rather than the one after the projection head.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "The paper tackles an important issue: that of training more generic representation. I appreciate the idea of mixing invariance and equivariance for this objective.\n\nI also find the idea of representing equivariance using angle preservation as a neat idea.", "weaknesses": "I think this work has some interesting intuitions and hope that if it is not accepted the authors will continue to make it stronger. While I like the idea of mixing invariance and equivariance I think this work could be strengthened by a more thoughtful use of the contrastive framework and a more solid experimental evaluation. Let me elaborate on both these points. Since the experimental evaluation is, in my opinion, a more impactful weakness I start there, and follow with the use of the contrastive framework. \n\nThe experimental protocols should be improved in order to highlight the benefit of CARE. Results are not clearly showing why one should adopt CARE, some key experiments for this work are missing (transfer learning), and other experiments do not seem comparable with previously presented results.  Specifically,\n\n- The experiments that use the Wahba’s Problem and the Relative Rotation Equivariance should contain some stronger baselines. Wouldn't the authors say that it is expected that CARE achieves better results in terms of equivariance compared to SimCLR (that does not try to achieve orthogonal equivariance)? SimCLR could be a lower bound but 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. \n\n- The results on linear probing seems to report results on a not very common setting (training on CIFAR-10, 100, STL, ImageNet100, as opposed to training on ImageNet or ImageNetTiny and testing a linear probe on those datasets). For example see table Table 3 of BYOL, or Table 8 of SimCLR , or Table 3 (RRC column) of Duet. All these results were obtained training on ImageNet and testing on other datasets (including Cifar10 and Cifar100) and report much stronger results. I believe these results (i.e. Transfer Learning) are crucial for this work because they could support the claim that by using equivariance one can learn more generic and transferable features. Training and testing on the same dataset it is less interesting in my opinion. \n\n- Even for the few results that can be found about training SimCLR on CIFAR-10 (for example) the results reported seem below previously reported numbers. The original SimCLR paper reports results on CIFAR 10 in the Appendix and the accuracy is 94% (compared to 90.98%). My experience with contrastive frameworks is that once the pre-training pipeline is optimized and the training scales up (in terms of datasets, batch size, and epochs) the initial advantages one might observe by tweaking the model reduce until they (very often) disappear altogether. It would be great to show that this does not happen with equivariance.\n\n- On the protein point clouds experiments I find unclear to assess what is the desired (correct) trajectory. Perhaps this is due to my lack of experience with this protein task but the manuscript could explain more clearly what is the expected result (Figure 3). Also, it is known that SimCLR requires large batches (hence large datasets), I am not familiar with this dataset but what I wonder if its size makes it a suitable application for pre-training with SimCLR. Maybe a better test would be to pre-train on a larger dataset and fine tune on this dataset?\n\n- Similarly, the qualitative results on the image retrieval tasks are difficult to assess. I thought the results from SimCLR were better. This task should be explained more clearly and it would be even better to employ some quantitative evaluation (if possible). I also do not fully understand what is the role of the “input” compared to that of the “query”. \n\nIn terms of the contrastive learning framework, if I understand correctly all three terms of the proposed algorithm act on the same representation space (after projection head). This however introduces an ambiguity: is the objective to induce invariance or equivariance? While it is true that invariance satisfies equivariance it is also true that satisfying equivariance through invariance leads to the loss of information typical of the invariance-inducing algorithms (as the authors correctly mention in the paper). It is also why the common contrastive learning algorithms based on invariance impose this constraint after the projection head but use the representation before (where \"hopefully\" invariance is not achieved strongly). Therefore, while the equivariance idea is great and the use of angles is an interesting idea, its application seems to have \"just\" a regulation effect (especially given that the \\lambda value is so small). A more interesting (albeit I recognize this comment falls in the realm of speculations) use of the contrastive framework would be to impose equivariance on the embedding before the projection head (the embedding that are actually used for downstream tasks) and invariance after the projection head. By doing so, one can hope to take full advantage of the equivariant properties while still leveraging the constructive learning framework and the benefit of the invariance loss.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (poster)