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 proposes a molecule representation learning method which is guided by chemical reactions.  In particular, it leverages chemical reaction equations by forcing the sum of reactant embeddings and the sum of product embed- dings to be equal for each chemical equation.  This idea is simple and useful, sharing the spirit of Word2Vec and TransE.