Ualign: pushing the limit of template-free retrosynthesis prediction with unsupervised SMILES alignment

Abstract

Motivation

Retrosynthesis planning poses a formidable challenge in the organic chemical industry, particularly in pharmaceuticals. Single-step retrosynthesis prediction, a crucial step in the planning process, has witnessed a surge in interest in recent years due to advancements in AI for science. Various deep learning-based methods have been proposed for this task in recent years, incorporating diverse levels of additional chemical knowledge dependency.

Results

This paper introduces UAlign, a template-free graph-to-sequence pipeline for retrosynthesis prediction. By combining graph neural networks and Transformers, our method can more effectively leverage the inherent graph structure of molecules. Based on the fact that the majority of molecule structures remain unchanged during a chemical reaction, we propose a simple yet effective SMILES alignment technique to facilitate the reuse of unchanged structures for reactant generation. Extensive experiments show that our method substantially outperforms state-of-the-art template-free and semi-template-based approaches. Importantly, our template-free method achieves effectiveness comparable to, or even surpasses, established powerful template-based methods.

Scientific contribution

We present a novel graph-to-sequence template-free retrosynthesis prediction pipeline that overcomes the limitations of Transformer-based methods in molecular representation learning and insufficient utilization of chemical information. We propose an unsupervised learning mechanism for establishing product-atom correspondence with reactant SMILES tokens, achieving even better results than supervised SMILES alignment methods. Extensive experiments demonstrate that UAlign significantly outperforms state-of-the-art template-free methods and rivals or surpasses template-based approaches, with up to 5% (top-5) and 5.4% (top-10) increased accuracy over the strongest baseline.

Introduction

Retrosynthesis prediction is a crucial task in organic chemistry, aiding in finding efficient synthetic pathways from target molecules to accessible starting materials. Despite significant advancements in chemical synthesis technology, it still remains a challenge in industries like pharmaceuticals. The extensive search space and the incomplete understanding of chemical reaction mechanisms make retrosynthesis prediction difficult, even for experienced chemists. To address this issue, computer-assisted synthetic planning (CASP) has gained increasing attention in recent years, starting from the seminal work by Corey. This paper focuses on single-step retrosynthesis prediction, which is the fundamental step in CASP. It aims to predict the reactants that can lead to a given product molecule through a single reaction step.

Various deep-learning-based single-step retrosynthesis prediction methods have been proposed in recent years. These methods can be broadly classified into three groups based on their dependency on additional chemical knowledge: template-based, semi-template-based and template-free methods. Template-based methods [5, 7, 9, 44] require an extra database of reaction templates. They frame the retrosynthesis prediction as a classification or retrieval problem for reaction templates suitable for the given product molecule to be synthesized. Among these solutions, Retrosim [7] utilizes molecular similarity to rank reaction templates; LocalRetro [5] and GLN [9] use graph neural networks to model the relationship between reaction templates and molecules to predict the most suitable reaction template; RetroKNN [44] further improves upon LocalRetro by addressing the issue of data imbalance using K-nearest neighbors (KNN). Template-based methods have strong interpretability and can accurately predict reactants. However, these methods are often unable to cover all cases and suffer from poor scalability due to limitations imposed by the template database.

To overcome the limitations faced by template-based methods, researchers have turned to generative models. Semi-template-based methods incorporate chemical knowledge into generative models with the help of chemical toolkits like RDKit [19], breaking free from the limitations imposed by reaction templates. The key idea of most semi-template-based methods [6, 32, 33, 41, 45] is to first convert the product into synthons based on reaction center identification and then complete the synthons into reactants. Graph neural networks are commonly used for synthon prediction, followed by leaving group attachment [6, 33], conditional graph generation [32], or SMILES (Simplified Molecular Input Line Entry System) generation [45] for reactant completion. Apart from all above, RetroPrime [41] utilizes two independent Transformers to accomplish synthon prediction and reactant generation as separate tasks.

Semi-template-based methods to a certain extent are more in line with chemical intuition. However, these methods increase the complexity of inference and training as they break down retrosynthesis into two subtasks. Failures in synthon prediction directly affect subsequent reactant completion and overall performance. Besides, methods based on leaving group necessitates an extra leaving group database. This requirement, akin to template-based approaches, imposes limitations on the model’s scalability.

As generative models, Template-free methods opt to generate reactants directly from the given products. In comparison to generating graph structures, SMILES provides a way to represent molecules as strings. Taking advantage of this, most template-free methods [18, 31, 34, 40, 49] use Transformer models to translate between product SMILES and reactants SMILES. In particular, Graph2SMILES [35] replaces the Transformer encoder with a graph neural network, resulting in a permutation-invariant pipeline. There are also methods [27, 47] formulates the generation of reactants as a series of graph generation or editing operation and solve it auto-regressively. Existing template-free methods generally follows an auto-regressive generation strategy and use beam search for the generation process. Consequently, preserving a level of diversity in the resultant outputs has emerged as a critical consideration for template-free methods [39]. Due to the use of SMILES as input and output, most of template-free methods often overlook the rich topological and chemical bond information present in molecular graphs. Moreover, as reactants molecules need to be generated from scratch, template-free methods frequently suffer from validity issues and fail to leverage an important property of retrosynthesis prediction, i.e., the presence of many common substructures between products and reactants.

