Holographic-(V)AE: an end-to-end SO(3)-Equivariant (Variational) Autoencoder in Fourier Space

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.

HERMES: Holographic Equivariant neuRal network model for Mutational Effect and Stability prediction

“…HERMES architecture. We build upon improvements to the HCNN architecture presented in [2], leading to a model that is ∼2.75x faster, more accurate, with comparable number of parameters (∼3.5M) as the architecture in [1] (Figure 1A). In short, atomic neighborhoods -i.e., featurized point clouds -are first projected onto the orthonormal Zernike Fourier Basis, centered at the (masked) central residue's Cα.…”

“…Finally, the embedding is fed to an MLP to generate the desired predictions. We refer the reader to [2] for details of the architecture, and for a mathematical introduction to building SO(3)-equivariant models in Fourier space. We implement HERMES using e3nn [31].…”

“…Among the recently-developed self-supervised models is Holographic Convolutional Neural Network (HCNN) [1]. operation and foregoing of invariant skip connections makes this architecture 2.75x faster and with better performance on the pre-training task than the architecture in [1]; we refer the reader to [2] for details of the architecture. (B) Pre-training procedure.…”