KGCL: Knowledge-Enhanced Graph Contrastive Learning for Retrosynthesis Prediction Based on Molecular Graph Editing

1 Introduction

Retrosynthesis prediction is an indispensable strategy for designing drugs, from small molecules to complex natural products[Gothard et al., 2012]. AI-based models for retrosynthesis prediction has achieved great progresses, which are divided into three classes: template-based, template-free, and semi-template-based methods. Template-based methods cannot predict reactions beyond the available templates, despite their strong interpretability [Segler and Waller, 2017; Dai et al., 2019; Chen and Jung, 2021; Xie et al., 2023]. Template-free methods do not rely on predefined templates [Tetko et al., 2020; Seo et al., 2021], but the reactants generated by these methods are probably chemically invalid and lack interpretability[Wan et al., 2022; Tu and Coley, 2022; Yao et al., 2024; Han et al., 2024]. To balance the dependence on templates and interpretability of the prediction process, semi-template-based methods were proposed, which perform retrosynthesis in two sequential steps: (1) breaking the molecule by identifying a reaction center, and (2) transforming the resulting fragments into potential reactants[Sacha et al., 2021].

∗Corresponding Author

The retrosynthesis prediction method based on molecular graph editing, which is inspired by the arrow-pushing method in reaction mechanism description, is one of the most representative semi-template-based strategies. As shown in Figure 1, the molecular graph editing methods formulate retrosynthesis as a sequential modification process on the product graph, which is guided by simplified reaction mechanisms. However, existing methods typically treat individual atoms in molecules as independent functional units when learning molecular graph representations. They ignore functional groups, which are substituents or moieties in molecules that cause the molecule's characteristic chemical reactions, thus limiting the predictive performance of the models[Zhong et al., 2023].

To address this issue, this paper incorporates the functional group knowledge from the chemical knowledge graph into the embeddings of molecular graphs. In details, we first extract functional group knowledge related to the target molecule from the knowledge graph, and then employ an attention mechanism to integrate the functional group knowledge into the atomic representation of the molecular graph. Such that the learned molecular embeddings become more discriminative and interpretable.

Furthermore, to enhance the accuracy of molecular graph representations, this paper introduces a contrastive learning strategy based on molecular graphs. The primary challenge is how to construct the rational positive and negative samples for the target molecular graph. Traditional random graph augmentation strategies are not suitable due to unique chemical constraints [Fang et al., 2023]. To address this issue, we use graph edit operations predicted by the model as perturbations to generate positive and negative samples for the target molecular graph.

The main contributions are summarized as follows:

• We propose a strategy to incorporate the chemical functional group knowledge from knowledge graph into retrosynthesis prediction based on molecular graph editing. To the best of our knowledge, this is the first study to apply the chemical knowledge graph to retrosynthesis prediction within the context of molecular graph editing.
• We present a new molecular graph contrastive learning method. This method generates positive and negative samples for the target molecular graph using molecular graph edit operations predicted by the model, thereby preserving the semantic information of molecular graphs to the greatest extent.
• Experiments on two benchmark datasets demonstrate that our proposed model KGCL consistently outperforms the compared baseline methods and achieve new state-of-the-art performance among the semi-templatebased methods.

2 RelatedWork

2.1 Knowledge-aware Semi-template-based Retrosynthesis Prediction

A series of semi-template-based retrosynthesis models were proposed in [Shi et al., 2020; Somnath et al., 2021; Wang et al., 2021; Sacha et al., 2021; Zhong et al., 2023]. Although achieved some progresses, these models did not consider the chemical knowledge when learning molecular representation. Thus the prediction performance of these models needs to be improved. Afterward, the models in [Chen et al., 2023; Liu et al., 2024] embedded the chemical knowledge, such as reaction center types and chemical synthesis rules, into the retrosynthesis systems to improve the accuracy and interpretability. But in these models, the knowledge was predefined and lacked flexibility. In the present paper, the knowledge is from the chemical knowledge graph, not from predefined rules. The knowledge graph has achieved great success to predict molecular properties[Fang et al., 2023], but not yet for retrosynthesis predictions. In the two scenarios, the fusion strategies of the chemical domain knowledge are different.

2.2 Molecular Graph Contrastive Learning

