Original: H OLOGRAPHIC -(V)AE: AN END - TO - END SO(3)- E QUIVARIANT (V ARIATIONAL ) A UTOENCODER IN F OURIER S PACE
Revised: H OLOGRAPHIC -(V)AE: AN END - TO - END SO(3)- E QUIVARIANT (V ARIATIONAL ) A UTOENCODER IN F OURIER S PACE

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: Group-equivariant neural networks have emerged as a data-efficient approach to solve classification and regression tasks, while respecting the relevant symmetries of the data.
Revised: Group-equivariant neural networks have emerged as a data-efficient approach to solve classification and regression tasks, while respecting the relevant symmetries of the data.

Original: However, little work has been done to extend this paradigm to the unsupervised and generative domains.
Revised: However, little work has been done to extend this paradigm to the unsupervised and generative domains.

Original: Here, we present
Revised: Here, we present

Original: Holographic -(V)AE
Revised: Holographic -(V)AE

Original: (H- (V)AE), a fully end-to-end SO(3)-equivariant (variational) autoencoder in Fourier space, suitable for unsupervised learning and generation of data distributed around a specified origin.
Revised: (H- (V)AE), a fully end-to-end SO(3)-equivariant (variational) autoencoder in Fourier space, suitable for unsupervised learning and generation of data distributed around a specified origin.

Original: H-(V)AE is trained to reconstruct the spherical Fourier encoding of data, learning in the process a latent space with a maximally informative invariant embedding alongside an equivariant frame describing the orientation of the data.
Revised: H-(V)AE is trained to reconstruct the spherical Fourier encoding of data, learning in the process a latent space with a maximally informative invariant embedding alongside an equivariant frame describing the orientation of the data.

Original: We extensively test the performance of H-(V)AE on diverse datasets and show that its latent space efficiently encodes the categorical features of spherical images and structural features of protein atomic environments.
Revised: We extensively test the performance of H-(V)AE on diverse datasets and show that its latent space efficiently encodes the categorical features of spherical images and structural features of protein atomic environments.

Original: Our work can further be seen as a case study for equivariant modeling of a data distribution by reconstructing its Fourier encoding.
Revised: Our work can further be seen as a case study for equivariant modeling of a data distribution by reconstructing its Fourier encoding.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: In supervised learning, the success of state-of-the-art algorithms is often attributed to respecting known inductive biases of the function they are trying to approximate.
Revised: In supervised learning, the success of state-of-the-art algorithms is often attributed to respecting known inductive biases of the function they are trying to approximate.

Original: One such bias is the invariance of the function to certain transformations of the input.
Revised: One such bias is the invariance of the function to certain transformations of the input.

Original: For example, image classification is translationally invariant.
Revised: For example, image classification is translationally invariant.

Original: To achieve such invariance, conventional techniques use data augmentation to train an algorithm on many transformed forms of the data.
Revised: To achieve such invariance, conventional techniques use data augmentation to train an algorithm on many transformed forms of the data.

Original: However, this solution is only approximate and increases training time significantly, up to prohibitive scales for high-dimensional and continuous transformations, e.g. rotations of an object in 3D. Alternatively, one could use invariant features of the data (e.g. pairwise distance between different features) as input to train any machine learning algorithm (Capecchi et al., 2020; Uhrin, 2021).
Revised: However, this solution is only approximate and increases training time significantly, up to prohibitive scales for high-dimensional and continuous transformations ( ∼ 500 augmentations are required to learn 3D rotation-invariant patterns (Geiger & Smidt, 2022)). Alternatively, one could use invariant features of the data (e.g. pairwise distance between different features) as input to train any machine learning algorithm (Capecchi et al., 2020; Uhrin, 2021).

Original: However, the choice of these invariants is arbitrary and the resulting network could lack in expressiveness.
Revised: However, the choice of these invariants is arbitrary and the resulting network could lack in expressiveness.

Original: Recent advances have developed neural network architectures that are equivariant under actions of different symmetry groups.
Revised: Recent advances have developed neural network architectures that are equivariant under actions of different symmetry groups.

Original: These networks can systematically treat and interpret various transformation in data, and learn models that are agnostic to these transformations.
Revised: These networks can systematically treat and interpret various transformation in data, and learn models that are agnostic to these transformations.

Original: For example, models equivariant to euclidean transformations have recently advanced the state-of-the-art on tasks over 3D point cloud data (Liao & Smidt, 2022; Musaelian et al., 2022; Brandstetter et al., 2022).
Revised: For example, models equivariant to euclidean transformations have recently advanced the state-of-the-art on tasks over 3D point cloud data (Liao & Smidt, 2022; Musaelian et al., 2022; Brandstetter et al., 2022).

Original: These models are more flexible and expressive compared to their purely invariant counterparts (Geiger &
Revised: These models are more flexible and expressive compared to their purely invariant counterparts (Geiger &

Original: Smidt, 2022), and exhibit high data efficiency.
Revised: Smidt, 2022), and exhibit high data efficiency.

Original: Extending such group invariant and equivariant paradigms to unsupervised learning could map out compact representations of data that are agnostic to a specified symmetry transformation (e.g. the global orientation of an object).
Revised: Extending such group invariant and equivariant paradigms to unsupervised learning could map out compact representations of data that are agnostic to a specified symmetry transformation (e.g. the global orientation of an object).

Original: In recent work Winter et al.
Revised: In recent work Winter et al.

Original: (2022) proposed a general mathematical framework for autoencoders that can be applied to data with arbitrary symmetry structures by learning an invariant latent space and an equivariant factor, related to the elements of the underlying symmetry group.
Revised: (2022) proposed a general mathematical framework for autoencoders that can be applied to data with arbitrary symmetry structures by learning an invariant latent space and an equivariant factor, related to the elements of the underlying symmetry group.

Original: Here, we focus on unsupervised learning that is equivariant to rotations around a specified origin in 3D, denoted by the group SO(3).
Revised: Here, we focus on unsupervised learning that is equivariant to rotations around a specified origin in 3D, denoted by the group SO(3).

Original: We encode the data in spherical Fourier space and construct holograms of the data that are conveniently structured for equivariant operations.
Revised: We encode the data in spherical Fourier space and construct holo- grams of the data that are conveniently structured for equivariant operations.

Original: These data holograms are inputs to our end-to-end SO(3)-equivariant (variational) autoencoder in spherical Fourier space, with a fully equivariant encoder-decoder architecture trained to reconstruct the Fourier coefficients of the input; we term this approach Holographic -(V)AE (H-(V)AE).
Revised: These data holograms are inputs to our end-to-end SO(3)-equivariant (variational) autoencoder in spherical Fourier space, with a fully equivariant encoder-decoder architecture trained to reconstruct the Fourier coefficients of the input; we term this approach Holographic -(V)AE (H-(V)AE).

Original: Similar to Winter et al.
Revised: Similar to Winter et al.

Original: , our network learns an SO(3)-equivariant latent space composed of a maximally informative set of invariants and an equivariant frame describing the orientation of the data.
Revised: , our network learns an SO(3)-equivariant latent space composed of a maximally informative set of invariants and an equivariant frame describing the orientation of the data.

Original: We extensively test the perfomance of H-(V)AE and demonstrate high accuracy in unsupervised classification and clustering tasks for spherical images and atomic point clouds within protein structures.
Revised: We extensively test the perfomance of H-(V)AE and demonstrate high accuracy in unsupervised classification and clustering tasks for spherical images and atomic point clouds within protein structures.

Original: The learned SO(3) invariant and equivariant representations would be useful for many real world applications in computer vision and structural biology.
Revised: The learned SO(3) invariant and equivariant representations would be useful for many real world applications in computer vision and structural biology.

Original:  B ACKGROUND 2. S PHERICAL HARMONICS AND IRREPS OF SO(3)
Revised:  B ACKGROUND 2. S PHERICAL HARMONICS AND IRREPS OF SO(3)

Original: We are interested in modeling  data that is distributed in 3D space and for which the global orientation of the data should not impact the inferred model (Einstein, 1916).
Revised: We are interested in modeling 3D data (i.e., functions in R 3 ), for which the global orientation of the data should not impact the inferred model (Einstein, 1916).

Original: We consider data that is distributed around a specified origin and can be expressed by the resulting spherical coordinates ( r, θ, ϕ ) around the origin.
Revised: We consider functions distributed around a specified origin, which we express by the resulting spherical coordinates ( r, θ, ϕ ) around the origin.

Original: In this case, the set of rotations about the origin define the 3D rotation group SO(3), and we will consider models that are rotationally equivariant under SO(3).
Revised: In this case, the set of rotations about the origin define the 3D rotation group SO(3), and we will consider models that are rotationally equivariant under SO(3).

Original: It is convenient to project data to spherical Fourier space to define equivariant transformations for rotations.
Revised: It is convenient to project data to spherical Fourier space to define equivariant transformations for rotations.

Original: To map a radially distributed function ρ ( r, θ, ϕ ) to a spherical Fourier space, we use the Zernike
Revised: To map a radially distributed function ρ ( r, θ, ϕ ) to a spherical Fourier space, we use the Zernike

Original: Fourier Transform (ZFT),
Revised: Fourier Transform (ZFT),

Original: [Equation]
Revised: [Equation]

Original: (1) where Y ℓm ( θ, ϕ ) is the spherical harmonics of degree ℓ and order m , where ℓ is a non-negative integer ( ℓ ≥ 0 ) and m is an integer within the interval − ℓ ≤ m ≤ ℓ .
Revised: where Y ℓm ( θ, ϕ ) is the spherical harmonics of degree ℓ and order m , where ℓ is a non-negative integer ( ℓ ≥ 0 ) and m is an integer within the interval − ℓ ≤ m ≤ ℓ .

Original: R nℓ ( r ) is the radial Zernicke polynomial in 3D (Eq.
Revised: R nℓ ( r ) is the radial Zernicke polynomial in 3D (Eq.

Original: A.7) with radial frequency n ≥ 0 and degree ℓ .
Revised: A.7) with radial frequency n ≥ 0 and degree ℓ .

Original: R nℓ ( r ) is non-zero only for even values of n − ℓ ≥ 0 .
Revised: R nℓ ( r ) is non-zero only for even values of n − ℓ ≥ 0 .

Original: Zernike polynomials - defined as the product
Revised: Zernike polynomials - defined as the product

Original: Y ℓm ( θ, ϕ ) R nℓ ( r ) - form a complete orthonormal basis in 3D, and therefore can be used to expand and retrieve 3D shapes, if large enough ℓ and n values are used; approximations that restrict the series to finite n and ℓ are often sufficient for shape retrieval, and hence, desirable algorithmically.
Revised: Y ℓm ( θ, ϕ ) R nℓ ( r ) - form a complete orthonormal basis in 3D, and therefore can be used to expand and retrieve 3D shapes, if large enough ℓ and n values are used; approximations that restrict the series to finite n and ℓ are often sufficient for shape retrieval, and hence, desirable algorithmically.

Original: Thus, in practice, we cap the resolution of the ZFT to a maximum degree L and a maximum radial frequency N .
Revised: Thus, in practice, we cap the resolution of the ZFT to a maximum degree L and a maximum radial frequency N .

Original: The operators that describe how spherical harmonics transform under rotations are called the Wigner
Revised: The operators that describe how spherical harmonics transform under rotations are called the Wigner

Original: D-matrices.
Revised: D-matrices.

Original: Notably, Wigner-D matrices are the irreducible representations (irreps) of SO(3), which implies that every element of the SO(3) group acting on any vector space can be represented as a direct sum of Wigner-D matrices.
Revised: Notably, Wigner-D matrices are the irreducible representations (irreps) of SO(3), which implies that every element of the SO(3) group acting on any vector space can be represented as a direct sum of Wigner-D matrices.

Original: As spherical harmonics form a basis for the irreps of SO(3), the SO(3) group acts on spherical
Revised: As spherical harmonics form a basis for the irreps of SO(3), the SO(3) group acts on spherical

Original: Fourier space via a direct sum of irreps.
Revised: Fourier space via a direct sum of irreps.

Original: Specifically, the ZFT encodes a data point into a tensor composed of a direct sum of features , each associated with a degree ℓ indicating the irrep that it transforms with under the action of SO(3).
Revised: Specifically, the ZFT encodes a data point into a tensor composed of a direct sum of features , each associated with a degree ℓ indicating the irrep that it transforms with under the action of SO(3).

Original: We refer to these tensors as SO(3)
Revised: We refer to these tensors as SO(3)

Original: -steerable tensors and to the vector spaces they occupy as SO(3) -steerable
Revised: -steerable tensors and to the vector spaces they occupy as SO(3) -steerable

Original: vector spaces , or simply steerable for short since we only deal with the SO(3) group in this work.
Revised: vector spaces , or simply steerable for short since we only deal with the SO(3) group in this work.

Original: We note that a tensor may contain multiple features of the same degree ℓ , which we generically refer to as distinct channels c .
Revised: We note that a tensor may contain multiple features of the same degree ℓ , which we generically refer to as distinct channels c .

Original: Throughout the paper, we refer to generic steerable tensors as h and index them by ℓ , m and c .
Revised: Throughout the paper, we refer to generic steerable tensors as h and index them by ℓ , m and c .

Original: We adopt the “hat” notation for individual entries to remind ourselves of the analogy with Fourier coefficients.
Revised: We adopt the “hat” notation for individual entries to remind ourselves of the analogy with Fourier coefficients.

Original: See
Revised: See

Original: Figure 1A for a graphical illustration of a tensor.
Revised: Figure 1A for a graphical illustration of a tensor.

Original: 2. M APPING BETWEEN SO(3)- STEERABLE VECTOR SPACES
Revised: 2. M APPING BETWEEN SO(3)- STEERABLE VECTOR SPACES

Original: Constructing equivariant operations equates to constructing maps between steerable vector spaces.
Revised: Constructing equivariant operations equates to constructing maps between steerable vector spaces.

Original: There are precise rules constraining the kinds of operations that guarantee a valid SO(3)-steerable output, the most important one being the Clebsch-Gordan (CG) tensor product ⊗ cg .
Revised: There are precise rules constraining the kinds of operations that guarantee a valid SO(3)-steerable output, the most important one being the Clebsch-Gordan (CG) tensor product ⊗ cg .

Original: The CG tensor product combines two features of degrees ℓ 1 and ℓ 2 to produce another feature of degree | ℓ 2 − ℓ 1 | ≤ ℓ 3 ≤ | ℓ 1 + ℓ 2 | .
Revised: The CG tensor product combines two features of degrees ℓ 1 and ℓ 2 to produce another feature of degree | ℓ 2 − ℓ 1 | ≤ ℓ 3 ≤ | ℓ 1 + ℓ 2 | .

