ProGAE: A Geometric Autoencoder-based Generative Model for  Disentangling Protein Conformational Space

1 INTRODUCTION . The complex and time-consuming calculations in molecular simulations have been significantly impacted by the application of machine learning techniques in recent years . In particular , deep learning has been applied to analysis and simulation of molecular trajectories to address diverse problems , such as estimating free energy surfaces , defining optimal reaction coordinates , constructing Markov State Models , and enhancing molecular sampling . For a comprehensive review of deep learning methods for analyzing and enhancing molecular simulations , see ( Noé et al. , 2020a ) and ( Noé et al. , 2020b ) . Specifically , there has been interest in modeling the underlying conformational space of proteins by using deep generative models , e.g . ( Ramaswamy et al. , 2020 ) and ( Bhowmik et al. , 2018 ; Guo et al. , 2020 ; Varolgüneş et al. , 2020 ) . This line of work has mainly attempted to respect the domain geometry by using convolutional AEs on features extracted from 3D structures . In parallel , learning directly from 3D structure has recently developed into an exciting and promising application area for deep learning . In this work , we learn the protein conformational space from a set of protein simulations using geometric deep learning . We also investigate how the geometry of a protein itself can assist learning and improve latent conformational space interpretability . Namely , we consider the influence of intrinsic and extrinsic geometry , where intrinsic geometry is independent of 3D embedding and extrinsic is not . Intrinsic geometric protein properties can be thought to be robust to conformation . To this end , we propose a Protein Geometric Autoencoder model , named ProGAE , to separately encode intrinsic and extrinsic protein geometries . The main contributions of this work are summarized : • Inspired by recent unsupervised geometric disentanglement learning works ( Tatro et al. , 2020 ; Wu et al. , 2019 ; Yang et al. , 2020 ) , we propose a novel geometric autoencoder named ProGAE that directly learns from 3D protein structures via separately encoding intrinsic and extrinsic geometries into disjoint latent spaces used to generate protein structures . • We further propose a novel formulation , in which network intrinsic input is taken as the Cα-Cα pseudo-bond distances , and the extrinsic input is the backbone bond orientations . • Analysis shows that the learned extrinsic geometric latent space can be used for drug classification and drug property prediction , where the drug is bound to the given protein . • We find that the intrinsic geometric latent space , even with small variation in the intrinsic input signal , is important for reducing geometric errors in reconstructed proteins . • We also demonstrate that the learned ProGAE can be transferred to a trajectory of the protein in a different state or a trajectory of a different protein all-together . 1.1 RELATED WORK . Recently , a body of work has used deep learning to learn from protein structures ( Graves et al. , 2020 ; Jing et al. , 2020 ; Klicpera et al. , 2020 ) . For example , Gainza et al . ( 2019 ) uses geometric deep learning to predict docking sites for protein interactions . Ingraham et al . ( 2019a ) solves the inverse folding problem using a graph transformer on the protein backbone . Degiacomi ( 2019 ) uses an AE to generate candidate proteins for docking . Hermosilla et al . ( 2020 ) leverages the notion of intrinsic and extrinsic geometry to define an architecture for a fold classification task . Additionally , there has been focus on directly learning the temporal aspects of molecular dynamics from simulation trajectories , which is not directly related to the current work . Please see Appendix A.1 for a detailed discussion . There is an existing body of recent works that use AE-based approaches for either analyzing and/or generating structures from from the latent space ( Bhowmik et al. , 2018 ; Guo et al. , 2020 ; Ramaswamy et al. , 2020 ; Varolgüneş et al. , 2020 ) , which are most closely related to this work . ( Bhowmik et al. , 2018 ) and ( Guo et al. , 2020 ) aim at learning from and generating protein contact maps , while ProGAE directly deals with 3D structures . Therefore a direct comparison of ProGAE with these methods is not possible . Ramaswamy et al . ( 2019 ) uses a 1D CNN autoencoder trained on backbone coordinates and uses a loss objective comprised of geometric MSE error and physicsbased ( bond length , bond angle , etc . ) error . Due to the unavailability of code or pre-trained model , we were unable to perform a direct comparison . Varolgüneş et al . ( 2020 ) uses a VAE with a Gaussian Mixture Prior for performing clustering of high-dimensional input configurations in the learned latent space . While the method works well on toy models and a standard Alanine Dipeptide benchmark , its performance drops as the size of the protein system grows to 15 amino acids , which is approximately an order smaller than the protein systems studied here . Also , their approach is likely not going to scale well to larger systems due to the use of fully-connected layers in the encoder . These mentioned works have not considered explicit disentangling of intrinsic and extrinsic geometries . To our knowledge , this work is the first to propose an autoencoder for the unsupervised modeling of the geometric disentanglement of protein conformational space captured in molecular simulations . This representation provides better interpretability of the latent space , in terms of the physico-chemical and geometric attributes , results in more geometrically accurate protein conformations , as well as scales and transfers well to larger protein systems . 2 PROGAE FOR PROTEIN CONFORMATIONAL SPACE . First , we introduce the input signal for our novel geometric autoencoder , ProGAE . We then discuss how ProGAE utilizes this signal to generate the conformational space of a protein . Geometric Features of Protein as Network Input ProGAE functions by separately encoding intrinsic and extrinsic geometry with the goal of achieving better latent space interpretability . We clarify these geometric notions . Mathematically , we can consider a manifold ( i.e . surface ) independent of its embedding in Euclidean space . Properties that do not depend on this embedding are known as intrinsic geometric properties , while properties that do are referred to as extrinsic . As an example , given two atoms of a protein , the intrinsic distance between them is the minimum sum of bond lengths in the bond path connecting them , whereas the extrinsic distance is their Euclidean distance in R3 . For an in-depth review of geometry , we refer the reader to ( Do Carmo , 2016 ) . As we will train ProGAE to learn the conformational space of a given protein , the protein primary structure is implicit . Then in treating it as a geometric object , we view the protein at the level of its backbone , which specifies its shape . Given primary structure , reconstructing the protein backbone is sufficient for reconstructing the entire protein . Of importance in the backbone are the Cα atoms , which are the centers of amino acids in the protein . Then a coarse-level description of the backbone is the Cα atoms connected linearly in terms of the protein sequence . This is known as the trace of the protein . We will use the backbone and trace as domains on which to define our signals . Both the protein backbone and its trace can be viewed as polygonal chain in Euclidean space . They are depicted in Figure 1 with their geometric features as network input . We can see that a polygonal chain can be determined up to translation given both the length and orientation of its line segments . Then it follows that the protein backbone can be determined given the length and orientation of its bonds . Here the length of these bonds is intrinsic while the orientation is extrinsic . Thus , to decouple the intrinsic and extrinsic geometry , we can consider encoding these signals . The length of covalent bonds undergo very little change during a simulation performed using an empirical force-field , like the simulations considered in this work . A standard deviation of less than 0.059Å from target bond lengths is common in PDB structures ( Jaskolski et al. , 2007 ) . To this end , we instead consider intrinsic geometry at a coarse level , so that the resulting signal has more variability . Specifically , we use length of the Cα-Cα pseudobonds in the trace as a representative of the intrinsic protein geometry , where as backbone bond orientations capture extrinsic geometry . We model the backbone by the graph , Gb = ( Vb , Eb ) , and the backbone trace by the graph , Gt = ( Vt , Et ) . Then our intrinsic and extrinsic signals , Int : Et → R and Ext : Eb → R3 are defined : Int ( Eij ) = ‖Eij‖2 , Eij ∈ Et , Ext ( Eij ) = sgn ( j − i ) Eij ‖Eij‖ , Eij ∈ Eb . ( 1 ) Network Architecture With the network inputs defined , we discuss the architecture of ProGAE . The core idea is to create an intrinsic latent space , LI ∈ Rni , and an extrinsic latent space , LE ∈ Rne , via separately encoding the intrinsic and extrinsic signals . Consequently , our network contains two encoders , Enci and Ence where : Enci ◦ Int ( Et ) ∈ LI , Ence ◦ Ext ( Eb ) ∈ LE . ( 2 ) We then jointly decode these latent vectors to recover the coordinates of the atoms in the protein backbone . Thus , we formally define the decoder : Dec : LI × LE → R|Vb|×3 . ( 3 ) This high level structure of ProGAE is depicted in Figure 1 . We provide additional details on the encoders and decoders . As these edge-based signals are defined on a geometric domain , it is sensible to learn feature representations using geometric convolution that respects the geometry of the data . The intrinsic encoder is simple , as the signal is defined on the backbone trace , which corresponds to a set of discrete curves . Here each curve corresponds to a protein fragment . Then we define Enci to be a series of 1D convolutions operating on each protein fragment . Each convolution is taken to have a kernel size of 3 and a stride of 2 , being followed with batch normalization layers and ReLU . In contrast , the extrinsic encoder operates on the backbone , which we associate with a graph . So the layers of graph attention networks ( GATs ) introduced in ( Veličković et al. , 2017 ) are a natural tool to use , albeit with some modification . Since the input signal is defined only on the edges of the graph , Eb , we define a signal on the graph vertices , Vb , as the average value of its incident edges , f0 ( vi ) : = 1 | { j ; Ei , · ∈ Eb } | ∑ j ; Ei , ·∈Eb Ext ( Eij ) , vi ∈ Vb . ( 4 ) Then the first layer of the extrinsic encoder uses the edge-convolution operator of ( Gong & Cheng , 2019 ) to map this graph signal to a signal defined exclusively on the graph vertices , Vb . The rest of the encoder contains successive graph attention layers with sparsity defined by a given neighborhood radius . At each layer , the signal is downsampled by a factor of two based on farthest point sampling . Given L layers , this defines a sequence of graphs , { Gb , i } Li=0 , with increasing decimation . As with Enci , each layer is followed with batch normalization and ReLU . Summarily , for l = 1 , 2 , ... , L , fl = σ ◦BN ◦GAT ( dl−1 ) where dl−1 = DS ( fl−1 ; 2 ) , f1 ( vi ) : = GAT ( f0 ( Vb ) , Ext ( Eb ) ) . ( 5 ) Global average pooling is applied to the encoder outputs to introduce invariance to size of Vt and Vb . Dense layers then map each result to their respective latent spaces , LI and LE . The Tanh function is applied to bound the latent space . This produces the intrinsic and extrinsic latent codes , zi and ze . The latent code z is taken as the concatenation of the two latent codes , [ zi , ze ] . A dense layer maps z to the a signal defined on the most decimated backbone graph , Gb , L . The structure of the decoder , Dec , is then analogous to Ence , though the convolutions are transposed . The output of Dec is the point cloud , P̂ , corresponding to the predicted coordinates of the backbone atoms , Vb ≈ P . Loss Function The first term in the loss function is a basic reconstruction loss , where P and P̂ are taken to be the true and predicted coordinates of the protein backbone atoms . Namely , we evaluate their difference using Smooth-L1 loss . This loss is defined , with δ = 2 , as SmoothL1 ( x , y ) : = # x∑ i=1 zi , where zi = min ( δ2 2 ( xi − yi ) 2 , δ|xi − yi| − 1 2 ) , ( 6 ) This loss function modifies L2 loss to be more robust to outliers ( Girshick , 2015 ) . As the reconstruction loss depends on the embedding of the protein in Euclidean space , it may not best measure if intrinsic geometry is faithfully reconstructed . To address this , we consider two encoded proteins with latent codes , [ zi,1 , ze,1 ] and [ zi,2 , ze,2 ] . Then we form a new latent variable , ẑi = ( 1− β ) zi,1 + βzi,2 , ẑe = ze,1 , β ∼ U [ 0 , 1 ] . ( 7 ) Each of these latent variable decodes to some point cloud P̂ . We let Int ( Êt , β ) , Int ( Êt,1 ) , and Int ( Êt,2 ) be the lengths of the Cα-Cα pseudobonds of the generated proteins from the interpolated latent code and the two given latent codes . We then introduce a bond length penalty given by , R ( P̂1 , P̂2 ) = Eβ ||Int ( Êt , β ) − ( ( 1− β ) Int ( Êt,1 ) + βInt ( Êt,2 ) ) ||1 , β ∈ U [ 0 , 1 ] . ( 8 ) This penalty can be viewed as promoting faithful reconstruction of the pseudobond length between Cα atoms , as well as a smooth interpolation of these lengths along paths in LI , that is independent of LE . This penalty is analogous to the metric preservation regularizer introduced in ( Cosmo et al. , 2020 ) for 3D meshes . Thus , the loss function L for ProGAE is , L ( ( P̂1 , P̂2 ) , ( P1 , P2 ) ) : = 2∑ i=1 SmoothL1 ( P̂i , Pi ) + λRR ( P̂1 , P̂2 ) . ( 9 )

Summary: The paper introduces a geometric variational autoencoder for capturing protein structural ensembles, disentangling intrinsic and extrinsic geometry into separate latent spaces. The model is shown to accurately reconstruct protein structure, and the difference between the intrinsic and extrinsic latent spaces are explored. Finally, the model is tested in a transfer-learning setting, where it displays encouraging results.