In this paper, we focus on the template-free generative approach for retrosynthesis prediction. Existing sequence-to-sequence methods have limitations in extracting robust molecular representations. They overlook the abundance of topological information and chemical bonds, and lack the ability to utilize atom descriptors as rich as those in graph-based methods. Furthermore, template-free methods overlook the fact that the molecular graph topology remains largely unaltered from reactants to products during chemical reactions, as they generate reactants from scratch. While there are methods that attempt to solve this problem using supervised SMILES alignment, they require complex data annotation and impact model training. Given these limitations, the following question naturally arises:

Can we effectively leverage the structural information of product molecules using a much simpler approach?.

To address these issues and further enhance template-free methods, we propose a novel graph-to-sequence pipeline called UAlign. Our approach employs a specifically designed graph neural network as an encoder, incorporating information from chemical bonds during message passing to create more powerful embeddings for the decoder. We introduce an unsupervised SMILES alignment mechanism that establishes associations between product atoms and reactant SMILES tokens which reduces the complexity of SMILES generation and enables the model to focus on learning chemical knowledge. Our model outperforms existing template-free methods by a large margin and demonstrates comparable performance against template-based methods.

Methods

We introduce UAlign, a novel single-step retrosynthesis prediction model based on an encoder-decoder architecture, as demonstrated in Fig. 1. It’s a fully template-free method without any molecule editing operation using RDKit [19]. We propose a specially designed variant of Graph Attention Network, which incorporates the information of chemical bonds to enhance the capability of capturing the structural characteristics of molecules.

Fig. 1

Overview of UAlign: Given a product molecule graph P and one of its DFS order $O_P$, the graph is first fed into the graph neural network called $\hbox {EGAT}^+$ to obtain node features H. Then the positional encoding is added to H according to the given DFS order $O_P$ to generate the order-aware node features $\hat{H}$. Finally the decoder takes $\hat{H}$ as input and generate the SMILES of reactants auto-regressively

Preliminary

A molecule can be represented as a graph, denoted by $G=(V,E)$, where V represents the atoms and E represents the chemical bonds. The SMILES representation of a molecule can be obtained by performing a depth-first search (DFS) starting from any arbitrary atom in the molecule graph. Given a molecule graph $G=(V, E)$, we can generate multiple DFS orders and each DFS order corresponds to a SMILES representation of the graph. Denoted the set of all possible DFS orders as $\mathcal {D}(V) \subseteq \mathcal {P}(V)$, $\mathcal {P}(V)$ represents all permutations of the set of atoms V. For each DFS order $O\in \mathcal {D}(V)$, we denoted its corresponding SMILES as Smiles(G, O), which lists all atoms in the order dictated by O. To facilitate our subsequent elaboration, we refer to the position of an atom a in the order O as its rank, denoted as rank(a, O). The atom with the minimal rank given order O is then defined as the root atom, denoted as root(G, O).

Model architecture

In this section, we provide an overview of the model’s architectural design and the rationale behind it. Our model adopts an encoder-decoder framework, as depicted in Fig. 1. The encoder is tasked with extracting molecular representations from the input products and supplying these as inputs to the decoder, which then generates a combination of reactants. We utilize graph neural networks for our encoder, which encodes the nodes through an iterative message-passing mechanism to derive node features. During each round of message-passing, the network collects and aggregates information from a node’s neighbors, thereby updating the node features. This design effectively integrates the topological information of the graph structure into the node features, naturally adapting to the task of molecular representation learning. The superiority of graph neural networks in this domain has been substantiated by a plethora of studies [15,16,17].

Moving on to the decoder, it is constructed based on the transformer decoder architecture [37]. Generating a graph poses a unique challenge due to the lack of inherent order among the graph’s nodes and the necessity to predict an adjacency matrix that is quadratic in relation to the number of nodes. However, the adoption of SMILES [42] circumvents this complexity by converting the molecular generation problem into a more tractable text sequence generation task. This transformation is crucial because it ensures that the length of the output sequence is linearly proportional to the number of atoms involved. And the framework of text sequence generation has been extensively applied in other domains, such as natural language processing, providing a robust foundation upon which our molecular generation model is built [12, 21, 22, 25]. The transformer decoder, equipped with a cross-attention mechanism, is adept at sequential generation conditioned on a given input, making it an ideal choice for our model’s decoder. The subsequent sections will delve into the intricate design of the encoder, decoder, and other components of our model.

$\hbox {EGAT}^+$

Chemical bonds play a significant role in determining the properties of molecules and contain valuable information. Previous studies [13, 26, 46] have demonstrated that incorporating edge information into graph neural networks can greatly enhance their ability to represent molecular structures. To fully leverage the information brought by chemical bonds, we propose a modified version of the Graph Attention Network (GAT) [38] called $\hbox {EGAT}^+$.