Original: Let h ℓ ∈ R 2 ℓ +1 be a generic degree ℓ feature, with individual components ˆ h ℓm for − ℓ ≤ m ≤ ℓ .
Revised: Let h ℓ ∈ R 2 ℓ +1 be a generic degree ℓ feature, with individual components ˆ h ℓm for − ℓ ≤ m ≤ ℓ .

Original: Then, the CG tensor product is given by:
Revised: Then, the CG tensor product is given by:

Original: [Equation]
Revised: [Equation]

Original: ) where C ( ℓ 3 m 3 ) ( ℓ 1 m 1 )( ℓ 2 m 2 ) are the Clebsch-Gordan coefficients (Tung, 1985).
Revised: (2) where C ( ℓ 3 m 3 ) ( ℓ 1 m 1 )( ℓ 2 m 2 ) are the Clebsch-Gordan coefficients (Tung, 1985).

Original:  H OLOGRAPHIC -(V)AE 3. SO(3) EQUIVARIANT LAYERS
Revised:  H OLOGRAPHIC -(V)AE 3. SO(3) EQUIVARIANT LAYERS

Original: Linearity (Lin).
Revised: Linearity (Lin).

Original: We construct linear layers acting on steerable tensors by learning degree-specific linear operations.
Revised: We construct linear layers acting on steerable tensors by learning degree-specific linear operations.

Original: Specifically, we learn weight matrices specific to each degree ℓ , and use them to map across degree- ℓ feature spaces by learning linear combinations of degree- ℓ features in the input tensor (Section A.2.1).
Revised: Specifically, we learn weight matrices specific to each degree ℓ , and use them to map across degree- ℓ feature spaces by learning linear combinations of degree- ℓ features in the input tensor (Section A.2.1).

Original: Efficient Tensor Product nonlinearity (ETP).
Revised: Efficient Tensor Product nonlinearity (ETP).

Original: We utilize the CG tensor product to inject nonlinearity in the network in an equivariant way, as was originally prescribed by Kondor et al.
Revised: We utilize the CG tensor product to inject nonlinearity in the network in an equivariant way, as was originally prescribed by Kondor et al.

Original: (2018) for SO(3)-equivariant convolutional neural networks (CNNs).
Revised: (2018) for SO(3)-equivariant convolutional neural networks (CNNs).

Original: This type of nonlinearity enables information flow between features of different degrees, which is necessary for constructing expressive models, and injects square nonlinearity.
Revised: This type of nonlinearity enables information flow between features of different degrees, which is necessary for constructing expressive models, and injects square nonlinearity.

Original: To significantly reduce the computational and memory costs of the tensor products, we leverage some of the modifications proposed by Cobb et al.
Revised: To significantly reduce the computational and memory costs of the tensor products, we leverage some of the modifications proposed by Cobb et al.

Original: Specifically, we compute tensor products channel-wise, i.e., only between features belonging to the same channel, and we limit the connections between features of different degrees to what Cobb et al.
Revised: Specifically, we compute tensor products channel-wise, i.e., only between features belonging to the same channel, and we limit the connections between features of different degrees to what Cobb et al.

Original: (2021) calls the MST subset.
Revised: (2021) calls the MST subset.

Original: Notably, the channel-wise computation constrains the input tensors to have the same number of channels for each feature.
Revised: Notably, the channel-wise computation constrains the input tensors to have the same number of channels for each feature.

Original: We found these modifications to be necessary to efficiently work with data encoded in large number of channels C and with large maximum degree L .
Revised: We found these modifications to be necessary to efficiently work with data encoded in large number of channels C and with large maximum degree L .

Original: See Section A.2.2 for details.
Revised: See Section A.2.2 for details, and Table A.7 for an ablation study showing the improvement in parameter efficiency provided by the ETP.

Original: Batch Norm (BN).
Revised: Batch Norm (BN).

Original: We normalize intermediate tensor representations degree-wise and channel-wise by the batch-averaged norms of the features, as initially proposed by Kondor et al.
Revised: We normalize intermediate tensor representations degree-wise and channel-wise by the batch-averaged norms of the features, as initially proposed by Kondor et al.

Original: (2018), and do it before the ETP; see Figure 1B and Section A.2.3 for details.
Revised: (2018), and do it before the ETP; see Figure 1B and Section A.2.3 for details.

Original: We found the use of this layer to speed up model convergence (Fig A.1).
Revised: We found the use of this layer to speed up model convergence (Fig A.2).

Original: Signal Norm (SN).
Revised: Signal Norm (SN).

Original: It is necessary to normalize activations computed by the CG tensor product to avoid their explosion.
Revised: It is necessary to normalize activations computed by the CG tensor product to avoid their explosion.

Original: We found using batch norm alone often caused activations to explode in the decoder during evaluation.
Revised: We found using batch norm alone often caused activations to explode in the decoder during evaluation.

Original: Thus, we introduce Signal Norm, whereby we divide each steerable tensor by its total norm,  and apply a degree-specific affine transformation ( w ℓ ) for added flexibility.
Revised: Thus, we introduce Signal Norm, whereby we divide each steerable tensor by its total norm, defined as the sum of the norms of each of the tensor’s features, and apply a degree-specific affine transformation ( w ℓ ) for added flexibility.

Original: Formally, the total norm for an individual tensor h is computed as:
Revised: Formally, the total norm for an individual tensor h is computed as:

Original: [Equation]
Revised: [Equation]

Original: and the features are updated as ˆ h cℓm =
Revised: and the features are updated as ˆ h cℓm =

Original: ˆ h cℓm w ℓ / √ N tot .
Revised: ˆ h cℓm w ℓ / √ N tot .

Original: We note that each pre-affine normalized tensor has a total norm of 1 , thus constraining the values of the individual features.
Revised: We note that each pre-affine normalized tensor has a total norm of 1 , thus constraining the values of the individual features.

Revised: Signal Norm can be seen as a form of Layer Norm that respects SO(3) equivariance (Ba et al., 2016).

Original: Clebsch-Gordan block (CG bl.)
Revised: Clebsch-Gordan block (CG bl.)

Original: We construct equivariant blocks, which we term Clebsch-Gordan blocks, by stacking together the equivariant layers introduced above, as shown in Figure 1B.
Revised: We construct equivariant blocks, which we term Clebsch-Gordan blocks, by stacking together the equivariant layers introduced above, as shown in Figure 1B.

Original: Each block can take a steerable tensor of arbitrary size composed of multiple features of arbitrary degrees, and can output a steerable tensor of arbitrary size with multiple features of arbitrary degrees, by following the sparsity of the CG tensor product ( | ℓ 2 − ℓ 1 | ≤ ℓ 3 ≤ | ℓ 1 + ℓ 2 | ).
Revised: Each block can take a steerable tensor of arbitrary size composed of multiple features of arbitrary degrees, and can output a steerable tensor of arbitrary size with multiple features of arbitrary degrees, by following the sparsity of the CG tensor product ( | ℓ 2 − ℓ 1 | ≤ ℓ 3 ≤ | ℓ 1 + ℓ 2 | ).

Original: Crucially, if the maximum feature degree in the input tensor is ℓ max , then the maximum feature degree that can be generated is 2 ℓ max , achieved by combining two features of degree ℓ max .
Revised: Crucially, if the maximum feature degree in the input tensor is ℓ max , then the maximum feature degree that can be generated is 2 ℓ max , achieved by combining two features of degree ℓ max .

Original: We employ additive skip connections, zero-padded when appropriate, to favor better gradient flow.
Revised: We employ additive skip connections, zero-padded when appropriate, to favor better gradient flow.

Original: 3. D ATA NORMALIZATION
Revised: 3. D ATA NORMALIZATION

Original: As per standard machine learning practice (Shanker et al., 1996), we normalize the data.
Revised: As per standard machine learning practice (Shanker et al., 1996), we normalize the data.

Original: We do this by dividing each tensor by the average square-root total norm of the training tensors, analogously to the Signal Norm.
Revised: We do this by dividing each tensor by the average square-root total norm of the training tensors, analogously to the Signal Norm.

Original: This strategy puts the target values on a similar scale as the normalized activations learned by the network, which we speculate to favor gradient flow.
Revised: This strategy puts the target values on a similar scale as the normalized activations learned by the network, which we speculate to favor gradient flow.

Original: 3. M ODEL ARCHITECTURE
Revised: 3. M ODEL ARCHITECTURE

Revised: H-(V)AE has a fully rotationally equivariant architecture, and learns a disentangled latent space consisting of a maximally informative invariant ( ℓ = 0 )

Revised: component z of arbitrary size, as well as three orthonormal vectors ( ℓ = 1 ), which represent the global 3D orientation of the object and reflect the coordinate frame of the input tensor.

Revised: Crucially, the disentangled nature of the latent space is respected at all stages of training, and is guaranteed by the model’s rotational equivariance.

Original: H-(V)AE’s architecture is shown in Figure 1C.
Revised: The architecture is shown in Figure 1C.

Original: The encoder takes as input a steerable tensor with maximum degree ℓ max =
Revised: The encoder takes as input a steerable tensor with maximum degree ℓ max =

Original: L and, via a stack of Clebsch-Gordan blocks, iteratively and equivariantly transfers information from higher degrees to lower ones, down to the final encoder layer with ℓ max = 1 .
Revised: L and, via a stack of Clebsch-Gordan blocks, iteratively and equivariantly transfers information from higher degrees to lower ones, down to the final encoder layer with ℓ max = 1 .

Original: The latent space learned by the encoder consists of a maximally informative invariant ( ℓ = 0 ) component of arbitrary size z , as well as three orthonormal vectors ( ℓ = 1 ) which represent the global 3D orientations, reflecting the coordinate frame of the input tensor.

Original: The frame is constructed by learning two vectors and usingthe Gram-Schmidtmethod to find the corresponding orthonormal bases (Schmidt, 1907).
Revised: The frame is constructed by learning two vectors and using Gram-Schmidt to find the corresponding orthonormal basis (Schmidt, 1907).

Original: The third orthonomal basis vector is then calculated as the cross product of the first two.
Revised: The third orthonomal basis vector is then calculated as the cross product of the first two.

Original: The decoder learns to reconstruct the input from the invariants and the frame, iteratively increasing the maximum degree ℓ max of the intermediate representations by leveraging the CG Tensor Product within the Clebsch-Gordan blocks.
Revised: The decoder learns to reconstruct the input from z and the frame, iteratively increasing the maximum degree ℓ max of the intermediate representations by leveraging the CG Tensor Product within the Clebsch-Gordan blocks.

Original: An interesting question to ask is: what does the trained decoder’s output look like if the frame is held constant (e.g. equal to the identity matrix)?
Revised: An interesting question to ask is: what does the trained decoder’s output look like if the frame is held constant (e.g. equal to the identity matrix)?

Original: We experimentally find that the reconstructed elements tend to be aligned with each other and hypothesize that the model is implicitly learning to maximize the overlap between training elements, providing empirical evidence in the
Revised: We experimentally find that the reconstructed elements tend to be aligned with each other and hypothesize that the model is implicitly learning to maximize the overlap between training elements, providing empirical evidence in the

Original: Appendix (Fig. A.2).
Revised: Appendix (Fig. A.3).

Original: We call this frame the canonical frame following the analogy with the canonical elements in Winter et al.
Revised: We call this frame the canonical frame following the analogy with the canonical elements in Winter et al.

Original: We note that it is possible to rotate original elements to the canonical frame thanks to the equivalence between the frame we learn and the rotation matrix within our implementation.
Revised: We note that it is possible to rotate original elements to the canonical frame thanks to the equivalence between the frame we learn and the rotation matrix within our implementation.

Original: Within the decoder, the maximum degree ℓ max ,b that can be outputted by each block b is constrained by the sparsity of the CG tensor product.
Revised: Within the decoder, the maximum degree ℓ max ,b that can be outputted by each block b is constrained by the sparsity of the CG tensor product.

Original: Specifically, ℓ max ,b ≤ 2 b where b ranges from 1 (first block) to B (last block).
Revised: Specifically, ℓ max ,b ≤ 2 b where b ranges from 1 (first block) to B (last block).

Original: Since we need to reconstruct features up to degree L in the decoder, we arrive at a lower bound for the number of blocks in the decoder set by ℓ max ,B ≥ L , or B ≥ log 2 L .
Revised: Since we need to reconstruct features up to degree L in the decoder, we arrive at a lower bound for the number of blocks in the decoder set by ℓ max ,B ≥ L , or B ≥ log 2 L .

Original: In our experiments, we set ℓ max ,b = min { 2 b , L } and do not let ℓ max ,b exceed the input’s maximum degree L .
Revised: In our experiments, we set ℓ max ,b = min { 2 b , L } and do not let ℓ max ,b exceed the input’s maximum degree L .

Original: Relaxing this condition might increase the expressive power of the network but at a significant increase in runtime and memory.
Revised: Relaxing this condition might increase the expressive power of the network but at a significant increase in runtime and memory.

Original: We leave the analysis of this trade-off to future work.
Revised: We leave the analysis of this trade-off to future work.

Original: For the encoder and the decoder to have similar expressive power, we construct them to be symmetric with respect to the latent space (Fig. 1C).
Revised: For the encoder and the decoder to have similar expressive power, we construct them to be symmetric with respect to the latent space (Fig. 1C).

Original: Optionally, we apply a linearity at the beginning of the encoder and at the end of the decoder; this is required for input data that does not have the same number of channels per degree since the ETP operates channel-wise.
Revised: Optionally, we apply a linearity at the beginning of the encoder and at the end of the decoder; this is required for input data that does not have the same number of channels per degree since the ETP operates channel-wise.

Revised: We empirically verify the equivariance of our model up to numerical errors in Table A.8.

Original: For H-VAE, we parameterize the invariant part of the latent space by an isotropic Gaussian, i.e., we learn two sets of size z invariants, corresponding to means and standard deviations.
Revised: For H-VAE, we parameterize the invariant part of the latent space by an isotropic Gaussian, i.e., we learn two sets of size z invariants, corresponding to means and standard deviations.

Original: 3. T RAINING OBJECTIVE
Revised: 3. T RAINING OBJECTIVE

