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 show the potential utility of H-(V)AE on structural biology tasks. Specifically, we train H-(V)AE on protein structure microenvironments, and show that its latent space can be used to extract compact embeddings of local structural features which, paired with a Random Forest Regressor, enable state-of-the-art predictions of protein-ligand binding affinity.