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: This paper proposes a method to generate molecular graphs with optimized properties. Molecular graphs are constructed by the iterative addition of molecular fragments in a deep reinforcement learning framework. The method is benchmarked against a set of baselines in a task to generate molecules that maximizes the docking score to a protein (NSP15) from the SARS-CoV-2 virus, and shows good performance.