Original: We train H-(V)AE to minimize the reconstruction loss L rec between the original and the reconstructed tensors, and, for H-VAE only, to minimize the KL-divergence of the posterior invariant latent space distribution q ( z | x ) from the selected prior p ( z ) (Kingma & Welling, 2013):
Revised: We train H-(V)AE to minimize the reconstruction loss L rec between the original and the reconstructed tensors, and, for H-VAE only, to minimize the KL-divergence of the posterior invariant latent space distribution q ( z | x ) from the selected prior p ( z ) (Kingma & Welling, 2013):

Original: [Equation]
Revised: [Equation]

Original: We use mean square error (MSE) for L rec , which as we show in Sec.
Revised: We use mean square error (MSE) for L rec , which as we show in Sec.

Original: A.2.4, respects the necessary property of SO(3) pairwise invariance, ensuring that the model remains rotationally equivariant.
Revised: A.2.4, respects the necessary property of SO(3) pairwise invariance, ensuring that the model remains rotationally equivariant.

Original: We cast an isotropic normal prior to the invariant latent space: p ( z ) =

Original: N ( 0 , I ) .

Original: Hyperparameters α and β control the trade-off between reconstruction accuracy and latent space regularization Higgins et al.

Original: We find it practical to scale the reconstruction loss by a dataset-specific scalar α since the MSE loss varies in average magnitude across datasets.

Original: When training H-VAE, we find it beneficial to keep β = 0 for a few epochs ( E rec ) so that the model can learn to perform meaningful reconstructions, and then linearly increasing it to the desired value for

Original: E warmup epochs to add structure to the latent space, an approach first used by Bowman et al.

Original: ℎ" !!"
Revised: ℎ" !!"

Original: ℎ" #$#" ℎ" #!"
Revised: ℎ" #$#" ℎ" #!"

Original: ℎ" ##" ℎ" !!
Revised: ℎ" ##" ℎ" !!

Original: % ℎ" #$#% ℎ" #!%
Revised: % ℎ" #$#% ℎ" #!%

Original: ℎ" ##% ℎ" !!
Revised: ℎ" ##% ℎ" !!

Original: & ℎ" #$#& ℎ" #!& ℎ" ##& ℎ" !!
Revised: & ℎ" #$#& ℎ" #!& ℎ" ##& ℎ" !!

Original: # ℎ" #$## ℎ" #!
Revised: # ℎ" #$## ℎ" #!

Original: # ℎ" ### 4 Channels ℓ =ℓ =
Revised: # ℎ" ### 4 Channels ℓ =ℓ =

Original: A
Revised: A

Original: CGbl.
Revised: CGbl.

Original: Gram-Schmidt
Revised: Gram-Schmidt

Original: Input tensor ℓ !"
Revised: Input tensor ℓ !"

Original: # = 𝐿
Revised: # = 𝐿

Original: Lin ℓ = Lin ℓ =
Revised: Lin ℓ = Lin ℓ =

Original: Lin ℓ = Lin ℓ =
Revised: Lin ℓ = Lin ℓ =

Original: Reconstructed tensor ℓ !"
Revised: Reconstructed tensor ℓ !"

Original: # = 𝐿
Revised: # = 𝐿

Original: Lin
Revised: Lin

Original: CGbl. CGbl.
Revised: CGbl. CGbl.

Original: CGbl.
Revised: CGbl.

Original: Lin
Revised: Lin

Original: Encoder
Revised: Encoder

Original: Decoder Latent space
Revised: Decoder Latent space

Original: [Figure]
Revised: [Figure]

Original: B
Revised: B

Original: Clebsch-Gordanblock
Revised: Clebsch-Gordanblock

Original: BN
Revised: BN

Original: ETP
Revised: ETP

Original: SN
Revised: SN

Original: Lin ⊕
Revised: Lin ⊕

