A 3D Generative Model for Structure-Based Drug Design

1 Introduction . Designing molecules that bind to a specific protein binding site , also known as structure-based drug design , is one of the most challenging tasks in drug discovery [ 2 ] . Searching for suitable molecule candidates in silico usually involves massive computational efforts because of the enormous space of synthetically feasible chemicals [ 22 ] and conformational degree of freedom of both compound and protein structures [ 11 ] . In recent years , we have witnessed the success of machine learning approaches to problems in drug design , especially on molecule generation . Most of these approaches use deep generative models to propose drug candidates by learning the underlying distribution of desirable molecules . However , most of such methods are generally SMILES/string-based [ 10 , 17 ] or graph-based [ 18 , 19 , 13 , 14 ] . They are limited by the lack of spatial information and unable to perceive how molecules interact with proteins in 3D space . Hence , these methods are not applicable to generating molecules that fit to a specific protein structure which is also known as the drug target . Another line of work studies 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . generating molecules directly in 3D space [ 8 , 28 , 29 , 20 , 30 , 15 ] . Most of them [ 8 , 28 , 29 ] can only handle very small organic molecules , not sufficient to generate drug-scale molecules which usually contain dozens of heavy atoms . [ 20 ] proposes to generate voxelized molecular images and use a post-processing algorithm to reconstruct molecular structures . Though this method could produce drug-scale molecules for specific protein pockets , the quality of the sampling is heavily limited by voxelization . Therefore , generating high-quality drug molecules for specific 3D protein binding sites remains challenging . In this work , we propose a 3D generative model to approach this task . Specifically , we aim at modeling the distribution of atom occurrence in the 3D space of the binding site . Formally , given a binding site C as input , we model the distribution p ( e , r|C ) , where r ∈ R3 is an arbitrary 3D coordinate and e is atom type . To realize this distribution , we design a neural network architecture which takes as input a query 3D coordinate r , conditional on the 3D context C , and outputs the probability of r being occupied by an atom of a particular chemical element . In order to ensure the distribution is equivariant to C ’ s rotation and translation , we utilize rotationally invariant graph neural networks to perceive the context of each query coordinate . Despite having a neural network to model the distribution of atom occurrence p ( e , r|C ) , how to generate valid and diverse molecules still remains technically challenging , mainly for the following two reasons : First , simply drawing i.i.d . samples from the distribution p ( e , r|C ) does not yield valid molecules because atoms within a molecule are not independent of each other . Second , a desirable sampling algorithm should capture the multi-modality of the feasible chemical space , i.e . it should be able to generate a diverse set of desired molecules given a specific binding context . To tackle the challenge , we propose an auto-regressive sampling algorithm . In specific , we start with a context consisting of only protein atoms . Then , we iteratively sample one atom from the distribution at each step and add it to the context to be used in the next step , until there is no room for new atoms . Compared to other recent methods [ 20 , 23 ] , our auto-regressive algorithm is simpler and more advantageous . It does not rely on post-processing algorithms to infer atom placements from density . More importantly , it is capable of multi-modal sampling by the nature of auto-regressive , avoiding additional latent variables via VAEs [ 16 ] or GANs [ 9 ] which would bring about extra architectural complexity and training difficulty . We conduct extensive experiments to evaluate our approach . Quantitative and qualitative results show that : ( 1 ) our method is able to generate diverse drug-like molecules that have high binding affinity to specific targets based on 3D structures of protein binding sites ; ( 2 ) our method is able to generate molecules with fairly high drug-likeness score ( QED ) [ 4 ] and synthetic accessibility score ( SA ) [ 6 ] even if the model is not specifically optimized for them ; ( 3 ) in addition to molecule generation , the proposed method is also applicable to other relevant tasks such as linker design . 2 Related Work . SMILES-Based and Graph-Based Molecule Generation Deep generative models have been prevalent in molecule design . The overall idea is to use deep generative models to propose molecule candidates by learning the underlying distribution of desirable molecules . Existing works can be roughly divided into two classes — string-based and graph-based . String-based methods represent molecules as linear strings , e.g . SMILES strings [ 34 ] , making a wide range of language modeling tools readily applicable . For example , [ 5 , 10 , 26 ] utilize recurrent neural networks to learn a language model of SMILES strings . However , string-based representations fail to capture molecular similarities , making it a sub-optimal representation for molecules [ 13 ] . In contrast , graph representations are more natural , and graph-based approaches have drawn great attention . The majority of graph-based models generate molecules in an auto-regressive fashion , i.e. , adding atoms or fragments sequentially , which could be implemented based upon VAEs [ 13 ] , normalizing flows [ 27 ] , reinforcement learning [ 35 , 14 ] , etc . Despite the progress made in string-based and graph-based approaches , they are limited by the lack of spatial information and thus unable to be directly applied to structure-based drug design tasks [ 2 ] . Specifically , as 1D/2D-based methods , they are unable to perceive how molecules interact with their target proteins exactly in 3D space . Molecule Generation in 3D Space There has been another line of methods that generate molecules directly in 3D space . [ 8 ] proposes an auto-regressive model which takes a partially generated molecule as input and outputs the next atom ’ s chemical element and the distances to previous atoms and places the atoms in the 3D space according to the distance constraints . [ 28 , 29 ] approach this task via reinforcement learning by generating 3D molecules in a sequential way . Different from the previous method [ 8 ] , they mainly rely on a reward function derived from the potential energy function of atomic systems . These works could generate realistic 3D molecules . However , they can only handle small organic molecules , not sufficient to generate drug-scale molecules which usually contain dozens of heavy atoms . [ 20 , 23 ] propose a non-autoregressive approach to 3D molecular generation which is able to generate drug-scale molecules . It represents molecules as 3D images by voxelizing molecules onto 3D meshgrids . In this way , the molecular generation problem is transformed into an image generation problem , making it possible to leverage sophisticated image generation techniques . In specific , it employs convolutional neural network-based VAEs [ 16 ] or GANs [ 9 ] to generate such molecular images . It also attempts to fuse the binding site structures into the generative network , enabling the model to generate molecules for designated binding targets . In order to reconstruct the molecular structures from images , it leverages a post-processing algorithm to search for atom placements that best fit the image . In comparison to previous methods which can only generate small 3D molecules , this method can generate drug-scale 3D molecules . However , the quality of its generated molecules is not satisfying because of the following major limitations . First , it is hardly scalable to large binding pockets , as the number of voxels grows cubically to the size of the binding site . Second , the resolution of the 3D molecular images is another bottleneck that significantly limits the precision due to the same scalability issue . Last , conventional CNNs are not rotation-equivariant , which is crucial for modeling molecular systems [ 25 ] . 3 Method . Our goal is to generate a set of atoms that is able to form a valid drug-like molecule fitting to a specific binding site . To this end , we first present a 3D generative model in Section 3.1 that predicts the probability of atom occurrence in 3D space of the binding site . Second , we present in Section 3.2 the auto-regressive sampling algorithm for generating valid and multi-modal molecules from the model . Finally , in Section 3.3 , we derive the training objective , by which the model learns to predict where should be placed and atoms and what type of atom should be placed . 3.1 3D Generative Model Design . A binding site can be defined as a set of atoms C = { ( ai , ri ) } Nbi=1 , where Nb is the number of atoms in the binding site , ai is the i-th atom ’ s attributes such as chemical element , belonging amino acid , etc. , and ri is its 3D coordinate . To generate atoms in the binding site , we consider modeling the probability of atom occurring at some position r in the site . Formally , this is to model the density p ( e|r , C ) , where r ∈ R3 is an arbitrary 3D coordinate , and e ∈ E = { H , C , O , . . . } is the chemical element . Intuitively , this density can be interpreted as a classifier that takes as input a 3D coordinate r conditional on C and predicts the probability of r being occupied by an atom of type e. To model p ( e|r , C ) , we devise a model consisting of two parts : Context Encoder learns the representation of each atom in the context C via graph neural networks . Spatial Classifier takes as input a query position r , then aggregates the representation of contextual atoms nearby it , and finally predicts p ( e|r , C ) . The implementation of these two parts is detailed as follows . Context Encoder The purpose of the context encoder is to extract information-rich representations for each atom in C. We assume a desirable representation should satisfy two properties : ( 1 ) contextawareness : the representation of an atom should not only encode the property of the atom itself , but also encode its context . ( 2 ) rotational and translational invariance : since the physical and biological properties of the system do not change according to rigid transforms , the representations that reflect these properties should be invariant to rigid transforms as well . To this end , we employ rotationally and translationally invariant graph neural networks [ 25 ] as the backbone of the context encoder , described as follows . First of all , since there is generally no natural topology in C , we construct a k-nearest-neighbor graph based on inter-atomic distances , denoted as G = 〈C , A〉 , where A is the adjacency matrix . We also denote the k-NN neighborhood of atom i as Nk ( ri ) for convenience . The context encoder will take G as input and output structure-aware node embeddings . The first layer of the encoder is a linear layer . It maps atomic attributes { ai } to initial embeddings { h ( 0 ) i } . Then , these embeddings along with the graph structure A are fed into L message passing layers . Specifically , the formula of message passing takes the form : h ( ` +1 ) i = σ W ` 0h ( ` ) i + ∑ j∈Nk ( ri ) W ` 1w ( dij ) W ` 2h ( ` ) j , ( 1 ) where w ( · ) is a weight network and dij denotes the distance between atom i and atom j . The formula is similar to continuous filter convolution [ 25 ] . Note that , the weight of message from j to i depends only on dij , ensuring its invariance to rotation and translation . Finally , we obtain { h ( L ) i } a set of embeddings for each atom in C. Spatial Classifier The spatial classifier takes as input a query position r ∈ R3 and predicts the type of atom occupying r. In order to make successful predictions , the model should be able to perceive the context around r. Therefore , the first step of this part is to aggregate atom embeddings from the context encoder : v = ∑ j∈Nk ( r ) W0waggr ( ‖r − rj‖ ) W1h ( L ) j , ( 2 ) where Nk ( r ) is the k-nearest neighborhood of r. Note that we weight different embedding using the weight network waggr ( · ) according to distances because it is necessary to distinguish the contribution of different atoms in the context . Finally , in order to predict p ( e|r , C ) , the aggregated feature v is then passed to a classical multi-layer perceptron classifier : c = MLP ( v ) , ( 3 ) where c is the non-normalized probability of chemical elements . The estimated probability of position r being occupied by atom of type e is : p ( e|r , C ) = exp ( c [ e ] ) 1 + ∑ e′∈E exp ( c [ e ′ ] ) , ( 4 ) where E is the set of possible chemical elements . Unlike typical classifiers that apply softmax to c , we make use of the extra degree of freedom by adding 1 to the denominator , so that the probability of “ nothing ” can be expressed as : p ( Nothing|r , C ) = 1 1 + ∑ exp ( c [ e′ ] ) ) . ( 5 )

Summary: This paper tackles the problem of generating molecules that bind to a specific protein. To make use of the spatial information to target the proteins exactly in 3D space, the paper proposes an autoregressive model where the atoms are sampled sequentially from a learned distribution that estimates the probability of atoms' occurrence in 3D space. In fact,   the paper models the probability of a position in 3D space being occupied by a certain type of atom given as input the 3D coordinate and the binding site.  the model encodes the context which is the binding site and concatenates it with the input query position r  and finally predicts the type atom that could occur at the position r.