Original: R OBUST
Revised: R OBUST

Original: S ELF -S UPERVISED
Revised: S ELF -S UPERVISED

Original: L EARNING
Revised: L EARNING

Original: WITH
Revised: WITH

Original: L IE
Revised: L IE

Original: G ROUPS
Revised: G ROUPS

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Deep learning has led to remarkable advances in computer vision.
Revised: Deep learning has led to remarkable advances in computer vision.

Original: Even so, today’s best models are brittle when presented with variations that differ even slightly from those seen during training.
Revised: Even so, today’s best models are brittle when presented with variations that differ even slightly from those seen during training.

Original: Minor shifts in the pose, color, or illumination of an object can lead to catastrophic misclassiﬁcations.
Revised: Minor shifts in the pose, color, or illumination of an object can lead to catastrophic misclassifications.

Original: State-of-the art models struggle to understand how a set of variations can affect different objects.
Revised: State-of-the art models struggle to understand how a set of variations can affect different objects.

Original: We propose a framework for instilling a notion of how objects vary in more realistic settings.
Revised: We propose a framework for instilling a notion of how objects vary in more realistic settings.

Original: Our approach applies the formalism of Lie groups to capture continuous transformations to improve models’ robustness to distributional shifts.
Revised: Our approach applies the formalism of Lie groups to capture continuous transformations to improve models’ robustness to distributional shifts.

Original: We apply our framework on top of state-of-the-art self-supervised learning (SSL) models, ﬁnding that explicitly modeling transformations with Lie groups leads to substantial performance gains of greater than 10% for MAE on both known instances seen in typical poses now presented in new poses, and on unknown instances in any pose.
Revised: We apply our framework on top of state-of-the-art self-supervised learning (SSL) models, finding that explicitly modeling transformations with Lie groups leads to substantial performance gains of greater than 10% for MAE on both known instances seen in typical poses now presented in new poses, and on unknown instances in any pose.

Original: We also apply our approach to ImageNet, ﬁnding that the Lie operator improves performance by almost 4%.
Revised: We also apply our approach to ImageNet, finding that the Lie operator improves performance by almost 4%.

Original: These results demonstrate the promise of learning transformations to improve model robustness 1 .
Revised: These results demonstrate the promise of learning transformations to improve model robustness 1 .

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: State-of-the-art models have proven adept at modeling a number of complex tasks, but they struggle when presented with inputs different from those seen during training.
Revised: State-of-the-art models have proven adept at modeling a number of complex tasks, but they struggle when presented with inputs different from those seen during training.

Original: For example, while classiﬁcation models are very good at recognizing buses in the upright position, they fail catastrophically when presented with an upside-down bus since such images are generally not included in standard training sets (Alcorn et al., 2019).
Revised: For example, while classification models are very good at recognizing buses in the upright position, they fail catastrophically when presented with an upside-down bus since such images are generally not included in standard training sets (Alcorn et al., 2019).

Original: This can be problematic for deployed systems as models are required to generalize to settings not seen during training (“out-of-distribution (OOD) generalization”).
Revised: This can be problematic for deployed systems as models are required to generalize to settings not seen during training (“out-of-distribution (OOD) generalization”).

Original: One potential explanation for this failure of OOD generalization is that models exploit any and all correlations between inputs and targets.
Revised: One potential explanation for this failure of OOD generalization is that models exploit any and all correlations between inputs and targets.

Original: Consequently, models rely on heuristics that while effective during training, may fail to generalize, leading to a form of “supervision collapse” (Jo & Bengio, 2017; Ilyas et al., 2019; Doersch et al., 2020; Geirhos et al., 2020a).
Revised: Consequently, models rely on heuristics that while effective during training, may fail to generalize, leading to a form of “supervision collapse” (Jo & Bengio, 2017; Ilyas et al., 2019; Doersch et al., 2020; Geirhos et al., 2020a).

Original: However, a number of models trained without supervision (self-supervised) have recently been proposed, many of which exhibit improved, but still limited OOD robustness (Chen et al., 2020;
Revised: However, a number of models trained without supervision (self-supervised) have recently been proposed, many of which exhibit improved, but still limited OOD robustness (Chen et al., 2020;

Original: Hendrycks et al., 2019; Geirhos et al., 2020b).
Revised: Hendrycks et al., 2019; Geirhos et al., 2020b).

Original: The most common approach to this problem is to reduce the distribution shift by augmenting training data.
Revised: The most common approach to this problem is to reduce the distribution shift by augmenting training data.

Original: Augmentations are also key for a number of contrastive self-supervised approaches, such as
Revised: Augmentations are also key for a number of contrastive self-supervised approaches, such as

Original: SimCLR (Chen et al., 2020).
Revised: SimCLR (Chen et al., 2020).

Original: While this approach can be effective, it has a number of disadvantages.
Revised: While this approach can be effective, it has a number of disadvantages.

Original: First, for image data, augmentations can only be applied to the pixels themselves, making it easy to, for example, rotate the entire image, but very difﬁcult to rotate a single object within the image.
Revised: First, for image data, augmentations are most often applied in pixel space, with exceptions e.g. Verma et al. This makes it easy to, for example, rotate the entire image, but very difficult to rotate a single object within the image.

Original: Since many of the variations seen in real data cannot be approximated by pixel-level augmentations, this can be quite limiting in practice.
Revised: Since many of the variations seen in real data cannot be approximated by pixel-level augmentations, this can be quite limiting in practice.

Original: Second, similar to adversarial training (Madry et al., 2017; Kurakin et al., 2016), while augmentation can improve performance on known objects, it often fails to generalize to novel objects (Alcorn et al., 2019).
Revised: Second, similar to adversarial training (Madry et al., 2017; Kurakin et al., 2016), while augmentation can improve performance on known objects, it often fails to generalize to novel objects (Alcorn et al., 2019).

Original: Third, augmenting to enable generalization for one form of variation can often harm the performance on other forms of variation (Geirhos et al., 2018; Engstrom et al., 2019), and is not guaranteed to provide the expected invariance to variations (Bouchacourt et al., 2021b).
Revised: Third, augmenting to enable generalization for one form of variation can often harm the performance on other forms of variation (Geirhos et al., 2018; Engstrom et al., 2019), and is not guaranteed to provide the expected invariance to variations (Bouchacourt et al., 2021b).

Original: Finally, enforcing invariance is not guaranteed to provide the correct robustness that generalizes to new instances (as discussed in Section 2).
Revised: Finally, enforcing invariance is not guaranteed to provide the correct robustness that generalizes to new instances (as discussed in Section 2).

Original: Lie Operator
Revised: Lie Operator

Original: Evaluate Robustness
Revised: Evaluate Robustness

Revised: Natural variation

Revised: Structure variation as Transformations

Revised: Generalize variation across instances

Original: [Figure]
Revised: [Figure]

Revised: Lie Operator