Original: 3. R ECONSTRUCTION ASSESSMENT VIA COSINE LOSS
Revised: We cast an isotropic normal prior to the invariant latent space: p ( z ) = N ( Hyperparameters α and β control the trade-off between reconstruction accuracy and latent space regularization Higgins et al. We find it practical to scale the reconstruction loss by a dataset-specific scalar α since the MSE loss varies in average magnitude across datasets.

Revised: When training H-VAE, we find it beneficial to keep β = 0 for a few epochs ( E rec ) so that the model can learn to perform meaningful reconstructions, and then linearly increasing it to the desired value for

Revised: E warmup epochs to add structure to the latent space, an approach first used by Bowman et al.

Original: As  MSE values are on different scales across tensors with features of different degrees, it is not possible to use MSE as a universal measure for “goodness” of the reconstructions.
Revised: As the scale of MSE depends on the characteristics of the data, e.g. the size of the tensors representing the data and their irreps (Fig. A.4), it is difficult to contextualize MSE values across datasets. It would be desirable to have a dimensionless metric that measures absolute “goodness” of reconstructions that is comparable across datasets.

Original: Instead, we propose the metric Cosine loss, which is a normalized dot product generalized to operate on pairs of steerable tensors (akin to cosine similarity), and modified to be interpreted as a loss:[Equation] with
Revised: For this purpose, we propose the metric Cosine loss which is a normalized dot product generalized to operate on pairs of steerable tensors (akin to cosine similarity), and modified to be interpreted as a loss:

Original: [Equation]
Revised: [Equation]

Original: The Cosine loss is interpretable across different datasets: it is pairwise invariant, has a minimum of zero for perfect reconstructions, and has an average value of 1.0 for a pair of random tensors of any size, with a smaller variance for larger tensors (Sec. A.2.5).
Revised: The Cosine loss is interpretable across different datasets: it is pairwise invariant, has a minimum of zero for perfect reconstructions, and has an average value of 1.0 for a pair of random tensors of any size, with a smaller variance for larger tensors (Sec. A.4).

Original: The interpretability of Cosine loss comes at the price of ignoring the relative norms of the features that are being compared, making the measure unable to reconstruct norms and thus not suitable as a training objective.
Revised: The interpretability of Cosine loss comes at the price of ignoring the relative norms of the features that are being compared, making the measure unable to reconstruct norms and thus not suitable as a training objective.

Original: However, as norms are easier to reconstruct than directions, we still find the Cosine loss useful as a noisy estimate of the model’s reconstruction ability.
Revised: However, as norms are easier to reconstruct than directions, we still find the Cosine loss useful as a noisy estimate of the model’s reconstruction ability.

Original: Furthermore, Cosine loss tends to correlate well with MSE, especially in the mid-to-low reconstruction quality regime (SpearmanR = 0 .
Revised: Furthermore, Cosine loss correlates almost perfectly with MSE, especially in the mid-to-low reconstruction quality regime (SpearmanR = 0 .

Original: 99 , Fig.
Revised: 99 , Fig.

Original: A.3).
Revised: A. and Table A.6).

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

Original: Group-equivariant neural networks.
Revised: Group-equivariant neural networks.

Original: Group-equivariant neural networks have improved the stateof-the-art on many supervised tasks, thanks to their data efficiency (Hutchinson et al., 2021;
Revised: Group-equivariant neural networks have improved the stateof-the-art on many supervised tasks, thanks to their data efficiency (Hutchinson et al., 2021;

Original: Bekkers et al., 2018; Romero & Cordonnier, 2021).
Revised: Bekkers et al., 2018; Romero & Cordonnier, 2021).

Original: Related to our work are 3D Euclidean neural networks, which are equivariant to (subsets of) the 3D Euclidean group: (Weiler et al., 2018;
Revised: Related to our work are 3D Euclidean neural networks, which are equivariant to (subsets of) the 3D Euclidean group: (Weiler et al., 2018;

Original: Brandstetter et al., 2022; Thomas et al., 2018; Batzner et al., 2022; Musaelian et al., 2022; Satorras et al., 2022; Fuchs et al., 2020; Liao & Smidt, 2022) and often use spherical harmonics and tensor products to construct
Revised: Brandstetter et al., 2022; Thomas et al., 2018; Batzner et al., 2022; Musaelian et al., 2022; Satorras et al., 2022; Fuchs et al., 2020; Liao & Smidt, 2022) and often use spherical harmonics and tensor products to construct

Original: SO(3) equivariant layers.
Revised: SO(3) equivariant layers.

Original: Seminal work on SO(3) equivariance has been conducted for spherical images (Cohen et al., 2018; Esteves et al., 2020); we were inspired by the fully Fourier approach of Kondor et al.
Revised: Seminal work on SO(3) equivariance has been conducted for spherical images (Cohen et al., 2018; Esteves et al., 2020); we were inspired by the fully Fourier approach of Kondor et al.

Original: (2018), and leveraged operations proposed by Cobb et al.
Revised: (2018), and leveraged operations proposed by Cobb et al.

Original: Equivariant representations of atomic systems.
Revised: Equivariant representations of atomic systems.

Original: There is a diverse body of literature on constructing representations of atomic systems that are invariant/equivariant to euclidean symmetries, lever- aging Fourier transforms and CG tensor products (Drautz, 2019; Musil et al., 2021).
Revised: There is a diverse body of literature on constructing representations of atomic systems that are invariant/equivariant to euclidean symmetries, leveraging Fourier transforms and CG tensor products (Drautz, 2019; Musil et al., 2021).

Original: Notably, Uhrin (2021) constructs SO(3)-invariant representations of point clouds using the ZFT and CG tensor product iterations.
Revised: Notably, Uhrin (2021) constructs SO(3)-invariant representations of point clouds using the ZFT and CG tensor product iterations.

Original: Our work can be seen as a data-driven instance of this framework, where we learn a compact invariant and equivariant latent space from data.
Revised: Our work can be seen as a data-driven instance of this framework, where we learn a compact invariant and equivariant latent space from data.

Original: Invariant autoencoders.
Revised: Invariant autoencoders.

Original: Several works attempt to learn representations that are invariant to certain classes of transformations.
Revised: Several works attempt to learn representations that are invariant to certain classes of transformations.

Original: Shu et al.
Revised: Shu et al.

Original: (2018) and Koneripalli et al.
Revised: (2018) and Koneripalli et al.

Original: (2020) learn general “shape” embeddings by learning a separate “deformation” embedding.
Revised: (2020) learn general “shape” embeddings by learning a separate “deformation” embedding.

Original: However, their networks are not explicitly equivariant to the transformations of interest.
Revised: However, their networks are not explicitly equivariant to the transformations of interest.

Original: Other work proposes to learn an exactly invariant embedding alongside an approximate (but not equivariant) group action to align the input and the reconstructed data.
Revised: Other work proposes to learn an exactly invariant embedding alongside an approximate (but not equivariant) group action to align the input and the reconstructed data.

Original: For example, Mehr et al.
Revised: For example, Mehr et al.

Original: (2018) learns in quotient space by sampling the group’s orbit and computing the reconstruction loss by taking the infimum over the group.
Revised: (2018) learns in quotient space by sampling the group’s orbit and computing the reconstruction loss by taking the infimum over the group.

Original: This approach is best suited for discrete and finite groups, and it is computationally expensive as it is akin to data augmentation.
Revised: This approach is best suited for discrete and finite groups, and it is computationally expensive as it is akin to data augmentation.

Original: Lohit & Trivedi (2020) construct an SO(3)-invariant autoencoder for spherical signals by learning an invariant latent space and minimizing a loss which first finds the rotation that best aligns the true and reconstructed signals, introducing an added optimization step and reconstructing the signals only up to a global rotation.
Revised: Lohit & Trivedi (2020) construct an SO(3)-invariant autoencoder for spherical signals by learning an invariant latent space and minimizing a loss which first finds the rotation that best aligns the true and reconstructed signals, introducing an added optimization steppotentially very expensive for 3D data - and reconstructing the signals only up to a global rotation.

Revised: To our knowledge, the method proposed by Lohit & Trivedi (2020) is the only approach to date for unsupervised learning of non-discretized SO(3)-invariant representations.

Original: In contrast, our approach  is fully equivariant, only requiring simple MSE for reconstruction, and reconstructing data in its original orientation.
Revised: However, the rotational invariance is manually imposed by Lohit & Trivedi (2020), which is distinct from our approach that is fully equivariant and only requires simple MSE for reconstruction of data in its original orientation.

Original: Group-equivariant autoencoders.
Revised: Group-equivariant autoencoders.

Original: A small body of work focuses on developing equivariant autoencoders.
Revised: A small body of work focuses on developing equivariant autoencoders.

Original: Several methods construct data and group-specific architectures to auto-encode data equivariantly, learning an equivariant representation in the process (Hinton et al., 2011; Kosiorek et al., 2019).
Revised: Several methods construct data and group-specific architectures to auto-encode data equivariantly, learning an equivariant representation in the process (Hinton et al., 2011; Kosiorek et al., 2019).

Original: Others use supervision to extract class-invariant and class-equivariant representations (Feige, 2022).
Revised: Others use supervision to extract class-invariant and class-equivariant representations (Feige, 2022).

Original: A recent theoretical work proposes to train an encoder that encodes elements into an invariant embedding and an equivariant group action, then using a standard decoder that uses the invariants to reconstruct the elements in a canonical form, and finally applying the learned group action to recover the data’s original form (Winter et al., 2022).
Revised: A recent theoretical work proposes to train an encoder that encodes elements into an invariant embedding and an equivariant group action, then using a standard decoder that uses the invariants to reconstruct the elements in a canonical form, and finally applying the learned group action to recover the data’s original form (Winter et al., 2022).

Original: Our method in SO(3) is closely related to this work, with the crucial difference that we use an equivariant decoder and that we learn to reconstruct the Fourier encoding of data.
Revised: Our method in SO(3) is closely related to this work, with the crucial difference that we use an equivariant decoder and that we learn to reconstruct the Fourier encoding of data.

Original: Furthermore, our model is variational in the invariant latent space.
Revised: Furthermore, our model is variational in the invariant latent space.

Original:  E XPERIMENTS 5. R OTATED MNIST ON THE SPHERE
Revised:  E XPERIMENTS 5. R OTATED MNIST ON THE SPHERE

Original: We extensively test the performance of H-(V)AE on the MNIST-on-the-sphere dataset (Deng, 2012).
Revised: We extensively test the performance of H-(V)AE on the MNIST-on-the-sphere dataset (Deng, 2012).

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

Original: we generate a discrete unit sphere using the Driscoll-Healey (DH) method with a bandwidth (bw) of 30, and project the MNIST dataset onto the lower hemisphere.
Revised: we generate a discrete unit sphere using the Driscoll-Healey (DH) method with a bandwidth (bw) of 30, and project the MNIST dataset onto the lower hemisphere.

Original: We consider two variants, NR/R and R/R, differing in whether the training/test images have been randomly rotated (R) or not (NR).
Revised: We consider two variants, NR/R and R/R, differing in whether the training/test images have been randomly rotated (R) or not (NR).

Original: For each dataset, we map the images to steerable tensors via the
Revised: For each dataset, we map the images to steerable tensors via the

Original: Zernike Fourier Transform (ZFT) with L = 10 , and a constant radial function R n ℓ = 1 , resulting in tensors with 121 coefficients. 
Revised: Zernike Fourier Transform (ZFT) with L = 10 , and a constant radial function R nℓ = 1 , resulting in tensors with 121 coefficients. [Figure]

Original: We train 8 models with combinations of the following features: training mode (NR vs. R), invariant latent space size z (16 vs. 120), and model type (AE vs. VAE).
Revised: We train 8 models with combinations of the following features: training mode (NR vs. R), invariant latent space size z (16 vs. 120), and model type (AE vs. VAE).

Original: In all cases the model architecture follows from Fig.
Revised: In all cases the model architecture follows from Fig.

Original: 1C; see Section A.4.2 for details.
Revised: 1C; see Section A.6.2 for details.

Original: All 8 models achieve very low reconstruction loss (Table 1) with no significant difference between training modes, indicating that the models successfully leverage SO(3)-equivariance to generalize to unseen orientations.
Revised: All 8 models achieve very low reconstruction loss (Table 1) with no significant difference between training modes, indicating that the models successfully leverage SO(3)-equivariance to generalize to unseen orientations.

Original: Predictably, AE models have lower reconstruction loss than VAE models, and so do models with a larger latent space.
Revised: Predictably, AE models have lower reconstruction loss than VAE models, and so do models with a larger latent space.

Original: Nonetheless, H-VAE achieves reliable reconstructions, as shown in Fig.
Revised: Nonetheless, H-VAE achieves reliable reconstructions, as shown in Fig.

Original: 2A and
Revised: 2A and

Original: Table 1.
Revised: Table 1.

Original: All 8 models produce an invariant latent space that naturally clusters by digit identity.
Revised: All 8 models produce an invariant latent space that naturally clusters by digit identity.

Original: We show this qualitatively for one of the models in Fig.
Revised: We show this qualitatively for one of the models in Fig.

Original: 2B, and quantitatively by clustering the data via K-Means with 10 centroids and computing standard clustering metrics of Purity (Aldenderfer &
Revised: 2B, and quantitatively by clustering the data via K-Means with 10 centroids and computing standard clustering metrics of Purity (Aldenderfer &

Original: Blashfield, 1984) and V-measure (Rosenberg & Hirschberg, 2007) in Table 1.
Revised: Blashfield, 1984) and V-measure (Rosenberg & Hirschberg, 2007) in Table 1.

Original: All 8 models achieve much better clustering metrics than Rot-Inv AE (Lohit & Trivedi, 2020), with VAE models consistently outperforming AE models.
Revised: All 8 models achieve much better clustering metrics than Rot-Inv AE (Lohit & Trivedi, 2020), with VAE models consistently outperforming AE models.

Original: We also train a linear classifier (LC) to predict digit identity from invariant latent space descriptors, achieving comparable accuracy to Rot-Inv AE with the same latent space size.
Revised: We also train a linear classifier (LC) to predict digit identity from invariant latent space descriptors, achieving comparable accuracy to Rot-Inv AE with the same latent space size.

Original: We do not observe any difference between VAE[Figure] and AE models in terms of classification accuracy.
Revised: We do not observe any difference between VAE and AE models in terms of classification accuracy.

Original: Using a KNN classifier instead of LC further improves performance (Table A.1).
Revised: Using a KNN classifier instead of LC further improves performance (Table A.2).

Original: As H-VAE is a generative model, we generate random spherical images by sampling invariant latent embeddings from the prior distribution, and observing diversity in digit type and style (Fig.
Revised: As H-VAE is a generative model, we generate random spherical images by sampling invariant latent embeddings from the prior distribution, and observing diversity in digit type and style (Fig.

Original: 2C and
Revised: 2C and

Original: Fig.
Revised: Fig.

Original: A.4).
Revised: A.5).

Original: We further assess the quality of the invariant latent space by generating images via linear interpolation of the invariant embeddings associated with two test images.
Revised: We further assess the quality of the invariant latent space by generating images via linear interpolation of the invariant embeddings associated with two test images.

Original: The interpolated images present spatially consistent transitions (Fig.
Revised: The interpolated images present spatially consistent transitions (Fig.

Original: 2D and Fig.
Revised: 2D and Fig.

Original: A.5), which is a sign of a semantically well-structured latent space.
Revised: A.6), which is a sign of a semantically well-structured latent space.

Original: To understand the meaning of the learned frames, we visualize the sum of images in the canonical frame (i.e. the identity matrix).
Revised: To understand the meaning of the learned frames, we visualize the sum of images in the canonical frame (i.e. the identity matrix).

Original: We hypothesize that H-(V)AE learns to optimally overlap the training images when in the same frame.
Revised: We hypothesize that H-(V)AE learns to optimally overlap the training images when in the same frame.

Original: Indeed, by visualizing the sum of the training images in the canonical frame, we can verify that images are well aligned within the same digit type and with varying degrees across different digit types depending on the content of training data (Fig. A.2). 5. S
Revised: Indeed, by visualizing the sum of the training images in the canonical frame, we can verify that images are well aligned within the same digit type and with varying degrees across different digit types depending on the content of training data (Fig. A.3).

Original: [Equation]
Revised: [Table]

Original: The Shrec17 dataset consists of 51k  3D models belonging to 55 object classes, with a 70/10/ train/valid/test split (noa).
Revised: The Shrec17 dataset consists of 51k colorless 3D models belonging to 55 object classes, with a 70/10/20 train/valid/test split (noa).

Original: We use the variant of the dataset where each datapoint is perturbed by random rotations.
Revised: We use the variant of the dataset where each model is perturbed by random rotations.

Revised: Converting 3D shapes into spherical images preserves topological surface information, while significantly simplifying the representation.

Original: We follow Cohen et al.
Revised: We follow Cohen et al.

Original: and project the data onto a DH spherical grid with a bandwidth of 90 via ray-casting, generating spherical images with 6 channels.
Revised: and project surface information from each model onto an enclosing DH spherical grid with a bandwidth of 90 via a ray-casting scheme, generating spherical images with 6 channels.

Original: We then apply the
Revised: We then apply the

Original: ZFT with L = 14 and a constant radial function R nℓ = 1 to each channel individually, resulting in a tensor with 1350 coefficients.
Revised: ZFT with L = 14 and a constant radial function R nℓ = 1 to each channel individually, resulting in a tensor with 1350 coefficients.

Original: We train an AE and a VAE model (Sec. A.4.3) and evaluate them similarly to the MNIST dataset and compute the Shrec17 retrieval metrics via the latent space linear classifier’s predictions (Table 1).
Revised: We train an AE and a VAE model (Sec. A.6.3) and evaluate them similarly to the MNIST dataset and compute the Shrec17 retrieval metrics via the latent space linear classifier’s predictions.

Original: H-AE achieves the best classification and retrieval results for autoencoder-based models (Lohit & Trivedi, 2020), and is competitive with supervised models(Esteves et al., 2020; Cobb et al., 2021) despite the lower grid bandwidth and the small latent space.
Revised: H-AE achieves the best classification and retrieval results for autoencoderbased models, and is competitive with supervised models despite the lower grid bandwidth and the small latent space (Table 1).

Original: Using KNN classification instead of a linear classifier further improves performance (Table A.2).
Revised: Using KNN classification instead of a linear classifier further improves performance (Table A.3).

Original: H-VAE achieves slightly worse classification results but better clustering metrics compared to H-AE.
Revised: H-VAE achieves slightly worse classification results but better clustering metrics compared to H-AE.

Original: While reconstruction loss is low, there is still significant margin of improvement.
Revised: While reconstruction loss is low, there is still significant margin of improvement.

Original: We partially attribute Lohit & Trivedi (2020)’s low scores to the low grid bandwidth.
Revised: We partially attribute Lohit & Trivedi (2020)’s low scores to the low grid bandwidth.

Original: However, we note that, unlike Lohit & Trivedi [Table] [Figure] (2020), the size and runtime of our method does not scale with grid bandwidth, since the size of the reconstructed tensor learned by our method does not depend on it.
Revised: However, we note that the size and runtime of our method does not scale with grid bandwidth, since the size of the reconstructed tensor learned by our method does not depend on it.

Original: 5. T OY AMINO ACIDS
Revised: 5. P ROTEIN NEIGHBORHOODS

Original: We train H-(V)AE on point clouds of amino acids extracted from structures in the Protein Data Bank (PDB) (Berman et al., 2000).

Original: We collect atomic point clouds of 50k residues from PDB evenly distributed across residue types and apply a 40/40/20 train/valid/test split.

Original: Residues of the same type have different conformations and naturally have noisy coordinates, making this problem a natural benchmark for our method.

Original: We consider atom-type-specific clouds (C, O, N and S; we exclude H) centered at the residue’s C α and compute the ZFT (Eq. 1) with L = 4 and N = 20 within a radius of 10 ˚A from the residue’s C α , and concatenate features of the same degree, resulting in a tensor with 940 coefficients. We train several H-AE and H-VAE models, all with z = 2 ; see Sec. A.4.4 for details.
Revised: A

Original: We consistently find that the latent space clusters by amino acid conformations (Fig. 3), with sharper cluster separations as more training data is added (Fig. A. and A.8).

Original: We find that test reconstruction loss decreases with more training data but the reconstruction is accurate even with little training data (from 0.153 Cosine loss with 400 training residues to 0.034 with 20,000); Table A.3. A similar trend is observed for KNN-based classification accuracy of residues (from 0.842 with 400 training residues to 0.972 with 20,000); (Table A.3). Notably, an untrained model, while achieving random reconstruction loss, still produces an informative invariant latent space (0.629 residue type accuracy), suggesting that the forced SO(3)-invariance grants a “warm start” to the encoder. We do not find significant improvements in latent space classification by training with a variational objective, and present ablation results in Table A.4.
Revised: B

Original: [Figure] We test H-(V)AE on amore challenging point cloud dataset, comprised of spherical atomic environments surrounding a residue within a protein structure, which we term protein neighborhoods .
Revised: We test H-(V)AE on a challenging point cloud dataset, comprised of spherical atomic environments surrounding a residue within a protein structure, which we term protein neighborhoods .

Original: We use
Revised: We use

Original: ProteinNet (AlQuraishi, 2019) splits to avoid any redundancy in training and test sets (Sec. A.4.5).
Revised: ProteinNet (AlQuraishi, 2019) splits to avoid any redundancy in training and test sets (Sec. A.6.5).

Original: We consider all atomic neighborhoods within a radius of 12.5 ˚A surrounding central residue types CYS,
Revised: We consider all atomic neighborhoods within a radius of 12.5 ˚A surrounding central residue types CYS,

Original: GLU or HIS (Fig. 4A).
Revised: GLU or HIS (Fig. 3A).

Original: We only consider backbone atoms (C α , C, O and N).
Revised: We only consider backbone atoms (C α , C, O and N).

Original: For each neighborhood centered at the central residue’s
Revised: For each neighborhood centered at the central residue’s

Original: Cα , we compute the ZFT (Eq. 1) with L = 4 and N = 26 .
Revised: Cα , we compute the ZFT (Eq. 1) with L = 4 and N = 26 .

Original: We then concatenate features of the same degree resulting in a tensor with 1240 coefficients.
Revised: We then concatenate features of the same degree resulting in a tensor with 1240 coefficients.

Original: The final dataset contains 303k/75k/4.4k (train/valid/test) tensors.
Revised: The final dataset contains 303k/75k/4.4k (train/valid/test) tensors.

Original: We train an H-AE model with an invariant latent space size z of 64 (see Sec. A.4.5).
Revised: We train an H- AE model with an invariant latent space size z of 64 (see Sec. A.6.5).

Original: While the test Cosine loss has still margins of improvement ( ∼ 0 . 161 ) the latent space is meaningfully organized according to relevant neighborhood statistics, such as the secondary structure of the central residue ( ∼ 87% KNN classification accuracy - Fig.
Revised: While the test Cosine loss has still margins of improvement ( ∼ 0 . 161 ) the latent space is meaningfully organized according to relevant neighborhood statistics, such as the secondary structure of the central residue ( ∼ 87% KNN classification accuracy - Fig.

Original: 4B), number of atoms in the neighborhood (Fig. 
Revised: 3B), number of atoms in the neighborhood (Fig. A.10), and neighborhood’s average Solvent Accessible Surface Area (SASA), which is a proxy value for how buried the residue is in the protein (Fig. A.11). We further benchmark H-(V)AE via extensive ablation experiments on a simpler point cloud dataset consisting of individual amino acids (Sec. A.3).  U SING AN UNCONSTRAINED DECODER It is possible to merge use our SO(3)-equivariant encoder and fourier-space formulation with Winter et al.

Original: A.9), and neighborhood’s average Solvent Accessible Surface Area (SASA), whichis a proxy value for how buried the residue is in the protein (Fig.
Revised: (2022)’s framework, in which the learned rotation (i.e. the frame) is applied to the output space of an unconstrained decoder instead of being fed as input to an equivariant decoder.

Original: A.10).
Revised: The resulting model is theoretically equivalent to ours with regards to respecting equivariance between input and output. We implement and test a version of this model on the Rotated MNIST-on-the-sphere dataset, and show that it is slightly worse but comparable to using our fully-equivariant decoder in terms of speed (1.3x slower) and performance (0.037 vs. 0.025 Cosine loss, Table A.1). Given the similarity in performance, we favor the simplicity and elegance of our equivariant decoder, which is constructed to be symmetric to the encoder, thus automatically endowing it with similar representational power and without the need to tune an architecture made with different base components. Implementation details and extended discussion on this matter can be found in Section A.4.

Original:  P ROPERTIES AND L IMITATIONS
Revised:  P ROPERTIES AND L IMITATIONS

Original: The information content of the latent space is limited by the maximum degree L and radial frequency
Revised: The information content of the latent space is limited by the maximum degree L and radial frequency

Original: N set as cutoff for the spherical Fourier series.
Revised: N set as cutoff for the spherical Fourier series.

Original: Such truncation is necessary for computational feasibility, but it limits resolution, evident when reconstructing a signal.
Revised: Such truncation is necessary for computational feasibility, but it limits resolution, evident when reconstructing a signal.

Original: Another related limitation is the higher inaccuracy of H-(V)AE in reconstructing features associated with larger degrees (Fig. A.11), resulting in a failure to learn fine-grained details of data.
Revised: Another related limitation is the higher inaccuracy of H-(V)AE in reconstructing features associated with larger degrees (Fig. A.12), resulting in a failure to learn fine-grained details of data.

Original: This failure could explain the network’s troubles on MNIST in differentiating between 4’s and 9’s, which tend to be hand-drawn similarly.
Revised: This failure could explain the network’s troubles on MNIST in differentiating between 4’s and 9’s, which tend to be hand-drawn similarly.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: In this work, we develop the first end-to-end SO(3)-equivariant (V)AE, suitable for data distributed around a center.
Revised: In this work, we develop the first end-to-end SO(3)-equivariant (V)AE, suitable for data distributed around a center.

Original: The model learns an invariant embedding describing the data in a “canonical” orientation alongside an equivariant frame describing the data’s original orientation relative to the canonical one.
Revised: The model learns an invariant embedding describing the data in a “canonical” orientation alongside an equivariant frame describing the data’s original orientation relative to the canonical one.

Original: We use the learned invariants to achieve state-of-the-art unsupervised clustering and classification results on various spherical image datasets, and atomic environments within protein structures.
Revised: We use the learned invariants to achieve state-of-the-art unsupervised clustering and classification results on various spherical image datasets, and atomic environments within protein structures.

Original: By making our model variational in its invariant latent space, we enhanced the quality of clustering quality and made the model generative.
Revised: By making our model variational in its invariant latent space, we enhanced the quality of clustering quality and made the model generative.

Original: Our model is defined fully in spherical Fourier space, and thus, can reach a desired expressiveness without a need for excessive computational resources.
Revised: Our model is defined fully in spherical Fourier space, and thus, can reach a desired expressiveness without a need for excessive computational resources.

Original: The example of protein neighborhoods demonstrates that our model can be used to efficiently learn
Revised: The example of protein neighborhoods demonstrates that our model can be used to efficiently learn

Original: SO(3) invariant and equivariant representations of high-dimensional data, such as atomic environments.
Revised: SO(3) invariant and equivariant representations of high-dimensional data, such as atomic environments.

Original: Specifically, it can be used to learn compact and informative representations of atomic environments as building blocks of larger protein units. While in this paper we limit our analysis to backbone atoms within a small protein subspace, our method can be extended to all-atom representations of the entire protein universe. 
Revised: While in this paper we limit our analysis to backbone atoms within a small protein subspace, our method can be extended to all-atom representations of the entire protein universe. Going forward, we expect our method will be useful in devising models for larger units from representations learned on smaller building blocks.

Revised: For example, embeddings learned by H-(V)AE representing protein neighborhoods can be used to coarse-grain full atom representations of protein structures to facilitate structure-based predictions.

Revised: A similar approach, albeit in the sequence domain, has been used by Omega-fold (Wu et al., 2022), where representations for amino acids in a protein from a separate language model are used as inputs to a neural network to learn an MSAfree model for protein folding.

Revised: Beyond proteins, our approach could be used for other hierarchical models in large 3D structures, for which respecting rotational symmetry is desirable.

Original: A A PPENDIX A. E XPANDED BACKGROUND ON SO(3)- EQUIVARIANCE A.1. I NVARIANCE AND E QUIVARIANCE
Revised: A A PPENDIX A. E XPANDED BACKGROUND ON SO(3)- EQUIVARIANCE A.1. I NVARIANCE AND E QUIVARIANCE

Original: Let f : X → Y be a function between two vector spaces and G a group, where G acts on X and via representation D X and on Y via representation D Y .
Revised: Let f : X → Y be a function between two vector spaces and G a group, where G acts on X and via representation D X and on Y via representation D Y .

Original: Then, f is said to be G -equivariant iff f
Revised: Then, f is said to be G -equivariant iff f

Original: ( D X ( g ) x ) =
Revised: ( D X ( g ) x ) =

Original: D Y ( g ) f ( x ) , ∀ x ∈ X ∧ ∀ g ∈ G .
Revised: D Y ( g ) f ( x ) , ∀ x ∈ X ∧ ∀ g ∈ G .

Original: We note that invariance is a special case of equivariance where D Y ( g )
Revised: We note that invariance is a special case of equivariance where D Y ( g )

Original: =
Revised: =

Original: I , ∀ g ∈ G .
Revised: I , ∀ g ∈ G .

Original: A.1. G ROUP REPRESENTATIONS AND THE IRREPS OF SO(3)
Revised: A.1. G ROUP REPRESENTATIONS AND THE IRREPS OF SO(3)

Original: Groups can concretely act on distinct vector spaces via distinct group representations.
Revised: Groups can concretely act on distinct vector spaces via distinct group representations.

Original: Formally, a group representation defines a set of invertible matrices
Revised: Formally, a group representation defines a set of invertible matrices

Original: D X ( g ) parameterized by group elements g ∈ G , which act on vector space X .
Revised: D X ( g ) parameterized by group elements g ∈ G , which act on vector space X .

Original: As an example, two vector spaces that transform differently under the 3D rotation group SO(3)- and thus have different group representations - are scalars, which do not change under the action of SO(3), and 3D vectors, which rotate according to the familiar 3D rotation matrices.
Revised: As an example, two vector spaces that transform differently under the 3D rotation group SO(3)- and thus have different group representations - are scalars, which do not change under the action of SO(3), and 3D vectors, which rotate according to the familiar 3D rotation matrices.

Original: A special kind of representation for any group are the irreducible representations (irreps) which are provably the “smallest” nontrivial (i.e., they have no nontrivial group-invariant subspaces) representations.
Revised: A special kind of representation for any group are the irreducible representations (irreps) which are provably the “smallest” nontrivial (i.e., they have no nontrivial group-invariant subspaces) representations.

Original: The irreps of a group are special because it can be proven that any finite-dimensional unitary group representation can be decomposed into a direct sum of irreps (Tung, 1985).
Revised: The irreps of a group are special because it can be proven that any finite-dimensional unitary group representation can be decomposed into a direct sum of irreps (Tung, 1985).

Original: This applies to
Revised: This applies to

Original: SO(3) as well, whose irreps are the Wigner-D matrices, which are (2 ℓ + 1 × 2 ℓ + 1) -dimensional matrices, each acting on a (2 ℓ + 1) -dimensional vector space:
Revised: SO(3) as well, whose irreps are the Wigner-D matrices, which are (2 ℓ + 1 × 2 ℓ + 1) -dimensional matrices, each acting on a (2 ℓ + 1) -dimensional vector space:

Original: [Equation]
Revised: [Equation]

Original: Therefore, every element of the SO(3) group acting on any vector space can be represented as a direct sum of
Revised: Therefore, every element of the SO(3) group acting on any vector space can be represented as a direct sum of

Original: Wigner-D matrices
Revised: Wigner-D matrices

Original: A.1. S TEERABLE FEATURES
Revised: A.1. S TEERABLE FEATURES

Original: A G-steerable vector is a vector x ∈ X that under some transformation group G , transforms via matrix-vector multiplication D X ( g ) x ; here, D X ( g ) is the group representation of g ∈ G .
Revised: A G-steerable vector is a vector x ∈ X that under some transformation group G , transforms via matrix-vector multiplication D X ( g ) x ; here, D X ( g ) is the group representation of g ∈ G .

Original: For example, a vector in 3D Euclidean space is SO(3)-steerable since it rotates via matrix-vector multiplication using a rotation matrix.
Revised: For example, a vector in 3D Euclidean space is SO(3)-steerable since it rotates via matrix-vector multiplication using a rotation matrix.

Original: However, we can generalize 3D rotations to arbitrary vector spaces by employing the irreps of SO(3).
Revised: However, we can generalize 3D rotations to arbitrary vector spaces by employing the irreps of SO(3).

Original: We start by defining a degree- ℓ feature as a vector that is SO(3)-steerable by the ℓ th Wigner-D matrix
Revised: We start by defining a degree- ℓ feature as a vector that is SO(3)-steerable by the ℓ th Wigner-D matrix

Original: D ℓ .
Revised: D ℓ .

Original: Given the properties of irreps, we can represent any SO(3)-steerable vector as the direct sum of two or more independent degree- ℓ features, e.g. x = x ℓ 1 ⊕ x ℓ 2 ⊕ ...
Revised: Given the properties of irreps, we can represent any SO(3)-steerable vector as the direct sum of two or more independent degree- ℓ features, e.g. x = x ℓ 1 ⊕ x ℓ 2 ⊕ ...

Original: ⊕ x ℓ n .
Revised: ⊕ x ℓ n .

Original: The resulting vector, which we refer to as a tensor to indicate that it is composed of multiple individually-steerable vectors, is SO(3)-steerable via the direct sum of Wigner-D matrices of corresponding degrees.
Revised: The resulting vector, which we refer to as a tensor to indicate that it is composed of multiple individually-steerable vectors, is SO(3)-steerable via the direct sum of Wigner-D matrices of corresponding degrees.

Original: This tensor is a block-diagonal matrix with the Wigner-D matrices along the diagonal: D ( g ) =
Revised: This tensor is a block-diagonal matrix with the Wigner-D matrices along the diagonal: D ( g ) =

Original: D ℓ 1 ( g ) ⊕
Revised: D ℓ 1 ( g ) ⊕

Original: D ℓ 2 ( g ) ⊕ ... ⊕ D ℓ n ( g ) .
Revised: D ℓ 2 ( g ) ⊕ ... ⊕ D ℓ n ( g ) .

Original: A.1. S PHERICAL HARMONICS AND THE S PHERICAL F OURIER T RANSFORM
Revised: A.1. S PHERICAL HARMONICS AND THE S PHERICAL F OURIER T RANSFORM

Original: Spherical harmonics are a class of functions that form a complete and orthonormal basis for functions f ( θ, ϕ ) defined on a unit sphere; θ and ϕ are the azimuthal and the polar angles in the spherical coordinate system.
Revised: Spherical harmonics are a class of functions that form a complete and orthonormal basis for functions f ( θ, ϕ ) defined on a unit sphere; θ and ϕ are the azimuthal and the polar angles in the spherical coordinate system.

Original: In their complex form, spherical harmonics are defined as,
Revised: In their complex form, spherical harmonics are defined as,

Original: [Equation]
Revised: [Equation]

Original: where ℓ is a non-negative integer ( 0 ≤ ℓ ) and m is an integer within the interval − ℓ ≤ m ≤ ℓ .
Revised: where ℓ is a non-negative integer ( 0 ≤ ℓ ) and m is an integer within the interval − ℓ ≤ m ≤ ℓ .

Original: P mℓ (cos θ ) is the Legendre polynomial of degree ℓ and order m , which together with the complex exponential e imϕ define sinusoidal functions over the angles θ and ϕ in the spherical coordinate system.
Revised: P mℓ (cos θ ) is the Legendre polynomial of degree ℓ and order m , which together with the complex exponential e imϕ define sinusoidal functions over the angles θ and ϕ in the spherical coordinate system.

Original: Spherical harmonics are used to describe angular momentum in quantum mechanics.
Revised: Spherical harmonics are used to describe angular momentum in quantum mechanics.

Original: Notably, spherical harmonics also form a basis for the irreps of SO(3), i.e., the Wigner-D matrices.
Revised: Notably, spherical harmonics also form a basis for the irreps of SO(3), i.e., the Wigner-D matrices.

Original: Specifically, the SO(3) group acts on the ℓ th spherical harmonic via the ℓ th Wigner-D matrix:
Revised: Specifically, the SO(3) group acts on the ℓ th spherical harmonic via the ℓ th Wigner-D matrix:

Original: [Equation]
Revised: [Equation]

Original: Therefore, any data encoded in a spherical harmonics basis is acted upon by the SO(3) group via a direct sum of the Wigner-D matrices corresponding to the basis functions being used.
Revised: Therefore, any data encoded in a spherical harmonics basis is acted upon by the SO(3) group via a direct sum of the Wigner-D matrices corresponding to the basis functions being used.

Original: Using our nomenclature, any such data encoding constitutes a steerable tensor.
Revised: Using our nomenclature, any such data encoding constitutes a steerable tensor.

Original: We can thus map any function f ( θ, ϕ ) defined on a sphere into a steerable tensor using the Spherical Fourier Transform (SFT):
Revised: We can thus map any function f ( θ, ϕ ) defined on a sphere into a steerable tensor using the Spherical Fourier Transform (SFT):

Original: [Equation]
Revised: [Equation]

Original: The signal can be reconstructed in the real space using the corresponding inverse Fourier transform.
Revised: The signal can be reconstructed in the real space using the corresponding inverse Fourier transform.

Original: For computational purposes, we truncate Fourier expansions at a maximum angular frequency L , which results in an approximate reconstruction of the signal ˜ f
Revised: For computational purposes, we truncate Fourier expansions at a maximum angular frequency L , which results in an approximate reconstruction of the signal ˜ f

Original: ( θ, φ ) up to the angular resolution allowed by L ,
Revised: ( θ, φ ) up to the angular resolution allowed by L ,

Original: [Equation]
Revised: [Equation]

Original: Here, ˆ f ℓm are the functions’ Spherical Fourier coefficients.
Revised: Here, ˆ f ℓm are the functions’ Spherical Fourier coefficients.

Original: A.1. Z ERNIKE POLYNOMIALS AND THE Z ERNIKE F OURIER TRANSFORM
Revised: A.1. Z ERNIKE POLYNOMIALS AND THE Z ERNIKE F OURIER TRANSFORM

Original: To encode a function ρ ( r, θ, ϕ ) with both radial and angular components, we use Zernike Fourier transform,
Revised: To encode a function ρ ( r, θ, ϕ ) with both radial and angular components, we use Zernike Fourier transform,

Original: [Equation]
Revised: [Equation]

Original: where R nℓ ( r ) is the radial Zernike polynomial in 3D defined as,
Revised: where R nℓ ( r ) is the radial Zernike polynomial in 3D defined as,

Original: [Equation]
Revised: [Equation]

Original: Here, 2 F 1 ( · ) is an ordinary hypergeometric function, and n is a non-negative integer representing a radial frequency, controlling the radial resolution of the coefficients.
Revised: Here, 2 F 1 ( · ) is an ordinary hypergeometric function, and n is a non-negative integer representing a radial frequency, controlling the radial resolution of the coefficients.

Original: R nℓ ( r ) is non-zero only for even values of n − ℓ ≥ 0 .
Revised: R nℓ ( r ) is non-zero only for even values of n − ℓ ≥ 0 .

Original: Zernike polynomials form a complete orthonormal basis in 3D, and therefore, can be used within a Fourier transform to expand and retrieve any 3D shape, if large enough ℓ and n coefficient are used.
Revised: Zernike polynomials form a complete orthonormal basis in 3D, and therefore, can be used within a Fourier transform to expand and retrieve any 3D shape, if large enough ℓ and n coefficient are used.

Original: We refer to the Fourier transform of Eq.
Revised: We refer to the Fourier transform of Eq.

Original: A.6 as the Zernike
Revised: A.6 as the Zernike

Original: Fourier Trasform (ZFT).
Revised: Fourier Trasform (ZFT).

Original: To represent point clouds, a common choice for the function ρ ( r ) ≡ ρ ( r, θ, ϕ ) is the sum of
Revised: To represent point clouds, a common choice for the function ρ ( r ) ≡ ρ ( r, θ, ϕ ) is the sum of

Original: Dirac-
Revised: Dirac-

Original: δ functions centered at each point:
Revised: δ functions centered at each point:

Original: [Equation]
Revised: [Equation]

Original: This choice is powerful because the forward transform has a closed-form solution that does not require a discretization of 3D space for numerical computation.
Revised: This choice is powerful because the forward transform has a closed-form solution that does not require a discretization of 3D space for numerical computation.

Original: Specifically, the ZFT of a point cloud follows:
Revised: Specifically, the ZFT of a point cloud follows:

Original: [Equation]
Revised: [Equation]

Original: Similar to SFT, we can reconstruct the data using inverse ZFT and define approximations by truncating the angular and radial frequencies at L and N , respectively,
Revised: Similar to SFT, we can reconstruct the data using inverse ZFT and define approximations by truncating the angular and radial frequencies at L and N , respectively,

Original: [Equation]
Revised: [Equation]

Revised: The use of other radial bases is possible within our framework, as long as they are complete.

Revised: Orthonormality is also desirable as it ensures that each basis encodes different information, resulting in a more efficient encoding of the coefficients.

Revised: We use Zernike polynomials following Boyd & Yu (2011), which demonstrates that encoding with Zernike polynomials result in a faster convergence compared to the radial basis functions localized at different radii, as well as most other orthogonal harmonic bases, with the exception of Logan-Shepp.

Original:  A. D ETAILS OF H-(V)AE COMPONENTS A.2. L INEARITY
Revised: “Faster convergence” indicates that fewer frequencies are required to encode the same information. Uhrin (2021) also uses Zernike to construct invariant descriptors of atomic environments. Other equivariant methods use Bessel functions (Musaelian et al., 2022), though, according to Boyd & Yu (2011), Zernike encoding results in faster convergence. A. D ETAILS OF H-(V)AE COMPONENTS A.2. L INEARITY

Original: Let us consider a feature h ℓ containing C features of the same degree ℓ .
Revised: Let us consider a feature h ℓ containing C features of the same degree ℓ .

Original: h ℓ can be represented as a
Revised: h ℓ can be represented as a

Original: C × (2 ℓ + 1) matrix where each row constitutes an individual feature.
Revised: C × (2 ℓ + 1) matrix where each row constitutes an individual feature.

Original: Then, we learn weight matrix
Revised: Then, we learn weight matrix

Original: W ℓ ∈ R C × K that linearly maps h ℓ to h ℓ ∈ R K × (2 ℓ +1) :
Revised: W ℓ ∈ R C × K that linearly maps h ℓ to h ℓ ∈ R K × (2 ℓ +1) :

Original: [Equation]
Revised: [Equation]

Original: A.2. E FFICIENT T ENSOR P RODUCT (ETP)
Revised: A.2. E FFICIENT T ENSOR P RODUCT (ETP)

Original: Channel-wise tensor product nonlinearity.
Revised: Channel-wise tensor product nonlinearity.

Original: We effectively compute C tensor products, each between features belonging to the same channel c , and concatenate all output features of the same degree.
Revised: We effectively compute C tensor products, each between features belonging to the same channel c , and concatenate all output features of the same degree.

Original: In other words, features belonging to different channels are not mixed in the nonlinearity; the across-channel mixing is instead done in the linear layer.
Revised: In other words, features belonging to different channels are not mixed in the nonlinearity; the across-channel mixing is instead done in the linear layer.

Original: This procedure reduces the computational time and the output size of the nonlinear layers with respect to the number of channels C , from O ( C 2 ) for a “fully-connected” tensor product down to O ( C ) .
Revised: This procedure reduces the computational time and the output size of the nonlinear layers with respect to the number of channels C , from O ( C 2 ) for a “fully-connected” tensor product down to O ( C ) .

Original: The number of learnable parameters in a linear layer are proportional to the size of the output space in the preceding nonlinear layer.
Revised: The number of learnable parameters in a linear layer are proportional to the size of the output space in the preceding nonlinear layer.

Original: Therefore, reducing the size of the nonlinear output substantially reduces the complexity of the model and the number of model parameters.
Revised: Therefore, reducing the size of the nonlinear output substantially reduces the complexity of the model and the number of model parameters.

Original: This procedure also forces the input tensor to have the same number of channels for all degrees.
Revised: This procedure also forces the input tensor to have the same number of channels for all degrees.

Original: We refer the reader to Cobb et al.
Revised: We refer the reader to Cobb et al

Original: (2021) for further details and for a nice visualization of this procedure.

Original:  Spanning Tree (MST) subset for degree mixing
Revised: Minimum Spanning Tree (MST) subset for degree mixing.

Original: To compute features of the same degree ℓ 3 using the CG Tensor Product, pairs of features of varying degrees may be used, up to the rules of the CG Tensor Product.
Revised: To compute features of the same degree ℓ 3 using the CG Tensor Product, pairs of features of varying degrees may be used, up to the rules of the CG Tensor Product.

Original: Specifically, pairs of features with any degree pair ( ℓ 1 , ℓ 2 ) may be used to produce a feature of degree ℓ 3 as long as | ℓ 1 − ℓ 2 | ≤ ℓ 3 ≤ ℓ 1 + ℓ 2 .
Revised: Specifically, pairs of features with any degree pair ( ℓ 1 , ℓ 2 ) may be used to produce a feature of degree ℓ 3 as long as | ℓ 1 − ℓ 2 | ≤ ℓ 3 ≤ ℓ 1 + ℓ 2 .

Original: Features of the same degree are then concatenated to produce the final equivariant (steerable) output tensor.
Revised: Features of the same degree are then concatenated to produce the final equivariant (steerable) output tensor.

Original: Since each produced feature (often referred to as a “fragment” in the literature Kondor et al.
Revised: Since each produced feature (often referred to as a “fragment” in the literature Kondor et al.

Original: Cobb et al.
Revised: Cobb et al.

Original: (2021)) is independently equivariant, computing only a subset of them still results in an equivariant output, albeit with lower representational power.
Revised: (2021)) is independently equivariant, computing only a subset of them still results in an equivariant output, albeit with lower representational power.

Original: Reducing the number of computed fragments is desirable since their computation cannot be easily parallelized.
Revised: Reducing the number of computed fragments is desirable since their computation cannot be easily parallelized.

Original: In other words, to reduce complexity we should identify a small subset of fragments that can still offer sufficient representational power.
Revised: In other words, to reduce complexity we should identify a small subset of fragments that can still offer sufficient representational power.

Original: In this work we adopt the “MST subset” solution proposed by Cobb et al.
Revised: In this work we adopt the “MST subset” solution proposed by Cobb et al.

Original: (2021), which adopts the following strategy: when computing features of the same degree ℓ 3 , exclude the degree pair ( ℓ 0 , ℓ 2 ) if the ( ℓ 0 , ℓ 1 ) and the ( ℓ 1 , ℓ 2 ) pairs have already been computed.
Revised: (2021), which adopts the following strategy: when computing features of the same degree ℓ 3 , exclude the degree pair ( ℓ 0 , ℓ 2 ) if the ( ℓ 0 , ℓ 1 ) and the ( ℓ 1 , ℓ 2 ) pairs have already been computed.

Original: The underlying assumption behind this solution is that the last two pairs already contain some information about the first pair, thus making its computation redundant.
Revised: The underlying assumption behind this solution is that the last two pairs already contain some information about the first pair, thus making its computation redundant.

Original: The resulting subset of pairs can be efficiently computed via the Minimum Spanning Tree of the graph describing the possible pairs used to generate features of a single degree ℓ 3 , given the maximum desired degree ℓ max .
Revised: The resulting subset of pairs can be efficiently computed via the Minimum Spanning Tree of the graph describing the possible pairs used to generate features of a single degree ℓ 3 , given the maximum desired degree ℓ max .

Original: As multiple such trees exist, we choose the one minimizing the computational complexity by weighting each edge (i.e. each pair) in the graph accordingly (edge ( ℓ 1 , ℓ 2 ) gets weight (2 ℓ 1 + 1)(2 ℓ 2 + 1) ).
Revised: As multiple such trees exist, we choose the one minimizing the computational complexity by weighting each edge (i.e. each pair) in the graph accordingly (edge ( ℓ 1 , ℓ 2 ) gets weight (2 ℓ 1 + 1)(2 ℓ 2 + 1) ).

Original: The subset is also augmented to contain all the pairs with same degree to inject more nonlinearity.
Revised: The subset is also augmented to contain all the pairs with same degree to inject more nonlinearity.

Original: This procedure reduces the complexity in number of pairs with respect to ℓ max from O ( ℓ 2 max ) - when all possible pairs are used - down to O ( ℓ max ) .
Revised: This procedure reduces the complexity in number of pairs with respect to ℓ max from O ( ℓ 2 max ) - when all possible pairs are used - down to O ( ℓ max ) .

Original: We refer the reader to Cobb et al.
Revised: We refer the reader to Cobb et al.

Original: (2021) for more details and for a nice visualization.
Revised: (2021) for more details and for a nice visualization.

Original: A.2. B ATCH NORM
Revised: A.2. B ATCH NORM

Original: Let us consider a batch of steerable tensors h which we index by batch b , degree ℓ , order m and channel c .
Revised: Let us consider a batch of steerable tensors h which we index by batch b , degree ℓ , order m and channel c .

Original: During training, we compute a batch-averaged norm for each degree ℓ and each channel c as,
Revised: During training, we compute a batch-averaged norm for each degree ℓ and each channel c as,

Original: [Equation]
Revised: [Equation]

Original: Similar to standard batch normalization, we also keep a running estimate of the training norms
Revised: Similar to standard batch normalization, we also keep a running estimate of the training norms

Original: N c,tr ( i ) ℓ using momentum ξ , set to 0 .
Revised: N c,tr ( i ) ℓ using momentum ξ , set to 0 .

Original: 1 in all our experiments:
Revised: 1 in all our experiments:

Original: [Equation]
Revised: [Equation]

Original: A.13)