Inspired by the success of contrastive learning in image and language domains, it was introduced to molecular learning and was shown to be effective [Chen et al., 2024]. Some researchers generated positive and negative samples by random perturbations for data augmentation [Luo et al., 2023; Zheng et al., 2023; Sun et al., 2022]. However, these methods inevitably introduce variance in critical semantic information, which can mislead contrastive learning [You et al., 2020;

Sun et al., 2021a]. To address this issue, some works predefined some domain knowledge and then utilized them to generate two augmented views [Sun et al., 2021a; Kim et al., 2023]. In the present paper, we need not predefine any domain knowledge, but directly utilize the graph edits predicted by the model to construct positive and negative samples.

3 ProposedMethod

The overall framework of the KGCL model is shown in Figure 2, and the details are in the following subsections.

3.1 ProblemDefinition

A molecule M is represented as a graph G = (A, B) in retrosynthesis prediction based on molecular graph editing. A is the set of atoms and B is the set of bonds in the molecule, where each atom corresponds to a node and each bond corresponds to an edge. The graph edit set includes an atomlevel edit set Ea, a bond-level edit set Eb and a moleculelevel edit set Eg. Specifically, Ea includes (1) atom modification and (2) attaching a leaving group to the atom, Eb includes (1) bond modification and (2) bond deletion, and edits in Eg determine whether the graph editing process has been completed. The task of the retrosynthesis prediction model is, given a target product molecule Mp and its graph Gp, to predict a sequence of graph edits (e1u , · · · , elu , · · · , eLu ) with length L and u ∈ {a, b, g} (i.e., elu ∈ Ea ∪ Eb ∪ Eg), which modifies Gp sequentially, until the graphs Greactant of the reactant molecules, which could potentially synthesize the product, are obtained. This process can be formally described as follows:

where R denotes the set of all possible reactant graphs. The priority order of these operations is the same as in [Zhong et al., 2023].

3.2 Knowledge-basedMolecular Graph Enhancement

In KGCL, the i-th atom and the undirected bond between the i-th and j-th atoms are initially encoded as hi = (hi,1, · · · hi,na ) with dimension na and hij = (hij,1, · · · , hij,nb ) with dimension nb, respectively, for 1 ≤ i, j ≤ |A|. The initial features of atoms and bonds are predefined in terms of chemical properties. More details can be found in [Zhong et al., 2023].

A functional group is a group of atoms in a molecule with distinctive chemical properties, regardless of the other atoms in the molecule. To enhance the accuracy of molecular learning, we integrate the functional group knowledge into the representations of molecular graphs.

This study only focuses on the functional groups that are actually present in the target molecule to reduce the noises introduced by irrelevant functional groups. Specifically, given a product molecule Mp and the predefined set of functional groups F = {f1, · · · , fi, · · · , fn}, where n is the number of predefine functional groups, we first determine whether each fi is a substructure of the molecule Mp using the substructure matching algorithm. Therefore, we obtain the set of functional groups within Mp, denoted as Fp = {fp1 , · · · , fpj , · · · , fpm }, where m is the number of the functional groups in Mp. Subsequently, the embeddings of these functional groups Hfunc = [hTfunc,1 , · · · , hTfunc,j , · · · , hTfunc,m]T , where hfunc,j is the embedding vector of fpj , are retrieved from the chemical knowledge graph [Fang et al., 2023]. Finally, these functional group embeddings are incorporated into the representation of the molecular graph using an attention mechanism.

Specifically, the embeddings of the atoms in the original molecular graph are used as queries to attend to all functional groups. The knowledge-enhanced embedding ˜hi for the i-th atom is computed as follows:

where df is the dimension of the functional group embedding. Next, the embeddings of bonds are enhanced and injected direction information as follows:

where ⊕ denotes the concatenation operation. Finally, the knowledge-enhanced molecular graph is denoted as

3.3 The Molecular Graph Encoder

We use the directed message passing neural network (DMPNN) [Gilmer et al., 2017] to encode the molecular graph. At the (t + 1)-th iteration, hidden representation ˜h(t+1) ij of each edge is updated based on messages m(t+1) ij according to

where ˜hik is the initial feature of the edge from the node i to the k in the knowledge-enhanced molecular graph, Winput is the learnable weight matrix, N(i) is the neighbors of the i-th vertex, and mij is the message feature vector from the node i to the node j. Mt and Ut are the message and edge update function:

After T iterations, the embedding of the i-th vertex is updated

where σr is the SELU activation function, Woutput is the learnable weight matrix and c is the bias of the fully connected network.

3.4 The Generation of Edit Sequences

