T-Rex: Text-assisted Retrosynthesis Prediction

We first denoted a retrosynthesis pair as $R=(G_{p},\{G_{r_{c}}\}_{c=1}^{C_{r}})$. In reaction center identification, we used a binary matrix $\mathbf{B}\in\{0,1\}^{n\times n}$ as the labels for reaction centers, where $\mathbf{B}_{ij}=1$ recognized the chemical bond between $i^{th}$ node and $j^{th}$ node as the reaction center. Therefore, the center identification task was formulated as a binary prediction task. The input of this task is the structure information of the product molecules, e.g., the SMILES sequence, and the output is the probability of each bond being identified as the reaction center.

To obtain a high-quality molecule structure representation, we integrate the textual information and graphic information together. We first get the text description of the products by ChatGPT, which generates descriptions based on the IUPAC name of the product. The IUPAC name represents a unique representation of a chemical structure, based on a set of mapping rules between structures and language phrases, and is characterized in a more natural language-like format as compared to SMILES. Therefore, the IUPAC name can serve as a bridge between the chemical molecules and large-scale language models. Here we use the PubChemPy Pub package to generate IUPAC name for each molecule in our dataset. We then generate the text description by the following prompt:

Please delineate the structural features, functional aspects, and applicable implementations of the molecule {{ NAME }}, commencing with the introduction:"The molecule is {{ NAME }}". Reasoning the most plausible type for synthesizing this molecule in the final step, and offer a rationale for your choice.

where the {{ NAME }} item is replaced by the IUPAC name. For generated text description, we trained a BERT model to get the textual embedding $\mathbf{H}_{t}$. Then we also utilize a L-layer R-GCN introduced in Sec. 2 for computing graphic embedding vector $\mathbf{H}_{g}^{L}$:

$\displaystyle\mathbf{h}^{L}=\text{R-GCN}_{1}(G_{p}),\mathbf{H}_{g}^{L}=\text{Sum}(\mathbf{h}^{L}).$

(2)

After that, we obtain the embedding for each edge by cross concatenation:

$\displaystyle\mathbf{e}_{ij}=\mathbf{h}_{i}^{L}||\mathbf{h}_{j}^{L}||\mathbf{V}_{ij}||\mathbf{H}_{g}^{L}||\mathbf{H}_{t}$

(3)

where $||$ stands for vector concatenation, $\mathbf{h}_{i}^{L}\in R^{s}$ denotes the $i^{th}$ row of $h^{L}$ representing the features of $i^{th}$ node and $\mathbf{V}_{ij}=\mathbf{V}_{[i,j,:]}\in\{0,1\}^{b}$ is the edge type between $i^{th}$ node and $j^{th}$ node. When accounting for the scenario with reaction type given, we can incorporate this information by just concatenating its embedding with $\mathbf{e}_{ij}$. The final probability of the chemical bond being the reaction center is computed by:

$\displaystyle\mathbf{r}_{ij}=\sigma(\text{FF}_{1}(\mathbf{e}_{ij})).$

(4)

Here $\text{FF}_{1}$ denotes a feedforward network and $\sigma$ is the Sigmoid activation function. Subsequently, we decided the reaction center bond to be the one that connects the $i^{th}$ node and $j^{th}$ node with the highest $\mathbf{e}_{ij}$ value. The model was trained using the cross entropy loss

	$\displaystyle\mathcal{L}_{1}=-\sum_{x\in X}$	$\displaystyle\sum_{i\neq j}\lambda\mathbf{B}_{ij}\log(\mathbf{r}_{ij})$
		$\displaystyle+(1-\mathbf{B}_{ij})\log(1-\mathbf{r}_{ij}),$		(5)

where X is the training dataset and $\lambda\in[1,+\infty)$ is a weighting hyper-parameter employed to balance the extremely small logits of $\mathbf{r}_{ij}$.

After center identification, we then cut off the chemical bond of reaction centers to get the synthons where each synthon is a connected subgraph. Then we generate the reactants for each synthon using a synthon completion model. We simply utilize the same structure in Shi et al. (2020) to predict the candidate reactant pairs by beam search. We collected a list of most possible candidate pairs to train the model in re-ranking stage.