Original: We then update the features of the steerable tensor using the real batch-averaged norms during training, and the running batch-averaged norms during testing, together with a learned affine transformation:
Revised: We then update the features of the steerable tensor using the real batch-averaged norms during training, and the running batch-averaged norms during testing, together with a learned affine transformation:

Original: [Equation]
Revised: [Equation]

Original: A.2. P AIRWISE INVARIANT RECONSTRUCTION LOSS
Revised: A.2. P AIRWISE INVARIANT RECONSTRUCTION LOSS

Original: To reconstruct a signal within an equivariant model it is desirable to have a pairwise invariant reconstruction loss, i.e., a loss L rec such that L rec ( x , y ) =
Revised: To reconstruct a signal within an equivariant model it is desirable to have a pairwise invariant reconstruction loss, i.e., a loss L rec such that L rec ( x , y ) =

Original: L rec ( D ( g ) x , D ( g ) y ) where D is the representation of the group element g acting on the space that x and y inhabit (e.g. a rotation matrix if x and y are vectors in Euclidean 3D space, or a degree- ℓ wigner-D matrix if x and y are degree- ℓ vectors).
Revised: L rec ( D ( g ) x , D ( g ) y ) where D is the representation of the group element g acting on the space that x and y inhabit (e.g. a rotation matrix if x and y are vectors in Euclidean 3D space, or a degree- ℓ wigner-D matrix if x and y are degree- ℓ vectors).