The edit sequence (e1u , · · · , elu , · · · , eLu ) for the target molecular graph is generated in an autoregressive manner [Zhong et al., 2023], in which three MLPs are used to predict the atomlevel, the bond-level and the graph-level edits, respectively. At the l-th step, to strengthen the relationship between the current editing operation and previous ones, the atom embeddings obtained from the D-MPNN encoder are updated by:

where ˜h(0) i = ˜hi defined in (12), Wp and Wc are the learnable weight matrices, and σr is the SELU activation function. After the atom features are updated, the bond features are recomputed using the features of the atoms at its two ends, and the molecular features are obtained by summing the hidden representation of all its atoms as follows:

where ⊕ denotes the concatenation operation.

At last, the logits o(l) (a,i) , o(l) (b,ij) and o(l) g corresponding to the atom-level edit a ∈ Ea, bond-level edit b ∈ Eb, and graphlevel edit at the l-th step are predicted as follows:

where Wi is the weight matrix, and ci is the corresponding bias for 1 ≤ i ≤ 6, and σr is the SELU activation function.

3.5 Molecular Graph Contrastive Learning Based on Edit Sequences

In order to enhance the representation learning ability of the encoder, we propose a novel molecular graph contrastive learning strategy. One main challenge in contrastive learning is to generate rational positive and negative samples:

• The anchor sample should be semantically similar to the positive sample while far from the negative samples;
• In our context, the molecules corresponding to the positive and negative samples should satisfy the chemical property constraints.

To address this issue, we leverage predicted edit sequences as augmentation operations to construct positive and negative samples. Specifically, let Sbatch = {M1, · · · ,Mi, · · · ,MN} denote the set of product molecules in the current training batch. For each molecule Mi, the predicted edit sequence is (ei1 , ei2 , · · · , eiLi ), where Li is the length of current edit sequence for Mi. This sequence generates a set of intermediates S+ i = {Mli = eil ◦ · · · ◦ ei1 (Mi)|1 ≤ l ≤ Li}. We randomly select one M+i from S+ i as the positive sample of Mi. The set of negative samples for Mi is defined as Si = ∪j̸=i{M1j ,M2j , · · · ,MLj j }, which includes intermediates generated from all other molecules in the batch.

In this study, we adopt a variant of InfoNCE loss as the contrastive loss, which assigns dynamic Gaussian weights to negative samples[Wu et al., 2024]. The contrastive loss function is defined as:

where ˜h ˜ Gi , ˜h˜ G+ i and ˜h˜ Gj are molecular graph embedding for the molecule Mi, its positive sample M + i and its negative sample M - j as defined in (15), τ is temperature coefficient and wj is the weight of the negative sample M - j , which is dynamically computed based on the similarity between the anchor sample Mi and the negative sample M - j as follows:

where µ and σ, two hyperparameters, are the mean and variance, respectively. Samples closer to µ have larger weights. The smaller the value of σ is, the more pronounced the weight differences between samples are.

3.6 Overall Loss Function for Model Training

We utilize the following cross-entropy loss, referred to as the edit prediction loss, to measure the difference between the ground truth labels and the predicted labels:

where L is the length of the graph edit sequence, Gm is an intermediate in the current batch during the training process. y(l) i , y(l) ij and y(l) ˜ G are the ground-truth labels for the atom edit a ∈ Ea, the bond edit b ∈ Eb and the molecular graph edit at the l-th step. o(l) (a,i) , o(l) (b,ij) and o(l) ˜ G are the output scores of the model for the corresponding edits.

The overall loss is defined as the weighted combination of the edit prediction loss and the contrastive learning loss:

where λ is a hyperparameter that controls the weight of the contrastive loss.

4 Experiments

4.1 Experimental Setup

Datasets. Our experiments are conducted on public benchmark datasets USPTO-50K [Schneider et al., 2016] and USPTO-FULL [Dai et al., 2019]. The USPTO-50K is a highquality dataset that contains 50,016 reactions with the correct atom-mapping. It is divided into the training, validation, and test sets according to the partitioning scheme as [Coley et al., 2017], with each set containing 40k, 5k, and 5k reactions, respectively. Compared with the USPTO-50K dataset, the USPTO-FULL dataset has higher coverage and diversity, containing 1000k unfiltered chemical reactions. We adopt the same splitting method as [Dai et al., 2019] and divide it into 800k, 100k, and 100k reactions as the training set, validation set, and test set respectively. For fair comparisons, the extraction of graph editing operations and other data preprocessing are the same as those in [Zhong et al., 2023].

