Holographic-(V)AE: an end-to-end SO(3)-Equivariant (Variational) Autoencoder in Fourier Space
Authors: Gian Marco Visani
Michael Neal Pun
Armita Nourmohammad
Venue: Submitted to ICLR 2023 (2023)

Abstract: 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. However, little work has been done to extend this paradigm to the unsupervised and generative domains. Here, we present Holographic-(V)AE (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. 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. 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. Our work can further be seen as a case study for equivariant modeling of a data distribution by reconstructing its Fourier encoding.

Under review as a conference paper at ICLR 2023
HOLOGRAPHIC-(V)A:: AN END-TO-END SO(3)- EQUIVARIANT (VARIATIONAL) AUTOENCODER IN FOURIER SPACE
Anonymous authors Paper under double-blind review
ABSTRACT
Group-equivariant neural networks have emerged as data-efficient approach to solve classification and regression tasks, while respecting the relevant symmetries of the data. However, little work has been done to extend this paradigm to the unsupervised and generative domains. Here, we present Holographic-(V)a (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 specified origin. H-(V)AE is trained to reconstruct the spherical Fourier encod- ing 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. We extensively test the performance of H- on diverse datasets and show that its latent space efficiently encodes the categorical features of spheri- cal images and structural features of protein atomic environments. Our work can further be seen case study for equivariant modeling of a data distribution by reconstructing its Fourier encoding.
(H-
INTRODUCTION
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. One such bias the variance of the function to certain transformations of the input. For example, image classification is translationally invariant. To achieve such invariance, conventional techniques use data augmentation to train an algorithm many transformed forms of the data. However, this solution is only approx- imate and increases training time significantly, up to prohibitive scales for high-dimensional and continuous transformations (~-500 augmentations
required to learn 3D rotation- inviriinnn pat- terns (Geiger & Smidt )022).. Alternatively, one could use invariant features the data (e.g. pair- wise distance between different features) as input train any machine learning algorithm (Capecchi] et al. 2020 Uhrin 2021). However, the choice of these invariants is arbitrary and the resulting net- work could lack expressiveness.
Recent advances have developed neural network architectures that are equivariant under actions of different symmetry groups. These networks can systematically treat and interpret various transfor- mation in data, and learn models that are agnostic to these transformations. For example, models equivariant to euclidean transformations have recently advanced the state-of-the-an th-t on tasks over 3D point cloud data (Liao & Smidi [2022] Musaelian et al.) 2022 Brandstetter et al., (022). These models are more flexible and expressive compared to their purely invariant counterparts (Geiger & Smidt 2022). and exhibit high data efficiency.
Extending such group invariant and equivariant paradigms to unsupervised learning could map out compact representations of data that agnostic to specified symmetry transformation (e.g. the global orientation an object). In recent work |Winter et al. (2022) proposed a general mathemat- ical framework for autoencoders that can be applied data with arbitrary symmetry structures by learning an invariant latent space and an equivariant factor, related to the elements of the underlying symmetry group.
Here, we focus on unsupervised learning that is equivariant to rotations around a specified origin in 3D, denoted by the group SO(3). We encode the data in spherical Fourier space and construct holo-

---

Under review as a conference paper at ICLR 2023
grams of the data that are conveniently structured for equivariant operations. These data holograms are inputs to our end-to-end SO(3)-equivarian) (variational) autoencoder in spherical Fourier space, with fully equivariant encoder-decoder architecture trained to reconstruct the Fourier coefficients of the input; we term this approach Holographic-( AV (H-(V)AE). Similar to Winter et al. (2022). network learns -Quivaaquiaariant latent space composed of maximally informative set of invariants and an equivariant frame describing the orientation of the data.
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 struc- tures, The learned SO(3) invariant and equivariant representations would be useful for many real world applications in computer vision and structural biology.
2
BACKGROUND
2.1
SPHERICAL HARMONICS AND IRREPS OF SO(3)
We are interested in modeling 3D data (i.e., functions in RR²) for which the global orientation of the data should not impact the inferred model (Einstein 1916). We consider functions distributed around a specified origin, which we express by the resulting spherical coordinates (r, 0, o) around the origin. In this case, the set of rotations about the origin define the 3D rotation group SO(3), and will consider models that are rotationally equivariant under SO(3).
is convenient to project data to spherical Fourier space to define equivariant transformations for rotations.
To map a radially distributed function p(r, 0. Φ) to a spherical Fourier space, we use the Zernike Fourier Transform (ZFT),
Dhm / plr, Vem (0, dV
(1)
where (0, Φ) is the spherical harmonics of degree l and order m, where é is a non-negative integer (f 2 0) and m an integer within the interval -f S m S f. R2 (r) is the radial Zernicke polynomial in 3D (Eq. A.7p with radial frequency n 2 0 and degree l. RT(rr) is non-zero only for even values of n f 2 Zernike polynomials defined as the product Yon (f, (r) form a complete orthonormal basis in 3D, and therefore can be used to expand and retrieve 3D shapes, if large enough l and n values are used; approximations that restrict the series to finite n and & are often sufficient for shape retrieval, and hence, desirable algorithmically. Thus, in practice, we cap the resolution of the ZFT to a maximum degree and a maximum radial frequency N.
The operators that describe how spherical harmonies transform under rotations called the Wigner D-matrices. Notably, Wigner-D matrices are the irreducible representations (irreps) of SO(3), which implies that every element of the SO(3) group acting any vector space can be represented direct sum of Wigner-D matrices.
As spherical harmonies form a basis for the irreps of SO(3), the SO(3) group acts on spherical Fourier space via a direct sum of irres.. Specifically, the ZFT encodes a data point into a lensor composed of a direct sum of features, each associated with degree & indicating the irrep that transforms with under the action of SO(3). We refer to these tensors as SO(3)-steerable tensors and to the vector spaces they occupy as SO(3)-steerable vector spaces, or simply steerable for short since we only deal with the SO(3) group in this work. We note that tensor may contain multiple features of the same degree f, which we generically refer to as distinct channels c. Throughout the paper, we refer to generic steerable tensors as h and index them by f, and We adopt the "hat" notation for individual entries to remind ourselves of the analogy with Fourier coefficients. See Figure TA for graphical illustration of a tensor.
2.2 MAPPING BETWEEN SO(3) SSEERAEE VECTOR SPACES
Constructing equivariant operations equates to constructing maps between steerable vector spaces. There are precise rules constraining the kinds of operations that guarantee a valid SO(3) output, the most important one being the Clebsch-Gordan (CG) tensor product mag- The CG tensor
2

---

Under review as a conference paper at ICLR 2023
product combines two features of degrees & and lg to produce another feature of degree 162 16/- f3 1 Ez]. Let hg IR2 be generic degree feature, with individual components hem for - S m I l. Then, the CG tensor product is given by:
hegma = (tes, Bog h
(ltma) (fema))(2ma)) herm heyms
(2)
172. E1
where ((fama))
(eam))(tma)) are the Clebsch-Gordan coefficients Tung 1985).
3 HOLOGRAPHIC-(V)AE
3.1 SO(3) EQUIVARIANT LAYERS
Linearity (Lin). construct linear layers acting on steerable tensors by learning degree-specific linear operations. Specifically, we learn weight matrices specific to each degree E, and them to map across degree- feature spaces by learning linear combinations of degree-l features in the input tensor (Section A.2.I.
Efficient Tensor Product nonlinearity (ETP). We utilize the CG tensor product to inject nonlinearity in the network in an equivariant way, as was originally prescribed by Kondor et al. (0188) for SO(3)-equivarian convolutional neural networks (CNNs). This type of nonlinearity enables information flow between features of different degrees, which is necessary for constructing expressive models, and injects square nonlinearity. To significantly reduce the computational and memory costs of the tensor products, we leverage some of the modifications proposed by Cobb al.) (202T). 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. (202I) calls the MST subset. Notably, the channel-wise computation constrains the input tensors to have the same number of channels for each feature. 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. See Section A.2.2 for details, and Table for ablation study showing the improvement in parameter efficiency provided by the ETP.
Batch Norm (BN). We normalize intermediate tensor representations degree- wegree-wise wise and channel by the aatch-averaged norms of the features, as initially proposed by Kondor et al.) (2018), and do before the ETP; see Figure B and Section A.2.3 for details. We found the use of this layer to speed up model convergence (Fig A.2).
Signal Norm (SN). It is necessary to normalize activations computed by the CG tensor product to avoid their explosion. We found using batch norm alone often caused activations to explode in the decoder during evaluation. 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 's features, and apply a degree-specific affnn transformation (wy) for added flexibility. Formally, the total norm for an individual tensor h is computed as:
Ee EcEm=-(he (hgm))
No = Σ
(3)
2e +
and the features are updated as hgm hgg we/ /Neee- We note that each pre-affine normalized tensor has a total norm of 1, thus constraining the values of the individual features. Signal Norm can be seen as a form of Layer Norm that respects SO(3) equivariance (Ba et al. 2016).
Clebsch-Gordan block (CG bl.) We construct equivariant blocks, which we term Clebsch-Gordan blocks, by stacking together the equivariant layers introduced above, as shown in Figure BB. Each block can take a steerable tensor of arbitrary size composed of multiple features of arbitrary degrees, and can output steerable tensor of arbitrary size with multiple features of arbitrary degrees, by following the sparsity of the CG tensor product (l62 - L1 < &g < 1e lz)) Crucially, the maximum feature degree in the input tensor is Euax, then the maximum feature degree that can be generated is 2mmxx achieved by combining two features of degree lmax. We employ additive skip connections, zero-padded when appropriate, favor better gradient flow.
3

---