Our proposed model explicitly incorporates edge features, which represent the information derived from chemical bonds, into the message passing process. During each iteration of message passing, the $\hbox {EGAT}^+$ applies self-attention to each node and its one-hop neighbors to calculate attention coefficient according to both node features and edge features. It then aggregates the both node and edge features of these neighbors, considering the attention coefficients, to update the node features. Denote the node feature of atom u as $h^{(k)}_u$ and the edge feature between atom u and v as $e_{u,v}^{(k)}$ at k-th iteration of message passing. In math, the message passing mechanism can be written as

$$\begin{aligned} \begin{aligned} \tilde{e}_{u,v}^{(k)}&= \textrm{FFN}^{(k)}_e (e_{u,v} ^{(k)}),\\ \tilde{h}_{u}^{(k)}&= \textrm{FFN}^{(k)}_n (h_u^{(k)}),\\ c_{u,v}&= \textbf{a}^T [\tilde{h}_u^{(k)} \Vert \tilde{h}_v^{(k)}\Vert \tilde{e}_{u,v}^{(k)}],\\ \alpha _{u,v}&= \frac{\exp (\textrm{LeakyReLU}(c_{u,v}))}{\sum _{v'\in \mathcal {N}(u)\cup \{u\}} \exp (\textrm{LeakyReLU}(c_{u,v'}))},\\ h^{(k+1)}_u&= \sum _{v\in \mathcal {N}(u) \cup \{u\}} \alpha _{u,v} \left( \tilde{h}_{u}^{(k)} + \tilde{e}_{u,v} ^ {(k)}\right) ,\\ e^{(k+1)}_{u,v}&= \textrm{FFN}_m^{(k)}([h_{u}^{(k+1)}\Vert h_{v}^{(k+1)} \Vert e^{(k)}_{u,v}]), u\ne v, \end{aligned} \end{aligned}$$

(1)

where $\textrm{FFN}_m^{(k)}$, $\textrm{FFN}^{(k)}_e$ and $\textrm{FFN}^{(k)}_n$ are three different feed forward networks, $\textbf{a}$ is a learnable parameter, $\mathcal {N}(u)$ denotes the one-hop neighbors of node u and $\Vert$ denotes the concatenation operation. Since there are no chemical bonds with the same beginning and ending atoms, the $e^{(k)}_{u, u}$ is also set as a learnable parameter shared among all atoms. The residual connection and layer normalization [2] are applied to prevent over-smoothing while enlarging the receptive field of the model [43].

The initial node features $h_u^{(0)}$ and edge features $e_{u,v}^{(0)}$ are determined via several chemical property descriptors, whose details are shown in Supplementary Sec. 6. After K iterations of message passing, we can obtain the encoded features $h^{(K)}_u$ of all atoms and make up the output $H\in \mathbb {R}^{V_P \times d}$ of the encoder, where d denotes the embedding size.

SMILES alignment

For single-step retrosynthesis prediction, a significant proportion of structures are shared between product molecules and reactant molecules [40, 50]. However, SMILES-based methods often have to generate the reactant SMILES from scratch, even if most of the structures of reactants are the same as those of the products. This results in the underutilization of input information and becomes the bottleneck of template-free retrosynthesis prediction methods. There are methods [31, 40] addressing this issue through supervised SMILES alignment, which involves adding supervised information to establish the correspondence between input and output tokens through cross-attention over the input and predicted tokens. This supervised training approach not only requires complex data annotation algorithms but also limits the diversity of the model’s attention map, thereby further affecting the model’s performance. To address the above-mentioned issues, we propose the unsupervised SMILES alignment method as follows.

Assuming we can identify the location of each product atom in the reactants’ SMILES and provide it to the model, a natural correspondence can be established between the input and output atoms. However, during the inference process, revealing this information would lead to label leakage, which is not permitted. Therefore, we propose the following modification: when providing an order of product atoms, we expect the model to generate atom tokens in the reactants’ SMILES in this given order as closely as possible. By doing so, we can establish a correspondence between the product atoms and the reactants’ SMILES tokens using unsupervised methods without leaking any labels. We refer to this type of reactants’ SMILES, which aims to preserve the given order of atom tokens as much as possible, as order-preserving reactant SMILES. Note that SMILES represents atoms in a molecule according to a certain DFS order, the provided order should also be a DFS order for the product molecule.

The generation of order-preserving reactant SMILES will be introduced as follows. Given the product molecule $P=(V_P,E_P)$ with a DFS order $O_P\in \mathcal {D}(V_P)$ and the corresponding set of reactant molecules $\mathcal {R}=\{R_1, R_2,\ldots R_l\}$, for each reactant $R=(V_R, E_R)\in \mathcal {R}$, we can find a depth-first order $O_R \in \mathcal {D}(V_R)$ that has a nearly consistent atomic appearance sequence with $O_P$ as the product and reactants have similar structures. For convenience, we name such a order as $O_P$-corresponding order of R and denote it as $CO(R, O_P)$. Mathematically, it’s defined as