Implementation details. The KGCL model uses the Adam optimizer for gradient descent training. For the USPTO-50K dataset, the initial learning rate is set to 0.001 (0.0001 for the USPTO-FULL dataset), and a polynomial decay learning rate scheduler is applied. When the improvement in accuracy on the validation set is less than a predefined threshold 0.01 within 5 consecutive epochs, it is considered to have reached a plateau (or peak). At this point, the learning rate is reduced by a decay factor 0.8. The hidden dimension of the encoder is set to 256, the number of iterations of message passing is 10, and node embeddings are dropped with the probability 0.15. We use three MLPs with a hidden dimension 512 and a dropout rate of 0.2 to predict the initial edit scores. The model was trained for 200 epochs using a batch size of 256. All experiments were conducted on a single NVIDIA RTX A6000 GPU.

Evaluation and baselines. We employ the top-k exact match, round-trip and MaxFrag accuracy as the metrics to measure the retrosynthesis performance. The top-k exact match accuracy is calculated by comparing the standard Simplified Molecular Input Line Entry System (SMILES) of predicted reactants with the ground-truths in the dataset [Weininger, 1988]. We report the top-k accuracy where k=1, 3, 5, 10 and 50. The round-trip accuracy is computed by evaluating the similarity between the ground-truth product and the product generated by a forward-synthesis model using the predicted reactants [Schwaller et al., 2020]. We utilize Molecular Transformer (MT) [Schwaller et al., 2019] as a forward-synthesis prediction model to compute the roundtrip accuracy. The MaxFrag accuracy is defined as the exact match between the largest fragment of the predicted reactant and the ground-truth in the dataset [Tetko et al., 2020].

We compared our model with representative baseline models in recent years, include:

• template-based methods: Retrosim[Coley et al., 2017], Neuralsym [Segler and Waller, 2017], GLN [Dai et al., 2019], LocalRetro [Chen and Jung, 2021], RetroKNN[Xie et al., 2023];
• Tempate-free methods: Aug.Transformer [Tetko et al., 2020], GTA[Seo et al., 2021], Dual-TF [Sun
• et al., 2021b], Retroformer [Wan et al., 2022], Graph2SMILES [Tu and Coley, 2022], NAG2G [Yao et al., 2024];
• Semi-template-based methods: MEGAN [Sacha et al., 2021], RetroPrime [Wang et al., 2021], GraphRetro [Somnath et al., 2021], G2Retro [Chen et al., 2023], Graph2Edits [Zhong et al., 2023], MARS [Liu et al., 2024].

4.2 Main Results

The comparison of model performance, in terms of the topk exact match accuracy, on the USPTO-50K dataset is presented in Tables 1 and 2. When the reaction class is unknown, the KGCL model achieved a top-1 accuracy of 56.3%. In more detail, KGCL achieves state-of-the-art performance among all semi-template-based methods for all values of k and surpasses the strong baseline Graph2Edits by margins of 1.2% and 1.9% in top-1 and top-10 accuracy, respectively. In addition, the performance of KGCL outperforms those of all template-free methods and is better than those of most template-based methods. Even compared with RetroKNN, a strong baseline for template-based methods, KGCL also shows strong competitiveness. When the reaction type is known, KGCL also outperforms all other semi-templatebased methods for all k values and improves the top-1 accuracy by 1.6% compared to the strong baseline Graph2Edits.

To evaluate the performance of our model in generating valid retrosynthetic suggestions and predicting the correct reaction class, we compared KGCL with baseline methods in terms of round-trip and MaxFrag accuracy. The comparative results on the USPTO-50K dataset are presented in Table 3. The top-1 round-trip accuracy of KGCL reached 89.1%, which is nearly comparable to the performance of LocalRetro and surpasses the strong baseline Graph2Edits by In order to verify the generalization ability and robustness of our method, we evaluated the model performance on the USPTO-FULL dataset, which has a larger number of reactions and more diverse reaction classes. The results are shown in Table 4. For fair comparison, we did not clean the USPTOFULL dataset, including the test set. Despite the large amount of noises in the original dataset, our method still achieve considerable performance. Moreover, for k=1, 3, 5, 10, our model achieves state-of-the-art performance among all semitemplate-based retrosynthesis methods. Additionally, KGCL also shows strong competitiveness compared with current leading template-free methods, which further highlights its potential for efficient search in more complex and diverse reaction spaces.

Table 3: Top-k round-trip and MaxFrag accuracy on USPTO-50K dataset without reaction classes.

Table 4: Top-k exact match accuracy on USPTO-FULL dataset without reaction classes. * indicates that invalid reactions in the test set were removed.