In the previous stage, we collected a lot of retrosynthesis pairs. However, it is notable that the generated reactants may mismatch the ground truth due to the limited capabilities GNN-based molecule representations. To determine the most possible reactants, we proposed the re-ranking stage by utilizing the capability of ChatGPT. Given a set of retrosynthesis pairs $S=\{R_{i}\}_{i=1}^{k}$ for the product $x$, where $k$ is the beam size in the synthon completion and $R_{i}=(G_{p},\{G_{r_{c}}\}_{c=1}^{C_{r_{i}}})$ is an individual retrosynthesis pair sharing the same product with other pairs in the set. The re-ranking model is a classifier which identified the candidate reactants from several candidates which have been predicted as the retrosynthesis path with a high probability. The input of the re-ranking clasifier is the same as the first stage model, i.e., the concatenation of text embeddings and molecular graph embeddings. The textual description $T^{i}$ of the reactant candidate $R_{i}$ was generated by ChatGPT using a different prompt:

Please delineate the structural features, functional aspects, and applicable implementations of the molecules {{ NAME }} and possible reactants {{ REACTANT1 }} and {{ REACTANT2 }} to synthesize it. You should generate the descriptions respectively as above example. These descriptions are linked by " [SEP] ", and each commences with the introduction:"The molecule is …".

where the {{ NAME }}, {{ REACTANT1 }} and {{ REACTANT2 }} are the IUPAC names of the product and the two reactants. In the case when there only exists one reactant, we repeat the IUPAC name twice to make it compatible with our paradigm. Then we used a BERT model as the encoder of the textual descriptions:

$\displaystyle\mathbf{H}_{t}^{i}=\text{BERT}_{2}(T^{i}).$

(6)

To embed the molecular graph information, we also utilized the R-GCN structure as the encoder:

$\displaystyle\mathbf{h}^{L,i}=\text{R-GCN}_{2}(\{G_{r_{c}}\}_{c=1}^{C_{r_{i}}}),\mathbf{H}_{g}^{i}=\text{Sum}(\mathbf{h}^{L,i}).$

(7)

Here we exclude the graph of product since its information has been already extracted in the previous stage. Afterwards, the embedding for the retrosynthesis pair is obtained by concatenating the textual and graphic embedding vectors:

$\displaystyle\mathbf{H}^{i}=\mathbf{H}_{g}^{i}||\mathbf{H}_{t}^{i}.$

(8)

Finally, we get a probability score for $R_{i}$ being recognized by the classifier:

$\displaystyle a^{i}=\text{Softmax}(\text{FF}_{2}(\mathbf{H}^{i})),$

(9)

where $\text{FF}_{2}$ is another feedforward network and $a^{i}=[a_{-}^{i},a_{+}^{i}]$ is the computed logits from the classifier. The confidence score, i.e., the probability for each candidate is then denoted as $a_{+}^{i}$.