Under review as a conference paper at ICLR 2023
3.2 DATA NORMALIZATION
As per standard machine learning practice (Shanker et al. 1996). we normalize the data. We do this by dividing each tensor by the average square-root total norm of the training tensors. analogously to the Signal Norm. 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.
3.3 MODEL ARCHITECTURE
H-(V)AB has a fully rotationally equivariant architecture, and learns a disentangled latent space consisting of a maximally informative invariant (e 0) component z of arbitrary size, as well as three orthonormal vectors (E = 1), which represent the global 3D orientation of the object and reflect the coordinate frame of the input tensor. Crucially, the disentangled nature of the latent space is respected at all stages of training, and is guaranteed by the model's 's rotational equivariance. The architecture is shown Figure [ec. The encoder takes as imput steerable tensor with maximum degree max I. and, via a stack of Clebsch-Gordan blocks, iteratively and equivariantly transfers information from higher degrees lower ones, down to the final encoder layer with fmax 1. The frame is constructed by learning two vectors and using Gram-Schmide to find the corresponding orthonormal basis (Schmidt, [907). The third orthonomal basis vector is then calculated as the cross product of the first two. The decoder learns to reconstruct the input from and the frame, iteratively increasing the maximum degree L of the intermediate representations by leveraging the CG Tensor Product within the Clebsch-Gordan blocks.
An interesting question ask is: what does the trained decoder's output look like if the frame is held constant (e.g. equal to the identity matrix)? We experimentally find that the recon- structed elements tend to be aligned with each other and hypothesize that the model implicitly learning to maximize the overlap between training elements, providing empirical evidence in the Appendix (Fig. A.3). We call this frame the canonical frame following the analogy with the canonical elements in Winter et al.] (2022). 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.
Within the decoder, the maximum degree fmax, that be outputted by each block b is constrained by the sparsity of the CG tensor product. Specifically, b S 2 where b ranges from (first block) to (last block). Since we need to reconstruct features up to degree L in the decoder, arrive at lower bound for the number of blocks in the decoder set by L 2 L, or B > log, L. In experiments, we set & min(2', 2 L}} and do not let L exceed the input's maximum degree T.. Relaxing this condition might increase the expressive power of the network but at significant increase in runtime and memory. We leave the analysis of this trade-off to future work. For the encoder and the decoder to have similar expressive power, we construct them to be symmetric with respect the latent space (Fig. [TC). 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. We empirically verify the equivariance of our model up to numerical errors in Table A.8
For H-VAE, we parameterize the invariant part of the latent space by an isotropic Gaussian, i.., we learn two sets of size z invariants, corresponding to means and standard deviations.
3.4 TRAINING OBJECTIVE
We train H-(V)AE to minimize the reconstruction loss Cree between the original and the recon- structed tensors, and, for H-VAE only, to minimize the KL-divergence of the posterior invariant latent space distribution (<>(<) from the selected prior p(z) (Kingma & Welling 2013):
f(x, x²) = a (1r, x²) + 3 [Dx
We use mean square eπror (MSE) for Lee, which show in Sec. A.2.4 respects the neces- sary property of SO(3) pairwise invariance, ensuring that the model remains rotationally equivariant.
4

---

Under review as a conference paper at ICLR 2023
A
B
Clobsch- Gordan block
BN
ETP
SN
Channels
Input tensor
Reconstructed tensor fu &
Lin
Lin
CG
CG bl.
CG bl
CG bl.
Gram- Gram--Shhinn Schmidl
space
Decoder
Figure 1: Schematic of the Network architecture. A: Schematic of a steerable tensor with &mxx and 4 channels per feature degree. We choose a pyramidal representation that naturally follows the expansion in size of features of higher degree. B: Schematic of Clebsch-Gordan Block (CG bl.), with batch norm (BN), efficient tensor product (ETP), and signal norm (SN), and Linear (Lin) operations. C: Schematic of the H-AE architecture. We color-code features of different degrees in the input and in the latent space for clarity. The H- VAE schematic differs only in the latent space, where two sets of invariants learned (means and standard deviations of isotropic Gaussian distribution).
We cast an isotropic normal prior the invariant latent space: p(z) I). Hyper- parameters and 3 control the trade-off between reconstruction accuracy and latent space regularization Higgins et al. (022).. We find practical scale the reconstruction loss by dateset- scalar a since the MSE loss varies in average magnitude across datasets. When training H-VAE, we find it beneficial to keep 3 for a few epochs (Eree) so that the model can learn to perform meaningful reconstructions, and then linearly increasing it to the desired value for Ewarmup epochs to add structure to the latent space, an approach first used by Bowman et al. (2116).
3.5 RECONSTRUCTION ASSESSMENT VIA COSINE LOSS
As the scale of MSE depends on the characteristics of the data, e.g. the size of the tensors represent- ing the data and their irreps (Fig. A.4), difficult to contextualize MSE values across datasets. It would be desirable to have dimensionless metric that measures absolute "goodness" of recon- structions that is comparable across datasets. For this purpose, we propose the metric Cosine loss which is normalized dot product generalized operate on pairs of steerable tensors (akin to cosine similarity), and modified to be interpreted as a loss:
Cosine(ar, y) =1 ( with
a ④ y
Σ (xe Beg De )=
[(x
(5)
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 pair of random tensors of any size, with smaller variance for larger tensors (Sec. A.4). 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 norms are easier to reconstruct than directions, we still find the Cosine loss useful as noisy estimate of the model's reconstruction ability. Furthermore, Cosine loss correlates almost perfectly with MSE, especially in the mid-to-low reconstruction quality regime (SpearmanR 0.99, Fig. A.4 and Table A.6).
training objective. However,
5

---

Under review as a conference paper at ICLR 2023
4 RELATED WORK
Group-equivariant neural networks. Group-equivarian neural networks have improved the state- of-the-art on many supervised tasks, thanks to their data efficiency (Hutchinson et al. 2021] Bekkers)
ict al. 20188: Romero & Cordonnier (021). Related to our work are 3D Euclidean neural networks,
which are equivariant to (subsets of) the 3D Euclidean group: (Weiler al. 2018 Brandstetter et al.
2022 Thomas et al. 2018 Batzner et al., 2022 Musaelian et al., [022 Satorras et al. [022] Fuchs
al. 2020 Liao & Smidt 2022 and often use spherical harmonics and tensor products to construct
SO(3) equivariant layers.
Seminal work on SO(3) equivariance has been conducted for spherical
images (Cohen al., 2018 Esteves al.. 2020); we were inspired by the fully Fourier approach of Kondor et al. 2018), and leveraged operations proposed by Cobb et al. (202T)
Equivariant representations of atomic systems. There is diverse body of literature on construct- ing representations of atomic systems that are invarian/equivaarin to euclidean symmetries, lever- aging Fourier transforms and CG tensor products (Drautz, 2019; Musil et al. 202T). Notably, Chrin 202T constructs SO(3)-invariant representations of point clouds using the ZFT and CG tensor prod- uct iterations. 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.
Invariant autoencoders. Several works attempt to learn representations that are invariant to certain classes of transformations. [Shu et al. (2018) and Koneripalli et al. (2020) learn general "shape" embeddings by learning a separate "deformation" embedding. However, their networks are not explicitly equivariant to the transformations of interest. Other work proposes to learn exactly invariant embedding alongside an approximate (but not equivariant) group action to align the input and the reconstructed data. For example, Mehr al. (2018) learns in quotient space by sampling the group's 's orbit and computing the reconstruction loss by taking the infimum over the group. This approach is best suited for discrete and finite groups, and it is computationally expensive as it is akin to data augmentation. [ohiit & Trivedi (2020) construct an SO(3)-invaiiant autoencoder for spherical signals by learning an invariant latent space and minimizing loss which first finds the rotation that best aligns the true and reconstructed signals, introducing an added optimization step- potentially very expensive for 3D data and reconstructing the signals only up to a global rotation. 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. However, the rotational invariance is manually imposed by Lohit & Trivedi] (2020), which is distinct from our approach that fully equivariant and only requires simple MSE for reconstruction of data its original orientation.
Group-equivariant autoencoders. A small body of work focuses on developing equivariant au- toencoders. 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 eet al. 2019). Others use supervision to extract class-invariant and class-equivariant representa- tions Feige 2022). A recent theoretical work proposes to train an encoder that encodes elements into invariant embedding and an equivariant group action, then using 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 al. (022). 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. Furthermore, our model is variational in the invariant latent space.
5 EXPERIMENTS
5.1 ROTATED MNIST ON THE SPHERE
We extensively test the performance of H-(V)AE on the MNIST-on-the-spheen dataset (Deng) 2012). Following Cohen et al. (2018) we generate a discrete unit sphere using the Driscoll-Healey (DH) method with bandwidth (bw) of 30, and project the MNIST dataset onto the lower hemisphere. We consider two variants, NR/R and R/R, differing in whether the training/test images have been randomly rotated (R) or not (NR). For each dataset, we map the images to steerable tensors via the Zernike Fourier Transform (ZFT) with L = 10, and a constant radial function R = 1, resulting in tensors with 121 coefficients.

---

Under review as a conference paper at ICLR 2023
Original
Reconstructed
Q
0
J
C
2
2
2
/
Start
Target
2
8
R
Figure 2: H-VAE on MINST-on-tien MiNSS-on-the-spherr. Evaluation on rotated digits for an H-VAE trained on non-rotated digits with 2 = 16. A: Original and reconstructed images in the canonical frame after inverse transform from Fourier space. The images are projected onto a plane. Distortions at the edges and flipping are side-effects of the projection. B: visualization of the latent space via 2D UMAP (Mcinnes et al., 2020). Data points are colored by digit identity. C: Shown are cherry-picked images generated by feeding the decoder invariant embeddings sampled from the prior distribution and the canonical frame. D: Example image trajectory by linearly interpolating through the learned invariant latent space. We interpolate between the learned invariant embeddings of the Start and the Target images. Then, we feed each embedding to the decoder alongside the canonical frame.
We train models with combinations of the following features: training mode (NR vs. R), invariant latent space size (16 120), and model type (AE vs. VAE). In all cases the model architecture follows from Fig. [[: see Section A...2 for details. All models achieve very low reconstruction loss (Table D with no significant difference between training modes, indicating that the models successfully leverage -quivvrauve to generalize to unseen orientations. Predictably, AE models have lower reconstruction loss than VAE models, and do models with larger latent space. Nonetheless, H-VAE achieves reliable reconstructions, as shown in Fig. [2A and Table I All 8 models produce an invariant latent space that naturally clusters by digit identity. We show this qualitatively for one of the models in Fig. B and quantitatively by clustering the data via K-Means with 10 centroids and computing standard clustering metrics of Purity (Aldenderfer & Blashfield 1984) and V-measure [Rosenberg & Hirschberg 2007) in Table
All 8 models achieve much better clustering metrics than Rot-Inv AE (Lohit & Trivedi 2020). with VAE models consistently outperforming AE models. We also train a linear classifier (LC) to predict digit identity from invariant latent space descriptors, achieving comparable accuracy to Rnt-Inv AE with the same latent space size. We do not observe any difference between VAE and AE models in terms of classification accuracy. Using KNN classifier instead of LC further improves performance (Table A.2).
As H-VAE is generative model, generate random spherical images by sampling invariant latent embeddings from the prior distribution, and observing diversity in digit type and style (Fig. [2 and Fig. A.5). 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. The interpolated images present spatially consistent transitions (Fig. 2 and Fig. A.6), which is sign of semantically well-structured latent space.
To understand the meaning of the learned frames, we visualize the sum images in the canonical frame (i.e. the identity matrix). We hypothesize that H-(V)AE learns optimally overlap the training images when in the same frame. Indeed, by visualizing the sum the training images verify that images 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)
in the canonical frame,
7

---

Under review as a conference paper at ICLR 2023
Table 1: Performance metries on MNIST-on-the-spher and Shrec17. Reconstruction loss, clus- tering metrics, classification accuracy in the latent space using linear classifier, and retrieval metrics (Shrec17 only) are shown. We only report scores presented the corresponding papers of origin.
Dataset
Type
Method
bw
Costne
Purity
V-meas.
Class.
P@N RØN
FION
MAP
NDCG
Supervised
Cinb et ml 2021 NR/
30
0.993
NIOX
120
30
0.30
0.35
0.891
FFAL
30
R-R (Quss)
30
0.881
H1 -LL NR/R (Ours)
16
30
0.025
0.62
0.51
0.820
MNIST
Unsupervised
H-AF. R:R (Qurs)
16
30
0.024
0.65
0.52
0.833
H. VAE NKV/ (Our)
120
30
0.037
0.70
0.59
0.583
HA- R/R (Curs)
120
30
0.037
0.65
0.53
0.884
H-VAF NR/R (Ours)
30
VAE R/R (Churs)
30
0.57
Supervised
Esseves al 2020
1228
128
0.717
0.719
0.737
0.710
0.708
0.685
0.679
0.758
Triscn
30
0.578
0.351
0.361
0.215
Unspervised
90
0.130
0.50
0.41
0.684
0.548
0.569
0.545
0.500
0.397
H-VAR (Ours)
90
0.151
0.52
0.42
0.631
0.512
0.537
0.512
0.463
0.568
5.2 SHREC17
The Shrec17 dataset consists of 51k colorless 3D models belonging to 55 object classes, with a 70/10/20 train/valid/test split (noa). We use the variant of the dataset where each model is perturbed by random rotations. Converting 3D shapes into spherical images preserves topological surface information, while significantly simplifying the representation. We follow Cohen et al. (2018) and project surface information from each model onto an enclosing DH spherical grid with a bandwidth of 90 via ray-casting scheme, generating spherical images with 6 channels. We then apply the ZFT with I, = 14 and a constant radial function R2 = 1 to each channel individually, resulting in tensor with 1350 coefficients. We train AE and VAE model (Sec. (A.6.3) and evaluate them similarly to the MNIST dataset and compute the Shrec 17 retrieval metrics via the latent space linear classifier's 's predictions. H-AE achieves the best classification and retrieval results for autoencoder- based models, and is competitive with supervised models despite the lower grid bandwidth and the small latent space (Table []. Using KNN classification instead of linear classifier further improves performance (Table A.3). H-VAE achieves slightly worse classification results but better clustering metries compared to H-AE While reconstruction loss low, there still significant margin of improvement. We partially attribute Lohit & Trivedi (2020)ss 's low scores to the low grid bandwidth. 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 method does not depend on it.
5.3 PROTEIN NEIGHBORHOODS
We test H-(V)AE on a challenging point cloud dataset, comprised of spher- ical atomic environments surrounding a residue within protein structure, which we term protein neighborhoods. We use ProteinNet (AIQuraishi] (019) splits avoid any redundancy in training and test sets (Sec. A.6.5)). We consider all atomic neighborhoods within a radius of 12.5Á surrounding central residue types CYS. GLU or HIS (Fig. (A). We only consider backbone atoms (Co, C, and N). For each neighborhood centered at the central residue's Ca, we compute the ZFT (Eq. [ with L = and N = 26. We then con- catenate features of the same degree re- sulting in tensor with 1240 coefficients. The final dataset contains 303k/75k/4,44 (train/valid/estt) tensors. We train H-
A
B
A sheet
a-helix
loop
Figure 3: H-AE on protein neighborhoods. A: An exam- ple protein neighborhood of backbone atoms: Ca (yel- low), C (gray), N (blue) and O (red). The secondary structure of the central residue is shown overlaid in yel- low. B: 2D UMAP visualization of the invariant latent space learned by H-AE, colored by secondary structure the central residue.
8

---

Under review as a conference paper at ICLR 2023
AE model with an invariant latent space size of 64 (see Sec. A.6.5). While the test Cosine loss has still margins of improvement (~ 0.161) the latent space is meaningfully organized according to relevant neighborhood statistics, such the secondary structure of the central residue (~ 87% KNN classification accuracy Fig. 3), number of atoms in the neighborhood (Fig. A.10), and neighbor- hood's average Solvent Accessible Surface Area (SASA), which is proxy value for how buried the residue is in the protein (Fig. [A.II. We further benchmark H-(V)AE via extensive ablation experiments on simpler point cloud dataset consisting of individual amino acids (Sec. A.3
USING AN UNCONSTRAINED DECODER
It is possible to merge use our SO(3) -equivaanv encoder and fourier-space formulation with Winter al. (2022)'s 's framework, in which the learned rotation (i.e. the frame) applied to the output space of an unconstrained decoder instead of being fed as input to an equivariant decoder. The resulting model 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-spheer 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 AAII) Given the sim- ilarity in performance, we favor the simplicity and elegance of our equivariant decoder, which is constructed be symmetric the encoder, thus automatically endowing it with similar represen- tational 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
7 PROPERTIES AND LIMITATIONS
The information content of the latent space is limited by the maximum degree L and radial frequency N set as cutoff for the spherical Fourier series. Such truncation is necessary for computational feasibility, but it limits resolution, evident when reconstructing a signal.
Another related limitation 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. 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.
8 CONCLUSION
In this work, we develop the first end-to-end SO(3)-cquivariant (V)AE, suitable for data distributed around center. The model learns invariant embedding deseribing the data in "canonical" orientation alongside an equivariant frame describing the data's original orientation relative to the canonical one.
We use the learned invariants to achieve state-of-the-art unsupervised clustering and classification results various spherical image datasets, and atomic environments within protein structures. By making our model variational in its invariant latent space, we enhanced the quality of clustering quality and made the model generative. Our model is defined fully in spherical Fourier space, and thus, can reach a desired expressiveness without a need for excessive computational resources.
The example of protein neighborhoods demonstrates that our model can be used to efficiently learn SO(3) invariant and equivariant representations of high-dimensional data, such as atomic environ- ments. 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 rep- resentations learned on smaller building blocks. 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. A similar approach, albeit in the sequence do- main, has been used by Omega-fold (Wu et al., 2022). where representations for amino acids in a protein from separate language model are used as inputs to neural network to learn MSA- free model for protein folding. Beyond proteins, our approach could be used for other hierarchical models in large 3D structures, for which respecting rotational symmetry is desirable.

---

Under review as a conference paper at ICLR 2023
REFERENCES
SHREC 2017: Large-scale 3D) Shape Retrieval from ShapeNet Core55.
URL. https://
shapenet. CS www.staadd........ edu/ shrec17/
Mark Aldenderfer and Roger Blashfield. Cluster Analysis. SAGE Publications, Inc., 2455 Teller
Road, Thousand Oaks Califoria 91320 United States of America, 1984. ISBN 978-0-8039-2341
978-1-4129-836448 doi: 2418836444 / sagepub.
com/ hlom/bookccluee clabokrrnlstttr-aaaaa lssis
Mohammed AlQuraishi. ProteinNet: a standardized data set for machine learning of protein structure. BMC Bioinformatics, 20(1):311, June 2019. ISSN 1471-2105. doi: 10.1186/
$12859-019-2932-0. /ttttss::://// /tw. http///o..ooogg..o//
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E. Hinton. Layer Normalization, July 2016. URL www.////g/b////eeea abs www.ab/ss.. 060./1641770000 arXiv: 06450 [es,, stat].
http://arxiv. org.:iii.orrrraaoo
Simon Batzner, Albert Musaclian, Lixin Sun, Mario Geiger, Jonathan P. Mailoa, Mordechai Ko- mbluth, Nicola Molinari, Tess E. Smidt, and Boris Kozinsky. E(3)-cquivariant graph neural networks for data-efficient and accurate interatomic potentials. Nature Communications, 13 (1):2453, May 2022. ISSN 2041-1723. doi: 10. 067110001/111477022222 29939-5 URL httpst:::
7/www. www.ratt........o.o com attaau/e..ccoorrtt..ooeeeeeeee-.
Number: 1 Publisher: Nature
Publishing Group.
Erik J. Bekkers, Maxime W. Lafarge, Mitko Veta, Koen AJ Eppenhof, Josien PW Pluim, and Remco Duits. Translaaasnttion Covariant Convolutional Networks for Medical Image Analysis, June 2018.
wttp C3393 arXiv: 11t::1800400393 [cs, math].
Helen M. Berman, John Westbrook, Zukang Feng. Gary Gilliland, N. Bhat, Helge Weissig, Ilya N. Shindyalov, and Philip E. Bourne. The Protein Data Bank. Nucleic Acids Research, 28(1):235- 242, January 2000. ISSN 0305-1048. doi: 10. 10.10/9/nr/228. 1ataaa/28........ 1.235. URI. rcsb..rrrg
Samuel R. Bowman, Luke Vilnis, Oriol Vinyals, Andrew M. Dai, Rafal Jozefowicz, and Samy Bengio. Generating Sentences from Continuous Space, May 2016. URL http::/arxiv
oxq/abs/1511 06349 arXiv:1511.0644 [cs].
John P. Boyd and Fu Yu. Comparing seven spectral methods for interpolation and for solv- ing the Poisson equation in disk: Zernike polynomials, Logan-Shepp ridge polynomials, Chebyshev-Fourey Series, cylindrical Robert functions, Bessel-Fourie expansions, square-to- disk conformal mapping and radial basis functions. Journal of Computational Physics, 230(4):
1408-1438, February 2011. 1. ISSN 0021-9991.
doi: 10. jo.10016/.jcp,220001111 7/www. www.iciinnnn wcweecediiiccccccooccoooocccc.cc rect. cwm www.//cceneii haamm//cinnce//rree cttp://ciiaeneeeaaaa.eee//////.. 0T111//000000099100000000
URL https:
Johannes Brandstetter, Rob Hesselink, Elise der Pol, Erik J. Bekkers, and Max Welling. Ge- ometric and Physical Quantities Improve E(3) Equivariant Message Passing, March 2022. URL
http::/arxivv orgg abs 02905 arXiv:2110.0290 [cs, stat].
Alice Capecchi, Daniel Probst, and Jean-Louis Reymond. One molecular fingerprint to rule them all: drugs, biomolecules, and the metabolome. Journal of Cheminformatics, :333,) June 2020. ISSN 1758- 2946. doi: 01111113311-000--012
URL htttp:://0cl..rg111...1
113321-020-00420
Sidhartha Chaudhury, Sergey Lyskov, and Jeffrey J. Gray. PyRosetta: a script-based interface for implementing molecular modeling algorithms using Rosetta. Bioinformatics, 26(5):689-
691, March 2010. ISSN 1367-4803. doi: 10. 10009//oooooomaais//t000.0 URL https:
7/doi. www.oogg........... 10. 1tt////ooofoorratcc/////oooooo/o
Oliver J. Cobb, Christopher G. R. Wallis, Augustine N. Mavor-Parker Augustin Marignier, Matthew A. Price, Mayeul d' Avezac, and Jason D. McEwen. Efficient Generalized Spherical CNNs, March 2021. URL [astro-ph].
/2100 11661116 161v:2010..1661
Taco S. Cohen, Mario Geiger, Jonas Koehler, and Max Welling. Spherical CNNs, February 2018. URL httt:://ariiv wwwporrr///b////o 10130 1801 10130 [es, stat).
10

---

Under review as a conference paper at ICLR 2023
Li Deng. The MNIST Database of Handwritten Digit Images for Machine Learning Research [Best of the Web]. IEEE Signal Processing Magazine, 29(6):141 141-144. November 2012. ISSN 1558-
0792. doi: 10. 10.11099/MSP22112222111 1477. Conference Name: IEEE Signal Processing Magazine.
Ralf Drautz. Atomic cluster expansion for accurate and transferable interatomic potentials. Physical Review B, 99())014104, January 2019, doi: 10.1 10.1103/hysRev..990 1et.0.1999.11144. URLhttps::
link. aank.aps.orr. http://///.../.1........///...// Society.
99, 014104
Publisher: American Physical
Albert Einstein. Die Grundlage der allgemeinen Relativitätsshcorii. Annalen der Physik, (49):770~ 822, 1916.
Carlos Esteves, Ameesh Makadia, and Kostas Daniilidis. 2020.
Spherical CNNs, October 10731 arXiv:2006.100 10731 [es].
Ilya Feige. Invariant-equivariann representtinno learning for multi-class data. February 2022. URL httpps://ppeneviier htt. www.nntt/ooommdddiiiimmm e4o0099
Fabian B. Fuchs, Daniel E. Worrall, Volker Fischer, and Max Welling. SEG))Transtfrmen::: 3D Roto-Translation Equivariant Attention Networks, November 2020. URL http://arxiv.
org/abs/ 2006 10503 arXiv:2006.1050) [cs, stat].
Mario Geiger and Tess Smidt. e3nn: Euclidean Neural Networks, July 2022. URL http://
arxiv.org/ abs /2007 09453 arXiv:2207.09453 [cs].
Irina Higgins, Loic Matthey, Arka Pal, Christopher Burgess, Xavier Glorot, Matthew Botvinick, Shakir Mohamed, and Alexander Lerchner. beta- VAE: Learning Basic Visual Concepts with a
Constrained Variational Framework. July 2022. /RL/atpss//ppenneviee...ee wttp:/iiwewwee/frrnmmeee.
id-Sy2fzU9gu
Geoffrey E. Hinton, Alex Krizhevsky, and Sida D. Wang. Transforming Auto-Encoders. In Timo Honkela, Wlodzislaw Duch, Mark Girolami, and Samuel Kaski (eds.). Artificial Neural Networks and Machine Learning ICANN 20/1, Lecture Notes in Computer Science, pp. 44-51, Berlin, Heidelberg, 2011. Springer. ISBN 978-3-642-21735-77 doi: 10.1007/978-3 1177783--444 6-2-3----4222-3-7-11
Michael Hutchinson, Charline Le Lan, Sheheryar Zaidi, Emilien Dupont, Yee Whye Teh, and Hyunjik Kim. LieTransformer: Equivariant self-attention for Lie Groups, June 2021. URL http:://axxi org/ absa 201.2 10885 arXiv:2012.1088. 10885 [cs, stat].
Diederik P. Kingma and Jimmy Ba. Adam: Method for Stochastic Optimization, January 2017.
URL http:t/t//// /twpvvaaavvv/aaa///i//eaaaaaaava org/ abs/ abs/aaass/11/ 11t./1412..0 6980 urXiv:1412.6980 [cs].
Diederik P. Kingma and Max Welling. Auto-Encoding Variational Bayes. December 2013. doi:
10.48550/arXiv. 1312.61 14. URL
/1312. 0111l///11111...1111111 6114v10
Risi Kondor, Zhen Lin, and Shubhendu Trivedi. Clebsch-Gordan Nets: a Fully Fourier Space Spheri- cal Convolutional Neural Network, November 2018. UR//ttpp://aaavvv.oo////bbb//b//
09231 N.Kiv:880.0092 [cs, stat].
Kaushik Koneripalli, Suhas Lohit, Rushil Anirudh, and Pavan Turaga. Rate-Invariant Autoencoding of Time-Series. In ICASSP 2020 2020 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 3732-3736, May 2020. doi: 10. 10.1 10.1199/CAAS999077620000 9053983. ISSN: 2379-190X.
Adam R. Kosiorek, Sara Sabour, Yee Whye Teh, and Geoffrey Hinton. Stacked Capsule Autoen- coders, December 2019. /BLLtt:://arxvv..rr///bbb [cs, stat].
06818 11Xiv:1906.0681s
Yi-Lun Liao and Tess Smidt. Equiformer: Equivariant Graph Attention Transformer for 3D Atom- istic Graphs, June 2022. URL http:::aaxw..oor/////a///// wttp://////.bb/././///... arXiv:2206. 1190 [physics].
Suhas Lohit and Shubhendu Trivedi. Cnnttiinnnnnaiiannttttoo for Signals on Spheres, De- cember 2020.
04474 2xX447012.04474 [es].

---

Under review as a conference paper at ICLR 2023
Leland McInnes, John Healy, and James Melville. UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction, September 2020. URL http:t/ http::///xxxvv.o/////// /bs/
1802. 03426 arXiv:1802.03426 [cs, stat].
Éloi Mehr, André Lieutier, Fernando Sanchez Bermudez, Vincent Guitteny, Nicolas Thome, and Matthieu Cord. Manifold Learning in Quotient Spaces. 2018 IEEE/CVF Conference on Com- puter Vision and Pattern Recognition, pp. 9165-9174, June 2018. doi: 10.1109/CVPR.2018
00955. ISSN: 2575-7075.
Albert Musaelian, Simon Batzner, Anders Johansson, Lixin Sun, Cameron Owen, Mordechai Kornbluth, and Boris Kozinsky. Learning Local Equivariant Representations for Large- Scale Atomistic Dynamics, April 2022. URL arXiv:2204.05249 [cond-mat, physics:physics].
Felix Musil, Andrea Grisafi, Albert P. Bartók, Christoph Ortner, Gábor Csányi, and Michele Ceriotti. Physics-Inspired Structural Representations for Molecules and Materials. Chemical Reviews, 121 (16):9759-9815, August 2021. ISSN 0009-2665. doi: 2.1021/acs.chemrent 1ecemmeev.c000211 URL http::::/do..org//// 102110111aae chemrev. 1c00021 Society.
Publisher: American Chemical
Fabian Pedregosa, Gaël Varoquaux, Alexandre Gramfort, Vincent Michel, Bertrand Thirion, Olivier Grisel, Mathieu Blondel, Peter Prettenhofer, Ron Weiss, Vincent Dubourg, Jake Vanderplas, Alexandre Passos, David Cournapeau, Matthieu Brucher, Matthieu Perrot, and Edouard Duch- esnay, Scikit-leam: Machine Learning Python. Journal of Machine Learning Research, ISSN 1533-7928.
12(85):2825 2830, 2011.
URL http: http://nni...coopppppe//////
pedregosalla. htm]
David W. Romero and Lean-Baptiste Cordonnier. Group Equivariant Stand-Alone Self-Attention For Vision, March 2021.
arXiv:2010.00977
[cs, stat].
Andrew Rosenberg and Julia Hirschberg. V-Measure: A Conditional Entropy- External Clus- ter Evaluation Measure. In Proceedingy of the 2007 Joint Conference on Empirical Methods in Natural Language Processing and Computational Natural Language Learning TEMNLL-CONLL pp. 410-420, Prague, Czech Republic, June 2007. Association for Computational Linguistics. URL https: http:/::///laaattooooo...oro 107-14333
Victor Garcia Satorras, Emiel Hoogeboom, and Max Welling. E(n) Equivariant Graph Neural Net- works, February 2022. URL http:://axx....rrrtt org. http//bb/220020 09844 arXiv:2102.09844
[es, stat].
Erhard Schmidt. Zur Theorie der linearen und nichtlinearen Integralgleichungen. Mathematische Annalen, 63(4):433-476 December 1907. ISSN 1432-1807. doi: 10.1007/BFO01449770 URL httpps://on
M. Shanker, M. Y. Hu, and M. S. Hung. Effect of data standardization on neural net- work training. Omega, 24(4):385-397, August 1996. ISSN 0305-0483. 1016/0000500443196000000000
doi: 10.
URL https: /ots://wwi /www. sci eww.eineeddr...... www.eddireet.commccc....... wttps/.oc.cc//////ccccccc/c articlee/p11/330008399000
Zhixin Shu, Mihir Sahasrabudhe, Alp Guler, Dimitris Samaras, Nikos Paragios, and lasonas Kokki- nos. Deforming Autoencoders: Unsupervised Disentangling of Shape and Appearance, June 2018. UR..tttp:/erxxv.crrr
/106 C6503 10Kiv:18060000
[cs].
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, and Patrick Riley. Tensor field networks: Rotation- and transsationn-equivariant neural networks for 3D point clouds, May 2018. URL UtL.attp//aarri...crr//////22222
11Xiv:1802.08219 [cs].
Wu-Ki Tung. Group Theory in Physics. 1985. ISBN 978-9971-966-57-77
Martin Uhrin. Through the eyes of a descriptor: Constructing complete, invertible descriptions of atomic environments. Physical Review B, 104(14):144110, October 2021. ISSN 2469-9950, 2469-9969. doi: 10. 10.110/hyssseB.... 1...Ree........111.. 144.144110. URL http: http:://arxi...rrr/bbbb///////// orxiv..rg//bss//10.. 09319 arXiv:2104.09319 [cond-mat].
12

---

Under review as a conference paper at ICLR 2023
Pascal Vincent, Hugo Larochelle, Yoshua Bengio, and Pierre- Manzagol. Extracting and composing robust features with denoising autoencoders. In Proceedings of the 25th international conference on Machine learning, ICML '08, pp. 1096-1 103, New York, NY, USA, July 2008, As- sociation for Computing Machinery. ISBN 978- 1--05566-800-44 doi: 10.1145/1
1 044455511916566 1390294]
Maurice Weiler, Mario Geiger, Max Welling, Wouter Boomsma, and Taco Cohen. 3D Steerable
CNNs: Learning Rotationally Equivariant Features in Volumetric Data, October 2018. URL
htttp:///arxiv.or//abb.. 010110254777777
arXiv:1807.022: [es, stat].
Robin Winter, Marco Bertolini, Tuan Le, Frank Noé, and Djork-Amé Clevert. Unsupervised
Learning of Group Invariant and Equivariant Representations, February 2022. URL http:
arXiv:2202.07559 [cs].
Ruidong Wu, Fan Ding, Rui Wang, Rui Shen, Xiwen Zhang, Shitong Luo, Chenpeng Su, Zuofan Wu, Qi Xie, Bonnie Berger, Jianzhu Ma. and Jian Peng. High-resolution de novo structure pre- diction from primary sequence, July 2022. URL https:://///...... www biorxiv. org/ content/
10. 10.1101/2022. 07. 21. 500999v1
Pages:
Section: New Results.
13

---

Under review as a conference paper at ICLR 2023
A APPENDIX
A.1 EXPANDED BACKGROUND ON EEUUAAROAN
A.1.1 INVARIANCE AND EQUIVARIANCE
Let ʃ - a function between two vector spaces and group, where acts on X and via representation Dx and on Y via representation Dy. Then, is said to be G-equivariant iff f(Dx(g)z) = Dy (g) (x), Vo E A Vg E 6. We note that invariance special case of equivariance where Dy (g) I Vg E 6.
A.1.2 GROUP REPRESENTATIIN AND THE IRREPS OF SO(3)
Groups can concretely act on distinct vector spaces via distinct group representations. Formally, group representation defines a set of invertible matrices Dx(g) parameterized by group elements B, which act on vector space X. As 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.
A special kind of representation for any group the irreducible representations (irreps) which provably the "smallest" nontrivial (i.e., they have no nontrivial group-invariant subspaces) represen- tations. The irreps of group are special because it can be proven that any finite-dimensional unitary group representation can be decomposed into direct sum of irreps (Tung 1985). This applies to SO(3) well, whose irreps are the Wigner-D matrices, which are (2 2 di-dimensional matrices, each acting on a (2l - 1) dimeenssio vector space:
Dz(g) for f 0, 1,2...
(A.1)
Therefore, every element of the SO(3) group acting any vector space can be represented direct sum of Wigner-D matrices.
A.1.3 STEERABLE FEATURES
G-steerable vector is vector E X that under some transformation group 6, transforms via matrix-vector multiplication Dx (g)a;; here, Dx is the group representation of E @. For example, vector in 3D Fuclidean space is Stee)-sterrbble since rotates via matrix-vector multi- plication using rotation matrix.
However, we can generalize 3D rotations to arbitrary vector spaces by employing the irreps of SO(3). We start by defining degree-E feature as a vector that is S0(3)-steerable by the fth Wigner-D matrix Dr. Given the properties of irreps. we can represent any stterrble vector as the direct sum of two vector, which we refer to as a tensor to indicate that it is composed of multiple individuall-sstere vectors, is Ste3)-steerable via the direct sum of Wigner-DD matrices of corresponding degrees. This tensor is block-diagonal matrix with the Wigner-D matrices along the diagonal: D(g) = Dr, (g) Di, (B) @ ... Ds, (g).
independent degree-E features, c.g.
It, I2,
@ Et., The resulting
A.1.4 SPHERICAL HARMONICS AND THE SPHERICAL FOURIER TRANSFORM
Spherical harmonics are class of functions that form complete and orthonormal basis for func- tions fe, Φ) defined on a unit sphere; and ó are the azimuthal and the polar angles in the spherical coordinate system. In their complex form, spherical harmonics are defined as,
Yem (0, 2n + 1 (n m)) im pm (cos
(A.2)
4m (n + m))
where E is a non-negative integer (0 < and m is an integer within the interval -( < m < l. P. 0) is the Legendre polynomial of degree l and order m, which together with the complex exponential imm define sinusoidal functions over the angles and ø in the spherical coordinate system. Spherical harmonics are used to describe angular momentum in quantum mechanics.
14

---

Under review as a conference paper at ICLR 2023
Notably, spherical harmonics also form basis for the irreps of SO(3), i.c, the Wigner-D matrices. Specifically, the SO(3) group acts on the eth spherical harmonic via the hth Wigner-D
matrix:
Yem (0.
(ESO
E Drt " (g) em (0.
(A.3)
Therefore, any data encoded in a spherical harmonics basis acted upon by the SO(3) group via direct sum of the Wigner-D matrices corresponding to the basis functions being used. Using our nomenclature, any such data encoding constitutes a steerable tensor. We can thus map any function f(0, defined on sphere into a steerable tensor using the Spherical Fourier Transform (SFT):
fem --m-J f
2m
TT f(0, 6) (0, 6) sin 0 d0 dø J."
(A.4)
The signal can be reconstructed in the real space using the corresponding inverse Fourier transform. For computational purposes, we truneate Fourier expansions at a maximum angular frequency L, which results in an approximate reconstruction of the signal f(0. :p) up to the angular resolution allowed by I,
f(0,0)) - Σ Σ Yem (0, Φ)
(A.5)
2= m=-t
Here, Jim are the functions' Spherical Fourier coefficients.
A.1.5 ZERNIKE POLYNOMIALS AND THE ZERNIKE FOURIER TRANSFORM
To encode a function p(r, 0) with both radial and angular components, we use Zernike Fourier transform,
2 /
p(r, 0, Y (mm (r) dV
(A.6)
where R7 (r)) is the radial Zemike polynomial 3D defined as,
|r| (- (---11 11-1 n
3 z 1rP))
R2 (r) = (-1) 2n +3
3
2 2
(A.7)
n-E
Here, (.) is an ordinary hypergeometric function, and n a non-negative integer representing radial frequency, controlling the radial resolution of the coefficients. RT(r) (r) is non-zero only for even values of n - l 2 O. Zernike polynomials form a complete orthonormal basis in 3D, and therefore, be used within Fourier transform expand and retrieve any 3D shape, if large enough l and n coefficient are used. We refer to the Fourier transform of Eq. A.6 as the Zernike Fourier Trasform (ZFT).
To represent point clouds, common choice for the function p(r) p(r,0, is the sum of Dirac-8 functions centered at each point:
p(r)
Σ 8(p(r:) p(r))
(A.8)
iEpoints
This choice is powerful because the forward transform has a closed-form solution that does not require discretization of 3D space for numerical computation. Specifically, the ZFT of a point cloud follows:
2
Σ (ra) (tm(:)))
(A.9)
sEpoints
Similar to SFT, we can reconstruct the data using inverse ZFT and define approximations by trun- cating the angular and radial frequencies at L and N, respectively,
p(r, 01r,0,5) =
Σ
Σ
2h R3 (r) (0.5)
(A.10)
2=0 m=-
The use of other radial bases is possible within our framework, as long as they are complete. Or- thonormality is also desirable it ensures that each basis encodes different information, resulting
15

---

Under review as a conference paper at ICLR 2023
in more efficient encoding of the coefficients. We use Zemike polynomials following [oyd & Yu (201T), which demonstrates that encoding with Zernike polynomials result in a faster con- vergence compared to the radial basis functions localized at different radii, well as most other orthogonal harmonic bases, with the exception of Logan-Shepp. "Faster convergence" indicates that fewer frequencies required to encode the same information. [Chrin (2021) also uses Zernike to construct invariant descriptors of atomic environments. Other equivariant methods use Bessel func- tions (Musaelian et al. 2022), though, according Boyd Yu (2011), Zernike encoding results in faster convergence.
A.2
DETAILS OF H-(V)AE COMPONENTS
A.2.1 LINEARITY
Let us consider a feature he containing C' features of the same degree f. he can be represented as a C' x (2(+1) matrix where each row constitutes an individual feature. Then, we learn weight matrix W. E RCCK that linearly maps hi to hig E R ((2+1));
h = W he
(A.11)
A.2.2 EFFICIENT TENSOR PRODUCT (ETP)
Channel-wise tensor product nonlinearity. We effectively compute C tensor products, each between features belonging to the same channel c. and concatenate all output features of the same degree. 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. This procedure reduces the computational time and the output size of the nonlinear layers with respect to the number of channels C, from (CC²) for a "fully-connected" tensor product down to OCC). The number of learnable parameters in linear layer are proportional to the size of the output space in the preceding nonlinear layer. Therefore, reducing the size of the nonlinear output substantially reduces the complexity of the model and the number of model parameters. This procedure also forces the input tensor to have the same number of channels for all degrees. We refer the reader to Cobb et al. 202T for further details and for nice visualization of this procedure.
Minimum Spanning Tree (MST) subset for degree mixing. To compute features of the same degree lg using the CG Tensor Product, pairs of features of varying degrees may be used, up to the rules of the CG Tensor Product. Specifically, pairs of features with any degree pair (f1, E2)
may be used to produce a feature of degree &a as long as |e, &2 < la S E1 + fz. Features of the same degree are then concatenated to produce the final equivariant (steerable) output tensor.
Since each produced feature (often referred to "fragment" in the literature Kondor et al. (2018); Cobb et al. (202T) is independently cquivariant, computing only a subset of them still results in an equivariant output, albeit with lower representational power. Reducing the number of computed fragments is desirable since their computation cannot be easily parallelized. In other words, to reduce complexity we should identify a small subset of fragments that can still offer sufficient representational power. In this work we adopt the "MST subset" solution proposed by Cobb et al.) (2021). which adopts the following strategy: when computing features of the same degree la- exclude the degree pair (fg, (2) if the (fo. E,) and the (41, 6x) pairs have already been computed. 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.
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 single degree és. given the maximum desired degree Lman- 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 (f1, 42) gets weight (2f, + 1)(2(2 + 1)). The subset is also augmented to contain all the pairs with same degree inject more nonlinearity. This procedure reduces the complexity in number of pairs with respect to fmax from ( when all possible pairs are used down to (ffaa)). We refer the reader to Cobb al. (02T) for more details and for nice visualization.
16

---

Under review as a conference paper at ICLR 2023
A.2.3 BATCH NORM
Let us consider batch of steerable tensors which index by batch b, degree f. order m and channel c. During training, we compute batch-averaged norm for each degree and each channel as,
(hem))
(A.12)
B
m=-
Similar to standard batch normalization, we also keep running estimate of the training norms Ne.e(()) using momentum ç, set 0.1 all our experiments:
No Ne = < + (1 ()
(A.13)
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 trans- formation:
high N
training
(A.14)
Em
hehh Ner)r())
&m
w
evaluation
(A.15)
A.2.4 PAIRWISE INVARIANT RECONSTRUCTION LOSS
To reconstruct signal within equivariant model it is desirable have pairwise invariant reconstruction loss, i.e., a loss Cree such that (r, y) Crec(D(g)et) Dgg)y) where D is the representation of the group element g acting on the space that æ and y inhabit (e.g. rotation matrix if æ and are vectors in Euclidean 3D space, or degree-E wigner-D matrix if x: and y are degree- vectors). This property is necessary for the model remain equivariant, i.e., given that the network is agnostic to the transformation of the input under group operation + D(g)æ by producing similarly transformed output y + (ggy,, we want reconstruction loss be agnostic to the same kind of transformation as well
The MSE loss is pairwise invariant for any degree- feature which SO(3) acts via the f's Wigner-D matrix. Consider two degree-l features æ; and yr acted upon by a Wigner-D matrix De(g) parameterized by rotation g (we drop the and indexing for clarity):
MSE Dy) (Dx Dy) (Dx Dy)
(D y))² (D(x -y)
(x y) D D(x - y)
(x: - y) (a -y)
since Wigner-D matrices are unitary
MSE(x,
(A.16)
Since the MSE loss is pairwise invariant for every pair of degree-l features, it is thus pairwise invariant for pairs of steerable tensors composed via direct products of steerable features.
A.2.5 COSINE
Proof that Cosine loss is pairwise invariant. The generalized dot product @ from Eq. 5 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. Therefore, the whole Cosine loss expression is pairwise invariant, since all of its components are pairwise invariant.
On the use case of Cosine loss. 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 pro- vides estimate for how much better the reconstructions can get if the network's hyperparameters
17

---

Under review as a conference paper at ICLR 2023
were to be further optimized. For example, looking at the Cosine loss in Table A.6 that our model trained on 1hrec17 (best Cosine = 0.00..1 0.130) is not as well optimized as our model trained on MNIST (best Cosine 0,017). Using MSE, the trend is reversed (1.8x10) -3 vs. 6.7x10 "), since the scale of MSE depends on the size of the irreps of the data (Fig. A.4
A.3 TOY AMINO-ACIDS
We train H-(V)AE on single amino acids represented as atomic point clouds, ex-
tracted from structures in the Protein Data Bank (PDB) (Berman et al.] 2000).
ALA
CYS
SER
We collect atomic point clouds of 50k residues from PDB evenly distributed
750
THh
across residue types and apply a 40/40/20
train/valid/test split. Residues of the same type have different conformations and nat- urally have noisy coordinates, making this problem a natural benchmark for our method.
GLM
25
GL
45
ANC
We consider atom-type-specific clouds (C, O, N and S: we exclude H) centered at the residue's s Ca and compute the ZFT (Eq. [ with L and N 20 within a radius of 10A from the residue's Co, and concate- nate features of the same degree, resulting in tensor with 940 coefficients. We train several H-AE and H-VAE models, all with
Figure A.1: H-VAE on amino acids. H-VAE was trained on 1,000 residues with 0.025 and = 2. The invariant latent space clusters by amino acid confor- mations. The highlighted clusters for PHE and TYR contain residue pairs with similar conformations; TYR and PHE differ by one oxygen at the end of their ben- zene rings. We compare conformations by plotting each residue in the standard backbone frame (right); and y axes are set by the orthonormalized Ca-N and Ca-C vectors, and a axis is their cross product.
2; see Sec. [c.6,[.6.]ffo details.
We consistently find that the latent space clusters by amino acid conformations (Fig.AA.I). with sharper cluster separations as more training data is added (Fig. A.88
and (.9). We find that test reconstruction loss decreases with more training data but the reconstruc- tion accurate even with little training data (from 0.153 Cosine loss with 400 training residues to 0,034 with 20,000); Table A.4 A similar trend is observed for dased classification accuracy of residues (from 0.842 with 400 training residues to 0.972 with 20,000); (Table A.4). Notably, an untrained model, while achieving random reconstruction loss, still produces an informative invari- ant 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.5
A.4 USING A NON-EQUTVARIANT DECODER: EXTENDED DISCUSSION
Winter al. 2022) propose to construct group-equivariant autoencoders by using 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. By contrast, for SO(3) we propose to use an equivariant decoder, whereby the learned group element is fed as input the decoder. Such "unconstrained decoder" procedure can in principle be merged with our equivariant encoder and Fourier-space approach in two ways. For each, we argue in favor of using our equivariant decoder.
1) Reconstructing the Fourier coefficients of the data. To apply the learned group ele- ment on the decoder's output, the Wigner-D matrices for the data's irreps need to be computed from the group element. Then, the Wigner-D matrices can be used to "rotate" the tensor. This has to be done on-the-fly, fly, and it can be done quickly using functions provided in the e3nn package (Geiger & Smidt 2022) and by smartly vectorizing operations. We implemented this procedure by using a simple Multi-Layer Perceptron with SiLU non- -ineeaiassss as a decoder. By using e3nn to compute Wigner-D D matrices in batches, and by clever construction of tensor multiplications such that runtime scales linearly with lax and is constant with regards number of channels and batch size,
18

---

Under review as a conference paper at ICLR 2023
achieve models that run with comparable speed those using our equivariant decoder, and have comparable performance on MNIST (Table A.I). Given the empirical similarities we observe, though a limited use case, favor the simplicity and elegance of equivariant decoder. "Simplicity" because we construct the decoder to be symmetric to the encoder, thus endowing it automatically with similar representational power and without the need tune architecture made with different base components. Furthermore, we highlight that our method generates intermediate equivariant representations in the decoder, rather than intermediate invariant representations. These intermediate equivariant representations may be of interest to study in and of themselves.
2) Reconstructing the data in real space. In this case, we do not have to compute Wigner- matrices on-the-fly, since the learned frame can be used directly in the output space as a rotation matrix. However, since the encoder only sees truncated Fourier representation of the data, which is by construction lossy, while the loss is computed over fine- fine-grained real-space, this model might be too difficult to train. 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. To avoid the denoising effect, could learn to reconstruct data in real space after an Inverse Fourier Transform (IFT). However, computing the IFT on-the-fly is very expensive and requires discretization of the input space, the point of being prohibitive for point clouds. This is not a bottleneck for the Forward Fourier Transform if the cloud is parameterized by Dirac-Delta distributions. i.e., for point clouds, as the integral can be computed exactly (Eq. A.9).
Table A.I: Performance comparison between our H-(V)AE and H-(V)AE with Winter et al. (2022)'s 's non-equivariant decoder formulation, on the LNNIST-on-th--pphee dataset. The non- equivariant decoders are constructed simple MLPs with SiLU non-linearities, with the following hidden layer sizes: [32,64,128, 12,,,6,,566,)) 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. All other training details are kept the same (Sec. A.6).
LC Class.
KNN Class. Ace.
Method
Speed
MSE
Cosine
Purity
V-meas.
Acc.
H-AE NR/R
1.0
9.3
0.025
0.62
0.51
0.820
0.862
H-AE unconst. decuder NR/R
16
1.3
1.3 10-3
0.037
0.61
0.48
0,802
0.856
H-VAE NR/R
1.0
1.4 10
0.057
0.67
0.54
0.812
0.848
H- H-VAF unconst. decoder NR/R
16
1.3
2.1 10
0.057
0.62
0.51
0.781
0.853
A.5 IMPLEMENTATION DETAILS
Without loss of generality, we use real spherical harmonics for implementation of H-(V)AE. We leverage e3nn Geiger& & Smidt (2022), using their computation of the real spherical harmonies and their Clebsch- coefficients.
In our code, we offer the option ically, at each block we allow the users to specify whether to compute the Tensor Product channel-wise fully-connected across channels, and whether compute using efficient or fully connected degree mixing.
the Full Tensor Product instead of the ETP. Specif-
A.6 EXPERIMENTAL DETAILS
A. ARCHITECTURE SPECIFICATION
We describe model architectures follows. We specify the number of blocks B, which is the same for the encoder and the decoder. We specify two lists, (i) DegreesList which contains the maximum degree Emar of the output of each block b, and (ii) ChannelsList, containing the channel sizes Cb, of each block b. These lists are in the order as they appear in the encoder, and are reversed for the decoder. When it applies, Omn- As noted in the main text, we use a fixed formula to determine k but we specify it for clarity.
specify the number of output channels of the initial linear projection
19

---

Under review as a conference paper at ICLR 2023
A.6.2 MNIST ON THE SPHERE
Model architectures. For models with invariant latent space size 16, we use 6 blocks, DigreesList [10, 10, 8, 4,2, 1] and ChannelsList = [16, 16, 16, 16, 16, 16). with total of 227k parameters.
For models with invariant latent space size = 120, we use 6 blocks, DegreesList [10, 10,8, 8,4,2,1] and Channelsl.ist = [16, 16, 16, 32, 64, 120), with total of 453k parame- ters.
Training details. We keep the learning schedule similar as possible for all models. We use a = 50. We train all models for 80 epochs using the Adam optimizer (Kingma & Ba 2017) with default parameters, batch size of 100, and initial learning rate of 0.001, which we decrease exponentially by one order of magnitude over 25 epochs. For VAE models, we use 3 = 0.2, Ersc = 25 and Ewarmpp 35. We utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models. Training took ~ 4.5 hours on single NVIDIA A40 GPU for each model.
A.6.3 SHREC17
Model architectures. Both AE and VAE models have 2 =40, blocks, DegreesList [14, 14, 14, 8, 3, 1], ChannelsList [12, 12, 12, 20, 24, 32, 40), Cinit = 12, with total of 518k parameters.
Training details. We keep the learning schedule as similar as possible for all models. We 1000. We train all models for 120 epochs using the Adam optimizer with default parame- ters, 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. For VAE models, we use B = 0.2, Free 25 and Ewammup = 10. We utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models. Training took 11 hours single NVIDIA A40 GPU for each model.
A.6.4 TOY AMINO ACIDS
Pre-processing of protein structures. We sample residues from the set of training structures pre- processed as described in Sec. A.6.5
Fourier projection. We set the maximum radial frequency to N 20 as it corresponds to radial resolution matching the minimum inter-atomic distances after rescaling the atomic neighborhoods of radius 10.0A to fit within a sphere of radius 1.0, necessary for Zernike transform.
The channel composition of the data tensors can be described in notation analogous to that used by e3nn but without parity specifications which specifies the number of channels C' for each feature of degree single units Cxê: 44x0 40x1 + 40x2 + 36x3 + 36x4.
Model architectures. All models have 2 = 2, 6 blocks, DegreesList = [4, 4, 4, 4, 1], ChannelsList = [60, 40, 24, 16, 16, 8). Cinit = 48, with total of 495k parameters. We note that the initial projection necessary since the number of channels differs across feature degrees in the data tensors.
Training details. We keep the learning schedule similar as possible for all models. We use a = 400. 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. For VAE models, we use E = 25 and Ewarmup = 10. We utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models. We vary the batch size according to the size of the training and the validation datasets. We use the following size-bttssze-bthhhee size) pairs: (400-4), (1,000-10), (2,000-20), (5,000-50), (20,000-20). Training took ~ 45 minutes on a single NVIDIA A40 GPU for each model.
Evaluation. 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. We keep relative proportions of residue types even in all datasets. We perform the data ablation with H-AE as well as H-VAE models with
20

---

Under review as a conference paper at ICLR 2023
3 = 0.025 and 0.1. We further perform a 3 ablation using the full (20,000) dataset, over the following choices of 3:: [O(AE), 0.025, 0.05, 0.1, 0.25, 0.5].
For robust results, train 3 versions of each model and compute averages of quantitative metrics of reconstruction loss and classification accuracy.
For fair comparison across models with varying amounts of training and validation data, we perform a 5-fold cross-validation-litc procedure over the 10k test residues, where the classifier is trained over 4 folds of the test data and evaluated on the fifth one. If validation data is needed for model selection (e.g. for LC), we use 10% of the training data.
A.6.5 PROTEIN NEIGHBORHHOODS
Pre-processing of protein structures. We model protein neighborhoods extracted from tertiary protein structures from the Protein Data Bank (PDB) (Berman et al. 22000). We use ProteinNet's splittings for training and validation sets to avoid redundancy, e.g. due similarities homologous protein domains (AIQuraishi 2019). Since PDB ids were only provided for the training and valida- tion sets, we used ProteinNet's training set as both our training and validation set and ProteinNet's validation set as testing set. Specifically, we make [80%, 20%) split in the ProteinNet's training data to define our training and validation sets. This splitting resulted in 10,957 training structures, 2,730 validation structures, and 212 testing structures. We use pyRosetta (Chaudhury ct al. 2010 to assign Solvent Accessible Surface Area (SASA) to every atom.
Projection details. 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 neighbor- hoods of radius 12.5A to fit within sphere of radius 1.0, necessary for Zernike transform.
Model architectures. All models have 2 64, 6 blocks, DegreesList [4.4.4,4,] 1], ChannelsList [128, 128, 96, 96, 64, 64]. Cunt 64, with a total of 3.5M parameters. We note that the initial projection is necessary since the number of channels differs across feature degrees in the data tensors.
Training details. We keep the learning schedule as similar as possible for all models. We use 400. We train all models for 120 epochs using the Adam optimizer with default parame- ters, batch size 256, and an initial learning rate of 0.002, which we decrease exponentially by one order of magnitude over 40 epochs. For VAE models, utilize the model with the lowest loss on validation data, only after the end of the warmup epochs for VAE models. Training took 10 hours single NVIDIA A40 GPU for each model.
Exec 40 and Ewamup 40. We
A.6.6 LATENT SPACE CLASSIFICATION
Linear classifier. We implement the linear classifier as one-layer fully connected neural network with input size equal to the invariant embedding of size and output size equal to the number of desired classes. 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. 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). At evaluation time, we select the class with the highest probability value. We use PyTorch for our implementation.
KNN Classitier We use the sklearn (Pedregosa al. 2011) implementation with default param- eters. At evaluation time, we select the class with the highest probability value.
A.6.7 CLUSTERING METRICS
Purity (Aldnnderfer & Blashfield, 1984). We first assign a class to each cluster based on the most prevalent class in it. Purity is computed as the of correctly classified items divided by the total number of items. Purity measures classification accuracy, and ranges between 0 (worst) and (best).
21

