RetCL: A Selection-based Approach for Retrosynthesis via Contrastive Learning

1 Introduction

Retrosynthesis [Corey, 1991], finding a synthetic route starting from commercially available reactants to synthesize a target product (see Figure 1a), is at the center of focus for discovering new materials in both academia and industry. It plays an essential role in practical applications by finding a new synthetic path, which can be more cost-effective or avoid patent infringement. However, retrosynthesis is a challenging task that requires searching over a vast number of molecules cient selection framework, named RETCL (retrosynthesis via contrastive learning). To this end, we design two effective selection scores in synthetic and retrosynthetic manners. To be specific, we use the cosine similarity between molecular embeddings of the product and the reactants computed by graph neural networks. For training the score functions, we also propose a novel contrastive learning scheme [Sohn, 2016; He et al., 2019; Chen et al., 2020] with hard negative mining [Harwood et al., 2017] to overcome a scalability issue while handling a large-scale candidate set.

∗This work was partially done while the first author visited Samsung Advanced Institute of Technology.

1 The supplementary material is available at arXiv:2105.00795.

To demonstrate the effectiveness of our RETCL, we conduct various experiments based on the USPTO database [Lowe, 2012] containing 1.8M chemical reactions in the US patent literature. Thanks to our prior knowledge on the candidate reactants, our method achieves 71.3% test accuracy and significantly outperforms the baselines without such prior knowledge. Furthermore, our algorithm demonstrates its superiority even when enhancing the baselines with candidate reactants, e.g., our algorithm improves upon the existing template-free approach [Chen et al., 2019] by 11.7%. We also evaluate the generalization ability of RETCL by testing USPTO-50k-trained models on the USPTO-full dataset; we obtain 39.9% test accuracy while the state-of-the-art template-based approach [Dai et al., 2019] achieves 26.7%.

We believe our scheme has the potential to improve further in the future, by utilizing (a) additional chemical knowledge such as atom-mapping or leaving groups [Shi et al., 2020; Somnath et al., 2020]; (b) various constrastive learning techniques in other domains, e.g., computer vision [He et al., 2019; Chen et al., 2020; H´enaff et al., 2019; Tian et al., 2019], audio processing [Oord et al., 2018], and reinforcement learning [Srinivas et al., 2020].

2 Selection-based Retrosynthesis via Contrastive Learning

2.1 Overview of RETCL

In this section, we propose a selection framework for retrosynthesis via contrastive learning, coined RETCL. Our framework is based on solving the retrosynthesis task as a selection problem over a candidate set of commercially available reactants given the target product. Especially, we design a selection procedure based on molecular embeddings computed by graph neural networks and train the networks via contrastive learning.

To this end, we define a chemical reaction R → P as a synthetic process of converting a reactant-set R = {R1, . . . ,Rn}, i.e., a set of reactant molecules, to a product molecule P (see Figure 1a). We aim to solve the problem of retrosynthesis by finding the reactant-set R from a candidate set C which can be synthesized to the target product P . Especially, we consider the case when the candidate set C consists of commercially available molecules. Throughout this paper, we say that the synthetic direction (fromR to P) is forward and the retrosynthetic direction (from P to R) is backward.

Note that our framework stands out from the existing works in terms of the candidate set C. To be specific, (a) templatefree approaches [Lin et al., 2019; Karpov et al., 2019; Shi et al., 2020] choose C as the whole space of (possibly unavailable) molecules; and (b) template-based approaches [Coley et al., 2017; Segler and Waller, 2017; Dai et al., 2019] choose C as possible reactants extracted from the known reaction templates. In comparison, our framework neither requires (a) search over the entire space of molecules, or (b) domain knowledge to extract the reaction templates.

We now briefly outline the RETCL framework. Our framework first searches the most likely reactant-sets R1, . . . ,RT ⊂ C in a sequential manner based on a backward selection score ψ(R|P,Rgiven), and then ranks the reactant-sets using ψ(R|P,Rgiven) and another forward score φ(P|R). For learning the score functions, we propose a novel contrastive learning scheme with hard negative mining for improving the selection qualities. We next provide detailed descriptions of the search procedure and the training scheme in Section 2.2 and 2.3, respectively.

2.2 Search Procedure with Graph Neural Networks

