Spatial Graph Attention and Curiosity-driven Policy for Antiviral Drug Discovery

1 INTRODUCTION . This work aims to address the challenge of establishing an automated process for the design of objects with connected components , such as molecules , that optimize specific properties . Achieving this goal is particularly desirable in drug development and materials science , where manual discovery remains a time-consuming and expensive process ( Hughes et al. , 2011 ; Schneider et al. , 2020 ) . However , there are two major difficulties that have long impeded rapid progress . Firstly , the chemical space is discrete and massive ( Polishchuk et al. , 2013 ) , presenting a complicated environment for an Artificial Intelligence ( AI ) approach to efficiently and effectively explore . Secondly , it is not trivial to compress such connected objects into feature representations that preserve most of the information , while also being highly computable for Deep Learning ( DL ) methods to exploit . We introduce Distilled Graph Attention Policy Networks ( DGAPN ) , a framework that advances prior work in addressing both of these challenges . We present a Reinforcement Learning ( RL ) architecture that is efficiently encouraged to take innovative actions with an environment that is able to construct a dynamic and chemically valid fragment-based action space . We also propose a hybrid Graph Neural Network ( GNN ) that comprehensively encodes graph objects ’ attributes and spatial structures in addition to adjacency structures . The following paragraphs discuss how we addressed limitations of prior work and its relevance to antiviral drug discovery . For more descriptions of key prior methodologies that we used as benchmarks in this paper , see Section 5 . Graph Representation Learning Despite their spatial efficiency , string representation of molecules acquired by the simplified molecular-input line-entry system ( SMILES ) ( Weininger , 1988 ) suffers from significant information loss and poor robustness ( Liu et al. , 2017 ) . Graph representations have become predominant and preferable for their ability to efficiently encode an object ’ s scaffold structure and attributes . Graph representations are particularly ideal for RL since intermediate representations can be decoded and evaluated for reward assignments . While GNNs such as Graph Convolutional Networks ( GCNs ) ( Kipf & Welling , 2016 ) and Graph Attention Networks ( GATs ) ( Veličković et al. , 2017 ) have demonstrated impressive performance on many DL tasks , further exploitation into richer information contained in graph-structured data is needed to faithfully represent the complexity of chemical space ( Morris et al. , 2019 ; Wang et al. , 2019 ; Chen et al. , 2020 ) . In this work , we made improvements to previous studies on attributes encoding and structural encoding . For structural encoding , previous studies have covered adjacency distance encoding ( Li et al. , 2020 ) , spatial cutoff ( Pei et al. , 2020 ) and coordinates encoding ( Schütt et al. , 2017 ; Danel et al. , 2020 ) . Our work presents an alternative approach to spatial structure encoding similar to Gilmer et al . ( 2017 ) which do not rely on node coordinates , but different in embedding and updating scheme . Distinct from Danel et al . ( 2020 ) and Chen & Chen ( 2021 ) , we extended attentional embedding to be edge-featured , while still node-centric for message passing efficiency . Reinforcement Learning A variety of graph generative models have been used in prior work , predominantly Variational Autoencoders ( VAEs ) ( Simonovsky & Komodakis , 2018 ; Samanta et al. , 2020 ; Liu et al. , 2018 ; Ma et al. , 2018 ; Jin et al. , 2018 ) and Generative Adversarial Networks ( GANs ) ( De Cao & Kipf , 2018 ) . While some of these have a recurrent structure ( Li et al. , 2018 ; You et al. , 2018b ) , RL and other search algorithms that interact dynamically with the environment excel in sequential generation due to their ability to resist overfitting on training data . Both policy learning ( You et al. , 2018a ) and value function learning ( Zhou et al. , 2019 ) have been adopted for molecule generation : however , they generate molecules node-by-node and edge-by-edge . In comparison , an action space consisting of molecular fragments , i.e. , a collection of chemically valid components and realizable synthesis paths , is favorable since different atom types and bonds are defined by the local molecular environment . Fragment-by-fragment sequential generation has been used in VAE ( Jin et al. , 2018 ) and search algorithms ( Jin et al. , 2020 ; Xie et al. , 2021 ) , but has not been utilized in a deep graph RL framework . In this work , we designed our environment with the Chemically Reasonable Mutations ( CReM ) ( Polishchuk , 2020 ) library to realize a valid fragment-based action space . Furthermore , we enhanced exploration by employing a simple and efficient technique , adapting Random Network Distillation ( RND ) ( Burda et al. , 2018 ) to GNNs and proposing surrogate innovation rewards for intermediate states . Antiviral Drug Discovery — A Timely Challenge The severity of the COVID-19 pandemic highlighted the major role of computational workflows to characterize the viral machinery and identify druggable targets for the rapid development of novel antivirals . Particularly , the synergistic use of DL methods and structural knowledge via molecular docking simulations is at the cutting edge of molecular biology — consolidating such integrative protocols to accelerate drug discovery is of paramount importance ( Yang et al. , 2021 ; Jeon & Kim , 2020 ; Thomas et al. , 2021 ) . Here we experimentally examined our architecture on the task of discovering novel inhibitors targeting the SARS-CoV-2 non-structural protein endoribonuclease ( NSP15 ) , which is critical for viral evasion of host defense systems ( Pillon et al. , 2021 ) . Structural information about the putative protein-ligand complexes was integrated into this framework with AutoDock-GPU ( Santos-Martins et al. , 2021 ) , which leverages the GPU resources from leadership-class computing facilities , including the Summit supercomputer , for high-throughput molecular docking ( Legrand et al. , 2020 ) . We show that our results outperformed state-of-the-art generation models in finding molecules with high affinity to the target and reasonable synthetic accessibility . 2 PROBLEM FORMULATIONS . Our goal is to establish a set of decision rules to generate graph-structured data that maximizes compound objectives under certain constraints . Similar to prior formulations , the generating process is defined as a time homogeneous Markov Decision Process ( MDP ) . We give a formal definition of this process in Appendix A . Under this setting , the action policies and state transition dynamics at step t can be factorized according to the Markov property : P ( at|s0 , a0 , s1 , a1 , . . . , st ) = P ( at|st ) : = π ( at|st ) ( 1 ) P ( st+1|s0 , a0 , s1 , a1 , . . . , st , at ) = P ( st+1|st , at ) : = ρ ( st+1|st , at ) ( 2 ) where { st , at } t are state-action sequences . A reward function r ( s , a ) is used to assess an action a taken at a given state s. The process terminates at an optional stopping time T and sT is then proposed as the final product of the current generating cycle . We aim to estimate the optimal policy π in terms of various objectives to be constructed later in the experiment section . 3 PROPOSED METHOD . 3.1 ENVIRONMENT SETTINGS . In the case of molecular graphs , single-atom or single-bond additions are often not realizable by known biochemical reactions . Rather than employing abstract architectures such as GANs to suggest synthetic accessibility , we use the chemical library CReM ( Polishchuk , 2020 ) to construct our environment such that all next possible molecules can be obtained by one step of interchanging chemical fragments with the current molecule . This explicit approach is considerably more reliable and interpretable compared to DL approaches . A detailed description of the CReM library can be found in Appendix B.1 . At each time step t , we use CReM to sample a set of valid molecules vt+1 as the candidates for the next state st+1 based on st . Under this setting , the transition dynamics are deterministic , the underlying set A of the action space can be defined as equal to S of the state space , and action at is induced by the direct selection of st+1 . With an abuse of notation , we let r ( st+1 ) : = r ( st , at ) . 3.2 SPATIAL GRAPH ATTENTION . We introduce a graph embedding mechanism called Spatial Graph Attention ( sGAT ) in an attempt to faithfully extract feature vectors ht ∈ Rdh representing graph-structured objects such as molecules . Two different types of information graphs constructed from a connected object are heterogeneous and thus handled differently in forward passes as described in the following sections . See Figure 1 for an overview . 3.2.1 ATTENTION ON ATTRIBUTION GRAPHS . The attribution graph of a molecule with n atoms and e bonds is given by the triple ( A , N , E ) , where A ∈ { 0 , 1 } n×n is the node adjacency matrix , N is the node attribution matrix of dimension n× dn and E is the edge attribution matrix of dimension e× de . Each entry aij ofA is 1 if a bond exists between atom i and j , and 0 otherwise . Each row vector ni of N is a concatenation of the properties of atom i , including its atomic number , mass , etc. , with the categorical properties being one-hot encoded . E is formed similar to N , but with bond attributes . We denote a row vector of E as eij if it corresponds to the bond between atom i and j . We proceed to define a multi-head forward propagation that handles these rich graph information : let hnk ∈ R1×dhn denote a given representation for nk , heij ∈ R1×dhe denote a representation for eij , then the m-th head attention αmij from node j to node i ( i 6= j ) is given by αmij = softmax j ( ⋃ k : aik=1 { σ ( [ hniWn , m ‖ heikWe , m ‖ hnkWn , m ] · attmT ) } ) ( 3 ) where softmax j is the softmax score of node j ; ‖ is column concatenation ; σ is some non-linear activation ; Wn , m ∈ Rdhn×dwn , We , m ∈ Rdhe×dwe are them-th head weight matrices for nodes and edges respectively ; attm ∈ R1× ( 2dwn+dwe ) is the m-th head attention weight . The representations after a feed-forward operation are consequently given as follow : h′ni = aggr1≤m≤nm σ ∑ j : aij=1 αmij · hnj + hni Wn , m ( 4 ) h′eij = aggr1≤m≤nm { σ ( [ hniWn , m ‖ heijWe , m ‖ hnjWn , m ] ·Wh , m ) } ( 5 ) where Wh , m ∈ R ( 2dwn+dwe ) ×dwe ; nm is the total number of attention heads and aggr denotes an aggregation method , most commonly mean , sum , or concat ( Hamilton et al. , 2017 ) . We note that we have not found significant difference across these methods and have used mean for all aggregations in our experiments . In principle , a single-head operation on nodes is essentially graph convolution with the adjacency matrix Â = Ã + I where Ã is attention-regularized according to ( 3 ) . This approach sufficiently embeds edge attributes while still being a node-centric convolution mechanism , for which efficient frameworks like Pytorch-Geometric ( Fey & Lenssen , 2019 ) have been well established .

Summary: With the goal of generating molecules that bind to functional sites of SARS-Cov-2 protein, the paper proposed a reinforcement learning model with a fragment-based framework for chemical structure modification. In the network part, spatial graph attention and spatial convolution are utilized to extract more structural information from the input graph into the representation of nodes or graphs. Based on the actor-critic algorithm, the reinforcement learning part is designed to find the state with the best docking score computed by docking programs and some novel approaches like PPO and RND are used to train the model more effectively. In the experiments, the model shows great performance while reducing the complexity of chemical synthesis meanwhile.