---

Under review as a conference paper at ICLR 2023
V-measure (Rosenberg & Hirschberg/ 2007). This common clustering metric strikes bal- ance between favoring homogeneous (high homogeneity score) and complete (high completeness score) clusters. Clusters defined as homogeneous when all elements in the same cluster belong to the same class (akin to a precision). Clusters are defined as complete when all elements belonging the same class put in the same cluster (akin to recall). The V-measure is computed as the harmonic mean of homogeneity and completeness in a given clustering.
A.6.8
ON THE COMPLEMENTARY NATURE CLASSIFICATION ACCURACY AND CLUSTERING METRICS
The clustering metrics "purity" and "V-measure" and the supervised metric "classification accu- racy" characterize different qualities of the latent space, and, while partly correlated, they are complementary each other.
Both classes of metries 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 result, these metrics focus on different features of the latent space. For example, the clustering metrics are largest when the test data naturally forms clusters with data points of the same label. While this case can result in 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 subset of the latent space features, instead of relying on all 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 a sign that class-related information is more evenly distributed across the latent space dimensions. Overall, the complementary aspect of these metries makes necessary use all of them when comparing the performance of different models in each task.
22

---

Under review as a conference paper at ICLR 2023
Table A.2: Evaluation of network performances for MNIST-on-the-spher using KNN classifier instead of linear classifier in the latent space. Results are significantly better than when using linear classifier for models with smaller 16) latent space, comparable for the other models.
Type
Method
Acc.
KNN Acc.
H-AE NR/R
120
0.877
H-AE R/R
120
0.881
0.886
H-AE NR/R
16
0.820
0.862
Unsupervised
H-AE R/R
16
0.833
0.876
H-VAE NR/R
120
30
0.883
0.879
VAV R/R
120
30
0,884
0.895
H-VAE NR/R
16
30
0.812
0.848
LA-E R/R
16
30
0.830
0.874
Table A.3: QEvaluation of network performances for Shree17 using a KNN classifier instead of linear classifier in the latent space. Results are better than when using linear classifier.
Class. Acc.
Type
Method
bw
P@N
R@N
FI@N
mAP
NDCG
Unsupervised L.C
H-AE
40
90
0,654
0.548
0.569
0.545
0.500
0.597
H-VAB
40
90
0.631
0.512
0.537
0.512
0.463
0.568
Unsupervised KNN
H-AE
40
90
0.672
0.560
0.572
0.555
0.501
0.599
H-V
40
90
0.658
0.541
0.558
0.539
0.487
0.591
Table A.4: Quantitative data ablation results on the Toy amino acids dataset. A random- classifier has an expected accuracy of 0.050.
H-AE
H-
0.255)
H-VAR (..1
train
MSE 10
Casine los
LCA.s.
KNN Aec.
MSE
Crodiee Joss
LC Arc.
KNN Act:
MSE
Casine lose
LC Acc.
KNN Aec.
10015
0.608
0.6.3
155 10
0.951
0.424
0.55h
1.5xI0
0.581
0.424
(1656
10--4
11:55
0.842
10-1
0.548
0.163
0.750
0.856
10-1
5,000
0.961
20,000
0.034
0.58
0.972
2.4
0.037
0.667
0.971
0.047
0.662
0.966
23

---

Under review as a conference paper at ICLR 2023
Table A.5: Quantitative data ablation results on the Toy amino acids dataset. Models were trained on the full dataset (# train 20,000).
MSE
Cosine
LC Acc.
KNN Acc.
(AE)
2.2 10 10
0.034
0.580
0.972
0.025
2.4 10
0.037
0.666
0.971
0.05
2.5 x 10 10
0.039
0.669
0.968
0.1
2.9 10
0,047
0.661
0,966
0.25
6.8 10
0.132
0.597
0.854
0.5
1.1 10-3
0.203
0.467
0.722
Table A.6: Test MSE and Cosine loss for H-(V)AE models trained on MNIST, Shrec17 and
Protein Neighborhoods. MSE and Cosine values are strongly correlated within datasets but not across datasets.
Dataset
Method
Z
bw
MSE
Cosine
H-AE NR/R
120
30
6.2 10
0.017
H-AE R/R
120
30
6.8 x 10-1
0.018
H-AE NR/R
16
30
9.3
10
0.025
H-AE R/R
16
30
8.9 10
0.024
MNIST
H-VAE NR/R
120
30
1.4 X 10
0.037
H VAE R/R
120
30
1.4 10
0.037
H-VAE NR/R
16
30
2.2 10
0.057
H-VAE R/R
16
30
2.1 10
0.055
H-AE
40
90
1,8 x 10
0.130
Shrec 17
H-VAE
40
90
2.2 10
0.151
Protein Neighborhoods
H-AE
40
90
6.1 10
0.161
Table A.7: Training speed and reconstruction ablations of H-(V)AE models with different Ten- sor Product rules. To make comparison fair, models were trained using the same training hy- perparameters as described in A.6 and all models of parameters. Speed was computed as training time and divided by the time of the model using ETP within each dataset. Models with the ETP consistently generates better reconstructions and are usually the fastest. The speed and performance gains of the ETP are most apparent on the Protein Neighborhoods task.
constructed to have comparable number
Dataset
Method
TP-type
Cinis
ChannelaList
DegreesList
Params
Speed
MSE
Cosine
H-AE
ETP
Noue
[66.16.66.16.6.100
[0.10.8422.1]
227k
10
MNIST
H.AF
Full-TP
None
[7.6.66616
[10,8,4.21]
229k
1.1
1.6
10
0.044
H-AE
Full-TP
Nooc
[5.5.5.5.716]]
[10.10.8.4...]]
227k
1.7
1,5
H.AF
FTP
12
(12.12.12.30:24333..000
(14.14,14.842.1
518k
1.0
1,8x100
0.130
H-AE
Rull-TP
None
[5.5.5.8.40]
[14.844221]
518k
0.9
1,9 10
0.137
H-AE
Full-TP
Noue
(4.3.3.6.6.6.4))
(14,14,14.842.1
513k
HEAE
ETP
64
1126.125.96996664..66
[44442.11
3.5M
L.
6.1 10
0.161
[17,24.64]
[42.1]
H-AE
Pull-TP
56
[17,17,64]
[42.
3.6M
2.8
7.5 10
0.207
24

---

Under review as a conference paper at ICLR 2023
for some of our trained H-(V)AE models. Errors computed over 2,000 randomly sampled spherical tensors, each with a randomly sampled rotation. Standard deviation is shown alongside the mean. We also show the mean and standard deviation of the absolute value of the output coefficients, to enable contextualization of the measured equiv- ariance error. The equivariance error due to numerical error (absolute difference in coefficients by rotating input vs. output tensor) is consistently three orders of magnitude lower than the typical absolute value of coefficients, indicating that equivariance preserved. The same trend occurs for untrained models (not shown here for simplicity).
Table A.8: Mean equivariance
Dataset
Method
z
# train
Equiv. Error
Abs. Value
H-AE NR/R
16
(2.3 + 2.7)x10-4
(1.0 0.9)x 10-1
MNIST
H-VAE NR/R
16
(2.5 1.9)x10-4
(1.3 1: 1.8)x 10
H-AE
40
(1,412.4)10
(0.4 1 1.I)xx0
Shrec17
H-VAE
40
31,4
(0.4 + 2.2) 10
H-AE
2
1,000
(4.7 4.7)
(2.6 4.1)x
H-VAE 0.025
2
1,000
(5.4 (.9)xx11
(2.8 1 4.0)x 10-2
Toy Aminoacids
H-AE
20,000
(9.3 1 6.4)x100
(3.5 4.7)x11
H- H-VAE 3 0.025
2
20,000
(9.0 + 4.7)x10-5
(3.0 4.3)x1000
Protein Neighborhoods
H-AE
64
(1.8 + 2.0)x10
(1.0 1.7)x10
H-AE
0.7
Yes Batch Norm
No Batch Norm
0.5
Loss
0.4
0.3
0.1
10
20
30
40
50
60
70
80
Epoch
Figure A.2: Training loss trace of H-AE, with and without Batch Norm, MNIST-on-the- sphere. Models were trained with the (NR/R; z = 16; AE) specification. The loss on validation data follows the same trend, but it is not shown for simplicity.
25

---

Under review as a conference paper at ICLR 2023
AE NR/R
AE R/R
VAE NRR/R
VAE R/R
J
Trained with all digits
A
J
V
J
J
Trained with Is and 7s only
Figure A.3: Empirical tests to show H-(V)AE learns canonical frame in the MNIST-on-the-
sphere. For each of the 4 models with 2 = 16, we train a version using only images containing Is visualize the sum of training images of digits I and 7,
and 7s. For each of the resulting 8 models, when rotated to the canonical frame. To achieve the canonical frame we rotate each image such that the frame learned by the encoder is the same for all images, and specifically, to the x 3 identity matrix. We compute the sums of images with the same digit, and overlay them with different colors for ease of visualization. We the hypothesis whether H-(V)AE learns frames that align the training images such that they maximally overlap; we do so in two ways.
make correspond
First, if the hypothesis were true, all canonical images of the same digit should maximally or near- maximally overlap- since they have very similar shape and thus, their overlays would look like a "smooth" version of that digit. Indeed, we find this statement be true for all models irrespective of their training strategy.
Second, we consider the alignment of images of different digits. We take and 7s as examples given their similarity in shape. If the hypothesis were true, models trained with only Is and 7s should align canonical 1s along the long side of canonical 7s; indeed we find this to be consistently the case. The same alignment between is and 7s. however, does not necessarily hold for models trained with all digits. This is because maximizing overlap across a set diverse shapes does not necessarily maximize the overlap within any independent pair of such shapes. Indeed, we find that canonical 1s and canonical 7s do not overlap optimally with each other for models trained with all digits.
We note that these tests do not provide proof, but rather empirical evidence of the characteristics of task.
frames learned by H-(V)AE on the
26

---

Under review as a conference paper at ICLR 2023
MNIST on the sphere
Shrec17
1.2
1.0
Cosine loss
0.6
0.2
SpearmanR: 0.82
SpearmanR: 0.95
(at low loss: 0.97)
(at low loss: 1.00)
0.0
0.2
1.0
0.0
0.4
0.6
0.8
1.0
1e-2
Toy amino-acids
Protein neighborhoods
1.2
1.
Cosine loss
0.8
0.4
0.2
SpearmanR:
SpearmanR:
(at low loss: 0.99)
(at low loss: 0.99)
0.5
1.5 MSE
0.6
1.0
1.4
1e-3
MSE
1e-3
Figure A.4: Correlation between Cosine loss and MSE values between pairs of random tensors. For each dataset, we sample a batch of = 1000 tensors with dataset-specific feature degrees and channel sizes, where each coefficient normalization step performed in the real experiment and normalize each tensor by the average total norm of the batch. We then generate a "noisy" version of each tensor by adding (normalized) Gaussian noise to each coefficient with standard deviation sampled from a uniform distribution between 0 and some maximum noise level (10 in these plots). This procedure results N pairs of tensors with varying degrees of similarity between them. We compute the MSE and Cosine loss for all N pairs of tensors and visualize them. The two loss values are well correlated rank as measured by Spearman Correlation. The correlation is significantly stronger in the regime of reconstruction loss below a Cosine loss of 0.5 (SpearmanR 0.99), a value well above the maximum Cosine loss achieved by H-(V)AE in all our experiments. All the p-vulues for the Spearman Correlations shown the plot are significant (< 0.05).
sampled from normal distribution. We mimic the
27

---

Under review as a conference paper at ICLR 2023
Q
C
R
R
Q
R
P
8
J
P
P
P
S
R
F
5
P
3
D
8
/
K
R
0
P
e
Figure A.5: Random samples generated by the (NR/R; z 16; VAE MNIST-on-the-sphere model. We sample invariant latent embeddings from the prior distribution (isotropic normal) and feed them the decoder alongside the canonical frame to generate tensors We then compute the inverse SFT to map the generated tensor to images in real space. The samples show a wide range of diversity in digit and style.
28

---

Under review as a conference paper at ICLR 2023
3
3
3
3
3
C
C
3
3
3
3
8
R
D
0
0
0
0
Q
6
9
9
3
4
D
0
0
0
8
8
S
S
F
/
/
/
/
P
P
P
P
P
P
/
/
4
Q
Q
6
Q
8
8
8
8
8
8
8
2
2
2
P
R
R
8
8
8
8
8
8
8
2
2
2
2
2
Figure A.6: Trajectories across the latent space for the (NR/R; z 16; VAE) MNIST-on-the- sphere model. We compute pairs of invariant latent embeddings using the model's encoder, and linearly interpolate between them through the latent space. We then feed the interpolated embed- dings into the decoder, together with the canonical frame, and compute the inverse SFT to get the image real space. The left and right columns show the original images (after forward and inverse SFT) rotated to be placed in the learned canonical frame, whereas the center rows show the inter- polated images. We can see that all trajectories are smooth, respecting spatial consistency, sign of well structured latent space.
Figure A.7: 2D visualization via UMAP of the invariant latent embeddings of 1hrec17 test data
learned by H-(V)AE. Left: H-AE, Right: H-VAE Points are colored by class (55 classes).
29

---

Under review as a conference paper at ICLR 2023
400
1,000
2,000
5,000
20,000
Figure A.8: Amino acid latent space learned by H-AE. Visualization of the test data's 's invariant la- tent space learned by H-AE trained with varying amounts of the training data. As more training data is added, the separation of clusters containing residues with most similar conformations becomes more distinct. Notably, even with no training data, conformation clusters can identified.
Figure A.9: Amino acid latent space learned by H-VAE. Visualization the test data's invariant latent space learned by H-VAE (3 0.025) trained with varying amounts of training data.
30

---

Under review as a conference paper at ICLR 2023
Number of atoms
250
200
150
100
50
Figure A. 10: Protein neighborhood latent space learned by H-AE annotated by the constituent atoms. 2D UMAP visualization of the 64-dimensional invariant latent space learned by H-AE, colored by number the atoms in the neighborhood shown. A clear gradient can be seen from the bottom of the plot to the top.
Average SASA
6
5
4
3
2
Figure A.11: Protein neighborhood latent space learned by H-AE annotated by SASA 2D UMAP visualization the 64-dimensional invariant latent space learned by H-AE, colored by the average Solvent Accessible Surface Area (SASA) computed for all atoms in each neighborhood is shown. A larger average SASA indicates that larger proportion of the neighborhood is at the pro- tein's surface. The surface neighborhoods are concentrated the lower-left side of the map, and the buried neighborhoods are concentrated at the top. Predictably, there is an inverse correlation between the number of atoms (Fig. A.10) and closeness to the surface, as do not represent protein's ss surrunnding environment.
31

---

Under review as a conference paper at ICLR 2023
1e-3
MNIST on the sphere
1e 4
Toy amino acids
2.0
1.5
2
0.5
0.0
Feature degree f
Feature degree /
Reconstruction loss as function of feature degree L. Test reconstruction loss (MSE) of H-VAE split by feature degree f, for the MNIST-on-the-spher (left) and Toy amino acids dataset (right). In both cases, features of larger degrees harder to reconstruct accurately. The increase in loss is more steep for smaller degrees.
Figure A.12:
J
P
2
Q
3
Figure A.13: Visual proof of the disentanglement in the latent space of MNIIS-on-the-ppeers For each row, the invariant embedding z is held fixed, and different frame (i.e., the rotation matrix) used. Frames are sampled randomly and differ across rows, with the exception the first column, which is always the identity frame. Then, and the frame are fed to the decoder and the Inverse Fourier Transform used to generate the reconstructed spherical image, which projected onto plane for the ease of visualization. Modulo the distortions given by projecting the image onto plane, it clear that the invariant embedding contains all semantic information, and the frame solely determines the orientation of the image.
32


Reviews:
[{"review_id": "QVdeGw_Gat", "reviewer": "Reviewer_Wgbf", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "The authors propose an SO(3) equivariant variational autoencoder, acting on steerable features. Using standard datasets, they examine the accuracy of reconstruction, and whether the latent representations of objects are useful for object classification and regression tasks.", "summary_of_the_review": "The model is new, and I believe it may prove to be useful, but I would like the usefulness to be shown empirically. I would like the authors to go beyond 'let's use spherical harmonics in a VAE' to 'here's what it can do that a simpler E(3) equivariant VAE can't do'.", "strength_and_weaknesses": "I think the idea of the paper is that using steerable features and equivariant models for unsupervised learning should yield good representations for downstream prediction tasks.  It is a reasonable hypothesis, but it wasn’t clear to me if the results really support it. Of the tasks presented, which ones are already better solved in simpler ways? (I’m not familiar with literature on unsupervised classification of 3D shapes).\n\nE(3) equivariant autoencoders already exist, for example in Satorras et al. (2022) https://arxiv.org/pdf/2102.09844.pdf, but I have not previously seen examples using steerable features.  Introducing spherical harmonics etc. adds complexity, which ideally would be justified by results.  Maybe the cited paper by Brandstetter et al. could be a source of ideas about when the extra complexity is worth it.", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "1: The contributions are neither significant nor novel."}, "extra_text": {}}, {"review_id": "F1occn1NwV", "reviewer": "Reviewer_FEZB", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This work presents H-VAE, an end-to-end SO(3)-equivariant VAE.  The input is a radial signal in $R^3$ called $\\rho(r,\\theta,\\phi)$ which is mapped via the Zernike Fourier Transform to a set of coefficients of Zernicke polynomials which are products of spherical harmonics and radial basis functions.  The SO(3)-equivariant encoder and decoder both operate in Zernike Fourier space, using block diagonal linear maps, tensor product, and Clebsch-Gordon decomposition to achieve equivariance.  The latent space is formed by a single element of SO(3) and an invariant vector.  The authors demonstrate in experiments that this results in the model learning a canonical frame for inputs.  That is, inputs are essentially rotated to be aligned to a canonical frame and the rotation is stored along with an invariant encoding of the canonically rotated input.  The authors demonstrate their method in 4 experimental domains, spherical MNIST, SHREC 17 3D model classification, Amino Acids, and Protein neighborhoods.  The evaluation shows the model achieves good clustering by class and meaningful alignments.  The model also has good reconstructions and can be sampled from to produce reasonable outputs.  ", "summary_of_the_review": "This paper contains a promising approach and is well-written but the motivation and experiments seem incomplete to me.  My score is tentative and I willing to adjust it based on discussion.  ", "strength_and_weaknesses": "## Strengths\n- The model is well-described and reasonable.  In contains block diagonal linear layers acting in Fourier space (as found in other SO(3)-equivariant methods, e.g. Thomas '18), tensor product (ETP) non-linearities (Kondor et al., '18) instead of radial non-linearities.  The authors introduce Signal norm to stabilize learning.  The latent space uses a single SO(3) element constructed using Gram-Schmidt and an invariant vector, a design inspired by Winters '22.  While most of these elements originate elsewhere I think they make for a reasonable design here which, so far as I can tell, is the first end-to-end SO(3)-equivariant VAE.     \n- The part of the experiments I found most interesting was the idea to insert the identity into the decoder and thus see the output in canonical orientation.  From the results, it does seem the model learns to align inputs leading to a disentanglement of frame and identity which is quite interesting and could have potential applications in 3D alignment or classification problems in which it is difficult to classify objects without disentangling pose.  \n\n## Weaknesses\n- The motivation seems a bit weak to me.  The experiments appear to use relatively small point clouds or spherical signals at a single radial distance.  The latent space embeddings do not seem to lead to better performance in downstream tasks.  They do seem better clustered, have nice interpolations, and can be easily sampled, but why is this desirable and is it not true of previous methods?  I would wager there are good answers to these questions, but I didn't get a strong impression from the paper.  In particular, it would nice to see applications in which orientation-identity disentanglement is critical.\n- The paper could use an ablation study.  How do features such as the Fourier space, the Zernicke basis, ETP, signal norm, and a small degree latent space contribute to success?  In particular, most SO(3) methods use simple radial basis functions localized at different radii.  Why is the Zernicke basis superior? Does orthogonality contribute so useful property?  \n- The paper could use more baselines.  There are many VAEs which could applied in the experimental settings as well as other generative models.  Could some of the claimed features (such as frame-aligned latent vectors), better clustering, etc. be realized by models with different generative paradigms or without equivariance or without equivariance but trained with data augmentation or using some form of canonicalization?  My hypothesis would be the proposed method is better at disentangling the frame, but this could be demonstrated.  \n- While the theory is sound, the equivariance of the model should be empirically confirmed to verify correct implementation and quantify any error arising from approximation.\n- I have a concern/question about the decoder and the size of the latent space.  Given that the input has features with high frequency $l_{max}$ it seems a bit strange to me go down to just an element of SO(3) and invariant features.  The reconstructions do seem to show this is adequate and since the encoder and decoder and complex non-linear functions, I don't think there is any theoretical problem, but it still strikes me as an unusual design choice.   Other works use higher-dimensional mixtures of higher frequency harmonics to encode latent representations of spherical signals.  I'm not sure this is the wrong choice, but I'd like to see the comparison.  \n- I'm not sure, but here is one version of my concern: the decoder takes $g \\in SO(3)$ and $z \\in \\mathbb{R}^n$ invariant, so $d(g,z) = g d(I, z)$.  However, this uses up the equivariance constraint, so that $d(I,z)$ is unconstrained.  It could thus be a simple MLP with g acting at the end and have the same equivariance property.  So what is gained by using the given architecture? \n\n\n## Questions\n- I'd like a proof signal norm is equivariant.  It seems like summing the components of the vectors would not be, but maybe I've misunderstood the notation.  \n- Is it possible to use a quantitative measure of for the quality of the random samples from the prior?  It may make it easier to compare to other work. \n- Is it not possible to fill in some missing values in table 1 in some way? Why do the supervised methods lack class. Acc.? ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "0Y-_SOignV", "reviewer": "Reviewer_nuL1", "summary": "", "questions": "", "limitations": "", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "This paper proposes an end-to-end $SO(3)$-equivariant (V)AE. The idea is to learn, in the spherical Fourier space, an dis-entangled embedding of the input datum -- one corresponding to the invariant embedding describing the datum in a \"canonical\" orientation, and the other being the equivariant frame describing the datum's orientation. Experiments have been conducted trying to verifying the authors' claim that the proposed model can achieve good unsupervised clustering and classification results on spherical images and protein structures.", "summary_of_the_review": "Although the idea is interesting, there is a lot of room for the paper to improve.", "strength_and_weaknesses": "**Strength**\n1. Inspired on the prior work by Winter et al. (2022), the idea of learning disentangled representations corresponding to the group-invariant embedding of the data as well as its equivariant frame describing its $SO(3)$ rotation is interesting. \n2. The authors have tried to conduct extensive experiments to verify the usage of the proposed model.\n\n**Weakness**\n1. The writing and structure of the paper can be significantly improved to make the paper clearer. For example, most of the technical terms introduced in Section 2 need further explanation for clarity, and many definitions and explanation of the appendix A.1 should be moved to the main text.\n2. Since the theoretical contribution of the paper is limited, one would expect the authors to better explain the architectural design of the proposed model. However, most of the implementation details are buried in the appendix, which makes Section 3 of the main text especially confusing.\n3. The authors claim that they achieved state-of-the-art performance in unsupervised clustering and classification. However, the only method to which the paper compared is an arxiv paper by Lohit & Trivedi, 2020. This makes the claim less compelling.\n4. Even though the authors explained why MSE is not a \"good\" measure of the reconstruction, the models are still trained using MSE. It is thus reasonable and fair to compare the results using MSE.\n5. It is interesting to see whether the model indeed learns disentangled representation. An easy experiment to conduct is to change the frame in the latent space with a fixed $z$. I am wondering whether this will generate a rotated copy.\n6. In table 1, the previous work by Lohit & Trivedi achieves the best result when measured in \"purity and V-means\", while it achieves the worst result when measured in \"classification accuracy\". This makes me wonder whether the valuation metric proposed by the authors are convincing.\n7. Since Shrec17 is 3D data to begin with, what is the rationale of projecting it to a sphere?", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}]

Author rebuttal:

Decision: Reject