4.3 Ablation Study

To investigate the contributions of different components, we conducted ablation experiments on the USPTO-50K dataset. The results are presented by comparing the top-1 exact match accuracy of different models on the USPTO-50K dataset without reaction classes, as shown in Table 5. Model I removes the functional group knowledge and contrastive learning strategy, and its accuracy is 55.1%. Model II adds the graph contrastive learning strategy into Model I, and the top1 accuracy increases by 0.5% to 55.6%, which proves the effectiveness of the graph contrastive learning strategy. Similarly, Model III incorporates the functional group knowledge into Model I, and the top-1 accuracy increases by 0.4%, reaching 55.5%. The results show the importance of the functional group knowledge for the retrosynthesis prediction task. Model IV is a complete KGCL model. Its top-1 accuracy is improved by 1.2% compared to Model I, which exceeds the cumulative gains achieved by using the functional group knowledge and the graph contrastive learning strategy independently. This demonstrates the superior performance achieved when these two strategies work together. Therefore, every component is indispensable.

4.4 Case Study

To better understand the impact of the functional group knowledge on retrosynthetic predictions, we randomly selected a product molecule from the test set of the USPTO-50K dataset. As shown in Figure 3, in the first step, Graph2Edits and our model predicted correctly. In the second step,

• (a) The retrosynthesis prediction process of Graph2Edits

(b) The retrosynthesis prediction process of our method

Figure 3: Investigation of importance of the functional group knowledge. Darker colors indicate higher weights of the functional groups.

our model still predicted correctly, but Graph2Edits made a wrong prediction. The following is a detailed analysis: All halogen elements possess oxidability due to the same outer electrons number, but with the increase of the electron layers, from F to I, electrophilicity is decreasing and nucleophilicity continue to increase. It is well known that iodine is not only a good nucleophile but also a good leaving group. Graph2Edits incorrectly predicted to attach the fluorine atom to the benzene ring in the second step. And our model identified the presence of phenyl and iodine functional groups in the intermediate. Thus it predicted to attach the iodine atom to the benzene ring. Our prediction is more reasonable: On one side, the iodobenzene radical intermediates has high reactivity and presents electron-deficient property for the iodine substituents, thus it is more suitable to react with the nucleophilic iodine atoms. On the other side, iodine atom has the maximum atomic radius among halogens, which makes the C-I bond easily broken. Therefore, the reactants predicted by our model can produce the target product more easily.

4.5 Visualization of Molecular Graph Embeddings

To evaluate the quality of the learned molecular embeddings, we visualize the embeddings of complex product molecules and their intermediates from different reaction classes, as shown in Figure 4. Specifically, we randomly selected a complex product molecule from each reaction class in the test set of the USPTO-50K dataset. The selected molecules have at least one of the following characteristics: multiple reaction centers, stereochemical features, or a long synthetic pathway [Herges, 1994; Zhong et al., 2023]. We then used Graph2Edits and KGCL with beam width = 10 and beam size = 50 to generate graph edit sequences for these molecules and obtain their intermediates. Subsequently, the high-dimensional feature vectors of these molecules and their intermediates were reduced to a two-dimensional embedding space using the t-distributed stochastic neighbor embedding (t-SNE) [Van der Maaten and Hinton, 2008] method.

It can be seen from Figure 4a that the embeddings of product molecules and their intermediates generated by Graph2Edits present a dispersed distribution, and the boundaries of the intermediate embeddings of different target molecules are not sufficiently distinct. In contrast, KGCL (Fig. 4b) is able to effectively distinguish intermediate embeddings of different target molecules, and all intermediate embeddings of the same target molecule are tightly clustered together. The visualization results show that our model is able to capture differences between different molecular structures and maintain the similarity between different intermediates of the same molecule, which provides strong support to predict the graph edit sequences.

5 Conclusion

In the present paper, we improved the accuracy of retrosynthesis prediction by introducing the functional knowledge and the contrastive learning strategy. The experimental results on two benchmark datasets, USPTO-50K and USPTO-FULL, verified the effectiveness of our model KGCL. The success of KGCL depends on merging the domain knowledge from the chemical knowledge graph into the deep learning model in a rational way. In the future, we will further improve the prediction accuracy by introducing more domain knowledge, such as the electronic effect and steric effect, into the retrosynthesis prediction model.

Acknowledgements

This study was supported in part by the National Natural Science Foundation of China (Grant No. 62472082, 62277009).