Towards understanding retrosynthesis by energy-based models

Retrosynthesis is a critical problem in organic chemistry and drug discovery [ 1–5 ] . As the reverse process of chemical synthesis [ 6 , 7 ] , retrosynthesis aims to find the set of reactants that can synthesize the provided target via chemical reactions ( Fig 1 ) . Since the search space of theoretically feasible reactant candidates is enormous , models should be designed carefully to have the expression power to learn complex chemical rules and maintain computational efficiency . Recent machine learning applications on retrosynthesis , including sequence- and graph-based models , have made significant progress [ 3 , 8 , 9 ] . Sequence-based models treat molecules as one-dimensional token sequences ( SMILES [ 10 ] , bottom of Fig 1 ) and formulate retrosynthesis as a sequence-to-sequence problem , where recent advances in neural machine translation [ 11 , 12 ] can be applied . In this principle , the LSTMbased encoder-decoder frameworks and , more recently , transformer-based approaches have achieved promising results [ 13 , 12 , 14 ] . On the other hand , graph-based models have a natural representation of humaninterpretable molecular graphs , where chemical rules are easily applied . Graph-based approaches that perform graph matching with templates ( e.g . chemical rules ) or reaction centers have reached encouraging results . Among those , G2Gs [ 15 ] , RetroXpert [ 16 ] and GraphRETRO [ 17 ] outperform template-based methods by inferring reaction centers in a supervised way . In this paper , we focus on one-step retrosynthesis , which is also the foundation of multi-step retrosynthesis [ 3 ] . Our goal here is to provide a unified view of both sequence- and graph-based retrosynthesis models using an energy-based model ( EBM ) framework . It is beneficial because : First , the model design with EBM is very flexible . Within this framework , both types of models can be formulated as different 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . EBM variants by instantiating the energy function into specific forms . Second , EBM provides principled ways for training models , including maximum likelihood estimator , pseudo-likelihood , etc . Third , a unified view is critical to provide insights into different EBM variants , as it is easy to extract commonalities and differences between EBM variants , understand strengths and limitations in model design , compare the complexity of learning or inference , and inspire novel EBM variants . To summarize our contributions : • We propose a unified energy-based model ( EBM ) framework that integrates sequence- and graph-based models for retrosynthesis . To our best knowledge , this is the first effort to unify and exploit inner connectivity between different models . • We perform rigorous evaluations by running tens of experiments on different model designs . Revealing the performance to the community contributes to the development of retrosynthesis models . • Inspired by such a unified framework , we propose a novel generalized dual EBM variant that performs consistent training over forward and backward prediction directions . This model improves the state-of-the-art by 4.3 % . 1 Energy-based model for Retrosynthesis . Algorithm 1 EBM framework [ Train Phase ] : Learning Input : Reactants X and products y . 1 . Parameterize X and y in Sequence or Graph format . 2 . Design Eθ { e.g . dual , perturbed , bidirectional , graph-based , etc } // Sec 2 3 . Select training loss to learn Eθ and obtain θ∗ // Sec 3 Return θ∗ [ Test Phase ] : Inference // Sec 4 Input : θ∗ , ytest , Proposal P . // Sec 4 4 . Obtain a list of X candidates by P . Ltest ← P ( ytest ) 5 . X∗ = argminX∈Ltest Eθ∗ ( X , ytest ) Return : X∗ Retrosynthesis is to predict a set of reactant molecules from a product molecule . We denote the product as y , and the set of reactants predicted for one-step retrosynthesis as X . The key for retrosynthesis is to model the conditional probability p ( X|y ) . EBM provides a common theoretical framework that can unify many retrosynthesis models , including but not limited to existing models . An EBM defines the distribution using an energy function [ 18 , 19 ] . Without loss of generality , we define the joint distribution of product and reactants as follows : pθ ( X , y ) = exp ( −Eθ ( X , y ) ) Z ( θ ) ( 1 ) where the partition function Z ( θ ) =∑ y ∑ X exp ( −Eθ ( X , y ) ) is a normalization constant to ensure a valid probability distribution . Since the design of Eθ is free of choice , EBMs can be used to unify many retrosynthesis models by instantiating the energy function E ( θ ) with various designs . Note there is a trade-off between model expression capacity and learning tractability . EBM is also easy to obtain conditioning with different partition functions . The forward prediction probability for reaction outcome prediction pθ ( y|X ) can be written as exp ( −Eθ ( X , y ) ) ∑ y′ exp ( −Eθ ( X , y′ ) ) with the same form of energy function . The proposed framework works as follows : Step 1 , design and train an energy function Eθ ( Sec 2 and Sec 3 ) , and Step 2 use Eθ for inference in retrosynthesis ( Sec 4 ) . See Fig 2 and Algorithm 1 . 2 Model Design . Based on how to parameterize reactant and product molecule X and y , the model designs can be divided into two categories : sequence-based and graph-based models . 2.1 Sequence-based Models . Here we describe several sequence-based parametriztion to instantiate our EBM framework , which use SMILES string as representations of molecules . We first define the sequence-based notations . Given a reactant molecule x , we denote its SMILES representation as s ( x ) . Superscript s ( x ) ( i ) denotes the character at i-th position of the SMILES string . For simplicity , we use x ( i ) when possible . Reactants of a chemical reaction are usually a collection of molecules : X = { x1 , x2 , .. , xj , .. , x|X| } , where xj is the j-th reactant molecule . The SMILES representation of a molecule set X , denoted as s ( X ) , is a concatenation of s ( x ) for every x in X with “ . ” in between : “ s ( x1 ) .s ( x2 ) ... s ( x|X| ) ” . We use X ( i ) as the short form of s ( X ) ( i ) to denote the i-th position of the concatenated SMILES . 2.1.1 Full energy-based model . We start by proposing a most flexible EBM that imposes the minimum restrictions on the design of Eθ . All the variants proposed in Sec 2.1 are special instantiations of this model ( e.g . by specifying different Eθ ) . The EBM is defined as follows : p ( X|y ) = exp ( −Eθ ( X , y ) ) ∑ X′∈P ( M ) exp ( −Eθ ( X ′ , y ) ) ( 2 ) ∝ exp ( −Eθ ( X , y ) ) ( 3 ) Here the energy function Eθ : P ( M ) ×M 7→ R takes a molecule set and a molecule as input , and outputs a scalar value . M defines the set of all possible molecules . P ( · ) represents the power set . P ( M ) denotes domain of reactant sets X . Due to the intractability of the partition function , training involves additional information e.g. , template or approximation of the partition ( See Sec 3 ) . 2.1.2 Ordered model . One design of energy function is factoring the input sequence in an autoregressive manner [ 12 , 20 ] . pθ ( X|y ) = exp ( |s ( X ) |∑ i=1 log pθ ( X ( i ) |X ( 1 : i−1 ) , y ) ) ( 4 ) = exp ( |s ( X ) |∑ i=1 log exp ( hθ ( X ( 1 : i−1 ) , y ) > e ( X ( i ) ) ) ∑ c∈S exp ( hθ ( X ( 1 : i−1 ) , y ) > e ( c ) ) ) ( 5 ) where pθ ( X ( i ) |X ( 1 : i−1 ) , y ) is parameterized by a transformer hθ ( p , q ) : S|p| × S|q| 7→ R|S| where S is vocabulary . e ( c ) is a one-hot vector with dimension c set to 1 . This choice of hθ ( p , q ) enables efficient computing of the partition function , as it outputs a vector with length equal to |S| to represent logits ( unnormalized log probability ) for each value in vocabulary . Here , maximum likelihood estimator ( MLE ) is feasible for training , as this factorization allows tractable partition function . 2.1.3 Dual model . A different design is to leverage on duality of retrosynthesis and reaction prediction . They are a pair of mutual reversible processes that factorize the joint distribution in different orders , where reaction prediction is “ forward direction ” – p ( y|X ) ) and retrosynthesis is the “ backward direction ” – p ( X|y ) . With additional prior modeling , the joint probability p ( X , y ) factorizes to either p ( X|y ) p ( y ) or p ( y|X ) p ( X ) . We propose a training framework that leverages on the duality of the forward and backward directions and performs consistent training between the two to bridge the divergence . Algorithm 2 Dual Model [ Train Phase ] : Learning : Input : Reactants X and product y . Let θ = { γ , α , η } Define Eθ as Eq ( 7 ) Eθ = log pγ ( X ) + log pα ( y|X ) + log pη ( X|y ) 1 . Train backward : η∗ = argminη Ldual = argmaxη Ê [ log pη ( X|y ) ] 2 . Train prior and forward : Plug in η∗ pmix ( X , y ) = 11+β p̂ ( X , y ) + β 1+β p̂ ( y ) pη∗ ( X|y ) γ∗ , α∗ = argminγ , α Ldual = Êpmix ( X , y ) [ log pγ ( X ) + log pα ( y|X ) ] [ Test Phase ] : Inference : Input : θ∗ = { γ∗ , α∗ , η∗ } , ytest , Proposal P. L← P ( ytest ) X∗ = argminX∈LEθ∗ ( X , y test ) Return X∗ The advantage of the duality of reversible processes has been demonstrated in other applications as well . He et al . [ 21 ] trained a reinforcement learning ( policy gradient ) model to achieve duality in natural language processing and improved performances . Wei et al . [ 22 ] treated code summary and code generation as a pair of dual tasks , and improved efficacy by imposing symmetry between attention weights of LSTM encoderdecoder in forward and backward directions . Despite of their encouraging results , these models are not ideal for stable and efficient training for retrosynthesis , as policy gradient methods suffer from high variance and LSTM has sub-optimal performance . Therefore we propose a novel training method that is simple yet efficient for retrosynthesis task . We impose duality constraints by training forward direction on a mixture of samples drawn from the backward and original dataset . To our best knowledge , we are the first to apply duality to retrosynthesis and to impose duality constraints by samples drawn from one direction . The EBM is defined : p ( X|y ) ∝ exp ( log pγ ( X ) + log pα ( y|X ) + log pη ( X|y ) ) ( 6 ) = exp ( −Eθ ( X , y ) ) ( 7 ) where prior p ( X ) , forward likelihood p ( y|X ) , and backward posterior P ( X|y ) are modeled as autoregressive models ( Sec 2.1.2 ) , parameterized by transformers with parameters γ , α , and η . Note energy function can be designed free of choice . The consistent training is achieved by minimizing the “ dual loss ” , where the duality constraints in the equation below are imposed to penalize KL divergence of the two directions , i.e. , KL ( backward|forward ) . For simplicity , we fix the backward probability in the dual loss , and therefore entropy H ( backward ) is dropped . γ∗ , α∗ , η∗ = arg min γ , α , η ` dual ( 8 ) ` dual = − ( Ê [ log pγ ( X ) + log pα ( y|X ) ] ︸ ︷︷ ︸ forward direction ( 9 ) + βÊyÊX|y [ log pγ ( X ) + log pα ( y|X ) ] ︸ ︷︷ ︸ duality constraints + Ê [ log pη ( X|y ) ] ︸ ︷︷ ︸ backward direction ) ( 10 ) = −Êpmix ( X , y ) [ log pγ ( X ) + log pα ( y|X ) ] − Ê [ log pη ( X|y ) ] ( 11 ) where Ê indicates expectation over empirical data distribution p̂ ( X , y ) . The duality constraints βÊyÊX|y [ log pγ ( X ) + log pα ( y|X ) ] is the expectation of the forward direction log pγ ( X ) + log pα ( y|X ) with respect to empirical backward data distribution ÊyÊX|y , where ÊyÊX|y are approximated by samples drawn from pη ( X|y ) , as y is given so p ( y ) = 1. β is scale parameter . In our implementation we use size k-beam search to draw samples efficiently . Combining “ forward ” and “ duality constraints ” terms ( Eq 11 ) , we can see that the first term of the dual loss is to train the forward direction on the mixture distribution of the original data and samples drawn from backward directions pmix ( X , y ) = 11+β p̂ ( X , y ) + β 1+β p̂ ( y ) pη ( X|y ) . Put every piece together ( Algorithm 2 and Fig 4 in Appendix ) . Here is our training procedure . Since we parameterize the three probabilities separately , the optimization of dual loss breaks into two steps : • Step 1 : Train backward . η does not depend on forward direction under empirical data distribution . η∗ = argminη Ldual = argmaxη Ê [ log pη ( X|y ) ] . η can be learned by MLE . • Step 2 : Train prior and forward . We plug η∗ into pmixη∗ ( X , y ) . γ∗ , α∗ = argminγ , α Ldual = argmaxγ , α Êpmix η∗ ( X , y ) [ log pγ ( X ) + log pα ( y|X ) ] . γ , α can be learned by MLE . We provide ablation study of each component of dual loss in Appendix 5.4 . The results show that each component in dual loss contribute to the final performance positively .

Summary: This paper presents a unified framework with energy-based models for retrosynthesis prediction. Different methods (sequence and graph methods) can be untied with different energy functions within the framework. A dual variant model is presented and it performs consistent forward and backward training. Experimental results demonstrate this consistent training improves the retrosynthesis prediction significantly.