We first introduce the search procedure of RETCL in detail. To find a reactant-set R = {R1, . . . , Rn}, we select each element Ri sequentially from the candidate set C based on the backward (retrosynthetic) selection score ψ(R|P,Rgiven). It represents a selection score of a reactant R given a target product P and a set of previously selected reactants Rgiven ⊂ C. Note that ψ is also capable of selecting a special reactant Rhalt to stop updating the reactant-set. Using beam search, we choose top T likely reactant-sets R1, . . . , RT .

Furthermore, we rank the chosen reactant-sets R1, . . . , RT based on the backward selection score ψ(R|P,Rgiven) and the forward (synthetic) score φ(P|R). The latter represents the synthesizability of R for P. Note that ψ(R|P,Rgiven) and φ(P|R) correspond to backward and forward directions of a chemical reaction R → P, respectively (see Section 2.1 and Figure 1a). Using both score functions, we define an overall score on a chemical reaction R → P as follows:

the space of permutations defined on the integers 1, . . . , n+1 satisfying π(n+1) = n+1. Based on score(P,R), we decide the rankings of R1, . . . ,RT for synthesizing the target product P. We note that the maxπ∈Π operator and the 1 n+2 term make the overall score (1) be independent of order and number of reactants, respectively. Figure 2 illustrates this search procedure of our framework.

Score design. We next elaborate our design choices for the score functions ψ and φ. We first observe that the molecular graph of the product P can be decomposed into subgraphs from each reactant of the reactant-set R, as illustrated in Figure 1a. Moreover, when selecting reactants sequentially, the structural information of the previously selected reactants Rgiven should be ignored to avoid duplicated selections. From these observations, we design the scores ψθ(R|P,Rgiven) and φ(P|R) as follows:

where d is the cosine similarity and fθ, gθ, hθ are embedding functions from a molecule to a fixed-sized vector with parameters θ. Note that one could think that fθ and gθ are query functions for a product and a reactant, respectively, while hθ is a key function for a molecule. Such a query-key separation allows the search procedure to be processed as an efficient matrix-vector multiplication. This computational efficiency is important in our selection-based setting because the number of candidates is often very large, e.g., |C| ≈ 6 × 105 for the USPTO dataset.

To parameterize the embedding functions fθ, gθ and hθ, we use the recently proposed graph neural network (GNN) architecture, structure2vec [Dai et al., 2016; Dai et al., 2019]. Incorporating reaction types. A human expert could have some prior information about a reaction type, e.g., carboncarbon bond formation, for the target product P. To utilize this prior knowledge, we add trainable embedding bias vectors u(t) and v(t) for each reaction type t into the query embeddings of ψ and φ, respectively. For example, φ(P|R) becomes CosSim(∑R∈R gθ(R) + v(t) , hθ(P)). The bias vectors are initialized by zero at beginning of training.

2.3 Training Scheme with Contrastive Learning

Finally, we describe our learning scheme for training the score functions defined in Section 2.1 and 2.2. We are inspired by how the score functions ψ(R|P,Rgiven) and φ(P |R) resemble the classification scores of selecting (a) the reactant R given the product P and the previously selected reactants Rgiven and (b) the product P given all of the selected reactants R, respectively. Based on this intuition, we consider two classification tasks with the following probabilities:

where τ is a hyperparameter for temperature scaling and C is the given candidate set of molecules. Note that we do not consider P and R ∈ R as available reactants and products for the classification tasks of p and q, respectively. This reflects our prior knowledge that the product P is always different from the reactants R in a chemical reaction. As a result, we arrive at the following losses defined on a reaction of the product P and the reactant-set R = {R1, . . . , Rn}:

where Rπ<i , Rn+1 and Π are the same as defined in (1). We note that minimizing the above losses increases the scores ψ(R|P,Rgiven) and φ(P|R) of the correct pairs of product and reactants, i.e., numerators, while decreasing that of wrong pairs, i.e., denominators. Such an objective is known as contrastive loss which has recently gained much attention in various domains [Sohn, 2016; He et al., 2019; Chen et al., 2020; Oord et al., 2018; Srinivas et al., 2020].

