Conformation-Guided Molecular Representation with Hamiltonian Neural Networks

1 INTRODUCTION . The past several years have seen a prevalence of the intersection between medical chemistry and deep learning . Remarkable progress has been made in various applications on small molecules , ranging from generation ( Jin et al. , 2018 ; You et al. , 2018 ) and property prediction ( Gilmer et al. , 2017 ; Cho & Choi , 2019 ; Klicpera et al. , 2020 ) to protein-ligand interaction analysis ( Lim et al. , 2019 ; Wang et al. , 2020 ) , yet all these tasks rely on well-designed numerical representations , or fingerprints , of molecules . These fingerprints encode molecular structures and serve as the indicators in downstream tasks . Early work of molecular fingerprints ( Morgan , 1965 ; Rogers & Hahn , 2010 ) started from encoding the two-dimensional ( 2D ) structures of molecules , i.e . the chemical bonds between atoms , often stored as atom-bond graphs . More recently , a trend of incorporating molecular geometry into the representations arose ( Axen et al. , 2017 ; Cho & Choi , 2019 ) . Molecular geometry refers to the conformation ( the three-dimensional ( 3D ) coordinations of atoms ) of a molecule , which contains widely interested chemical information such as bond lengths and angles , and thus stands vital for determining physical , chemical , and biomedical properties of the molecule . Whereas incorporating 3D geometry of molecules seems indeed beneficial , 3D fingerprints , especially in combination with deep learning , are still in infancy . The use of 3D fingerprints is limited by pragmatic considerations including i ) calculation costs , ii ) translational & rotational invariances , and iii ) the availability of conformations , especially considering the generated ligand candidates in drug discovery tasks . Furthermore , compared with current 3D algorithms , mature 2D fingerprints ( Rogers & Hahn , 2010 ; Gilmer et al. , 2017 ; Xiong et al. , 2020 ) are generally more popular with equivalent or even better performances in practice . For example , as a 2D approach , Attentive Fingerprints ( Attentive FP ) ( Xiong et al. , 2020 ) have become the de facto state-of-the-art approach . To push the boundaries of leveraging 3D geometries in molecular fingerprints , we propose HamNet ( Molecular Hamiltonian Networks ) . HamNet simulates the process of molecular dynamics ( MD ) ∗Equal Contribution . †Corresponding Author . to model the conformations of small molecules , based on which final fingerprints are calculated similarly to ( Xiong et al. , 2020 ) . To address the potential lack of labeled conformations , HamNet does not regard molecular conformations as all-time available inputs . Instead , A Hamiltonian engine is designed to reconstruct known conformations and generalize for unknown ones . Encoded from atom features , implicit positions and momentums of atoms interact in the engine following the discretized Hamiltonian Equations with learnable energy and dissipation functions . Final positions are supervised with real conformations , and further used as inputs to a Message-Passing Neural Network ( MPNN ) ( Gilmer et al. , 2017 ) to generate the fingerprints . Novel loss functions with translational & rotational invariances are proposed to supervise the Hamiltonian Engine , and the architecture of the Fingerprint Generator is elaborated to better incorporate the output quantities from the engine . We show via our conformation-reconstructing experiments that the proposed Hamiltonian Engine is eligible to better predict molecular conformations than conventional geometric approaches as well as common neural structures ( MPNNs ) . We also evaluate HamNet on several datasets with different targets collected in a standard molecular machine learning benchmark , MoleculeNet ( Wu et al. , 2017 ) , all following the same experimental setups . HamNet demonstrates state-of-the-art performances , outperforming baselines including both 2D and 3D approaches . 2 PRELIMINARIES . Notations . Given a molecule with n atoms , we use vi to denote the features of atom i , and eij that of the chemical bond between i and j ( if exists ) . Bold , upper-case letters denote matrices , and lower-case , vectors . All vectors in this paper are column vectors , and · > stands for the transpose operation . We use ⊕ for the concatenation operation of vectors . The positions and momentums of atom i are denoted as qi and pi , and the set of all positions of atoms in a molecule is denoted as Q = ( qi , · · · , qn ) > . N ( v ) refers to the neighborhood of node v in some graph . Graph Convolutional Networks ( GCNs ) . Given an attributed graph G = ( V , A , X ) , where V = { v1 , · · · , vn } is the set of vertices , A ∈ Rn×n the ( weighted ) adjacency matrix , and X ∈ Rn×d the attribute matrix , GCNs ( Kipf & Welling , 2017 ) calculate the hidden states of graph nodes as GCN ( L ) ( X ) ≡H ( L ) , H ( l+1 ) = σ ( ÂH ( l ) W ( l ) ) , H ( 0 ) = X , l = 0 , 1 , · · · , L− 1 . ( 1 ) Here , H = ( hv1 , · · · , hvn ) > are hidden representations of nodes , Â = D− 1 2AD− 1 2 is the normalized adjacency matrix , D with Dii = ∑ jAij is the diagonal matrix of node degrees , and W s are network parameters . Message-Passing Neural Networks ( MPNNs ) . MPNN ( Gilmer et al. , 2017 ) introduced a general framework of Graph Neural Networks ( GNNs ) . In the t-th layer of a typical MPNN , messages ( mt ) are generated between two connected nodes ( i , j ) based on the hidden representations of both nodes ( ht ) and the edge in-between . After that , nodes receive the messages and update their own hidden representations . A readout function is then defined over final node representations ( hT ) to derive graph-level representations . Denoted in formula , the calculation follows mt+1v = ∑ w∈N ( v ) Mt ( h t v , h t w , ev , w ) , h t+1 v = Ut ( h t v , m t+1 v ) , ŷ = R ( { hTv |v ∈ V } ) , ( 2 ) where Mt , Ut , R are the message , update and readout functions . Hamiltonian Equations . The Hamiltonian Equations depict Newton ’ s laws of motion in the form of first-order PDEs . Considering a system of n particles with positions ( q1 , · · · , qn ) and momentums ( p1 , · · · , pn ) , the dynamics of the system follow q̇i ≡ dqi dt = ∂H ∂pi , ṗi ≡ dpi dt = −∂H ∂qi , ( 3 ) where H is the Hamiltonian of the system , and equals to the total system energy . Generally , the Hamiltonian is composed of the kinetic energy of all particles and the potential energy as H = n∑ i=1 Ti + U . ( 4 ) Meanwhile , if dissipation exists in the system , the Hamiltonian Equations shall be adapted as q̇i = ∂H ∂pi , ṗi = − ( ∂H ∂qi + ∂Φ ∂q̇i ) = − ( ∂H ∂qi +mi ∂Φ ∂pi ) ( 5 ) where mi is the mass of the particle and Φ is the dissipation function which describes how the system energy is dissipated by the outer environment . 3 METHOD : HAMNET . Figure 1 shows the overall architecture of HamNet . HamNet consists of two modules : i ) a Hamiltonian Engine , where molecular conformations are reconstructed ; and ii ) a Fingerprint Generator , where final fingerprints are generated from atom & bond features and outputs from the Hamiltonian Engine . 3.1 HAMILTONIAN ENGINE . Discretized Hamiltonian Equations . The Hamiltonian Engine is designed to simulate the physical interactions between atoms in a molecule . To correctly incorporate the laws of motion , we discretize the Hamiltonian Equations with dissipation ( Equation 5 ) and model the energy and dissipation with learnable functions . At the t-th step ( t = 0 , 1 , · · · , T ) in the engine , for atom i , q ( t+1 ) i = q ( t ) i + η ∂H ( t ) ∂p ( t ) i , p ( t+1 ) i = p ( t ) i − η ( ∂H ( t ) ∂q ( t ) i +mi ∂Φ ( t ) ∂p ( t ) i ) . ( 6 ) Here , η is a hyperparameter of step size which controls the granularity of the discretization , mi is the ( normalized ) mass of atom i , andH ( t ) , Φ ( t ) are the learnable Hamiltonian and dissipation functions of q ( t ) and p ( t ) : ( superscripts skipped ) H = n∑ i=1 Ti ( pi ) + U ( qi , · · · , qn ) , Φ = n∑ i=1 φ ( pi ) . ( 7 ) It should be noted that in order to improve the expressional power of the network , we extend the concept of positions and momentums into implicit ones in a generalized df -dimensional space , i.e . q , p ∈ Rdf , df > 3 . We will discuss how to supervise these implicit quantities , and will show the influences of the dimensionality in the experimental results . Nonetheless , we parameterize the energy and dissipation function in a physically inspired manner : for the atom-wise kinetic energy , we generalize the definition of kinetic energy ( T = p 2 2m ) as the quadratic forms of the implicit momentums ; for the atom-wise dissipation , we adapt the Rayleigh ’ s dissipation function ( Φ = 12 ∑n i , j=1 cij q̇iq̇j ) into the generalized space . Denoted in formula , we have Ti ( pi ) = p > i W > T WTpi 2mi , φi ( pi ) = p > i W > φ Wφpi 2m2i . ( 8 ) For the potential energy , we simplify the Lennard-Jones potential ( U ( r ) = ( r−12 − r−6 ) ) with parameterized distances rs in the generalized space , that is , U = ∑ i 6=j uij , uij = r −4 ij − r −2 ij , r 2 ij ≡ r ( qi , qj ) = ( qi − qj ) > W > U WU ( qi − qj ) . ( 9 ) In both Equation 8 and Equation 9 , W s are network parameters . For the mass of the atoms , we empirically normalize the relative atomic mass Mi of atom i with mi = Mi/50 . One would also note that the parameters in each step of the engine remain the same , and thus the scale of parameters in the engine depends only on df and is irrelevant to the depth T . Initial positions and momentums . Graph-based neural networks are used to initialize these spatial quantities . Molecules are essentially graphs , while bonds between atoms can be of different types ( single , double , triple and aromatic ) . Instead of using channel-wise GCNs which assign different parameters for different types of edges ( Schlichtkrull et al. , 2018 ) , we calculate a bond-strength adjacency for the molecules : the strength of a bond depends on the atom-bond-atom tuple , i.e . Aij = sigmoid ( MLP ( ( vi ⊕ eij ⊕ vj ) ) ) if the bond exists , Aij = 0 otherwise . ( 10 ) With so-defined adjacency , we first encode the atoms with vanilla GCNs ( Equation 1 ) and concatenate all hidden layers following the DenseNet scheme ( Huang et al. , 2017 ) . Deep enough GCNs do capture entire molecular structures , however , atoms with identical chemical environment can not be distinguished , for example , carbon atoms in the benzene ring . This may not be a problem in conventional MPNNs , but atoms with coinciding positions are inacceptable in physics as well as the Hamiltonian engine . Therefore , we conduct an LSTM over the GCN outputs to generate unique positions for atoms . The order that the atoms appear in the LSTM conforms with the SMILES representations ( Weininger , 1988 ) of the molecule , which display atoms in the molecule in a specific topological order . Denoted in formula , initial positions and momentums of atoms are calculated as q̃i = L⊕ l=0 f ( l ) i , q ( 0 ) i = LSTMsi ( q̃s1 , q̃s2 , · · · , q̃sn ) ( 11 ) p̃i = L⊕ l=0 g ( l ) i , p ( 0 ) i = LSTMsi ( p̃s1 , p̃s2 , · · · , p̃sn ) ( 12 ) where f ( l ) i , g ( l ) i are hidden representations of atom i in the l-th GCN layer and sk is the atom at k-th position in the SMILES . As unique orders are assigned for atoms , their initial positions are thus unique . Conformation preserving . After the dynamical process in the Hamiltonian Engine , positions in the generalized Rdf space are transformed into real 3D space linearly , that is , Q̂ = QWtrans , Q̂ ∈ Rn×3 , Q = ( q1 , · · · , qn ) ∈ Rn×df . ( 13 ) Considering translational & rotational invariances , we do not require that Q̂ approximates the labeled 3D coordinations of atoms . Instead , three translational- and rotational-invariant metrics are proposed and used to supervise Q̂ : I ) Kabsch-RMSD ( K-RMSD ) . We use the Kabsch Algorithm ( Kabsch , 1976 ) to rotate and align the approximated atom positions ( Q̂ ) to the real ones ( QR ) , and then calculate the Root of Mean Squared Deviations ( RMSD ) of two conformations using atom mass as weights : Q̂K = Kabsch ( Q̂ ; QR ) , Lk−rmsd ( Q̂ , QR ) = √√√√∑ni=1mi × ∥∥q̂Ki − qRi ∥∥22∑n i=1mi . ( 14 ) One should note that the alignment Kabsch ( · ; · ) is calculated with SVD and is thus differentiable . II ) Distance Loss . Pair-wise distances between atoms are vital quantities in describing molecular conformations and enjoy desired invariances . Therefore , we propose the metric Ldist as L2dist ( Q̂ , Q R ) = 1 n2 n∑ i , j=1 ( ‖qi − qj‖22 − ∥∥qRi − qRj ∥∥22 ) 2 ( 15 ) III ) ADJ-k Loss . A drawback of the naive distance loss is that distances between far atoms are over-emphasized , leading to deformed local structures . Therefore , we further propose ADJ-k loss , where only distances between k-hop connected atoms are preserved under weights calculated from hop-distances . Denote the normalized , simple adjacency matrix as Ã,1 the ADJ-k loss is defined as L2adj−k ( Q̂ , Q R ) = 1 n n∑ i , j=1 Ãkij ( ‖qi − qj‖22 − ∥∥qRi − qRj ∥∥22 ) 2 ( 16 ) In implementation , we use a linear combination of K-RMSD and ADJ-3 losses to supervise the engine , i.e . LHE = Lk−rmsd + λLadj−3 , where λ is a hyperparameter .

Summary: The paper proposes a new method for generating fingerprints for small molecules. It is based on two components: a "Hamiltonian Engine" that runs a brief simulation, predicting the structure of the small molecule by minimizing a learned potential energy, and 2) a message passing algorithm that uses the predicted structure as input. The reported experimental results demonstrate state-of-the-art performance.