Deep Molecular Representation Learning via Fusing Physical and Chemical Information

1 Introduction . The intersection between deep learning and molecular science has recently caught the eye of researchers in both areas . Remarkable progress was made in applications including molecular property prediction [ 36 , 16 ] , molecular graph generation [ 37 , 12 , 28 ] , and virtual screening for drug discovery [ 8 , 31 ] , yet learning representations for molecules remains the first yet vital step . Molecular representation learning , or learning molecular fingerprints , aims to encode input notations of molecules into numerical vectors , which later serve as features for downstream tasks . Earlier deep molecular representation methods generally used off-the-shelf network architectures including message-passing neural networks [ 9 ] , graph attention networks [ 36 ] and Transformers [ 11 , 24 ] . These methods took either line notations ( e.g . SMILES3 ) or graph notations ( i.e . structural formulas ) of molecules as inputs , whereas physical and chemical essence of molecules was largely neglected . Notably , a trend of integrating 3D conformations ( i.e . the 3D Cartesian coordinates of atoms ) into molecular representations recently emerged [ 4 , 15 , 29 ] , while most of these methods assume the availability of labeled conformations of target molecules . In order to push the boundaries of molecular representation learning , we revisited molecules from both physical and chemical perspectives . Modern physicists generally regard molecules as particle systems that continuously move following the laws of ( quantum ) mechanics . The dominant conformations of molecules reflect the equilibriums of these micro mechanical systems , and are thus of wide interest . ∗Equal Contribution . †Corresponding Author . 3SMILES ( Simplified Molecular Input Line Entry Specification [ 34 ] ) is a widely used protocol that specifies ( non-unique ) line notations for molecules , CCO for ethanol , for example . 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Chemists , on the other hand , focus more on chemical bonds and functional groups of molecules , which denote the interactions of electrons and determine chemical / biomedical properties such as solubility and toxicity , etc . Nevertheless , physical and chemical information of molecules is not orthogonal . For example , torsions of bond lengths and angles greatly influence the dynamics of particle systems . Therefore , an ideal molecular representation is not only expected to capture both physical and chemical information , but also to appropriately fuse the two types of information . Based on the observations above , we propose PhysChem , a novel neural architecture that captures and fuses physical and chemical information of molecules . PhysChem is composed of two specialist networks , namely a physicist network ( PhysNet ) and a chemist network ( ChemNet ) , who understand molecules physically and chemically.4 PhysNet is a neural physical engine that learns dominant conformations of molecules via simulating molecular dynamics in a generalized space . In PhysNet , implicit positions and momenta of atoms are initialized by encoding input features . Forces between pairs of atoms are learned with neural networks , according to which the system moves following laws of classic mechanics . Final positions of atoms are supervised with labeled conformations under spatial-invariant losses . ChemNet utilizes a message-passing framework [ 9 ] to capture chemical characteristics of atoms and bonds . ChemNet generates messages from atom states and local geometries , and then updates the states of both atoms and bonds . Output molecular representations are merged from atomic states and supervised with labeled chemical / biomedical properties . Besides focusing on their own specialties , two networks also cooperate by sharing expertise : PhysNet consults the hidden representations of chemical bonds in ChemNet to generate torsion forces , whereas ChemNet leverages the local geometries of the intermediate conformations in PhysNet . Compared with existing methods , PhysChem adopts a more elaborated as well as interpretable architecture that incarnates physical and chemical understandings of molecules . Moreover , as PhysNet learns molecular conformations from scratch , PhysChem does not require labeled conformations of test molecules . This extends the applicability of PhysChem to situations where such labels are unavailable , for example , with neural-network-generated drug candidates . We evaluated PhysChem on several datasets in the MoleculeNet [ 35 ] benchmark , where PhysChem displayed state-of-the-art performances on both conformation learning and property prediction tasks . Results on cutting-edge datasets of SARS-CoV-2 further proved the effectiveness of PhysChem . 2 Related Work . Molecular Representation Learning Early molecular fingerprints commonly encoded line or graph notations of molecules with rule-based algorithms [ 19 , 22 , 25 ] . Along with the spurt of deep learning , deep molecular representations gradually prevailed [ 7 , 9 , 11 , 36 ] . More recently , researchers started to focus on incorporating 3D conformations of molecules into their representations [ 1 , 4 , 29 , 15 ] . Models that leveraged 3D geometries of molecules generally performed better than those that simply used graph notations , whereas most 3D models required labeled conformations of the target molecules . This limited the applicability of these models . Among previous studies , message-passing neural networks ( MPNNs ) proposed a universal framework of encoding molecular graphs , which assumed that nodes in graphs ( atoms in molecules ) passed messages to their neighbors , and then aggregated received messages to update their states . A general message-passing layer calculated mi = ∑ j∈N ( v ) M ( hi , hj , ei , j ) , hi ← U ( hi , mi ) , i ∈ V , ( 1 ) where i , j ∈ V were graph nodes , hs and es were states of nodes and edges , and M ( · ) , U ( · ) were the message and update functions . For graph-level tasks , MPNNs further defined a readout function which merged the states of all nodes into graph representations . In this work , ChemNet modifies the message-passing scheme of MPNNs : messages are extended to incorporate geometries of local structures , and states of both atoms and bonds are updated . See Section 3.4 for more details . Neural Physical Engines Recent studies showed that neural networks are capable of learning annotated ( or pseudo ) potentials and forces in particle systems , which made fast molecular simulations [ 38 , 17 ] and protein-folding tasks [ 26 ] possible . Notably , it was further shown in [ 16 , 27 ] that neural networks alone can simulate molecular dynamics for conformation prediction . As an instance , 4PhysNet and ChemNet are novel architectures , not to be confused with previous works with similar or identical names [ 33 , 32 , 20 ] . HamNet [ 16 ] proposed a neural physical engine that operated on a generalized space , where positions and momentums of atoms were defined as high-dimensional vectors . In the engine , atoms moved following Hamiltonian Equations with parameterized kinetic , potential and dissipation functions . PhysNet in our work is a similar engine . Nevertheless , instead of learning parameterized energies and calculating their negative derivatives as forces , we directly parameterize the forces between each pair of atoms . In addition , HamNet considered gravitations and repulsions of molecules based on implicit positions , while it ignored the effects of chemical interactions : for example , the types of chemical bonds were ignored in the energy functions . PhysChem fixes this issue via the cooperation mechanism between two specialist networks . Specifically , PhysNet takes chemical expertise ( the bond states ) from ChemNet and introduces torsion forces , i.e . forces that origin from torsions in chemical bonds , into the dynamics . See Section 3.3 for more details . Multi-Task Learning and Model Fusion The cooperation mechanism in PhysChem shares similar motivations with multi-task learning and model fusion . Multi-task learning [ 39 , 6 ] is now almost universally used in deep learning models . Representations are shared among a collection of related tasks in order to learn the common ideas . Model fusion , on the other hand , merges different models on identical tasks to improve performances [ 23 , 13 ] . Notably , these techniques have been previously applied to molecular tasks [ 30 , 13 , 31 ] . In PhysChem , conformation learning and property prediction tasks are jointly trained , with two specialist networks fused to achieve better performances . Nevertheless , the cooperation mechanism in PhysChem roots from observations in physics and chemistry , and enjoys better interpretability than straight-forward ensemble strategies . The advantages of the cooperation mechanism over multi-task strategies are also empirically shown in Section 4 . 3 Proposed Method : PhysChem . 3.1 Preliminaries . Notations In the statements and equations below , we use italic letters for scalars and indices , bold lower-case letters for ( column ) vectors , bold upper-case letters for matrices , calligraphic letters for sets , and normal letters for annotations . Common neural network layers are directly referred to : FC ( x ) denotes a fully-connected layer with input x ; GCN ( A , X ) denotes a ( vanilla ) graph convolutional network ( GCN ) [ 14 ] with adjacency matrix A and feature matrix X ; GRUcell ( s , x ) denotes a cell of the Gated Recurrent Units ( GRU ) [ 5 ] with recurrent state s and input signal x ; LSTM ( { xt } ) denotes a Long Short-Term Memory network [ 10 ] with input signals { xt } ; MaxPool ( { xt } ) denotes a max-pooling layer with inputs { xt } . Exact formulas are available in the appendix . ⊕ is the operator of concatenations . ∥·∥ denotes the l2 norm of the input vector . Problem Definition PhysChem considers molecular representation learning as a supervised learning task . It takes notations of molecules as inputs , conformations and chemical / biomedical properties as supervisions . PhysChem assumes a pre-conducted featurization process , after which a molecule can be denoted as an attributed graphM = ( V , E , n , m , Xv , Xe ) . Here , V is the set of n atoms , E ⊂ V × V is the set of m chemical bonds , Xv ∈ Rn×dv = ( xv1 , · · · , xvn ) ⊤ is the matrix of atomic features , and Xe ∈ Rm×de = ( xe1 , · · · , xem ) ⊤ that of bond features . Based on above inputs , PhysNet outputs the dominant conformations of molecules , and ChemNet outputs the representations as well as the predicted chemical / biomedical properties of molecules . Overview Figure 1 is a sketch of PhysChem . An initializer first encodes the inputs into initial atom and bond states ( v ( 0 ) and e ( 0 ) ) for ChemNet , along with the initial atomic positions and momenta ( q ( 0 ) and p ( 0 ) ) in PhysNet . Subsequently , L PhysNet blocks simulate neural molecular dynamics in the generalized space ; L ChemNet blocks conduct geometry-aware message-passing for atoms and bonds . Between each couple of blocks , implicit conformations ( qs ) and states of chemical bonds ( es ) are shared as expertise . At the top of PhysNet , a conformation decoder transforms implicit atomic positions into the 3D Euclidean space ; of ChemNet , a sequence of readout layers aggregate atom states into molecular representations , based on which molecular properties are predicted with task-specific layers . Notably , two specialist networks are jointly optimized in PhysChem . 3.2 Initializer . Given a molecular graphM = ( V , E , n , m , Xv , Xe ) , the initializer generates the initial values for both PhysNet and ChemNet variables . In the initializer , we first encode the input features into initial atom and bond states with fully connected layers , i.e . v ( 0 ) i = FC ( x v i ) , i ∈ V ; e ( 0 ) i , j = FC ( xei , j ) , ( i , j ) ∈ E . ( 2 ) We then adopt the initialization method in [ 16 ] to generate initial positions ( q ( 0 ) ) and momenta ( p ( 0 ) ) for atoms : a bond-strength adjacency matrix A ∈ Rn×n is estimated with sigmoid-activated FC layers on bond features , according to which a GCN captures the chemical environments of atoms ( as ṽ ) ; an LSTM then determines unique positions for atoms , especially for those with identical chemical environments ( carbons in benzene , for example ) . Denoted in formula , the initialization follows A ( i , j ) = { 0 , ( i , j ) /∈ E FCsigmoid ( xei , j ) , ( i , j ) ∈ E , Ṽ = GCN ( A , V ( 0 ) ) , ( 3 ) { ( q ( 0 ) i ⊕ p ( 0 ) i ) } = LSTM ( { ṽi } ) , i ∈ V ( 4 ) where V ( 0 ) = ( v ( 0 ) 1 , · · · , v ( 0 ) n ) ⊤ and Ṽ = ( ṽ1 , · · · , ṽn ) ⊤ are the atom states . The order of atoms in the LSTM is specified by the canonical SMILES of the molecule .

Summary: The authors are interested in the linked tasks of predicting (i) 3D conformations and (ii) chemical/biomedical properties of molecules. They show that by designing a model (PhysChem) that tackles both of these tasks simultaneously, they can often do better on each individual task than when these tasks are considered separately. This is empirically demonstrated on the QM{7,8,9} datasets, the MoleculeNet  Physical Chemistry Benchmark [32], as well as a SARS-CoV-2 dataset [3], where the authors also compare to a variety of older approaches.     ### Further Details on the Architecture To be more specific about the architecture, PhysChem consists of two networks (which share information at each layer and both take in as input the graph representation of a molecule): 1. PhysNet, is used to predict the conformation of a molecule and it does this by first predicting the initial positions of the atoms (using a graph convolutional network (GCN) and RNN) before then updating these positions using the equations of motion under the action of interatomic forces output by a fully connected NN. 2. ChemNet†, is used to predict chemical/biomedical properties of a molecule, and is a new variant of a message passing neural network [9].  † This is a new network, not to be confused with previous networks also called "ChemNet", e.g. Preuer, K., Renz, P., Unterthiner, T., Hochreiter, S. and Klambauer, G. (2018) ‘Fréchet ChemNet Distance: A Metric for Generative Models for Molecules in Drug Discovery’, Journal of chemical information and modeling, 58