Original: This property is necessary for the model to remain equivariant, i.e., given that the network is agnostic to the transformation of the input under group operation x
Revised: This property is necessary for the model to remain equivariant, i.e., given that the network is agnostic to the transformation of the input under group operation x

Original: → D ( g ) x by producing a similarly transformed output y → D ( g ) y , we want the reconstruction loss to be agnostic to the same kind of transformation as well
Revised: → D ( g ) x by producing a similarly transformed output y → D ( g ) y , we want the reconstruction loss to be agnostic to the same kind of transformation as well

Original: The MSE loss is pairwise invariant for any degree- ℓ feature on which SO(3) acts via the ℓ ’s
Revised: The MSE loss is pairwise invariant for any degree- ℓ feature on which SO(3) acts via the ℓ ’s

Original: Wigner-D matrix.
Revised: Wigner-D matrix.

Original: Consider two degree- ℓ features x ℓ and y ℓ acted upon by a Wigner-D matrix
Revised: Consider two degree- ℓ features x ℓ and y ℓ acted upon by a Wigner-D matrix

Original: D ℓ ( g ) parameterized by rotation g
Revised: D ℓ ( g ) parameterized by rotation g

Original: (we drop the g and ℓ indexing for clarity):
Revised: (we drop the g and ℓ indexing for clarity):

Original: [Equation]
Revised: [Equation]

Original: since Wigner-D matrices are unitary=
Revised: since Wigner-D matrices are unitary

Original: [Equation]
Revised: [Equation]

Original: (A.16)

Original: Since the MSE loss is pairwise invariant for every pair of degree- ℓ features, it is thus pairwise invariant for pairs of steerable tensors composed via direct products of steerable features.
Revised: Since the MSE loss is pairwise invariant for every pair of degree- ℓ features, it is thus pairwise invariant for pairs of steerable tensors composed via direct products of steerable features.

Original: A.2. C OSINE LOSS
Revised: A.2. C OSINE LOSS

Original: Proof that Cosine loss is pairwise invariant.
Revised: Proof that Cosine loss is pairwise invariant.

Original: The generalized dot product ⊙ from Eq.
Revised: The generalized dot product ⊙ from Eq.