In our implementation, we selected the top 3 reactant candidate from the first stage model as the inputs to the re-ranking model. And we label the candidates matching the ground truth as positive examples and others as negative examples. Notably, we observed a significant data distribution imbalance, i.e., two-thirds of instances have negative labels among the generated candidates. To augment the number of positive samples in the training set of the re-ranking model, we additionally add the ground truth reactants with positive labels. Then we train the re-ranking model as a binary classification task. The model is trained using cross entropy loss for the classification task. Additionally, we incorporate a contrastive learning loss term to optimize the molecule representations. Finally, the total loss is computed as follows:

	$\displaystyle\mathcal{L}_{2}=\sum_{S\in\mathbb{S}}\sum_{i}\text{CELoss}(a_{S}^{i},label_{S}^{i})$
	$\displaystyle-\alpha\sum_{j\neq i}\log\frac{\exp(\text{sim}(\mathbf{H}_{t}^{i},\mathbf{H}_{t}^{j})}{\sum_{m}\exp(\text{sim}(\mathbf{H}_{t}^{i},\mathbf{H}_{t}^{m}))},$		(10)

where $\text{sim}(\mathbf{u},\mathbf{v})=\frac{\mathbf{u}^{T}\mathbf{v}}{||\mathbf{u}||\cdot||\mathbf{v}||}$ is the cosine similarity function and $label_{S}^{i}$ is the label for the $i^{th}$ retrosynthesis pair in set $S$. Furthermore, the two losses are appropriately balanced via the utilization of a parameter denoted as $\alpha$.

During the inference phase, we compute the probability for each reactant pair among the top $k$ generated candidates. In the evaluation, we can calculated the top $k^{\prime}$ exact match accuracy by considering the reactant pairs with the largest $k^{\prime}$ scores. In our implementation, we evaluate top $k^{\prime}=1,3,5$ accuracy based on top $k=3,5,10$ candidates, respectively.

We first evaluated the retrosyntheis prediction performance using T-Rex. We calculated the top-k accuracy of the proposed approach on USPTO-50k and filtered USPTO-MIT datasets in both reaction class known and reaction class unknown settings. The results of T-Rex, G2Gs and G2Gs+text models when reaction class was not given is summarized in Fig. 3 and Appendix. A. We first noticed that G2Gs+Text outperformed the state-of-the-art template-free model G2Gs on Top 1 and Top 3 accuracy and achieved comparable performance on Top 5 and Top 10 accuracy on both datasets. This result demonstrates the effectiveness of incorporating text description generated by ChatGPT into the graph-to-graph retrosynthesis model. Furthermore, our T-Rex model exhibits superior performance on top-1 and top-3 accuracy on two datasets compared to the G2Gs+Text model, highlighting the benefits of the re-ranking stage in the T-Rex model.

We next examine the generalizability of our T-Rex model by evaluating the cross-dataset retrosythesis prediction. Specifically, we trained a model on USPTO-50k to predict the retrosythesis on USPTO-MIT Table 1. Since USPTO-MIT and USPTO-50k have disjoint reaction types, these two datasets present substantially different distributions. We found that T-Rex achieved the best results on Top 1 and Top 5 in this cross-dataset setting. On the reverse side, our method also obtained the best performance on Top 1 and Top 5 accuracy, reflecting the good generalizability of our T-REX model.

Next, we conducted a series of incremental experiments by training on the entire set of USPTO-50k and part of USPTO-MIT, and testing on the remaining USPTO-MIT. We summarized the performance in Fig. 4. T-Rex consistently outperformed other approaches with regard to different proportions of USPTO-MIT training set. This further confirms the strong generalizability of T-Rex and the benefit of using text for retrosynthesis prediction.

Finally, since USPTO-MIT does not have reaction type information, we exploited T-Rex to predict the reaction type. Specifically, we used USPTO-50k to train a reaction type classifier and then used this classifier to predict reaction type for each target product in a multi-class classification setting. We then incorporated this reaction type into our framework. We found that this can further improve the retrosynthesis prediction performance (Appendix 5.4), again demonstrating the ability of T-Rex to jointly learn from heterogeneous datasets.

After observing substantial improvement of T-Rex over the G2Gs+Text baseline, we then conducted ablation studies for different components in the re-ranking stage. In particular, we assess the variants without text embedding (denoted as T-Rex w/o text), without graph embedding (denoted as T-Rex w/o graph), without contrastive learning loss on USPTO-50k (denoted as T-Rex w/o c.l.l.) or with the texts generated by MolT5 (denoted as T-Rex(MolT5)). The results of top 1, top 3 and top 5 exact match accuracy are shown in Table 2.

We found that all these four components are important because the original T-Rex outperformed all these four variants. Moreover, the inclusion of text information contributes to performance improvements most substantially, as demonstrated by large drops in Top 3 and Top 5 accuracy by the variant that does not use the text information, necessitating the incorporation of text information into our framework. Moreover, the variant that uses MolT5 instead of ChatGPT to generate textual description performs substantially worse than T-Rex, indicating the importance of using high-quality text in our framework, further supporting the benefits of using text information for retrosynthesis prediction.

We also tried to improve the performance of our T-Rex model in the transferring test with the help of reaction type. Specifically, we started by training a reaction-type classifier on the USPTO-50k dataset. We then employed this classifier to predict the reaction types within the filtered USPTO-MIT dataset, even in cases where certain reactions cannot be definitively classified under any specific reaction type.

We have explored various potential classifier architectures, encompassing different combinations of graph features solely pertaining to the product (denoted as g.p.), graphic and textual information associated with the product (denoted as g.p.+t.p.), as well as graphic features of the reactants and product in conjunction with textual information specific to the product (denoted as g.p.+t.p.+g.r.). To make the experiments comparable, we only gave the reaction type information in the re-ranking stage. The top 1, 3, and 5 exact match accuracies are listed in Table 3.

We can see that with the help of predicted reaction types, the top 1 exact match accuracy can improve much, demonstrating its potential to incorporate information of new resources.

In our experiment setting, we used two types of BERT: BERT-small and PubMedBERT. And we also did the experiments with only one types of BERT for USPTO-50k, and the results are presented in Table 4. We can notice that the selection of different types of embedding models may also contribute to the performance of the retrosynthesis prediction task.

In Figure 3, we noticed T-Rex posed a great improvement than G2Gs on top-1 accuracy, but a smaller increase in top-3, 5 and 10. Therefore, we made deep investigation on the re-ranking stage in two aspects. Firstly, the increase in the top-1 accuracy results from the re-ranking stage. We have recorded the accuracies before and after re-ranking for each set of retrosynthesis pairs. For ground truth retrosynthesis pairs ranking first in the candidate ranking stage, 87.3% of them remain top-1 by the ranker in the re-ranking stage. And for those ranking 2 or 3 in the first stage which are ground truth pairs, 52.3% of them ranks top-1 in the second stage. Moreover, in the re-ranking stage, we have tried to involve only graph features, and the top-1 and top-3 accuracies are only 45.3% and 63.3%, respectively (Table 2), showing that the involvement of text information is still important for improving accuracies. All in all, we conclude that the dramatic improvement in top-1 accuracy may results from the architecture of ranker and the information from ChatGPT involved in the re-ranking stage.

To assess the generalizability of the T-Rex model, we computed accuracies across all 10 reaction types in one side of cross-dataset experiments (USPTO-MIT$\rightarrow$USPTO-50k). Notably, the top-1 accuracy across all reaction types reveals that the T-Rex model consistently outperforms G2Gs, demonstrating superior performance. And T-Rex outperforms G2Gs greatly in terms of accuracies for some rare reaction types including Oxidations, Functional group inter-conversion (FGI) as well as Reductions. Moreover, the top-1 accuracies of G2Gs and T-Rex for these specific reaction types within the test dataset along with the corresponding proportions of these types are presented in Table 5. T-Rex exhibits substantially higher accuracy than G2Gs for these specific types, showing its potential for generalization to rare reaction types.

Moreover, for top-5 and top-10 accuracies, we noticed that G2Gs can outperform T-Rex slightly for some reaction types. For example, for acylation and related processes, taking up 24.3% of the test dataset, the top-5 and top-10 accuracies for G2Gs are 64.8%, 70.2% while the top-5 and top-10 accuracies for T-Rex are 63.0%, 66.2%. However, for some rare reaction types, T-Rex instead outperforms G2Gs greatly. For example, for Functional group addition (FGA) reactions(taking up 0.52% in the test dataset), the top-5 and top-10 accuracies for G2Gs are both 0% while the top-5 and top-10 accuracies for T-Rex are both 12%; for FGI reactions (taking up 3.52% in the test dataset), the top-5 and top-10 accuracies for G2Gs are 26.7%, 34.3% while the top-5 and top-10 accuracies for T-Rex are 30.3%, 40.2%. From these types, we can still see T-Rex’s potential of generalization to rare reaction types.

We presented the t-SNE visualizations of the embeddings of G2Gs, G2Gs with text description(G2Gs+Text), and our T-Rex model in Fig. 5. Here, G2Gs and G2Gs+Text show the embeddings from the reaction center identification stage, while T-Rex shows the embeddings from the re-ranking stage, which are unavailable for the other two methods. We can see that the embeddings of both G2Gs and G2Gs+Text model do not show visible patterns while our T-Rex can show clear pattern for the retrosynthesis pairs, again demonstrating the effectiveness of the re-ranking stage. Furthermore, we visualized an example that the re-ranking stage successfully correct the mis-predicted reaction center by using the text information (Fig. 6). More examples are presented in Section E in the Appendix.