Unfortunately, the optimization of Lbackward and Lforward is intractable since the denominators of p(R|P,Rgiven, C) and q(P|R, C) require summation over the large set of candidate molecules C. To resolve this, for each mini-batch of reactions B sampled from the training dataset, we approximate C with CB = ⋃R→P∈B R ∪ {P} which is the set of all molecules in B. Then we arrive at the following training objective as L(B|θ) = 1 |B| ∑(R,P )∈B (Lbackward(P, R|θ, CB) + Lforward(P, R|θ, CB)).

Hard negative mining. In our setting, molecules in the candidate set CB are easily distinguishable. Hence, learning to discriminate between them is often not informative. To alleviate this issue, we replace the CB with its augmented version ˜ CB by adding hard negative samples, i.e., ˜ CB = CB ∪ ⋃M∈CB {Top-K nearest neighbors of M from C}, where K is a hyperparameter controlling hardness of the contrastive task. The nearest neighbors are defined with respect to the cosine similarity on {hθ(M)}M∈C. Since computing all embeddings {hθ(M)}M∈C for every iteration is timeconsuming, we update information of the nearest neighbors periodically. We found that this technique significantly improves the performance of RETCL (see Section 3.3).

3 Experiments

3.1 Experimental Setup

Dataset. We mainly evaluate our framework in USPTO-50k, which is a standard benchmark for the task of retrosynthesis. It contains 50k reactions of 10 reaction types derived from the US patent literature, and we divide it into training/validation/test splits following [Coley et al., 2017]. To apply our framework, we choose the candidate set of commercially available molecules C as the all reactants in the entire USPTO database as [Guo et al., 2020] did. This results in the candidate set with a size of 671,518. For the evaluation metric, we use the top-k exact match accuracy, which is widely used in the retrosynthesis literature. We also experiment with other USPTO benchmarks for more challenging tasks, e.g., generalization to unseen templates. We provide a more detailed description of the USPTO benchmarks in the supplementary material.

Hyperparameters. We use a single shared 5-layer structure2vec [Dai et al., 2016; Dai et al., 2019] architecture and three separate 2-layer residual blocks with an embedding size of 256. To obtain graph-level embedding vectors, we use sum pooling over mean pooling since it captures the size information of molecules. For contrastive learning, we use a temperature of τ = 0.1 and K = 4 nearest neighbors for hard negative mining. More details are provided in the supplementary material.

3.2 Single-step Retrosynthesis in USPTO-50k

Table 1 evaluates our RETCL and other baselines using the top-k exact match accuracy with k ∈ {1, 3, 5, 10, 20, 50}. We first note that our framework significantly outperforms a concurrent selection-based approach, 3 Bayesian-Retro [Guo et al., 2020], by 23.8% and 23.7% in terms of top-1 accuracy when reaction type is unknown and given, respectively. Furthermore, ours also outperforms template-based approaches utilizing the different knowledge, i.e., reaction templates instead of candidates, with a large margin, e.g., 18.8% over GLN [Dai et al., 2019] in terms of top-1 accuracy when reaction type is unknown.

3 Note that Bayesian-Retro [Guo et al., 2020] is not scalable to a large candidate set, e.g., 6 hours for testing one product.

Incorporating the knowledge of candidates into baselines. However, it is hard to fairly compare between methods operating under different assumptions. For example, templatebased approaches require the knowledge of reaction templates, while our selection-based approach requires that of available reactants. To alleviate such a concern, we incorporate our prior knowledge of candidates C into the baselines; we filter out reactants outside the candidates C from the predictions made by the baselines. As reported in Table 2, our framework still outperforms the template-free approaches with a large margin, e.g., Transformer [Chen et al., 2019] achieves 68.4% in the top-1 accuracy, while we achieve 78.9% when reaction type is given. Although GLN uses more knowledge than ours in this setting, its top-k accuracy is saturated to 93.3% which is the coverage of known templates, i.e., the upper bound of template-based approaches. However, our framework continues to increase the top-k accuracy as k increases, e.g., 97.5% in terms of top-200 accuracy.

3.3 Analysis and Ablation Study