Original: 5 is pairwise invariant in the same way that the dot product between two 3D vectors depends only on their relative orientations but not the global orientation of the whole two-vector system.
Revised: 5 is pairwise invariant in the same way that the dot product between two 3D vectors depends only on their relative orientations but not the global orientation of the whole two-vector system.

Original: Therefore, the whole
Revised: Therefore, the whole

Original: Cosine loss expression is pairwise invariant, since all of its components are pairwise invariant. 
Revised: Cosine loss expression is pairwise invariant, since all of its components are pairwise invariant. On the use case of Cosine loss.

Original: A. I MPLEMENTATION DETAILS
Revised: We introduce the Cosine loss as a measure of reconstruction that is both interpretable and comparable across datasets– the two characteristics that MSE does not have. A measure with these characteristics is practically useful for training of a networkd because it provides an estimate for how much better the reconstructions can get if the network’s hyperparameters were to be further optimized. For example, looking at the Cosine loss in Table A.6, we see that our model trained on Shrec17 (best Cosine = 0.130) is not as well optimized as our model trained on MNIST (best Cosine = 0.017).

Revised: Using MSE, the trend is reversed (1.8 × 10 − 3 vs. 6.7 × 10 − 3 ), since the scale of MSE depends on the size of the irreps of the data (Fig.

Revised: A.4).

Revised: [Figure]

Revised: [Figure]

Revised: We train H-(V)AE on single amino acids represented as atomic point clouds, extracted from structures in the Protein

Revised: Data Bank (PDB) (Berman et

Revised: We collect atomic point clouds of 50k residues from PDB evenly distributed across residue types and apply a 40/40/ train/valid/test split.

Revised: Residues of the same type have different conformations and naturally have noisy coordinates, making this problem a natural benchmark for our method.

Revised: We consider atom-type-specific clouds (C,

Revised: O, N and S; we exclude H) centered at the residue’s C α and compute the ZFT (Eq. 1) with L = 4 and N = 20 within a radius of 10 ˚A from the residue’s C α , and concatenate features of the same degree, resulting in a tensor with 940 coefficients.

Revised: We train several H-AE and H-VAE models, all with z = 2 ; see Sec.

Revised: A.6.4 for details.

Revised: We consistently find that the latent space clusters by amino acid conformations (Fig.

Revised: A.1), with sharper cluster separations as more training data is added (Fig. A. and A.9).

Revised: We find that test reconstruction loss decreases with more training data but the reconstruction is accurate even with little training data (from 0.153 Cosine loss with 400 training residues to 0.034 with 20,000); Table A.4.

Revised: A similar trend is observed for KNN-based classification accuracy of residues (from 0.842 with 400 training residues to 0.972 with 20,000); (Table A.4).

Revised: Notably, an untrained model, while achieving random reconstruction loss, still produces an informative invariant latent space (0.629 residue type accuracy), suggesting that the forced SO(3)-invariance grants a “warm start” to the encoder.

Revised: We do not find significant improvements in latent space classification by training with a variational objective, and present ablation results in Table A.5.

Revised: A. U SING A NON - EQUIVARIANT DECODER : EXTENDED DISCUSSION

Revised: Winter et al.

Revised: (2022) propose to construct group-equivariant autoencoders by using an equivariant encoder that learns an invariant embedding and a group element, and an unconstrained decoder which uses the invariants alone to reconstruct each datapoint in the ”canonical” form, before applying the learned group action in the output space.

Revised: By contrast, for SO(3) we propose to use an equivariant decoder, whereby the learned group element is fed as input to the decoder.

Revised: Such “unconstrained decoder” procedure can in principle be merged with our equivariant encoder and

Revised: Fourier-space approach in two ways.

Revised: For each, we argue in favor of using our equivariant decoder.

Revised: 1) Reconstructing the Fourier coefficients of the data.

Revised: To apply the learned group element on

Revised: the decoder’s output, the Wigner-D matrices for the data’s irreps need to be computed from the group element.

Revised: Then, the Wigner-D matrices can be used to “rotate” the tensor.

Revised: This has to be done on-the-fly, and it can be done quickly using functions provided in the e3nn package (Geiger & Smidt, 2022) and by smartly vectorizing operations.

Revised: We implemented this procedure by using a simple Multi-Layer Perceptron with SiLU non-linearities as a decoder.

Revised: By using e3nn to compute

Revised: Wigner-D matrices in batches, and by clever construction of tensor multiplications such that runtime scales linearly with ℓ max and is constant with regards to number of channels and batch size, we achieve models that run with comparable speed to those using our equivariant decoder, and have comparable performance on MNIST (Table A.1).

Revised: Given the empirical similarities we observe, though on a limited use case, we favor the simplicity and elegance of our equivariant decoder.

Revised: “Simplicity” because we construct the decoder to be symmetric to the encoder, thus endowing it automatically with similar representational power and without the need to tune an architecture made with different base components.

Revised: Furthermore, we highlight that our method generates intermediate equivariant representations in the decoder, rather than intermediate invariant representations.

Revised: These intermediate equivariant representations may be of interest to study in and of themselves.

Revised: 2) Reconstructing the data in real space.

Revised: In this case, we do not have to compute WignerD matrices on-the-fly, since the learned frame can be used directly in the output space as a rotation matrix.

Revised: However, since the encoder only sees a truncated Fourier representation of the data, which is by construction lossy, while the loss is computed over fine-grained real-space, this model might be too difficult to train.

Revised: We suspect this would make the model akin to a denoising autoencoder (Vincent et al., 2008) and it might be interesting to analyze, but that would be beyond the scope of this paper.

Revised: To avoid the denoising effect, we could learn to reconstruct data in real space after an Inverse Fourier Transform (IFT).

Revised: However, computing the IFT on-the-fly is very expensive and requires a discretization of the input space, to the point of being prohibitive for point clouds.

Revised: This is not a bottleneck for the Forward Fourier Transform if the cloud is parameterized by

Revised: Dirac-Delta distributions, i.e., for point clouds, as the integral can be computed exactly (Eq. A.9).

Revised: (2022)’s non-equivariant decoder formulation on the MNIST-on the sphere dataset constructed with SiLU non-linearities, with the following hidden layer sizes: [32,64,128,160,256].

Revised: We keep the number of parameters approximately the same to make model comparison fair, but we do not tune the architecture of the invariant decoders.

Revised: All other training details are kept the same (Sec. A.6).

Revised: [Table]

Revised: A. I MPLEMENTATION DETAILS

Original: Without loss of generality, we use real spherical harmonics for implementation of H-(V)AE.
Revised: Without loss of generality [Table] we use real spherical harmonics for implementation of H-(V)AE.

Original: We leverage e3nn Geiger & Smidt (2022), using their computation of the real spherical harmonics and their Clebsch-Gordan coefficients.
Revised: We leverage e3nn Geiger & Smidt (2022) [Table] using their computation of the real spherical harmonics and their Clebsch-Gordan coefficients.

Original: In our code, we offer the option to use the Full Tensor Product instead of the ETP.
Revised: In our code [Table] we offer the option to use the Full Tensor Product instead of the ETP.

Original: Specifically, at each block we allow the users to specify whether to compute the Tensor Product channel-wise or fully-connected across channels, and whether to compute using efficient or fully connected degree mixing.
Revised: Specifically, at each block we allow the users to specify whether to compute the Tensor Product channel-wise or fully-connected across channels, and whether to compute using efficient or fully connected degree mixing.

Original: However, for all the experiments in this paper, we use the ETP, which offers computational efficiency and sufficient expressivity.

Original: A. E XPERIMENTAL DETAILS A.4. A RCHITECTURE SPECIFICATION
Revised: A. E XPERIMENTAL DETAILS A.6. A RCHITECTURE SPECIFICATION

Original: We describe model architectures as follows.
Revised: We describe model architectures as follows.

Original: We specify the number of blocks B , which is the same for the encoder and the decoder.
Revised: We specify the number of blocks B , which is the same for the encoder and the decoder.

Original: We specify two lists, (i) DegreesList which contains the maximum degree ℓ max,b of the output of each block b , and (ii) ChannelsList, containing the channel sizes C b , of each block b .
Revised: We specify two lists, (i) DegreesList which contains the maximum degree ℓ max,b of the output of each block b [Table]

Original: These lists are in the order as they appear in the encoder, andare reversed for the decoder. When it applies, we specify the number of output channels of the initial linear projection C init . As noted in the main text, we use a fixed formula to determine ℓ max,b , but we specify it for clarity.
Revised: and

Original: A.4. MNIST ON THE SPHERE
Revised: A.6. MNIST ON THE SPHERE

Original: Model architectures.
Revised: Model architectures.

Original: For models with invariant latent space size z = 16 , we use 6 blocks,
Revised: For models with invariant latent space size z = 16 , we use 6 blocks,

Original: DegreesList =
Revised: DegreesList =

Original: [10 , 10 , 8 , 4 , 2 , 1] and ChannelsList =
Revised: [10 , 10 , 8 , 4 , 2 , 1] and ChannelsList =

Original: [16 , 16 , 16 , 16 , 16 , 16] , with a total of 227k parameters.
Revised: [16 , 16 , 16 , 16 , 16 , 16] , with a total of 227k parameters.

Original: Formodels with invariant latent space size z = 120 , we use 6 blocks, DegreesList =
Revised: Formodels with invariant latent space size z = 120 , we use 6 blocks, DegreesList =

Original: [10 , 10 , 8 , 4 , 2 , 1] and ChannelsList =
Revised: [10 , 10 , 8 , 4 , 2 , 1] and ChannelsList =

Original: [16 , 16 , 16 , 32 , 64 , 120] , with a total of 453k parameters.
Revised: [16 , 16 , 16 , 32 , 64 , 120] , with a total of 453k parameters.

Original: Training details.
Revised: Training details.

Original: We keep the learning schedule as similar as possible for all models.
Revised: We keep the learning schedule as similar as possible for all models.

Original: We use α = 50 .
Revised: We use α = 50 .

Original: We train all models for 80 epochs using the Adam optimizer (Kingma & Ba, 2017) with default parameters, a batch size of 100 , and an initial learning rate of 0 .
Revised: We train all models for 80 epochs using the Adam optimizer (Kingma & Ba, 2017) with default parameters, a batch size of 100 , and an initial learning rate of 0 .

Original: 001 , which we decrease exponentially by one order of magnitude over 25 epochs.
Revised: 001 , which we decrease exponentially by one order of magnitude over 25 epochs.

Original: For VAE models, we use β = 0 .
Revised: For VAE models, we use β = 0 .

Original: E rec = 25 and E warmup = 35 .
Revised: E rec = 25 and E warmup = 35 .

Original: We utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models.
Revised: We utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models.

Original: Training took ∼ 4 .
Revised: Training took ∼ 4 .

Original: 5 hours on a single NVIDIA
Revised: 5 hours on a single NVIDIA

Original: A40 GPU for each model.
Revised: A40 GPU for each model.

Original: A.4. S HREC
Revised: A.6. S HRECModel architectures.

Original: Model architectures.

Original: Both AE and VAE models have z = 40 , 7 blocks, DegreesList =
Revised: Both AE and VAE models have z = 40 , 7 blocks, DegreesList =

Original: [14 , 14 , 14 , 8 , 4 , 2 , 1] , ChannelsList =
Revised: [14 , 14 , 14 , 8 , 4 , 2 , 1] , ChannelsList =

Original: [12 , 12 , 12 , 20 , 24 , 32 , 40] , C init = 12 , with a total of 518k parameters.
Revised: [12 , 12 , 12 , 20 , 24 , 32 , 40] , C init = 12 , with a total of 518k parameters.

Original: Training details.
Revised: Training details.

Original: We keep the learning schedule as similar as possible for all models.
Revised: We keep the learning schedule as similar as possible for all models.

Original: We use α = 1000 .
Revised: We use α = 1000 .

Original: We train all models for 120 epochs using the Adam optimizer with default parameters, a batch size of 100 , and an initial learning rate of 0 . 0025 , which we decrease exponentially by two orders of magnitude over the entire 120 epochs.
Revised: We train all models for 120 epochs using the Adam optimizer with default parameters, a batch size of 100 , and an initial learning rate of 0 . 0025 , which we decrease exponentially by two orders of magnitude over the entire 120 epochs.

Original: For VAE models, we use β = 0 .
Revised: For VAE models, we use β = 0 .

Original: 2 , E rec =
Revised: 2 , E rec =

Original: and E warmup = 10 .
Revised: and E warmup = 10 .

Original: We utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models.
Revised: We utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models.

Original: Training took ∼ 11 hours on a single NVIDIA A40 GPU for each model.
Revised: Training took ∼ 11 hours on a single NVIDIA A40 GPU for each model.

Original: A.4. T OY AMINO ACIDS
Revised: A.6. T OY AMINO ACIDS

Original: Pre-processing of protein structures.
Revised: Pre-processing of protein structures.

Original: We sample residues from the set of training structures preprocessed as described in Sec.
Revised: We sample residues from the set of training structures preprocessed as described in Sec.

Original: sec:protein n eighborhoods d etails.
Revised: A.6.5.

Original: Fourier projection.
Revised: Fourier projection.

Original: We set the maximum radial frequency to N = 20 as it corresponds to a radial resolution matching the minimum inter-atomic distances after rescaling the atomic neighborhoods of radius 10 . 

Revised: We set the maximum radial frequency to N = 20 as it corresponds to a radial resolution matching the minimum inter-atomic distances after rescaling the atomic neighborhoods of radius 10 .

Original: ˚ A to fit within a sphere of radius 1.0, necessary for Zernike transform.
Revised: ˚ A to fit within a sphere of radius 1.0, necessary for Zernike transform.

Original: The channel composition of the data tensors can be described in a notation - analogous to that used by e3nn but without parity specifications - which specifies the number of channels C for each feature of degree ℓ in single units C x ℓ : 44x

Revised: The channel composition of the data tensors can be described in a notation - analogous to that used by e3nn but without parity specifications - which specifies the number of channels C for each feature of degree ℓ in single units C x ℓ : 44x

Original: + 40x

Revised: + 40x

Original: + 40x2 + 36x3 + 36x4.
Revised: + 40x2 + 36x3 + 36x4.

Original: Model architectures.
Revised: Model architectures.

Original: All models have z = 2 , 6 blocks, DegreesList =
Revised: All models have z = 2 , 6 blocks, DegreesList =

Original: ChannelsList =
Revised: ChannelsList =

Original: [60 , 40 , 24 , 16 , 16 , 8] , C init = 48 , with a total of 495k parameters.
Revised: [60 , 40 , 24 , 16 , 16 , 8] , C init = 48 , with a total of 495k parameters.