Original: exp (
Revised: Known instance [Figure] Unknown instance in typical [Figure]

Revised: MAE Lie +

Revised: [Figure]

Revised: [Figure]

Revised: [Figure]

Revised: ̂

Original: [Figure]
Revised: [Figure]

Revised: ( ̂ g

Original:  d 
Revised: [Figure] exp ( [Figure] d ̂ t d L d ) infer

Original: [Figure]
Revised: [Figure]

Revised: ̂ t

Original: [Figure]
Revised: [Figure]

Original: zz 
Revised: zz g

Revised: Known instances typical pose typical pose encourages z

Original: [Figure]
Revised: [Figure]

Original: encourages z
Revised: =

Original: [Figure]
Revised: [Figure]

Original: g
Revised: g

Original: [Figure]
Revised: [Figure]

Original:  in in representation space atypical [Figure]
Revised: seen in in representation space atypical pose

Original: SSL Loss
Revised: SSL Loss

Original: Lie Loss + Loss =
Revised: Lie Loss + Loss =

Original: Training Data
Revised: Training Data

Original: [Figure]
Revised: [Figure]

Original: |
Revised: |

Original: [Figure]
Revised: [Figure]

Original: Using this operator, we improve performance by >10% for MAE for both known instances in new poses and unknown instances in both typical
Revised: Using this operator, we improve performance by >10% for MAE for both known instances in new poses and unknown instances in both typical

Original: For these reasons, we choose to explicitly model the transformations of the data as transformations in the latent representation rather than trying to be invariant to it.
Revised: For these reasons, we choose to explicitly model the transformations of the data as transformations in the latent representation rather than trying to be invariant to it.

Original: To do so, we use the formalism of Lie groups.
Revised: To do so, we use the formalism of Lie groups.

Original: Informally, Lie groups are continuous groups described by a set of real parameters (Hall, 2003).
Revised: Informally, Lie groups are continuous groups described by a set of real parameters (Hall, 2003).

Original: While many continuous transformations form matrix Lie groups (e.g., rotations), they lack the typical structure of a vector space.
Revised: While many continuous transformations form matrix Lie groups (e.g., rotations), they lack the typical structure of a vector space.

Original: However, Lie group have a corresponding vector space, their Lie algebra, that can be described using basis matrices, allowing to describe the inﬁnite number elements of the group by a ﬁnite number of basis matrices.
Revised: However, Lie group have a corresponding vector space, their Lie algebra, that can be described using basis matrices, allowing to describe the infinite number elements of the group by a finite number of basis matrices.

Original: Our goal will be to learn such matrices to directly model the data variations.
Revised: Our goal will be to learn such matrices to directly model the data variations.

Original: To summarize, our approach structures the representation space to enable self-supervised models to generalize variation across objects.
Revised: To summarize, our approach structures the representation space to enable self-supervised models to generalize variation across objects.

Original: Since many naturally occurring transformations (e.g., pose, color, size, etc.) are continuous, we develop a theoretically-motivated operator, the Lie operator , that acts in representation space (see Fig. 1).
Revised: Since many naturally occurring transformations (e.g., pose, color, size, etc.) are continuous, we develop a theoretically-motivated operator, the Lie operator , that acts in representation space (see Fig. 1).

Original: Speciﬁcally, the Lie operator learns the continuous transformations observed in data as a vector space, using a set of basis matrices.
Revised: Specifically, the Lie operator learns the continuous transformations observed in data as a vector space, using a set of basis matrices.

Original: With this approach, we make the following contributions:
Revised: With this approach, we make the following contributions:

Original: We generate a novel dataset containing 3D objects in many different poses, allowing us to explicitly evaluate the ability of models to generalize to both known objects in unknown poses and to unknown objects in both known and unknown poses (Section 3).
Revised: We generate a novel dataset containing 3D objects in many different poses, allowing us to explicitly evaluate the ability of models to generalize to both known objects in unknown poses and to unknown objects in both known and unknown poses (Section 3).

Original: Using this dataset, we evaluate the generalization capabilities of a number of standard models, including ResNet-50, ViT, MLP-Mixer, SimCLR, CLIP, VICReg, and MAE ﬁnding that all state-of-the-art models perform relatively poorly in this setting (Section 3.2).
Revised: Using this dataset, we evaluate the generalization capabilities of a number of standard models, including ResNet-50, ViT, MLP-Mixer, SimCLR, CLIP, VICReg, and MAE finding that all state-of-the-art models perform relatively poorly in this setting (Section 3.2).

Original: We incorporate our proposed Lie operator in two recent SSL approaches: masked autoencoders (MAE, (He et al., 2021)), and Variance-Invariance-Covariance Regularization (VICReg (Bardes et al., 2021)) to directly model transformations in data (Section 2), resulting in substantial OOD performance gains of greater than 10% for MAE and of up to 8% for VICReg (Section 4.1).
Revised: We incorporate our proposed Lie operator in two recent SSL approaches: masked autoencoders (MAE, (He et al., 2021)), and Variance-Invariance-Covariance Regularization (VICReg (Bardes et al., 2021)) to directly model transformations in data (Section 2), resulting in substantial OOD performance gains of greater than 10% for MAE and of up to 8% for VICReg (Section 4.1).

Original: We also incorporate our Lie model in SimCLR (Chen et al., 2020)
Revised: We also incorporate our Lie model in SimCLR (Chen et al., 2020)

Original: (Appendix E).
Revised: (Appendix E).

Original: We run systemic ablations of each term of our learning objective in the MAE Lie model, showing the relevance of every component for best performance (Section 4.3).
Revised: We run systemic ablations of each term of our learning objective in the MAE Lie model, showing the relevance of every component for best performance (Section 4.3).

Original: We challenge our learned Lie operator by running evaluation on Imagenet as well as a rotated version of ImageNet, improving MAE performance by 3.94% and 2.62%, respectively.
Revised: We challenge our learned Lie operator by running evaluation on Imagenet as well as a rotated version of ImageNet, improving MAE performance by 3.94% and 2.62%, respectively.

Original: We also improve over standard MAE with MAE Lie when evaluated with ﬁnetuning on the iLab (Borji et al., 2016) dataset.
Revised: We also improve over standard MAE with MAE Lie when evaluated with finetuning on the iLab (Borji et al., 2016) dataset.

Original: These experiments show the applicability of our proposed Lie operator to realistic, challenging datasets (Section 5).
Revised: These experiments show the applicability of our proposed Lie operator to realistic, challenging datasets (Section 5).

Original:  M ETHODS
Revised:  M ETHODS

Original: Humans recognize even new objects in all sorts of poses by observing continuous changes in the world.
Revised: Humans recognize even new objects in all sorts of poses by observing continuous changes in the world.

Original: To mimic humans’ ability to generalize variation, we propose a method for learning continuous transformations from data.
Revised: To mimic humans’ ability to generalize variation, we propose a method for learning continuous transformations from data.

Original: We assume we have access to pairs of data, e.g., images, x , x 0 where x 0 is a transformed version of x corresponding to the same instance , e.g., a speciﬁc object, undergoing variations.
Revised: We assume we have access to pairs of data, e.g., images, x , x ′ where x ′ is a transformed version of x corresponding to the same instance , e.g., a specific object, undergoing variations.

Original: This is a reasonable form of supervision which one can easily acquire from video data for example, where pairs of frames would correspond to pairs of transformed data, and has been employed in previous works such as Ick & Lostanlen (2020);
Revised: This is a reasonable form of supervision which one can easily acquire from video data for example, where pairs of frames would correspond to pairs of transformed data, and has been employed in previous works such as Ick & Lostanlen (2020);

Original: Locatello et al.
Revised: Locatello et al.

Original: ; Connor et al.
Revised: ; Connor et al.

Original: (2021b); Bouchacourt et al.
Revised: (2021b); Bouchacourt et al.

Original: (2021a).
Revised: (2021a).

Original: Despite widespread use of such pairs in data augmentation, models still struggle to generalize across distribution shifts (Alcorn et al., 2019;
Revised: Despite widespread use of such pairs in data augmentation, models still struggle to generalize across distribution shifts (Alcorn et al., 2019;

Original: Geirhos et al., 2018; Engstrom et al., 2019).
Revised: Geirhos et al., 2018; Engstrom et al., 2019).

Original: To solve this brittleness, we take inspiration from previous work illustrating the possibility of representing variation via latent operators (Connor et al., 2021b;
Revised: To solve this brittleness, we take inspiration from previous work illustrating the possibility of representing variation via latent operators (Connor et al., 2021b;

Original: Bouchacourt et al., 2021a;
Revised: Bouchacourt et al., 2021a;

Original: Giannone et al., 2019).
Revised: Giannone et al., 2019).

Original: Acting in latent space avoids costly modiﬁcations of the full input dimensions and allows us to express complex operators using simpler latent operators.
Revised: Acting in latent space avoids costly modifications of the full input dimensions and allows us to express complex operators using simpler latent operators.

Original: We learn to represent data variation explicitly as group actions on the model’s representation space Z .
Revised: We learn to represent data variation explicitly as group actions on the model’s representation space Z .

Original: We use operators theoretically motivated by group theory, as groups enforce properties such as composition and inverses, allowing us to structure representations beyond that of an unconstrained network (such as an multi-layer perceptron, MLP) (see Appendix A).
Revised: We use operators theoretically motivated by group theory, as groups enforce properties such as composition and inverses, allowing us to structure representations beyond that of an unconstrained network (such as an multi-layer perceptron, MLP) (see Appendix A).

Original: Fortunately, many common transformations can be described as groups.
Revised: Fortunately, many common transformations can be described as groups.

Original: We choose matrix Lie groups to model continuous transformations, as they can be represented as matrices and described by real parameters.
Revised: We choose matrix Lie groups to model continuous transformations, as they can be represented as matrices and described by real parameters.

Original: However, they do not possess the basic properties of a vector space e.g., adding two rotation matrices does not produce a rotation matrix, without which back-propagation is not possible.
Revised: However, they do not possess the basic properties of a vector space e.g., adding two rotation matrices does not produce a rotation matrix, without which back-propagation is not possible.

Original: For each matrix Lie group G however, there is a corresponding Lie algebra g which is the tangent space of G at the identity.
Revised: For each matrix Lie group G however, there is a corresponding Lie algebra g which is the tangent space of G at the identity.

Original: A Lie algebra forms vector space where addition is closed and where the space can be described using basis matrices.
Revised: A Lie algebra forms vector space where addition is closed and where the space can be described using basis matrices.

Original: Under mild assumptions over G , any element of G can be obtained by taking the exponential matrix map: exp( t M ) where M is a matrix in g , and t is a real number.
Revised: Under mild assumptions over G , any element of G can be obtained by taking the exponential matrix map: exp( t M ) where M is a matrix in g , and t is a real number.

Original: Beneﬁting from the vector space structure of Lie algebras, we aim to learn a basis of the Lie algebra in order to represent the group acting on the representation space, using the exp map.
Revised: Benefiting from the vector space structure of Lie algebras, we aim to learn a basis of the Lie algebra in order to represent the group acting on the representation space, using the exp map.

Original: We refer the reader to Appendix A and Hall (2003); Stillwell (2008) for a more detailed presentation of matrix Lie groups and their algebra.
Revised: We refer the reader to Appendix A and Hall (2003); Stillwell (2008) for a more detailed presentation of matrix Lie groups and their algebra.

Original: 2. L EARNING TRANSFORMATIONS FROM ONE SAMPLE TO ANOTHER VIA L IE GROUPS
Revised: 2. L EARNING TRANSFORMATIONS FROM ONE SAMPLE TO ANOTHER VIA L IE GROUPS

Original: We assume that there exists a ground-truth group element, g , which describes the transformation in latent space between the embedding of x and x 0 .
Revised: We assume that there exists a ground-truth group element, g , which describes the transformation in latent space between the embedding of x and x ′ .

Original: Our goal is to approximate g using its Lie algebra.
Revised: Our goal is to approximate g using its Lie algebra.

Original: More formally, given a pair of samples (e.g., a pair of video frames), x and x 0 , and an encoder, f ( · ) , we compute the latent representations z and z g as z = f ( x ) and z g = f
Revised: More formally, given a pair of samples (e.g., a pair of video frames), x and x ′ , and an encoder, f ( · ) , we compute the latent representations z and z g as z = f ( x ) and z g = f

Original: ( x 0 ) .
Revised: ( x ′ ) .

Original: We therefore assume z g
Revised: We therefore assume z g

Original: = f ( x 0 )
Revised: = f ( x ′ )

Original: =
Revised: =

Original: g ( f ( x )) .
Revised: g ( f ( x )) .

Original: We aim to learn an operator ˆ g maximizing the similarity between z ˆ g = ˆ g ( f ( x )) and z g .
Revised: We aim to learn an operator ˆ g maximizing the similarity between z ˆ g = ˆ g ( f ( x )) and z g .

Original: We construct our operator by using the vector space structure of the group’s
Revised: We construct our operator by using the vector space structure of the group’s

Original: Lie algebra g : we learn a basis of g as a set of matrices L k 2 R m ⇥ m , k 2 { 1 , . . .
Revised: Lie algebra g : we learn a basis of g as a set of matrices L k ∈ R m × m , k ∈ { 1 , . . .

Original: , d } where d is the dimension of the Lie algebra and m is the dimension of the embedding space Z .
Revised: , d } m is the dimension of the embedding space Z and d is the dimension of the Lie algebra, i.e. the number of Lie generator matrices.

Revised: Intuitively, d expresses the degrees of freedom required to model the transformations.

Revised: Specifically, it corresponds to the dimension of the Lie algebra as a vector space 2 .

Original: We use this basis as a way to represent the group elements, and thus to model the variations that data samples undergo.
Revised: We use this basis as a way to represent the group elements, and thus to model the variations that data samples undergo.

Original: A matrix M in the Lie algebra g thus writes M = [Equation]
Revised: A matrix M in the Lie algebra g thus writes M = (cid:80) dk

Original:  some coordinate 
Revised: =1 t k L k for some coordinate vector t =

Original: [Equation]
Revised: ( t 1 , . . . , t d ) ∈ R d .

Original: Let us consider a pair z , z g with z g = g ( z ) .
Revised: Let us consider a pair z , z g with z g = g ( z ) .

Original: We assume the exponential map is surjective 2 , thus, there exists [Equation]
Revised: We assume the exponential map is surjective 3 , thus, there exists (cid:80) dk

Original:  such that
Revised: =1 t k L k ∈ g such that

Original: [Equation]
Revised: [Equation]

Original: where t = ( t 1 , . . .
Revised: where t = ( t 1 , . . .

Original: , t d ) are the coordinates in the Lie algebra of the group element g , and L is the 3D matrix in R d ⇥ m ⇥ m concatenating { L k , k 2 { 1 , . . .
Revised: , t d ) are the coordinates in the Lie algebra of the group element g , and L is the 3D matrix in R d × m × m concatenating { L k , k ∈ { 1 , . . .

Original: , d }} .
Revised: , d }} .

Original: Our goal is to infer ˆ g such that ˆ g ( z ) (the application of the inferred ˆ g to the representation z ) equals z g .
Revised: Our goal is to infer ˆ g such that ˆ g ( z ) (the application of the inferred ˆ g to the representation z ) equals z g .

Original: We denote z
Revised: We denote z

Original: ˆ g
Revised: ˆ g

Original: =
Revised: =

Original: ˆ g ( z ) .
Revised: ˆ g ( z ) .

Original: Speciﬁcally, our transformation inference procedure consists of
Revised: Specifically, our transformation inference procedure consists of

Original: Steps 1-4 in Algorithm 1.
Revised: Steps 1-4 in Algorithm 1.

Original: For step 1 in Algorithm 1, similar to Connor et al.
Revised: For step 1 in Algorithm 1, similar to Connor et al.

Original: (2021b)
Revised: (2021b)

Original: ; Bouchacourt et al.
Revised: ; Bouchacourt et al.

Original: (2021a), we use a pair of latent codes to infer the latent transformation.
Revised: (2021a), we use a pair of latent codes to infer the latent transformation.

Original: Additionally, we consider that we can access a scalar value, denoted  x,x 0 , measuring “how much” the two samples differ.
Revised: Additionally, we consider that we can access a scalar value, denoted δ x,x ′ , measuring “how much” the two samples differ.

Original: For example, considering frames from a given video, this would be the number of frames that separates the two original video frames x and x 0 , and hence comes for free 3 .The
Revised: For example, considering frames from a given video, this would be the number of frames that separates the two original video frames x and x ′ , and hence comes for free 4 .The

Original: coordinate vector corresponding to the pair ( z , z g ) is then inferred using a a multi-layer perceptron h 4 : 
Revised: coordinate vector corresponding to the pair ( z , z g ) is then inferred using a a multi-layer perceptron h 5 : [Equation]

Original: [Equation]
Revised: [Equation]

Original: When  x , x 0 is small, we assume the corresponding images and embedding differ by a transformation close to the identity which corresponds to an inﬁnitesimal ˆ t .
Revised: When δ x , x ′ is small, we assume the corresponding images and embedding differ by a transformation close to the identity which corresponds to an infinitesimal ˆ t .

Original: To illustrate this, we use a squared
Revised: To illustrate this, we use a squared

Original: L 2 -norm constraint on ˆ t .
Revised: L 2 -norm constraint on ˆ t .

Original: For smaller  x , x 0 (e.g., close-by frames in a video), larger similarity between z , z g is expected, in which case we want to encourage a smaller norm of ˆ t :
Revised: For smaller δ x , x ′ (e.g., close-by frames in a video), larger similarity between z , z g is expected, in which case we want to encourage a smaller norm of ˆ t :

Original: [Equation]
Revised: [Equation]

Original: where s :Z ⇥ Z ⇥
Revised: where s :

Original: [Equation] measures the similarity between z and z g .
Revised: Z × Z × R → [0 , 1] measures the similarity between z and z g .

Original: In our experiments, we simply use s ( z , z g ,  x , x 0 ) as a function of  x , x 0 and compute
Revised: In our experiments, we simply use s ( z , z g , δ x , x ′ ) as a function of δ x , x ′ and compute

Original: [Equation]
Revised: [Equation]

Original: When |  x , x 0 | decreases, s increases which strengthens the constraint in Eq.
Revised: When | δ x , x ′ | decreases, s increases which strengthens the constraint in Eq.

Original: 3: the inferred vector of coordinates ˆ t is therefore enforced to have small norm.
Revised: 3: the inferred vector of coordinates ˆ t is therefore enforced to have small norm.

Original:  L IE LEARNING OBJECTIVE
Revised: 2. L IE LEARNING OBJECTIVE

Original: We encourage z ˆ g to be close to z g , while preventing a collapse mode where every transformation would be inferred as identity.
Revised: We encourage z ˆ g to be close to z g , while preventing a collapse mode where every transformation would be inferred as identity.

Original: In practice, recall we obtain z g as the embedding f ( x 0 ) of a corresponding frame.
Revised: In practice, recall we obtain z g as the embedding f ( x ′ ) of a corresponding frame.

Original: Speciﬁcally, we use an InfoNCE loss to match pairs 5 :
Revised: Specifically, we use an InfoNCE loss to match pairs 6 :

Original: [Equation]
Revised: [Equation]

Original: where sim is a similarity function measuring how close two latent embeddings are (e.g., cosine similarity in our experiments), and ˜z belongs to set of negatives N for z g .
Revised: where sim is a similarity function measuring how close two latent embeddings are (e.g., cosine similarity in our experiments), and ˜z belongs to set of negatives N for z g .

Original: For a given z g of a given instance, the set N is comprised of latent codes of other instances.
Revised: For a given z g of a given instance, the set N is comprised of latent codes of other instances.

Original: We also include other instances’ transformed versions and their corresponding inferred transformed embeddings as negative samples.
Revised: We also include other instances’ transformed versions and their corresponding inferred transformed embeddings as negative samples.

Original: Crucially, we also include in the negatives set of z g (denoted N ) the frame encoding z in order to prevent the model from collapsing to a point where all the encodings of the same instance (i.e. any transformed version) would be mapped to the same point, and the inferred ˆ g would simply learned as identity.
Revised: Crucially, we also include in the negatives set of z g (denoted N ) the frame encoding z in order to prevent the model from collapsing to a point where all the encodings of the same instance (i.e. any transformed version) would be mapped to the same point, and the inferred ˆ g would simply learned as identity.

Original: Indeed, one can view transformed pairs simply as another form of data augmentation, such that enforcing that z g matches z
Revised: Indeed, one can view transformed pairs simply as another form of data augmentation, such that enforcing that z g matches z

Original: (i.e. implicitly enforcing invariance) would  bring robustness to the transformation g .
Revised: (i.e. implicitly enforcing invariance) would [Equation] bring robustness to the transformation g .

Original: However, we control for this in our experiments using the
Revised: However, we control for this in our experiments using the

Original: SimCLR model with our Lie operator (Appendix E, see SimCLR Frames baseline).
Revised: SimCLR model with our Lie operator (Appendix E, see SimCLR Frames baseline).

Original: We experiment[Equation] with a baseline that is simply encouraged to match pairs of rotated frames in the representation, and SimCLR Lie outperforms this baseline.
Revised: We experiment with a baseline that is simply encouraged to match pairs of rotated frames in the representation, and SimCLR Lie outperforms this baseline.

Original: Thus, this shows that viewing transformations as data augmentation is not sufﬁcient for robust generalization.
Revised: Thus, this shows that viewing transformations as data augmentation is not sufficient for robust generalization.

Original: To further encourage z g and z ˆ g to match, we add an Euclidean-type constraint:
Revised: To further encourage z g and z ˆ g to match, we add an Euclidean-type constraint:

Original: [Equation]
Revised: [Equation]

Original: The addition of the Euclidean loss helps to stabilize training, but using only the Euclidean loss (in place of our Lie loss) would not be enough to prevent the model from collapsing to mapping all transformed samples to the same point in latent space.
Revised: The addition of the Euclidean loss helps to stabilize training, but using only the Euclidean loss (in place of our Lie loss) would not be enough to prevent the model from collapsing to mapping all transformed samples to the same point in latent space.

Original: Indeed, in the case of the MAE model, we perform ablations of our Lie model in Section 4.3 and show that all terms in the training objective are important for best performance with the MAE Lie model.
Revised: Indeed, in the case of the MAE model, we perform ablations of our Lie model in Section 4.3 and show that all terms in the training objective are important for best performance with the MAE Lie model.

Original: 2. L EARNING A L IE OPERATOR IN SELF - SUPERVISED REPRESENTATIONS
Revised: 2. L EARNING A L IE OPERATOR IN SELF - SUPERVISED REPRESENTATIONS

Original: We use our Lie algebra model on top of two existing self-supervised (SSL) models, MAE and VICReg.
Revised: We use our Lie algebra model on top of two existing self-supervised (SSL) models, MAE and VICReg.

Original: In MAE, f is trained to reconstruct a masked input image, while VICReg is trained to match pairs of samples obtained from the same image via some form of data augmentations (e.g. two crops) and uses Variance-Invariance-Covariance regularization to prevent collapse.
Revised: In MAE, f is trained to reconstruct a masked input image, while VICReg is trained to match pairs of samples obtained from the same image via some form of data augmentations (e.g. two crops) and uses Variance-Invariance-Covariance regularization to prevent collapse.

Original: In Appendix E we also experiment applying our Lie operator on the SimCLR (Chen et al., 2020) model, showing gains in the linear setting in particular.
Revised: In Appendix E we also experiment applying our Lie operator on the SimCLR (Chen et al., 2020) model, showing gains in the linear setting in particular.

Original: Our loss will balance between two components: the original SSL model loss L ssl , applied to both samples of the pair separately, and the transformation loss L Lie .
Revised: Our loss will balance between two components: the original SSL model loss L ssl , applied to both samples of the pair separately, and the transformation loss L Lie .

Original: Combining these constraints gives us our ﬁnal loss, L :
Revised: Combining these constraints gives us our final loss, L :

Original: [Equation]
Revised: [Equation]

Original: The parameters  ssl and  Lie weight the two losses.
Revised: The parameters λ ssl and λ Lie weight the two losses.

Original: Note that L ssl can be different given the
Revised: Note that L ssl can be different given the

Original: SSL model backbone we use.
Revised: SSL model backbone we use.

Original: In our experiment we deploy our Lie operator on top of the
Revised: In our experiment we deploy our Lie operator on top of the

Original: Variance-Invariance-Covariance Regularization model (VICReg, (Bardes et al., 2021)), and of the very recent masked autoencoders model (MAE, (He et al., 2021)).
Revised: Variance-Invariance-Covariance Regularization model (VICReg, (Bardes et al., 2021)), and of the very recent masked autoencoders model (MAE, (He et al., 2021)).

Original: This loss is minimized with respect to the parameters of the encoder f , the parameters of the network h (multi-layer perceptron with two linear layers connected by a leaky ReLU) that infers the Lie algebra coordinates, and the Lie algebra basis matrices L k , k 2 { 1 , . . .
Revised: This loss is minimized with respect to the parameters of the encoder f , the parameters of the network h (multi-layer perceptron with two linear layers connected by a leaky ReLU) that infers the Lie algebra coordinates, and the Lie algebra basis matrices L k , k ∈ { 1 , . . .

Original: , d } which are used to compute ˆ g .
Revised: , d } which are used to compute ˆ g .

Original: The total loss is over all the instances ( N of them) and all possible pairs of transformed samples we have for each instance (if we have P pose changes, this will be P ( P  1) ).
Revised: The total loss is over all the instances ( N of them) and all possible pairs of transformed samples we have for each instance (if we have P pose changes, this will be P ( P − 1) ).

Original: During training, we estimate this loss for a minibatch by sampling instances ﬁrst, and sampling a pair of samples for each instance without replacement.
Revised: During training, we estimate this loss for a minibatch by sampling instances first, and sampling a pair of samples for each instance without replacement.

Original: The full procedure is illustrated in Fig.
Revised: The full procedure is illustrated in Fig.

Original: 1 and in
Revised: 1 and in

Original: Algorithm 1.
Revised: Algorithm 1.

Original:  E VALUATING GENERALIZATION ACROSS INSTANCES
Revised:  E VALUATING GENERALIZATION ACROSS INSTANCES

Original: Our goal is to evaluate how well models generalize variation across instances.
Revised: Our goal is to evaluate how well models generalize variation across instances.

Original: We focus on robustness to pose changes, a common, natural variation of objects.
Revised: We focus on robustness to pose changes, a common, natural variation of objects.

Original: We evaluate the learned representations’robustness via standard linear evaluation or supervised ﬁnetuning protocols. Speciﬁcally, we measure
Revised: We evaluate the learned representations’

Original: [Figure]
Revised: [Figure]

Revised: robustness via standard linear evaluation or supervised finetuning protocols.

Original: classiﬁcation top-1 accuracy across both known (seen during training) and unknown instances (not seen during training) in typical and new poses.
Revised: Specifically, we measure classification top-1 accuracy across both known (seen during training) and unknown instances (not seen during training) in typical and new poses.

Original: As only a portion of instances (less or equal to 50% in our experiments) are seen varying, motivated by works showing the beneﬁt of data reweighting for generalization challenges (Idrissi et al., 2022), we sample data at training so that the model sees as many images varying as non varying by upsampling the varying instances where needed.
Revised: As only a portion of instances (less or equal to 50% in our experiments) are seen varying, motivated by works showing the benefit of data reweighting for generalization challenges (Idrissi et al., 2022), we sample data at training so that the model sees as many images varying as non varying by upsampling the varying instances where needed.

Original: In order to perform well on known instances in new poses, models must learn to recognize an object regardless of its orientation and constitutes our weakest generalization test.
Revised: In order to perform well on known instances in new poses, models must learn to recognize an object regardless of its orientation and constitutes our weakest generalization test.

Original: In order to perform well on unknown instances, however, models must be able to generalize both across instances and poses, a substantially more difﬁcult task.
Revised: In order to perform well on unknown instances, however, models must be able to generalize both across instances and poses, a substantially more difficult task.

Original: 3. D ATA VARIATION
Revised: 3. D ATA VARIATION

Original: To conduct our evaluation, we develop a controlled dataset of common objects such as buses, cars, planes, etc.
Revised: To conduct our evaluation, we develop a controlled dataset of common objects such as buses, cars, planes, etc.

Original: based on realistic object models from Trimble Inc 6 .
Revised: based on realistic object models from Trimble Inc 7 .

Original: We render nearly 500,000 images of varying 2D and 3D poses uniformly distributed with a step size of 4 degrees.
Revised: We render nearly 500,000 images of varying 2D and 3D poses uniformly distributed with a step size of 4 degrees.

Original: To simulate natural changes in pose we organize pose changes temporally into frames (i.e. into images).
Revised: To simulate natural changes in pose we organize pose changes temporally into frames (i.e. into images).

Original: Each frame captures a 4 degree change in pose either in-plane rotation or across all three axes.
Revised: Each frame captures a 4 degree change in pose either in-plane rotation or across all three axes.

Original: In addition, we conduct experiments evaluating the Lie operator on the iLab 2M dataset following the preprocessing steps from Madan et al.
Revised: In addition, we conduct experiments evaluating the Lie operator on the iLab 2M dataset following the preprocessing steps from Madan et al.

Original: 3. G ENERALIZING POSE CHANGES IS CHALLENGING FOR S O TA VISION MODELS
Revised: 3. G ENERALIZING POSE CHANGES IS CHALLENGING FOR S O TA VISION MODELS

Original: We ﬁrst assess state-of-the-art vision models’ robustness to pose using a collection of supervised and self-supervised models (CLIP, ResNet-50, ViT-B/16, MLP Mixer, SimCLR, VICReg and MAE).
Revised: We first assess state-of-the-art vision models’ robustness to pose using a collection of supervised and self-supervised models (CLIP, ResNet-50, ViT-B/16, MLP Mixer, SimCLR, VICReg and MAE).

Original: All of these models rely on some form of data augmentation to provide invariance to nuisance features such as pose.
Revised: All of these models rely on some form of data augmentation to provide invariance to nuisance features such as pose.

Original: We use ImageNet (Deng et al., 2009) pretrained embeddings for our models (with the exception of CLIP, which is pretrained on a large corpus of image-text pairs)
Revised: We use ImageNet (Deng et al., 2009) pretrained embeddings for our models (with the exception of CLIP, which is pretrained on a large corpus of image-text pairs)

Original: (Radford et al.).
Revised: (Radford et al.).

Original: We only show supervised linear evaluation for VICReg as it is mainly explored in the linear evaluation setting for the fully supervised task in Bardes et al.
Revised: We only show supervised linear evaluation for VICReg as it is mainly explored in the linear evaluation setting for the fully supervised task in Bardes et al.

Original: While all models perform well on classiﬁcation for known instances in typical poses, the same models have a considerable generalization gap to varying poses.
Revised: While all models perform well on classification for known instances in typical poses, the same models have a considerable generalization gap to varying poses.

Original: In Figure 2, we show results for both linear evaluation and supervised ﬁnetuning classiﬁcation gaps relative to typical poses of known instances.
Revised: In Figure 2, we show results for both linear evaluation and supervised finetuning classification gaps relative to typical poses of known instances.

Original: Strikingly, all models we evaluated exhibited large generalization gaps of 30-70% suggesting that even modern self-supervised models struggle to generalize both to new poses and to new instances.
Revised: Strikingly, all models we evaluated exhibited large generalization gaps of 30-70% suggesting that even modern self-supervised models struggle to generalize both to new poses and to new instances.

Original: As such poor generalization can be problematic at best and dangerous at worst for deployed models, developing ways to improve generalization performance beyond simply augmenting the data is critical.
Revised: As such poor generalization can be problematic at best and dangerous at worst for deployed models, developing ways to improve generalization performance beyond simply augmenting the data is critical.

Original: In the next section, we introduce a new approach in which we encourage models not to be invariant to transformation, but rather to explicitly model and account for transformations.
Revised: In the next section, we introduce a new approach in which we encourage models not to be invariant to transformation, but rather to explicitly model and account for transformations.

Original: [Figure]
Revised: [Figure]

Original: 3. I NCORPORATING A LEARNED L IE OPERATOR IN SELF - SUPERVISED LEARNING MODELS
Revised: 3. I NCORPORATING A LEARNED L IE OPERATOR IN SELF - SUPERVISED LEARNING MODELS

Original: To improve robustness, we incorporate a self-supervised learning phase where we structure representations using our Lie operator.
Revised: To improve robustness, we incorporate a self-supervised learning phase where we structure representations using our Lie operator.

Original: During this phase, half of the instances are seen undergoing pose changes in 2D and 3D.
Revised: During this phase, half of the instances are seen undergoing pose changes in 2D and 3D.

Original: We use the angle difference between frames as an proxy for frame distance (denoted  in Section 2.1).
Revised: We use the angle difference between frames as an proxy for frame distance (denoted δ in Section 2.1).

Original: In the case of rotation around all axes, we use the mean of three angles.
Revised: In the case of rotation around all axes, we use the mean of three angles.

Original: We then evaluate the representations in the same manner as described above.
Revised: We then evaluate the representations in the same manner as described above.

Original: We implement our Lie operator in two SSL models: Masked Autoencoder (MAE) and VarianceInvariance-Covariance Regularization (VICReg).
Revised: We implement our Lie operator in two SSL models: Masked Autoencoder (MAE) and VarianceInvariance-Covariance Regularization (VICReg).

Original: VICReg is a SSL method, which encourages similar representations among two views (generated via data augmentation) of an image.
Revised: VICReg is a SSL method, which encourages similar representations among two views (generated via data augmentation) of an image.

Original: MAE is an SSL model whose learning objective is based on the reconstruction of a randomly masked input.
Revised: MAE is an SSL model whose learning objective is based on the reconstruction of a randomly masked input.

Original:  R ESULTS 4. SSL MODELS WITH L IE OPERATORS ARE MORE ROBUST ACROSS INSTANCES AND POSES
Revised:  R ESULTS 4. SSL MODELS WITH L IE OPERATORS ARE MORE ROBUST ACROSS INSTANCES AND POSES

Original: We evaluate the robustness beneﬁts of our Lie operator across several generalization settings.
Revised: We evaluate the robustness benefits of our Lie operator across several generalization settings.

Original: In
Revised: In

Original: Figure 3, we plot the performance of a MAE ( A ) and VICReg ( B ) models with and without our Lie operator.
Revised: Figure 3, we plot the performance of a MAE ( A ) and VICReg ( B ) models with and without our Lie operator.

Original: Encouragingly, we ﬁnd explicitly modeling transformations with Lie groups (referred to as
Revised: Encouragingly, we find explicitly modeling transformations with Lie groups (referred to as

Original: MAE Lie and VICReg Lie) improves robustness for all generalization settings.
Revised: MAE Lie and VICReg Lie) improves robustness for all generalization settings.

Original: For known instances (those seen during training), incorporating our Lie operator improves Top-1 classiﬁcation accuracy for new poses by 13 .
Revised: For known instances (those seen during training), incorporating our Lie operator improves Top-1 classification accuracy for new poses by 13 .

Original: 7% for MAE Lie and 7 .
Revised: 7% for MAE Lie and 7 .

Original: 7% for VICReg Lie.
Revised: 7% for VICReg Lie.

Original: In the more challenging settings of unknown instances which are never seen during training, we still observe performance gains for
Revised: In the more challenging settings of unknown instances which are never seen during training, we still observe performance gains for

Original: VICReg Lie of 0 .
Revised: VICReg Lie of 0 .

Original: 2% for new instances in typical and 1 .
Revised: 2% for new instances in typical and 1 .

Original: 2% for those in new poses.
Revised: 2% for those in new poses.

Original: For MAE Lie, we observe substantial performance gains of 12 . 4% for unknown instances in typical poses and 10 .
Revised: For MAE Lie, we observe substantial performance gains of 12 . 4% for unknown instances in typical poses and 10 .

Original: 3% for those in diverse poses.
Revised: 3% for those in diverse poses.

Original: We note both MAE and VICReg Lie bring even more substantial gains for settings when data diversity is limited (see Sec.
Revised: We note both MAE and VICReg Lie bring even more substantial gains for settings when data diversity is limited (see Sec.

Original: 4.2 Tables 1 and 2).
Revised: 4.2 Tables 1 and 2).

Original: 4. L IE OPERATOR ROBUSTNESS GAINS ARE SUSTAINED EVEN WITH LIMITED DIVERSE DATA
Revised: 4. L IE OPERATOR ROBUSTNESS GAINS ARE SUSTAINED EVEN WITH LIMITED DIVERSE DATA

Original: For our experiments so far, models have been exposed to a 50/50 mix of typical and atypical poses during supervised linear classiﬁcation and ﬁnetuning.
Revised: For our experiments so far, models have been exposed to a 50/50 mix of typical and atypical poses during supervised linear classification and finetuning.

Original: However, we want models to generalize across instances and to atypical poses even when presented with comparatively small amounts of diverse data.
Revised: However, we want models to generalize across instances and to atypical poses even when presented with comparatively small amounts of diverse data.

Original: We therefore evaluated our models by performing supervised linear classiﬁcation and ﬁnetuning with either 50%, 25%, or 5% of data featuring atypical poses (Tables 1 and 2).
Revised: We therefore evaluated our models by performing supervised linear classification and finetuning with either 50%, 25%, or 5% of data featuring atypical poses (Tables 1 and 2).

Original: Notably, we observed that while overall performance decreased with fewer diverse instances, as we might expect, the relative gains of the Lie models were largely consistent across settings, and if anything, were larger for settings with less diverse data.
Revised: Notably, we observed that while overall performance decreased with fewer diverse instances, as we might expect, the relative gains of the Lie models were largely consistent across settings, and if anything, were larger for settings with less diverse data.

Original: These results demonstrate the robustness of our approach to differing quantities of diverse data
Revised: These results demonstrate the robustness of our approach to differing quantities of diverse data

Original: [Table] [Table]
Revised: [Table] [Table]

Original: as mean ± standard error (+ absolute difference, x relative multiple) to the baseline model.
Revised: ± standard error (+ absolute difference, x relative multiple) to the baseline model.

Original: We do not experiment with VICReg ﬁnetuning as the original paper (Bardes et al., 2021) focuses on linear evaluation.
Revised: We do not experiment with VICReg finetuning as the original paper (Bardes et al., 2021) focuses on linear evaluation.

Original: [Table]
Revised: [Table]

Original: 4. A BLATIONS
Revised: 4. A BLATIONS

Original: In Table 2, compared to our MAE Lie, a model with only the SSL term (and no Euclidean or Lie loss), that is, the baseline MAE, suffers drops of 10 .
Revised: In Table 2, compared to our MAE Lie, a model with only the SSL term (and no Euclidean or Lie loss), that is, the baseline MAE, suffers drops of 10 .

Original: 3% , 9 . 4% , and
Revised: 3% , 9 . 4% , and

Original: 2% in top-1 accuracy on unknown instances, respectively for the 5% , 25% and 50% diversity settings.
Revised: 2% in top-1 accuracy on unknown instances, respectively for the 5% , 25% and 50% diversity settings.

Original: In addition, we ran systematic ablations for each loss component in our MAE Lie model.
Revised: In addition, we ran systematic ablations for each loss component in our MAE Lie model.

Original: We ﬁnd a consistent decrease in performance across all three diversity settings.
Revised: We find a consistent decrease in performance across all three diversity settings.

Original: Compared to our
Revised: Compared to our

Original: MAE Lie model, a model without the Euclidean loss term (  euc = 0 ) suffers drops of 10 . 5% , 9 .
Revised: MAE Lie model, a model without the Euclidean loss term ( λ euc = 0 ) suffers drops of 10 . 5% , 9 .

Original: 08% , and 3 . 9% in top-1 accuracy for unknown instances (when 50% , 25% , or 5% of instances vary in pose during training, respectively).
Revised: 08% , and 3 . 9% in top-1 accuracy for unknown instances (when 50% , 25% , or 5% of instances vary in pose during training, respectively).

Original: Similarly, compared to our MAE Lie, a model with only the Lie loss (i.e.  euc =  ssl = 0 ) suffers a drop of 5 .
Revised: Similarly, compared to our MAE Lie, a model with only the Lie loss (i.e. λ euc = λ ssl = 0 ) suffers a drop of 5 .

Original: 9% , 5 . 01% , and 1 . 99% in top-1 accuracy for unknown instances across the three diversity settings.
Revised: 9% , 5 . 01% , and 1 . 99% in top-1 accuracy for unknown instances across the three diversity settings.

Original: This result illustrates how each loss term is critical to the overall performance of our proposed MAE Lie model.
Revised: This result illustrates how each loss term is critical to the overall performance of our proposed MAE Lie model.

Original:  MAE L IE TRANSFER ON REAL OBJECTS DATASETS5. I MAGE N ET , R OT I MAGE N ET
Revised:  MAE L IE TRANSFER ON REAL OBJECTS DATASETS

Revised:  I MAGE N ET , R OT I MAGE N ET

Original: [Table]
Revised: [Table]

Original: In order to further demonstrate the beneﬁt of our proposed Lie operator to standard benchmarks, we perform evaluation on ImageNet (Deng et al., 2009) and some of its variants.
Revised: In order to further demonstrate the benefit of our proposed Lie operator to standard benchmarks, we perform evaluation on ImageNet (Deng et al., 2009) and some of its variants.

Original: Rot ImageNet consists of the same validation images on which we apply random rotations to assess how well models handle pose changes.
Revised: Rot ImageNet consists of the same validation images on which we apply random rotations to assess how well models handle pose changes.

Original: To do so, we ﬁnetune both the baseline MAE model and MAE Lie for 20k steps on
Revised: To do so, we finetune both the baseline MAE model and MAE Lie for 20k steps on

Original: ImageNet with a linear classiﬁcation head.
Revised: ImageNet with a linear classification head.

Original: Table 3 reports the performance on the ImageNet and Rot
Revised: Table 3 reports the performance on the ImageNet and Rot

Original: ImageNet validation sets.
Revised: ImageNet validation sets.

Original: We ﬁnd MAE Lie outperforms the baseline for rotated ImageNet validation images by a considerable +2 .
Revised: We find MAE Lie outperforms the baseline for rotated ImageNet validation images by a considerable +2 .

Original: 62% in top-1 accuracy and the unaltered ImageNet validation set (by
Revised: 62% in top-1 accuracy and the unaltered ImageNet validation set (by

Original: +3 .
Revised: +3 .

Original: 94% in top-1 accuracy) which contain pose along with other changing factors.
Revised: 94% in top-1 accuracy) which contain pose along with other changing factors.

Original: Details on the training can be found in Appendix B.1.3.
Revised: Details on the training can be found in Appendix B.1.3.

Original: 5. I L AB DATASET
Revised: 5. I L AB DATASET

Original: To demonstrate the applicability of our approach to other datasets, we also evaluate our MAE Lie model on the iLab 2M dataset.
Revised: To demonstrate the applicability of our approach to other datasets, we also evaluate our MAE Lie model on the iLab 2M dataset.

Original: The iLab 2M dataset (Borji et al., 2016) is a natural image dataset consisting of 2M images of toy objects, from multiple categories, undergoing variations in object instances, viewpoints, lighting conditions, and backgrounds.
Revised: The iLab 2M dataset (Borji et al., 2016) is a natural image dataset consisting of 2M images of toy objects, from multiple categories, undergoing variations in object instances, viewpoints, lighting conditions, and backgrounds.

Original: We follow the preprocessing steps in
Revised: We follow the preprocessing steps in

Original: Madan et al.
Revised: Madan et al.

Original: We partition the training samples into a validation set ( 10% ) and show results on the iLab test set.
Revised: We partition the training samples into a validation set ( 10% ) and show results on the iLab test set.

Original: We conduct supervised ﬁnetuning, the standard evaluation protocol for MAE.
Revised: We conduct supervised finetuning, the standard evaluation protocol for MAE.

Original: We observe MAE Lie yields nearly a 3% improvement on the test set ( 93 . 2% ± 0 .
Revised: We observe MAE Lie yields nearly a 3% improvement on the test set ( 93 . 2% ± 0 .

Original: 45% MAE Lie vs 90 .
Revised: 45% MAE Lie vs 90 .

Original: 6% ± 0 .
Revised: 6% ± 0 .

Original: 11% MAE) without re-training our Lie operator on the iLab dataset.
Revised: 11% MAE) without re-training our Lie operator on the iLab dataset.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Modeling variation Existing domain adaptation works (Zhou et al., 2020;
Revised: Modeling variation Existing domain adaptation works (Zhou et al., 2020;

Original: Robey et al., 2021; Nguyen et al., 2021) learn transformations that map one domain to the other.
Revised: Robey et al., 2021; Nguyen et al., 2021) learn transformations that map one domain to the other.

Original: In contrast to Ick & Lostanlen (2020); Hashimoto et al.
Revised: In contrast to Ick & Lostanlen (2020); Hashimoto et al.

Original: , we learn transformations in the latent space rather than image-space.
Revised: , we learn transformations in the latent space rather than image-space.

Original: Connor & Rozell (2020); Connor et al.
Revised: Connor & Rozell (2020); Connor et al.

Original: (2021a) learn transformations in the latent space but uses a decoder that performs image reconstruction, which is known to be hard to apply to complex, large scale image data.
Revised: (2021a) learn transformations in the latent space but uses a decoder that performs image reconstruction, which is known to be hard to apply to complex, large scale image data.

Original: In the context of classiﬁcation, Benton et al.
Revised: In the context of classification, Benton et al.

Original: (2020) use the Lie algebra formalism to learn augmentations, but they require a priori knowledge of the augmentations.
Revised: (2020) use the Lie algebra formalism to learn augmentations, but they require a priori knowledge of the augmentations.

Original: Tai et al.
Revised: Tai et al.

Original: (2019) provides invariance to continuous transformations, but they only experiment on small scale data (MNIST (LeCun et al., 2010) and SVHN (Netzer et al., 2011)).
Revised: (2019) provides invariance to continuous transformations, but they only experiment on small scale data (MNIST (LeCun et al., 2010) and SVHN (Netzer et al., 2011)).

Original: We also depart from their work as we explicitly focus on the challenge of robustness to distributional shifts.
Revised: We also depart from their work as we explicitly focus on the challenge of robustness to distributional shifts.

Original: Equivariant models (see Cohen & Welling (2016);
Revised: Equivariant models (see Cohen & Welling (2016);

Original: Cohen et al.
Revised: Cohen et al.

Original: (2020) among others) and disentanglement (Higgins et al., 2018;
Revised: (2020) among others) and disentanglement (Higgins et al., 2018;

Original: Giannone et al., 2019) also aim to structure latent representations, but are mainly applied to synthetic, small scale data.
Revised: Giannone et al., 2019) also aim to structure latent representations, but are mainly applied to synthetic, small scale data.

Original: Disentanglement also often requires costly pixel reconstruction.
Revised: Disentanglement also often requires costly pixel reconstruction.

Original: Spatial Transformer
Revised: Spatial Transformer

Original: Networks (Jaderberg et al.) learn an instance-dependent afﬁne transform for classiﬁcation, resulting in behavior such as zooming in on the foreground object.
Revised: Networks (Jaderberg et al.) learn an instance-dependent affine transform for classification, resulting in behavior such as zooming in on the foreground object.

Original: Our goal differs as we focus on the generalization capabilities of various models to pose changes.
Revised: Our goal differs as we focus on the generalization capabilities of various models to pose changes.

Original: Capsule networks (Sabour et al., 2017;
Revised: Capsule networks (Sabour et al., 2017;

Original: Mazzia et al., 2021) were designed to better generalize to novel views, however, it is difﬁcult for us to compare our approach to these as we require a strong pre-trained image model.
Revised: Mazzia et al., 2021) were designed to better generalize to novel views, however, it is difficult for us to compare our approach to these as we require a strong pre-trained image model.

Original: To our knowledge, no such model using capsule networks exists, making a fair comparison very difﬁcult.
Revised: To our knowledge, no such model using capsule networks exists, making a fair comparison very difficult.

Original: OOD Generalization failures Alcorn et al.
Revised: OOD Generalization failures Alcorn et al.

Original: Engstrom et al.
Revised: Engstrom et al.

Original: (2019) shed light on the failures of SoTA models on “unusual" examples that slightly vary from the training data, and Engstrom et al.
Revised: (2019) shed light on the failures of SoTA models on “unusual" examples that slightly vary from the training data, and Engstrom et al.

Original: (2019); Azulay & Weiss (2019) show that explicitly augmenting the data with such variations does not ﬁx the problem.
Revised: (2019); Azulay & Weiss (2019) show that explicitly augmenting the data with such variations does not fix the problem.

Original: Indeed, Bouchacourt et al.
Revised: Indeed, Bouchacourt et al.

Original: (2021b) found that data augmentation does not bring the expected invariance to transformations, and that invariance rather comes from the data itself.
Revised: (2021b) found that data augmentation does not bring the expected invariance to transformations, and that invariance rather comes from the data itself.

Original: Self-supervised models trained with InfoNCE theoretically recover the true data factors of variations (Zimmermann et al., 2021; von Kügelgen et al., 2021), but our experiments show that they do not generalize to distributional shifts.
Revised: Self-supervised models trained with InfoNCE theoretically recover the true data factors of variations (Zimmermann et al., 2021; von Kügelgen et al., 2021), but our experiments show that they do not generalize to distributional shifts.

Original:  D ISCUSSION
Revised:  D ISCUSSION

Original: We aim to improve the OOD generalization of SSL models by encouraging them to explicitly model transformations using Lie groups.
Revised: We aim to improve the OOD generalization of SSL models by encouraging them to explicitly model transformations using Lie groups.

Original: Using a novel dataset to evaluate OOD generalization, we found that applying our Lie operator to MAE and VICReg leads to markedly improved performance for both models, enabling better generalization both across poses and instances, both on our synthetic data and on ImageNet and its variations for MAE.
Revised: Using a novel dataset to evaluate OOD generalization, we found that applying our Lie operator to MAE and VICReg leads to markedly improved performance for both models, enabling better generalization both across poses and instances, both on our synthetic data and on ImageNet and its variations for MAE.

Original: Our results highlight that data augmentation alone is not enough to enable effective generalization and that structuring operators to explicitly account for transformations is a promising approach to improving robustness.
Revised: Our results highlight that data augmentation alone is not enough to enable effective generalization and that structuring operators to explicitly account for transformations is a promising approach to improving robustness.

Original: We focus on modeling pose changes, a common natural factor of variation in objects.
Revised: We focus on modeling pose changes, a common natural factor of variation in objects.

Original: However, Lie groups are capable of describing a variety of naturally occurring continuous transformations.
Revised: However, Lie groups are capable of describing a variety of naturally occurring continuous transformations.

Original: Future work could extend the application of Lie operators to these settings, possibly even to model multiple factors of variation at once.
Revised: Future work could extend the application of Lie operators to these settings, possibly even to model multiple factors of variation at once.

Original:  R EPRODUCIBILITY STATEMENT
Revised:  R EPRODUCIBILITY STATEMENT

Original: We made efforts to make all our work reproducible.
Revised: We made efforts to make all our work reproducible.

Original: In the main text, we describe every aspect of our learning objective and all the assumptions taken, with some background in Appendix A.
Revised: In the main text, we describe every aspect of our learning objective and all the assumptions taken, with some background in Appendix A.

Original: Regarding our experiments, all training details are in Appendix B or in main text, and datasets descriptions are in main text or
Revised: Regarding our experiments, all training details are in Appendix B or in main text, and datasets descriptions are in main text or

Original: Appendix C. We will release the code to reproduce all experiments and generate the dataset based on realistic object models from Trimble Inc (freely available under a General Model license) upon acceptance.
Revised: Appendix C. We will release the code to reproduce all experiments and generate the dataset based on realistic object models from Trimble Inc (freely available under a General Model license) upon acceptance.