Failure cases. Figure 3 shows examples of wrong predictions generated by our framework. We found that our predictions are still similar to the ground-truth (i.e., presented in USPTO-50k) reactants. For example, the third predictions of the examples A and B are partially correct; the larger reactant is correct while the smaller one is slightly different. In the example C, the ring at the center of the product is broken in the ground-truth reactants while our RETCL predicts non-broken reactants. Surprisingly, in a chemical database, Reaxys, we found a synthetic path starting from the second prediction 4 , as illustrated in Figure 4. The path is present in a chemical literature [Gonda and Novak, 2015]. These results show that our RETCL could provide meaningful information for retrosynthetic analysis in practice.

Nearest neighbors on molecular embeddings. For hard negative mining described in Section 2.3, it is required to find similar molecules using the cosine similarity on {hθ(M)}M∈C. As illustrated in Figure 5, hθ(M) is capable of capturing the molecular structures.

Effect of components. Table 4 shows the effect of components of our framework. First, we found that the hard negative mining as described in Section 2.3 increases the performance significantly. This is because there are many similar molecules in the candidate set C, thus a model could predict slightly different reactants without hard negative mining. We also demonstrate the effect of checking the synthesizablity of the predicted reactants with φ(P|R). As seen the fourth and fifth rows in Table 4, using φ(P|R) provides a 2.6% gain in terms of top-10 accuracy. Moreover, we empirically found that sum pooling for aggregating node embedding vectors is more effective than mean pooling. This is because the former can capture the size of molecules as the norm of vectors.

4 We note that our model is trained without considering reagents (e.g., sodium acetate or thiophene) as previous one-step models did.

3.4 More Challenging Retrosynthesis Tasks

Generalization to unseen templates. The advantage of our framework over the template-based approaches is the generalization ability to unseen reaction templates. To demonstrate it, we remove reactions of classes (i.e., reaction types) from 6 to 10 in training/validation splits of the USPTO-50k benchmark. Then the number of remaining reactions is 27k. In this case, the templates extracted from the modified dataset cannot be applied to the reactions of different classes. Hence the template-based approaches suffer from the generalization issue; for example, GLN [Dai et al., 2019] cannot provide correct predictions for reactions of unseen types as reported in Table 3, while our RETCL can provide correct answers.

We also conduct a more realistic experiment: testing on a larger dataset, the test split of USPTO-full dataset preprocessed by [Dai et al., 2019], using a model trained on a smaller dataset, USPTO-50k. We note that the number of reactions for training, 40k, is smaller than that of testing reactions, 100k. As reported in Table 5, our framework provides a consistent benefit over the template-based approaches. These results show that our strength of generalization ability.

Generalization to unseen candidates. The knowledge of the candidate set C could be updated after learning the RETCL framework. In this case, the set used in the test phase is larger than that in the training phase, i.e., Ctrain ⊊ Ctest. One can learn the framework once again, however someone wants to use it instantly without additional training. To validate that our framework can generalize to unseen candidates, we use a smaller candidate set during the training phase. To be specific, we use all molecules present in training and validation splits of USPTO-50k. In this case, the number of candidate reactants in the training phase is 91297. When testing, we use the original candidate set, in other words, |Ctest| = 671518. As reported in Table 6, this model achieves comparable performance to another model trained with a larger number of candidate reactants. This experiment demonstrates that our framework trained with a small corpora (e.g., USPTO-50k) can work with unseen candidates.

Multi-step retrosynthesis. In practice, retrosynthesis often requires multiple recursive steps since intermediate reactants could be unavailable. We remark that our RETCL can be applied into not only the singe-step retrosynthesis, but also the multi-step problem. We provide multi-step retrosynthesis experiments with our framework in the supplementary material. For example, we found that our RETCL can improve the quality (e.g., length) of discovered synthetic routes.

4 Conclusion

In this paper, we propose RETCL for solving retrosynthesis. To this end, we reformulate retrosynthesis as a selection problem of commercially available reactants, and propose a contrastive learning scheme with hard negative mining to train our RETCL. Extensive experiments show that our framework achieves outstanding performance for the USPTO benchmarks. Furthermore, we demonstrate the generalizability of RETCL to unseen reaction templates. We believe that extending our framework to multi-step retrosynthesis or combining with various contrastive learning techniques in other domains could be interesting future research directions.

Acknowledgments

This work was mainly supported by Samsung Electronics Co., Ltd (IO201211-08107-01) and Institute of Information & Communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No.2019-0-00075, Artificial Intelligence Graduate School Program (KAIST)).