$$\begin{aligned} CO(R, O_P) = \arg \min _{o\in \mathcal {D}(R)} \sum _{i\in O_P\cap o} \sum _{j\in O_P\cap o} inv(i, j, O_P, o), \end{aligned}$$

(2)

where the value of $inv(i, j, O_P, o)$ equals 1 if and only if $rank(i, O_P) < rank(j,O_P)$ and $rank(i,o) > rank(j, o)$, and equals 0 otherwise. We sort the reactants $\mathcal {R}$ according to $rank(root(R, CO(R, O_P)), O_P)$ in ascending order. Then we generate SMILES for each reactant molecule using its $O_P$-corresponding order and join them together using “.” to obtain order-preserving reactant SMILES.

For further discussion, we denote the order-preserving reactant SMILES given the reactant molecules $\mathcal {R}$ and a DFS order O of product as $OPSmiles(\mathcal {R}, O)$. An example of the process to generate order-preserving reactants SMILES is shown in Fig. 2. The detailed implementations are presented in Supplementary Sec. 5.1.

Fig. 2

An example of the process to generate order-preserving reactants SMILES. The atom mapping numbers shown on the figure are included only for clearer explanation and will be removed in our implementation to prevent any label leakage

Decoder

The decoder takes the node features $H\in \mathbb {R}^{V_P \times d}$ that are generated from the encoder, as well as the given DFS order $O_P$ for the product molecule graph as input. We use the vanilla Transformer decoder [37] as our decoder. As mentioned in "SMILES alignment" Section, the order information of product atoms are required for SMILES alignment. However, the Transformer decoder is permutation-invariant to memory [20, 35], meaning it is not sensitive to the order of the features from encoder. This implies that directly performing cross-attention over H may not effectively capture the relationship between product atoms and reactant SMILES tokens. To address this problem, we introduce position encoding to the node features based on the rank of each atom in the given DFS order $O_P$ to generate order-aware node features $\hat{H}$. Then given an input embedding sequence $Z\in \mathbb {R}^{m\times d}$, the Transformer decoder layer utilizes the order-aware node features as keys and values in all the cross-attention layers. This process ultimately generates the decoded embeddings $\hat{Z}\in \mathbb {R}^{m\times d}$. These embeddings are then fed into feed-forward layers $\mathrm {FFN_1}: \mathbb {R}^{d}\rightarrow \mathbb {R}^T$ to predict the tokens $\hat{T}$ that should be generated. In summary, the decoder can be mathematically expressed as

$$\begin{aligned} \begin{aligned} \hat{H}&= H + \textrm{PE}(O_P),\\ \hat{Z}&=\textrm{TransformerDecoder}(Z, \hat{H}),\\ \hat{T}&= \mathrm {FFN_1}(\hat{Z}). \end{aligned} \end{aligned}$$

(3)

Two-stage training

There is a significant distribution shift between graphs and SMILES representations. Moreover, our model is specifically designed to generate non-canonical SMILES, which may contain more complex patterns compared to canonical SMILES. To achieve this goal, we propose a two-stage training strategy in this paper. The first stage aims to align the distributions between two distinct modalities: SMILES and molecular graphs, while enabling the model to learn the patterns of non-canonical SMILES. Given a molecule graph M and one of its possible DFS orders $O_M$, the training task is to translate the graph into the corresponding SMILES representation based on the given order $O_M$. In detail, this is reached by training the model to generate $Smiles(M, O_M)$ given molecule M and DFS order $O_M$.

Once the first stage training converges, we proceed to the second stage, which focuses on retrosynthesis prediction. In this stage, the model is trained using the order-preserving reactant SMILES as targets. Given a product molecule graph P, a possible DFS order $O_P$, and a set of reactants $\mathcal {R}$, the model is expected to generate $OPSmiles(\mathcal {R}, O_P)$.

Data augmentation

Different from those Transformer-based methods [31, 34, 40] taking SMILES as input and canonical SMILES as target, our method takes a graph as input and is trained with non-canonical SMILES. That means the previous SMILES augmentation tricks are not suitable for us. Similar to [40], we choose to augment the training data on-the-fly.