Original: We note that the initial projection is necessary since the number of channels differs across feature degrees in the data tensors.
Revised: We note that the initial projection is necessary since the number of channels differs across feature degrees in the data tensors.

Original: Training details.
Revised: Training details.

Original: We keep the learning schedule as similar as possible for all models.
Revised: We keep the learning schedule as similar as possible for all models.

Original: We use α = 400 .
Revised: We use α = 400 .

Original: We train all models for 80 epochs using the Adam optimizer with default parameters and an initial learning rate of 0 . 005 , which we decrease exponentially by by one order of magnitude over 25 epochs.
Revised: We train all models for 80 epochs using the Adam optimizer with default parameters and an initial learning rate of 0 . 005 , which we decrease exponentially by by one order of magnitude over 25 epochs.

Original: For VAE models, we use E rec = 25 and E warmup = 10 .
Revised: For VAE models, we use E rec = 25 and E warmup = 10 .

Original: We utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models.
Revised: We utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models.

Original: We vary the batch size according to the size of the training and the validation datasets.
Revised: We vary the batch size according to the size of the training and the validation datasets.

Original: We use the following (dataset size-batch size) pairs: (400-4), (1,000-10), (2,000-20), (5,000-50), (20,000-20).
Revised: We use the following (dataset size-batch size) pairs: (400-4), (1,000-10), (2,000-20), (5,000-50), (20,000-20).

Original: Training took ∼ 45 minutes on a single NVIDIA A40 GPU for each model.
Revised: Training took ∼ 45 minutes on a single NVIDIA A40 GPU for each model

Original: Evaluation.
Revised: Evaluation.

Original: We perform our data ablations by considering training and validation datasets of the following sizes: 400, 1,000, 2,000, 5,000 and 20,000.
Revised: We perform our data ablations by considering training and validation datasets of the following sizes: 400, 1,000, 2,000, 5,000 and 20,000.

Original: We keep relative proportions of residue types even in all datasets.
Revised: We keep relative proportions of residue types even in all datasets.

Original: We perform the data ablation with H-AE as well as H-VAE models with β = 0 .
Revised: We perform the data ablation with H-AE as well as H-VAE models with β = 0 .

Original: 025 and 0 . 1 .
Revised: 025 and 0 . 1 .

Original: We further perform a β ablation using the full (20,000) dataset, over the following choices of β : [0 (AE) , 0 . 025 , 0 . 05 , 0 . 1 , 0 .
Revised: We further perform a β ablation using the full (20,000) dataset, over the following choices of β : [0 (AE) , 0 . 025 , 0 . 05 , 0 . 1 , 0 .

Original: For robust results, we train 3 versions of each model and compute averages of quantitative metrics of reconstruction loss and classification accuracy.
Revised: For robust results, we train 3 versions of each model and compute averages of quantitative metrics of reconstruction loss and classification accuracy.

Original: For a fair comparison across models with varying amounts of training and validation data, we perform a 5-fold cross-validation-like procedure over the 10k test residues, where the classifier is trained over 4 folds of the test data and evaluated on the fifth one.
Revised: For a fair comparison across models with varying amounts of training and validation data, we perform a 5-fold cross-validation-like procedure over the 10k test residues, where the classifier is trained over 4 folds of the test data and evaluated on the fifth one.

Original: If validation data is needed for model selection (e.g. for LC), we use 10% of the training data.
Revised: If validation data is needed for model selection (e.g. for LC), we use 10% of the training data.

Original: A.4. P ROTEIN N EIGHBORHOODS
Revised: A.6. P ROTEIN N EIGHBORHOODS

Original: Pre-processing of protein structures .
Revised: Pre-processing of protein structures .

Original: We model protein neighborhoods extracted from tertiary protein structures from the Protein Data Bank (PDB) (Berman et al., 2000).
Revised: We model protein neighborhoods extracted from tertiary protein structures from the Protein Data Bank (PDB) (Berman et al., 2000).

Original: We use ProteinNet’s splittings for training and validation sets to avoid redundancy, e.g. due to similarities in homologous protein domains (AlQuraishi, 2019).
Revised: We use ProteinNet’s splittings for training and validation sets to avoid redundancy, e.g. due to similarities in homologous protein domains (AlQuraishi, 2019).

Original: Since PDB ids were only provided for the training and validation sets, we used ProteinNet’s training set as both our training and validation set and ProteinNet’s validation set as our testing set.
Revised: Since PDB ids were only provided for the training and validation sets, we used ProteinNet’s training set as both our training and validation set and ProteinNet’s validation set as our testing set.

Original: Specifically, we make a [80% , 20%] split in the ProteinNet’s training data to define our training and validation sets.
Revised: Specifically, we make a [80% , 20%] split in the ProteinNet’s training data to define our training and validation sets.

Original: This splitting resulted in 10,957 training structures, 2,730 validation structures, and 212 testing structures.
Revised: This splitting resulted in 10,957 training structures, 2,730 validation structures, and 212 testing structures.

Original: We use pyRosetta (Chaudhury et al., 2010) to assign Solvent Accessible Surface Area (SASA) to every atom.
Revised: We use pyRosetta (Chaudhury et al., 2010) to assign Solvent Accessible Surface Area (SASA) to every atom.

Original: Projection details.
Revised: Projection details.

Original: We set the maximum radial frequency to N = 26 as it corresponds to a radial resolution matching the minimum inter-atomic distances after rescaling the atomic neighborhoods of radius 12 .
Revised: We set the maximum radial frequency to N = 26 as it corresponds to a radial resolution matching the minimum inter-atomic distances after rescaling the atomic neighborhoods of radius 12 .

Original: 5 ˚ A to fit within a sphere of radius 1.0, necessary for Zernike transform.
Revised: 5 ˚ A to fit within a sphere of radius 1.0, necessary for Zernike transform.

Original: Model architectures.
Revised: Model architectures.

Original: All models have z = 64 , 6 blocks, DegreesList =
Revised: All models have z = 64 , 6 blocks, DegreesList =

Original: ChannelsList =
Revised: ChannelsList =

Original: [128 , 128 , 96 , 96 , 64 , 64] , C init = 64 , with a total of 3.5M parameters.
Revised: [128 , 128 , 96 , 96 , 64 , 64] , C init = 64 , with a total of 3.5M parameters.

Original: We note that the initial projection is necessary since the number of channels differs across feature degrees in the data tensors.
Revised: We note that the initial projection is necessary since the number of channels differs across feature degrees in the data tensors.

Original: Training details.
Revised: Training details.

Original: We keep the learning schedule as similar as possible for all models.
Revised: We keep the learning schedule as similar as possible for all models.

Original: We use α = 400 .
Revised: We use α = 400 .

Original: We train all models for 120 epochs using the Adam optimizer with default parameters, a batch size of 256, and an initial learning rate of 0 . 002 , which we decrease exponentially by one order of magnitude over 40 epochs.
Revised: We train all models for 120 epochs using the Adam optimizer with default parameters, a batch size of 256, and an initial learning rate of 0 . 002 , which we decrease exponentially by one order of magnitude over 40 epochs.

Original: For VAE models, we use E rec =
Revised: For VAE models, we use E rec =

Original: 40 and E warmup = 40 .
Revised: 40 and E warmup = 40 .

Original: We utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models.
Revised: We utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models.

Original: Training took ∼ 10 hours on a single NVIDIA A40 GPU for each model.
Revised: Training took ∼ 10 hours on a single NVIDIA A40 GPU for each model.

Original: A.4. L ATENT SPACE CLASSIFICATIONLinear
Revised: A.6. L ATENT SPACE CLASSIFICATION

Original:  classifier.
Revised: Linear classifier.

Original: We implement the linear classifier as a one-layer fully connected neural network with input size equal to the invariant embedding of size z , and output size equal to the number of desired classes.
Revised: We implement the linear classifier as a one-layer fully connected neural network with input size equal to the invariant embedding of size z , and output size equal to the number of desired classes.

Original: We use cross entropy loss with logits as training objective, which we minimize for 250 epochs using the Adam optimizer with batch size 100, and initial learning rate of 0.01.
Revised: We use cross entropy loss with logits as training objective, which we minimize for 250 epochs using the Adam optimizer with batch size 100, and initial learning rate of 0.01.

Original: We reduce the learning rate by one order of magnitude every time the loss on validation data stops improving for 10 epochs (if validation data is not provided, the training data is used).
Revised: We reduce the learning rate by one order of magnitude every time the loss on validation data stops improving for 10 epochs (if validation data is not provided, the training data is used).

Original: At evaluation time, we select the class with the highest probability value.
Revised: At evaluation time, we select the class with the highest probability value.

Original: We use PyTorch for our implementation.
Revised: We use PyTorch for our implementation.

Original: KNN Classifier
Revised: KNN Classifier

Original: We use the sklearn (Pedregosa et al., 2011) implementation with default parameters.
Revised: We use the sklearn (Pedregosa et al., 2011) implementation with default parameters.

Original: At evaluation time, we select the class with the highest probability value.
Revised: At evaluation time, we select the class with the highest probability value.

Original: A.4. C LUSTERING METRICS
Revised: A.6. C LUSTERING METRICS

Original: Purity (Aldenderfer & Blashfield, 1984).
Revised: Purity (Aldenderfer & Blashfield, 1984).

Original: We first assign a class to each cluster based on the most prevalent class in it.
Revised: We first assign a class to each cluster based on the most prevalent class in it.

Original: Purity is computed as the sum of correctly classified items divided by the total number of items.
Revised: Purity is computed as the sum of correctly classified items divided by the total number of items.

Original: Purity measures classification accuracy, and ranges between 0 (worst) and 1 (best).
Revised: Purity measures classification accuracy, and ranges between 0 (worst) and 1 (best).

Original: V-measure (Rosenberg & Hirschberg, 2007).
Revised: V-measure (Rosenberg & Hirschberg, 2007).

Original: This common clustering metric strikes a balance between favoring homogeneous (high homogeneity score) and complete (high completeness score) clusters.
Revised: This common clustering metric strikes a balance between favoring homogeneous (high homogeneity score) and complete (high completeness score) clusters.

Original: Clusters are defined as homogeneous when all elements in the same cluster belong to the same class (akin to a precision).
Revised: Clusters are defined as homogeneous when all elements in the same cluster belong to the same class (akin to a precision).

Original: Clusters are defined as complete when all elements belonging to the same class are put in the same cluster (akin to a recall).
Revised: Clusters are defined as complete when all elements belonging to the same class are put in the same cluster (akin to a recall).

Original: The V-measure is computed as the harmonic mean of homogeneity and completeness in a given clustering.
Revised: The V-measure is computed as the harmonic mean of homogeneity and completeness in a given clustering.

Original: [Table]
Revised: The clustering metrics “purity” and “V-measure” and the supervised metric “classification accuracy” characterize different qualities of the latent space, and, while partly correlated, they are complementary to each other. Both classes of metrics are computed by comparing the ground truth labels to the predicted labels, and they mainly differ by how the predicted labels are assigned; the clustering metrics use an unsupervised clustering algorithm, while the classification metric uses a supervised classification algorithm to do so. As a result, these metrics focus on different features of the latent space.

Original: [Table]
Revised: For example, the clustering metrics are largest when the test data naturally forms clusters with all data points of the same label. While this case can result in a high supervised classification accuracy, clustering is not a necessary condition for high classification accuracy. Indeed, the supervised signal could make the predicted labels depend more heavily on a subset of the latent space features, instead of relying on all of them equally, which is what the clustering algorithm naturally does. Therefore, it is reasonable to conclude that having higher clustering metrics and a lower classification accuracy is a sign that class-related information is more evenly distributed across the latent space dimensions.

Revised: Overall, the complementary aspect of these metrics makes it necessary to use all of them when comparing the performance of different models in each task.

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

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

Revised: 0 ± 4.3) × 10 −

Original: [Figure]
Revised: [Figure]

Original:  AE- NR/R
Revised: Protein Neighborhoods H [Figure] AE

Original: AE - R/R
Revised: [Figure] (1.8 ± 2.0) × 10 − [Figure] (1.0 ± 1.7) × 10 −

Original:  VAE - NR/R
Revised: [Figure] AE - NR/R AE - R/R VAE - NR/R

Original: VAE - R/R
Revised: VAE - R/R

Original: [Figure]
Revised: [Figure]

Original: First, if the hypothesis were true, all canonical images of the same digit should maximally or nearmaximally overlap - since they have very similar shape - and thus, their overlays would look like a “smooth” version of that digit.
Revised: First, if the hypothesis were true, all canonical images of the same digit should maximally or nearmaximally overlap - since they have very similar shape - and thus, their overlays would look like a “smooth” version of that digit.

Original: Indeed, we find this statement to be true for all models irrespective of their training strategy
Revised: Indeed, we find this statement to be true for all models irrespective of their training strategy

Original: Second, we consider the alignment of images of different digits.
Revised: Second, we consider the alignment of images of different digits.

Original: We take 1s and 7s as examples given their similarity in shape.
Revised: We take 1s and 7s as examples given their similarity in shape.

Original: If the hypothesis were true, models trained with only 1s and 7s should align canonical 1s along the long side of canonical 7s; indeed we find this to be consistently the case.
Revised: If the hypothesis were true, models trained with only 1s and 7s should align canonical 1s along the long side of canonical 7s; indeed we find this to be consistently the case.

Original: The same alignment between 1s and 7s, however, does not necessarily hold for models trained with all digits.
Revised: The same alignment between 1s and 7s, however, does not necessarily hold for models trained with all digits.

Original: This is because maximizing overlap across a set of diverse shapes does not necessarily maximize the overlap within any independent pair of such shapes.
Revised: This is because maximizing overlap across a set of diverse shapes does not necessarily maximize the overlap within any independent pair of such shapes.

Original: Indeed, we find that canonical 1s and canonical 7s do not overlap optimally with each other for models trained with all digits.
Revised: Indeed, we find that canonical 1s and canonical 7s do not overlap optimally with each other for models trained with all digits.

Original: Wenote that these tests do not provide proof, but rather empirical evidence of the characteristics of frames learned by H-(V)AE on the MNIST-on-the-sphere task. 
Revised: Wenote that these tests do not provide proof, but rather empirical evidence of the characteristics of frames learned by H-(V)AE on the MNIST-on-the-sphere task. [Figure] S [Figure]

Original: [Figure]
Revised: [Figure]

Revised: 1e

Revised: Toy amino-acids