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: This paper presents a novel neural network module called Hamiltonian Neural Networks to learn representation of molecules. The module consists of two main components: 1) Hamiltonian Engine (HE), and 2) Fingerprint Generator (FG). The HE 1) takes a molecular graph with atom and bond features as inputs, and first generate "generalized" positions p and momentums q via GNN+LSTM. These ps and qs are fed into a discrete Hamiltonian system with dissipation, and produces (generalized) "conformations" of molecules. The FG also takes a molecular graph with atom and bond features + the generalized positions and momentums from 1) as inputs to generate the final vector embedding of the input molecule. HE is trained to fit the input 3D conformation with a weighted combination of two types lossses (K-RMSD + 5 * ADJ-k loss). Posing this "Hamiltonian" inductive bias to the model, the paper demonstrated that the prediction performance over multiple molecular tasks from MoleculeNet are all improved. Also, the analysis of HE module or the hyperparamter sensitivity analysis are provided.