For the first stage, at each iteration, for each molecule $M=(V_M, E_M)$, we have a 50% chance of using a random DFS order $O_M$ as the input for the model, and using the corresponding $Smiles(M, O_M)$ as the training target. For the other 50%, we randomly select another molecule $M'=(V_{M'}, E_{M'})$ from the dataset to form a new molecular graph $\tilde{M}=(V_M \cup V_{M'}, E_M \cup E_{M'})$, and find the DFS order $O_{\tilde{M}}$ that can generate canonical SMILES for $\tilde{M}$. $\tilde{M}$ and $O_{\tilde{M}}$ are then fed into the model and the target is set as the canonical SMILES of $\tilde{M}$. Such an augmentation method enables the model to output the atom tokens according to the given DFS order and be aware of different components within a graph.

For the second stage, at each iteration, for each product molecule $P=(V_P, E_P)$, we have a 50% probability of using a random DFS order as input, and for the remaining part, we use the DFS order capable of producing canonical SMILES for product as input. The target used for training is the order-preserving reactant SMILES generated based on the input DFS order. This data augmentation method allows the model to focus more on the DFS order for canonical product SMILES while also noticing the correspondence between product atoms and the output SMILES tokens.

Loss function

Both of the two stages of training can considered as a kind of translation between graphs and SMILES, thus we use the loss widely used for auto-regressive language generation models for training. Denote the training target as $T=\{t_1, t_2, \ldots , t_n\}$ and the output of the model $\hat{T}=\{\hat{t}_1, \hat{t}_2,\ldots ,\hat{t}_n\}$, the loss can be written as

$$\begin{aligned} \mathcal {L}=\sum _{i=1}^n l_{cls}(\hat{t}_i, t_i), \end{aligned}$$

(4)

where $l_{cls}(\cdot )$ is the classification loss.

Moreover, we have introduced numerous definitions to establish a foundation for understanding the model’s overall training process. To facilitate comprehension, we have encapsulated the entire training procedure in Algorithm 1.

Algorithm 1

The training procedure.

Results and discussion

In this section, we conduct extensive experiments to make a comprehensive evaluation of our proposed UAlign.

Evaluation protocol

Benchmark Datasets. We adopt three datasets for evaluation: (1) USPTO-50K consists of 50,016 atom-mapped reactions grouped into 10 different classes; (2) USPTO-FULL comprises 1,013,118 atom-mapped reactions without any reaction class information. (3) USPTO-MIT consists of 479,035 atom-mapped reactions without any reaction class information. To ensure a fair comparison, we adopt the same training/validation/test splits as those in a previous study [9] for USPTO-50K and USPTO-FULL datasets. The training/validation/test splits is aligned with the previous study [16]. The detailed data processing procedure and the statistical information of the processed dataset are presented in Supplementary Sec. 2.

Metrics. We utilize the following three evaluation metrics for evaluation: top-k accuracy, top-k SMILES validity, top-k round-trip accuracy and Computational Cost. The detailed definitions for the first three metrics are provided in Supplementary Sec. 3.

Performance comparison

Table 1 Top-k accuracy for retrosynthesis prediction on USPTO-50K

Table 2 Top-k accuracy for retrosynthesis prediction on USPTO-MIT

Table 3 Top-k accuracy for retrosynthesis prediction on USPTO-FULL

Table 4 Top-k SMILES validity for retrosynthesis prediction on USPTO-50K with reaction class unknown

Table 5 Top-k round-trip accuracy for retrosynthesis prediction on USPTO-50K with reaction class unknown

Table 6 Average inference time per sample on USPTO-50K dataset

Top-k Accuracy. We compare our model with existing single-step retrosynthesis prediction in terms of top-k accuracy on all the datasets. The results are summarized in Tables 1, 2 and 3. On the USPTO-50K dataset, our model achieves an average top-3 accuracy of 77.3%, average top-5 accuracy of 84.6% and average top-10 accuracy of 90.5% under the reaction class unknown setting, surpassing the SOTA template-free method by 3.2%, 4.0% and 4.9% respectively. And with reaction class given on USPTO-50K dataset, our model achieves an average top-3 accuracy of 86.7%, average top-5 accuracy of 91.5% and average top-10 accuracy of 95.0%, which exceeds the SOTA template-free method by 4.2%, 4.0% and 4.8% respectively. It is worth noting that, even when taking the standard deviation into consideration, the lower-bound performance of our model still surpasses all template-free methods in terms of all metrics. Moreover, our model outperforms all the semi-template-based methods with a noticeable margin. It’s also encouraging to see that our method, as a template-free method, achieves competitive or even superior performance against the powerful template-based methods such as LocalRetro under both settings of USPTO-50K dataset. On UPSTO-MIT dataset, our model achieves the top-1 accuracy of 59.9% and top-10 accuracy of 86.4%, which even outperforms the existing template-based SOTA method LocalRetro significantly. Additionally, our models achieved a top-1 accuracy of 50.4% on the USPTO-FULL dataset, which exceed that of the current SOTA model GTA by 3.8%. These findings sufficiently demonstrate the effectiveness of our method. The contribution of each proposed module will be further validated in "Ablation Study" section.

It is noteworthy that while template-based approaches have achieved remarkable performance on the USPTO-50K dataset, their reliance on external template libraries has emerged as a constraint as datasets grow in scale and complexity. This dependency leads to a substantial degradation in model performance. In contrast, template-free methods have demonstrated superior versatility and adaptability, qualities that render them especially appropriate for managing large-scale and intricate datasets.

Top-k SMILES Validity. We use vanilla Transformer, RetroPrime, Retroformer and Graph2SMILES as robust baselines to compare the validity of SMILES in our study. SMILES generation models for retrosynthesis tasks often encounter challenges with maintaining SMILES validity. We do not take the methods based on templates or molecule editing as baseline here because the validity of generated SMILES can be guaranteed by the templates or chemical toolkits. Unlike graph-based models, SMILES-based methods need to ensure that the generated content adheres to the parsing rules of SMILES, without leveraging chemical tools such as RDKit. Consequently, SMILES-based approaches are more susceptible to generating invalid SMILES compared to graph-based models. As shown in Table 4, our model demonstrates superior top-1 and top-5 molecule validity compared to other models, even without employing canonical SMILES as our training objective. This improvement could be attributed to the proposed two-stage training strategy and data augmentation, which assist the model in capturing various SMILES patterns effectively.

Top-k Round-Trip Accuracy. To assess the accuracy of our predicted synthesis plans, we utilize the Molecule Transformer [29] as the benchmark reaction prediction model and calculate the top-k round-trip accuracy. We take RetroPrime, Retroformer and Graph2SMILES as our strong SMILES-based baselines. We also use take graph-based method GraphRetro into comparison. The results are presented in Table 5. The results clearly indicate that our model outperforms all SMILES-based baselines by a considerable extent and even exceeds the well-established graph-based method, GraphRetro. This underscores the efficacy of our unsupervised SMILES alignment mechanism, which enables the model to efficiently leverage substructures from product molecules to construct reactants. This mechanism allows the model to focus more intently on learning reaction mechanisms, thereby yielding more plausible predictive outcomes. In summary, our model has exhibited a robust capacity for generating coherent and efficacious synthesis pathways, specifically tailored for advanced downstream applications such as multi-step retrosynthesis planning.

Computational Cost. The computational cost is a critical metric for single-step prediction models, particularly when these models are intended to be integrated with other searching algorithms for multi-step retrosynthesis planning and are expected to be invoked repeatedly. In our comparative analysis, we included SMILES-based baselines including RetroPrime, Retroformer, Graph2SMILES, and the graph-based baseline, GraphRetro. We performed inference on the test set of the USPTO-50K dataset using a single NVIDIA RTX 3090 graphics card. The average inference time per sample is detailed in Table 6.

As it is shown in the table, our model exhibits the second-fastest inference speed among the SMILES-based methods and shows a negligible difference in inference time when compared to the graph-based baseline, GraphRetro. The superior inference speed of Graph2SMILES is attributed to its use of an RNN as the decoder, which has a computational complexity linearly related to the length of the output sequence. In contrast, the other SMILES-based methods are constructed with a transformer decoder, resulting in a computational complexity that is quadratically related to the output sequence length. These results underscore the capability of our method to rapidly infer results based on input, positioning it favorably for integration with searching algorithms that require extensive exploration and trial-and-error in the construction of multi-step retrosynthesis planning systems.

It is noteworthy that our method does not implement batch-wise parallelism and only utilized 2GB of GPU memory during inference. There is significant room for optimization in our code, which could enable parallel inference for multiple samples, thereby achieving greater hardware utilization efficiency and faster average inference speeds.

Ablation study

Table 7 Effects of different modules on retrosynthesis performance in reaction class unknown setting of USPTO-50K dataset

We investigate the effects of different components in our proposed pipelines. The result is summarized in Table 7.

Two-Stage Training. We eliminate the initial training phase and directly train the model for the retrosynthesis prediction task. As indicated in Table 7, the two-stage training strategy has consistently led to enhancements in all evaluated metrics. This observation implies that the two-stage training strategy effectively enables the model to adeptly learn the intricacies of molecular SMILES representations, thereby yielding higher quality and more plausible retrosynthetic analysis outcomes.

Data Augmentation. We remove the data augmentation during the second training stage, which means training solely using the DFS order that can generate canonical SMILES. Table 7 demonstrates a significant decline in model performance across all metrics. This clearly demonstrates that our data augmentation significantly improves the model’s performance.

SMILES Alignment. In the training process, we remove all operations related to SMILES alignment. This includes the removal of the position encoding in Eq. 3, where the features H directly served as the input memory for the Transformer decoder. Since we eliminate the input related to the DFS order, the model was no longer trained using order-preserving reactants SMILES as the target but instead switched to canonical SMILES for product. Additionally, in this set of experiments, we remove the first training stage, which aligns the graph and SMILES modalities as the model architecture changes. The results are reported in Table 7, and they show a significant decline in performance compared to our full version, indicating that the proposed SMILES alignment algorithm is crucial for achieving excellent performance.

It is worth noting that even without data augmentation, two-stage training and SMILES alignment, our model still outperforms the vanilla Transformer by a large margin in terms of all metrics reported in the last line of Table 7. This indicates that graph-based molecular representation learning still has advantages over SMILES-based approaches, and our proposed $\hbox {EGAT}^+$ can extract effective molecular representations for downstream usage.

Case study (visualization of cross-attention mechanism in transformer with UAlign)

Fig. 3

Visualization of cross-attention over order-aware node features and the predicted tokens. The number on the y-axis is the map number of atoms in the product. The reactants atoms that not appear in product is colored red in the x-axis. $\circ$ represents the end token

We randomly select a case from the dataset and showcase the cross-attention map in Fig. 3. The cross-attention map indicates the correlation between reactant tokens and nodes in the input product graph. This map is obtained by averaging the attention coefficient from each attention head. From the figure, it is evident that the predicted tokens successfully locate their corresponding atoms in the product, which contributes to the accurate prediction. The SMILES alignment can also be observed to assist the model in correctly identifying the reaction center. In accordance with the figure, the bond between atom C:11 and N:9 breaks during the transformation into reactants. Our model effectively notices this occurrence and focus the attention of token $\hat{t}_7$ on the reaction centers C:11 and N:9. This strategic focusing successfully guides the completion of the reactants, ensuring that the leaving group is correctly attached to the appropriate atoms. Additionally, we note that the attention coefficient at token $\hat{t}_{14}$ is concentrated on atoms C:1 and N:9, which are the first atoms of each reactant molecule according to the given DFS order. This further indicates that our model is able to correctly identify the sites where the reaction occurs and accurately cleave the chemical bonds. Moreover, the attention of newly generated structures (i.e., tokens $\hat{t}_7$ to $\hat{t}_{13}$) is directed towards atoms C:1, C:11, and O:2, which correspond to the specific synthon they will attach to. This demonstrates that our model is able to generate appropriate functional groups based on the molecular structure information to form the reactants. All the aforementioned results illustrate that our proposed SMILES alignment method assists the model in comprehending molecular structural information and helps it to focus on learning chemical rules.

To further investigate the impact of the proposed SMILES alignment mechanism on model training, we visualize the cross-attention coefficients of different Transformer decoder layers. The visualization is provided in Supplementary Fig. 1 of Supplementary Information. From Supplementary Fig. 1, we can observe significant variations in the cross-attention across different layers. Additionally, the establishment of correspondence does not occur exclusively at certain layers, such as the first or last layer. This suggests that directly imposing supervised signals on the cross attention coefficient [31, 40] for SMILES alignment is not a wise approach, whether applied to all layers or only the last layer. This observation further corroborates our assertion in "SMILES Alignment" section that supervised SMILES alignment methods might diminish the diversity of cross-attention maps across different layers, consequently impairing the model’s capacity for representation. In contrast, unsupervised SMILES alignment methods do not exert such an adverse influence. This is why our unsupervised SMILES alignment mechanism achieves better results than supervised SMILES alignment.

Case study (multi-step retrosynthetic pathway planning)

Fig. 4

Multistep retrosynthesis predictions by our method. a Mitapivat b Pacritinib c Daprodust. The reaction centers and leaving groups are highlighted in different colors. The pathway pf molecules a and b come from literature, while the last one is verified by chemical experts

To explore the suitability of our model for multi-step retrosynthetic pathway planning, we select three distinct molecules as targets for synthetic route design, and the synthesis routes are obtained through iterative calls to our UAlign model, which is trained with the USPTO-FULL dataset. The predicted pathways are summarized in Fig. 4.

The first case study involves Mitapivat, a compound approved for the treatment of hereditary hemolytic anemias in February 2022 [1]. Our model successfully predicted the five-step synthetic route reported in [3], with each step consistently ranked within the top-2 predictions. The first step entails an amide coupling reaction, which our model placed at rank 2, yielding the reactants 1-(cyclopropylmethyl)piperazine (compound 2) and 4-(quinoline-8-sulfonamido)benzoic acid (compound 3). Notably, at the initial step, our model also proposed an alternative synthesis method utilizing the Borch Reductive Amination, which was ranked at the first and is consistent with the synthetic route delineated by Saunders et al. Subsequently, for the synthesis of 4-(quinoline-8-sulfonamido)benzoic acid, the model precisely executed a functional group protection strategy during the second step and accurately anticipated the subsequent formation of the sulfonamide, effectively deconstructing the target molecule into readily available precursors. For the synthesis of 1-(cyclopropylmethyl)piperazine, the model strategically protected the amine functional group with a tert-butyloxycarbonyl moiety at the outset and, in the ultimate step, prognosticated the N-alkylation reaction with a top-ranking accuracy. This example illustrates our model’s capability to uncover diverse reaction centers in molecular retrosynthetic design and to generate plausible reactant combinations based on these insights.

The second case under scrutiny is Pacritinib, an orally bioavailable and isoform selective JAK-2 inhibitor for the treatment of patients with myelofibrosis, which received FDA approval on February 28, 2022 [48]. As shown in Fig. 4(b), our model successfully delineates a eight-step synthesis, as described in the literature [4], tracing the synthetic pathway from commercially available 5-nitrosalicylaldehyde and 2,6-dichloropyrimidine to the final product. The initial step of the reverse synthesis is olefin metathesis, ranking the first in order of likelihood, followed by another rank-2 aromatic substitution of 4-(3-((allyloxy)methyl)phenyl)-2-chloropyrimidine (compound 12) and 3-((allyloxy)methyl)-4-(2-(pyrrolidin-1-yl)ethoxy)aniline (compound 13). Subsequently, synthesis of 4-(3-((allyloxy)methyl)phenyl)-2-chloropyrimidine was correctly identified via continuous allyl substitution and Suzuki cross-coupling reaction as the top and the second choices. The reverse synthesis of 3-((allyloxy)methyl)-4-(2-(pyrrolidin-1-yl)ethoxy)aniline was reduction of the nitro group, followed by another allyl substitution. In the final step, the model’s highest probability prediction was reduction of the aldehyde group, followed by a nucleophilic substitution. Despite the synthesis route involving a considerable number of steps and encompassing a variety of reaction types, our model successfully and accurately predicted each step within the top-2 choices. This accomplishment signifies the robustness and efficacy of our model in the context of retrosynthetic analysis.

The final case is Daprodustat, the first oral hypoxy-inducing factor prolyl hydroxylase inhibitor (HIF-PHI) for the treatment of renal anemia caused by chronic kidney disease (CKD) [11]. This novel compound received approval for market release from the FDA on the 1st of February, 2023 [48]. Our model predicted the three-step synthetic route. The first step reports the hydrolysis of ester at rank 3, which is aligned with the route provided by Duffy et al. Although next two steps provided by our method do not exist in the literature, there are all explainable. The synthesis of ethyl (1,3-dicyclohexyl-2,4,6-trioxohexahydropyrimidine-5-carbonyl)glycinate (compound 21) was identified via dehydration condensation of 1,3-dicyclohexyl-2,4,6-trioxohexahydropyrimidine-5-carboxylic acid (compound 23) and ethyl glycinate (compound 22) as the top choice, which avoided using toxic ethyl isocyanatoacetate reported in literature. In the final step, the model’s highest probability prediction was amidation of ester, resulting in cost-effective and readily accessible starting materials. This case demonstrates the robust extrapolative capacity of our model, highlighting its potential to generate synthetic routes that surpass those documented in the literature.

We also provide the results of multi-step retrosynthesis planning of two powerful baselines: SMILES-based method Retroformer [40] and graph-based method GraphRetro [33] in Supplementary Fig. 2 and Supplementary Fig. 3 respectively. The visualization reveals that while the Retroformer outperforms our method in the prediction of the synthetic route for Mitapivat, placing the literature pathway in a more advanced position, our model still accurately predicted each step of the literature-provided synthetic route within the top-2 choices. Conversely, when faced with compounds like Pacritinib, which has multiple potential reaction centers, the Retroformer exhibits disadvantages. This is evident in Supplementary Fig. 2 (b) from steps 4 to 7, where the literature-documented synthetic route is ranked beyond the third position by Retroformer. Additionally, it is observable that Retroformer lacks robust predictive power for complex reactions, such as those requiring ring formation. Supplementary Fig. 2 illustrates that Retroformer failed to successfully predict Pacritinib’s step-1 and was also unable to forecast both the literature pathways and the pathway validated by our chemical experts for Daprodust. Supplementary Fig. 3 further demonstrates that the performance of GraphRetro is marginally worse than that of Retroformer. Across the three presented cases, GraphRetro has not been able to successfully predict the synthetic routes. All the results above suggest that our model demonstrates a stronger capability in handling more complex molecules and reaction types compared to our baselines.

Discussion

We present UAlign, a novel graph-to-sequence pipeline that achieves state-of-the-art performance in the field of template-free methods. Our approach outperforms existing template-free and semi-template-based methods, while achieving comparable results to template-based methods. By utilizing a specially-designed graph neural network as the encoder, our model effectively leverages chemical and structural information from molecule graphs, resulting in powerful embedding for the decoder. Additionally, our proposed unsupervised SMILES alignment mechanism facilitates the reuse of shared substructures between reactants and products for generation, allowing the model to prioritize chemical knowledge even without complex data annotations. This significantly enhances the performance of the pipeline.

Despite achieving commendable performance, our work acknowledges areas for improvement. This work does not integrate much domain knowledge related to chemical reaction mechanisms in its design, which to some extent, compromises its interpretability. Similarly to most template-free methods, our work also faces challenges in generating diverse results. Additionally, we recognize a significant disparity between single-step retrosynthesis prediction and the complex reality of molecular synthesis route planning, underscoring the need for more realistic evaluation metrics to validate proposed models. To this end, we are charting a course for future exploration.

We remain committed to monitoring advancements in the understanding of chemical reaction mechanisms and intend to compile pertinent information from the field of chemical reactions, including kinetic and thermodynamic data, to construct a more interpretable single-step retrosynthesis prediction model. We also plan to build a multiple-step route design system for molecular synthesis, integrating UAlign with search algorithms and predictive models for reaction conditions. Based on such a system, we expect to work with chemists to synthesize complex molecules in the wet lab, assessing the capability of the retrosynthesis prediction model.