Chemical-Reaction-Aware Molecule Representation Learning

1 INTRODUCTION . How to represent molecules is a fundamental and crucial problem in chemistry . Chemists usually use IUPAC nomenclature , molecular formula , structural formula , skeletal formula , etc. , to represent molecules in chemistry literature.1 However , such representations are initially designed for human readers rather than computers . To facilitate machine learning algorithms understanding and making use of molecules , molecule representation learning ( MRL ) is proposed to map molecules into a lowdimensional real space and represent them as dense vectors . The learned vectors ( a.k.a . embeddings ) of molecules can benefit a wide range of downstream tasks , such as chemical reaction prediction ( Jin et al. , 2017 ; Segler & Waller , 2017 ) , molecule property prediction ( Zhang et al. , 2021 ) , molecule generation ( Mahmood et al. , 2021 ) , drug discovery ( Rathi et al. , 2019 ) , retrosynthesis planning ( Segler et al. , 2018 ) , chemical text mining ( Krallinger et al. , 2017 ) , and chemical knowledge graph modeling ( Bean et al. , 2017 ) . Researchers have proposed a great many MRL methods . A large portion of them , including MolBERT ( Fabian et al. , 2020 ) , ChemBERTa ( Chithrananda et al. , 2020 ) , SMILES-Transformer ( Honda et al. , 2019 ) , SMILES-BERT ( Wang et al. , 2019 ) , Molecule-Transformer ( Shin et al. , 2019 ) , and SABiLSTM ( Zheng et al. , 2019b ) , take SMILES2 strings as input and utilize natural language models , for example , Transformers ( Vaswani et al. , 2017 ) or BERT ( Devlin et al. , 2018 ) , as their base model . Despite the great power of such language models , they have difficulty dealing with SMILES input , 1For example , for glycerol , its IUPAC name , molecular formula , structural formula , and skeletal formula are propane-1,2,3-triol , C3H8O3 , CH2 OH CH OH CH2 OH , and OH OH OH , respectively . 2The Simplified Molecular-Input Line-Entry System ( SMILES ) is a specification in the form of a line notation for describing the structure of chemical species using short ASCII strings . For example , the SMILES string for glycerol is “ OCC ( O ) CO ” . because SMILES is 1D linearization of molecular structure , which makes it hard for language models to learn the original structural information of molecules simply based on “ slender ” strings ( see Section 4 for more discussion ) . Another line of MRL methods , instead , use graph neural networks ( GNNs ) ( Kipf & Welling , 2017 ) to process molecular graphs ( Merkwirth & Lengauer , 2005 ; Jin et al. , 2017 ; Gilmer et al. , 2017 ; Ishida et al. , 2021 ) . Though GNN-based methods are theoretically superior to SMILES-based methods in learning molecule structure , they are limited to designing fresh and delicate GNN architectures while ignoring the essence of MRL , which is generalization ability . Actually , we will show later that , there is no specific GNN that performs universally best in all downstream tasks of MRL , which inspires us to explore beyond GNN architectures . To address the limitations of existing work , in this paper , we propose using chemical reactions to assist learning molecule representations and improving their generalization ability . A chemical reaction is usually represented as a chemical equation in the form of symbols and formulae , wherein the reactant entities are given on the left-hand side and the product entities on the right-hand side . For example , the chemical equation of Fischer esterification of acetic acid and ethanol can be written as CH3COOH + C2H5OH → CH3COOC2H5 + H2O . A chemical reaction usually indicates a particular relation of equivalence between its reactants and products ( e.g. , in terms of conservation of mass and conservation of charge ) , and our idea is to preserve this equivalence in the molecule embedding space . Specifically , given the chemical reaction of Fischer esterification above , we hope that the equation hCH3COOH + hC2H5OH = hCH3COOC2H5 + hH2O also holds , where h ( · ) represents molecule embedding function . This simple constraint endows molecule embeddings with very nice properties : ( 1 ) Molecule embeddings are composable with respect to chemical reactions , which make the embedding space well-organized ( see Proposition 1 ) ; ( 2 ) More importantly , we will show later that , when the molecule encoder is a GNN with summation as the readout function , our model can automatically and implicitly learn reaction templates that summarize a group of chemical reactions within the same category ( see Proposition 2 ) . The ability of learning reaction templates is the key to improving the generalization ability of molecule representation , since the model can easily generalize its learned knowledge to a molecule that is unseen but belongs to the same category as or shares the similar structure with a known molecule . We show that the molecule embeddings learned by our proposed model , namely MolR ( chemicalreaction-aware molecule embeddings ) , is able to benefit a variety of downstream tasks , which makes it significantly distinct from all existing methods that are designed for only one downstream task . For example , MolR achieves 17.4 % absolute Hit @ 1 gain in reaction product prediction , 2.3 % absolute AUC gain on BBBP dataset in molecule property prediction , and 18.5 % relative RMSE gain in graph-edit-distance prediction , respectively , over the best baseline method . We also visualize the learned molecule embeddings and show that they are able to encode reaction templates as well as several key molecule attributes , e.g. , molecule size and the number of smallest rings . 2 THE PROPOSED METHOD . 2.1 STRUCTURAL MOLECULE ENCODER . A molecular graph is represented as G = ( V , E ) , where V = { a1 , · · · } is the set of non-hydrogen atoms and E = { b1 , · · · } is the set of bonds . Each atom ai has an initial feature vector xi encoding its properties . In this work , we use four types of atom properties : element type , charge , whether the atom is an aromatic ring , and the count of attached hydrogen atom ( s ) . Each type of atom properties is represented as a one-hot vector , and we add an additional “ unknown ” entry for each one-hot vector to handle unknown values during inference . The four one-hot vectors are concatenated as the initial atom feature . In addition , each bond bi has a bond type ( e.g. , single , double ) . Since the bond type can usually be inferred by the features of its two associated atoms and does not consistently improve the model performance according to our experiments , we do not explicitly take bond type as input . To learn structural representation of molecules , we choose GNNs , which utilize molecule structure and atom features to learn a representation vector for each atom and the entire molecule , as our base model . Typical GNNs follow a neighborhood aggregation strategy , which iteratively updates the representation of an atom by aggregating representations of its neighbors and itself . Formally , the k-th layer of a GNN is : hki = AGGREGATE ( { hk−1j } j∈N ( i ) ∪ { i } ) , k = 1 , · · · , K , ( 1 ) where hki is atom ai ’ s representation vector at the k-th layer ( h 0 i is initialized as ai ’ s initial feature xi ) , N ( i ) is the set of atoms directly connected to ai , and K is the number of GNN layers . The choice of AGGREGATE function is essential to designing GNNs , and a number of GNN architectures have been proposed . See Appendix A for a detailed introduction on GNN architectures . Finally , a readout function is used to aggregate all node representations output by the last GNN layer to obtain the entire molecule ’ s representation hG : hG = READOUT ( { hKi } ai∈V ) . ( 2 ) The READOUT function can be a simple permutation invariant function such as summation and mean , or a more sophisticated graph-level pooling algorithm ( Ying et al. , 2018 ; Zhang et al. , 2018 ) . An illustrative example of GNN encoder is shown in Figure 1a . 2.2 PRESERVING CHEMICAL REACTION EQUIVALENCE . A chemical reaction defines a particular relation “ → ” between reactant set R = { r1 , r2 , · · · } and product set P = { p1 , p2 , · · · } : r1 + r2 + · · · → p1 + p2 + · · · . ( 3 ) A chemical reaction usually represents a closed system where several physical quantities of the system retain constant before and after the reaction , such as mass , energy , charge , etc . Therefore , it describes a certain kind of equivalence between its reactants and products in the chemical reaction space . Our key idea is to preserve such equivalence in the molecule embedding space : ∑ r∈R hr = ∑ p∈P hp . ( 4 ) The above simple constraint is crucial to improving the quality of molecule embeddings . We first show , through the following proposition , that the chemical reaction relation “ → ” is an equivalence relation under the constraint of Eq . ( 4 ) : Proposition 1 Let M be the set of molecules , R ⊆M and P ⊆M be the reactant set and product set of a chemical reaction , respectively . If R → P ⇔ ∑ r∈R hr = ∑ p∈P hp for all chemical reactions , then “ → ” is an equivalence relation on 2M that satisfies the following three properties : ( 1 ) Reflexivity : A→ A , for all A ∈ 2M ; ( 2 ) Symmetry : A→ B ⇔ B → A , for all A , B ∈ 2M ; ( 3 ) Transitivity : If A→ B and B → C , then A→ C , for all A , B , C ∈ 2M . The proof of Proposition 1 is in Appendix B . One important corollary of Proposition 1 is that , the set of all subsets of M , i.e . 2M , is naturally split into equivalence classes based on the equivalence relation “ → ” . For all molecule sets within one equivalent class , the sum of embeddings of all molecules they consist of should be equal . For example , in organic synthesis , a target compound t may be made from three different sets of starting materials A , B , and C. Then the sets A , B , C as well as { t } belong to one equivalence class , and we have ∑ m∈A hm = ∑ m∈B hm = ∑ m∈C hm = ht . Note that the starting materials are usually small and basic molecules that frequently appear in a number of synthesis routes . Therefore , Eq . ( 4 ) forms a system of linear equations , wherein the chemical reaction equivalence imposes strong constraint on the embeddings of base molecules . As a result , the feasible solutions of molecule embeddings will be more robust , and the whole embedding space will be more organized . See the visualized result on molecule embedding space in Section 3.5 for more details . We can further show that , the constraint in Eq . ( 4 ) is also able to improve the generalization ability of molecule embeddings . To see this , we first define reaction center for a chemical reaction . The reaction center of R → P is defined as a subgraph of R consisting of atoms whose bonds have changed after reaction . For example , for the reaction in the upper part of Figure 1b , its reaction center is the two oxygen atoms marked in dark orange , since they are the only two atoms whose bonds have changed . Given the concept of reaction center , we have the following proposition : Proposition 2 Let R→ P be a chemical reaction where R is the reactant set and P is the product set , and C be its reaction center . Suppose that we use the GNN ( whose number of layers is K ) shown in Eqs . ( 1 ) and ( 2 ) as the molecule encoder , and set the READOUT function in Eq . ( 2 ) as summation . Then for an arbitrary atom a in one of the reactant whose final representation is hKa , the residual term ∑ r∈R hr − ∑ p∈P hp is a function of h K a if and only if the distance between atom a and reaction center C is less than K. The proof of Proposition 2 is in Appendix C. Proposition 2 indicates that the residual between reactant embedding and product embedding will fully and only depend on atoms that are less than K hops away from the reaction center . For example , as shown in Figure 1b , suppose that we use a 3- layer GNN to process Fischer esterification of propionic acid and propanol , then the residual between reactant embedding and product embedding will totally depend on the reaction center ( colored in orange ) as well as atoms whose distance from the reaction center is 1 or 2 ( colored in light orange ) . This implies that , if the GNN encoder has been well-optimized on this chemical equation and outputs perfect embeddings , i.e. , hCH3CH2COOH + hCH3CH2CH2OH = hCH3CH2COOCH2CH2CH3 + hH2O , then the equation hR1−CH2COOH + hR2−CH2CH2OH = hR1−CH2COOCH2CH2−R2 + hH2O will also hold for any functional group R1 and R2 , since the residual between the two sides of the equation does not depend on R1 or R2 that are more than 2 hops away from the reaction center . The induced general chemical reaction R1-CH2COOH + R2-CH2CH2OH→ R1-CH2COOCH2CH2-R2 + H2O is called a reaction template , which abstracts a group of chemical reactions within the same category . The learned reaction templates are essential to improving the generalization ability of our model , as the model can easily apply this knowledge to reactions that are unseen in training data but comply with a known reaction template ( e.g. , acetic acid plus propanol , butyric acid plus butanol ) . We will further show in Section 3.5 how reaction templates are encoded in molecule embeddings . Remarks . Below are some of our remarks to provide further understanding on the proposed model : First , compared with ( Jin et al. , 2017 ) that also learns reaction templates , our model does not need a complicated network to calculate attention scores , nor requires additional atom mapping information between reactants and products as input . Moreover , our model is theoretically able to learn a reaction template based on even only one reaction instance , which makes it particularly useful in few-shot learning ( Wang et al. , 2020 ) scenario . See Section 3.1 for the experimental result . Second , organic reactions are usually imbalanced and omit small and inorganic molecules to highlight the product of interest ( e.g. , H2O is usually omitted from Fischer esterification ) . Nonetheless , our model can still learn meaningful reaction templates as long as chemical reactions are written in a consistent manner ( e.g. , H2O is omitted for all Fischer esterification ) . Third , the number of GNN layers K can greatly impact the learned reaction templates according to Proposition 2 : A small K may not be enough to represent a meaningful reaction template ( e.g. , in Fischer esterification , the necessary carbonyl group “ C=O ” in carboxylic acid will not appear in the reaction template ifK < 3 ) , while a largeK may include unnecessary atoms for a reaction template and thus reduces its coverage ( e.g. , the Fischer esterification of formic acid HCOOH and methanol CH3OH is not covered by the reaction template shown in Figure 1b ) . The empirical impact of K is shown in Appendix D .

Summary: The paper is about learning a vector representation of molecules in a way that the learned representation preserves the equivalence of molecules with respect to chemical reactions. They do so by forcing the sum of reactant embeddings and the sum of product embeddings to be equal for each chemical equation.  They have also shown experimentally that such embeddings can improve the performance of downstream tasks such as chemical reaction prediction, molecule property prediction, and graph edit distance prediction problems.