RetroMoE: A Mixture-of-Experts Latent Translation Framework for Single-step Retrosynthesis

1 Introduction

Single-step retrosynthesis is a fundamental aspect of organic chemistry, especially vital for the pharmaceutical industry, as it enables the design of viable synthetic routes to create complex compounds. This process involves deducing the necessary reactants by working backwards from the final product, forming the basis for multi-step synthesis planning, where a complete route is constructed through a series of sequential single-step reactions.

Recently, single-step retrosynthesis has advanced significantly through computer-aided synthesis planning (CASP), particularly with the integration of machine learning. There are three primary research approaches in machine-learningbased retrosynthesis. The first approach is template-based methods [Segler and Waller, 2017; Dai et al., 2019; Chen and Jung, 2021; Xie et al., 2023], which involve searching a library to find the most relevant reaction templates that represent chemical reactions in synthesis. By applying these templates to product molecules, they facilitate straightforward reactant prediction. However, despite their state-of-the-art performance and interpretability, these methods struggle with generalization and scalability due to the limited template libraries. The second approach, semi-template-based methods [Shi et al., 2020; Chen et al., 2023; Somnath et al., 2021; Sacha et al., 2021; Liu et al., 2022; Zhong et al., 2023], either employ a twostage process (identifying reaction centers and completing synthons) or utilizes an auto-regressive graph editing technique. While these methods strike a balance between interpretability and generalization, they often suffer from cumulative errors due to their stage-wise processing nature. The third approach is template-free methods, which treat retrosynthesis as a sequence-to-sequence [Liu et al., 2017; Irwin et al., 2022; Kim et al., 2021] or graph-to-sequence [Tu and Coley, 2022; Wan et al., 2022] translation problem, converting products directly into reactants. Some methods [Kim et al., 2021; Igashov et al., 2023] also attempt to capture the diversity inherent in retrosynthesis procedures. However, despite their broader scope, these methods often exhibit lower performance compared to other approaches.

Figure 1: The retrosynthesis task involves identifying the reactants required to synthesize a given product molecule. Illustrated using both molecular graphs and SMILES strings, the task reveals that only minor modifications are typically needed in the molecule. This observation motivates the approach of projecting molecules into a latent space and formulating retrosynthesis as a latent translation problem from product to reactants.

In our study, we introduce a new template-free method, RetroMoE, which tackles the retrosynthesis task as a latent translation from products to reactants. Drawing from the molecule optimization literature [Du et al., ], we know that molecules with similar structures tend to cluster in the latent space. This insight aligns with the observation that only minor modifications are typically needed for reactants to synthesize products during chemical reactions, as illustrated in Fig. 1. Inspired by this, we aim to enhance the template-free method from a latent translation perspective. To achieve this, we introduce a non-symmetric variational autoencoder (VAE) model that constructs a latent space for both product and reactant molecules. We utilize a graph encoder to ensure the latent space accurately captures the structural information of molecules. For precise molecule prediction, we employ a transformer decoder. We then introduce a novel latent translation method that converts the product's latent embedding into the reactant's latent embedding. Given that retrosynthesis involves more complex chemical reactions than typical molecule optimization tasks, we employ a mixture-of-experts (MoE) network [Jacobs et al., 1991] for this translation. Finally, to ensure the precise prediction of reactants, we use a transformer decoder following the MoE network, further enhancing our model's accuracy.

The contributions of our work are summarized as follows:

• To the best of our knowledge, we are the first to formulate the retrosynthesis task as a latent translation problem.
• We propose a non-symmetric VAE model to construct the latent space and an MoE model to translate the product latent embedding to the reactant latent embedding.
• Our approach outperforms other template-free and semitemplate-based methods on the USPTO-50K and USPTOMIT datasets and is competitive with the template-based methods.

2 RelatedWork

2.1 Single-step Retrosynthesis

Several lines of research have explored learning methods for modeling single-step retrosynthesis, broadly categorized into three types: template-based, semi-template-based, and template-free. We introduce the semi-template-based methods in the supplementary material due to space limit.

2.2 Template-based Retrosynthesis

In retrosynthesis, a template is a structured representation of a chemical reaction, defining how a product subgraph transforms into one or more reactant subgraphs. It serves as a blueprint, detailing the rearrangement of specific atoms and bonds during the synthesis process.

Template-based models identify the most relevant reaction templates from a database. Notable methods include NeuralSym [Segler and Waller, 2017] and RetroSim [Coley et al., 2017], which match templates to products based on molecular similarity. MHNReact [Seidl et al., 2022] treats this as a content-based retrieval task using a modern Hopfield network. GLN [Dai et al., 2019] predicts joint conditional probabilities, integrating reactants into template relevance decisions. LocalRetro [Chen and Jung, 2021] emphasizes local templates, while RetroKNN [Xie et al., 2023] enhances this with k-nearest-neighbor (KNN) search. RetroComposer [Yan et al., 2022] composes new templates for retrosynthesis.

Despite their interpretability and strong performance, template-based methods struggle with generalization, as they cannot adapt to templates outside the database. They also face scalability issues, with performance deteriorating on larger datasets due to the increased likelihood of missing relevant templates during extensive searches.

2.3 Template-free Retrosynthesis

The template-free approach to retrosynthesis is first introduced by [Liu et al., 2017], framing the task as a sequenceto-sequence translation problem by processing the chemical language SMILES with natural language processing techniques. [Karpov et al., 2019] later incorporates the transformer architecture into retrosynthesis. To improve performance, AugTransformer[Tetko et al., 2020] uses extensive data augmentation, while Chemformer [Irwin et al., 2022] combines data augmentation with pre-trained models. GTA [Seo et al., 2021] optimizes training by designing attention masks to reduce transformer parameters. TiedTransformer [Kim et al., 2021] introduces a forward synthesis transformer to create a cycle-consistent framework, addressing diversity issues through latent variables, as also explored by [Chen et al., 2019; He et al., 2022]. Graph2Smiles [Tu and Coley, 2022] combines a graph encoder with a transformer decoder, while Retroformer [Wan et al., 2022] uses a graph transformer for reaction center detection and a transformer decoder for translating product SMILES to reactant SMILES. Recently, Retrobridge [Igashov et al., 2023] modeled retrosynthesis as Markov bridges between product and reactant distributions.

Our work aligns with the methods of [Kim et al., 2021; Chen et al., 2019; He et al., 2022; Igashov et al., 2023], but rather than fostering diversity through latent variables, we approach the task as a latent translation between products and reactants.

2.4 Molecule Optimization

Our work is closely related to the 1D/2D molecule optimization task, where the goal is to generate new molecules with desired properties by optimizing existing ones. HierVAE [Jin et al., 2020] employs a hierarchical variational autoencoder (VAE) to capture and model the complex hierarchical information of molecules during optimization. MSO [Winter et al., 2019] uses a sequence-level VAE, representing molecules as SMILES strings and reconstructing them from a learned latent space to explore and generate molecular structures. ChemSpace [Du et al., ] introduces the concept of identifying smooth latent directions that control molecular properties, emphasizing their role in interpreting structural changes in relation to property variations.

While these methods typically involve creating a latent space and using a single network to transform existing molecules into new ones, retrosynthesis presents unique challenges. Complex chemical reactions often involve multiple, rather than simple, transformations, requiring precise, rather than approximate, representation of reactants. To address these challenges, we propose a mixture-of-experts framework to effectively model diverse and complex reaction mechanisms.

Additionally, we introduce a non-symmetric graph-sequence VAE specifically tailored to the retrosynthesis task.

3 Method

3.1 ProblemFormulation

Given one product moleculeMP , the retrosynthesis task aims to predict a set of N reactant molecules {MRi }Ni=1 that can lead to MP . Furthermore, a molecule denoted as M, can be represented using two primary data formats, the SMILES string and the molecular graph.

For the SMILES format, the molecular structure S is expressed as a sequence of characters si, written as S := s1s2...sL, where L represents the total length of the string. Each character si represents a structural element, which could be an atom element, a chemical bond, a branching notation, etc. To represent multiple molecules, multiple SMILES strings can be concatenated using a period ".", resulting in a single, extended SMILES sequence.

Alternatively, a molecule can be conceptualized as a graph G = {V, E}. Here, V = {v1, .., vn} corresponds to the set of n atoms in the molecule, and E = {e1, .., em} represents the set of m bonds between these atoms. Each node (atom) in this graph is associated with a feature vector hi ∈ Rd that holds atomic properties such as aromaticity and electric charge. The complete atomic information is compiled into a node feature matrix H ∈ Rn×d . The adjacency matrix B ∈ Rn×n×c describes the topological relationships within M, with c denoting the number of bond types and Bijk indicating the presence or absence of a chemical bond of type k between atom i and atom j. To represent multiple molecules, a collection of molecules can be treated as a single disconnected graph, where each molecule forms an independent connected component within the graph. In our work, we use both molecular graph and SMILES string for our encoder and decoder, respectively.

3.2 Non-symmetric VAE

As noted in the molecule optimization literature [Du et al., ], molecules with similar structures tend to cluster together in the latent space, which corresponds to the observation that reactants typically undergo only minor modifications to synthesize products during chemical reactions. With this in mind, we aim to develop a latent space that enables us to reformulate the retrosynthesis task as a product-to-reactant latent translation task. To achieve this, we employ the Variational AutoEncoder (VAE) model [Kingma and Welling, 2013]. Below, we provide a brief overview of the VAE.

VAE Recap

VAEs estimate the Evidence Lower Bound (ELBO) on the log-likelihood p(x) of the input molecule x, using a proposal distribution q(z|x), where z represents the latent variables. The goal is to maximize the ELBO, defined as:

The first term of the ELBO is the reconstruction term, while the second term, the Kullback-Leibler (KL) divergence, measures the information loss when q(z|x) approximates p(z). In VAEs, p(z) typically follows a standard Gaussian distribution, N(0, I), suggesting that all latent dimensions should be independent or disentangled. Another important aspect to highlight is the reparameterization trick. In VAEs, z is expressed as a deterministic variable through z = f(ϵ, x), where ϵ is an auxiliary variable independently distributed as p(ϵ), and f(·) is the reconstruction function. The reparameterization is given by:

where ⊙ denotes the element-wise product. This trick facilitates the computation of gradients of the reconstruction function with reduced variance.

Non-symmetric Encoder-Decoder

In developing our VAE model, we employ a non-symmetric encoder-decoder framework. Specifically, we combine a graph-based encoder with a transformer-based decoder [Tu and Coley, 2022; Wan et al., 2022; Zeng et al., 2024], designed to efficiently encode molecular graphs into a latent space and subsequently decode them into SMILES strings. This approach is tailored to support the complex demands of the retrosynthesis task, ensuring accurate and meaningful translations from product to reactant representations.

As shown in Fig. 2, given a molecule graph G = {V, E} representing a product or reactant, we begin by converting it into two matrices: the node (atom) matrix H and the edge (bond) matrix B. These matrices are then processed by separate encoders: an atom encoder for H and a bond encoder for B, as they have different dimensions. Each encoder consists of a series of MLP layers that transform the respective matrices into atom features ˆH and bond features ˆB.

After processing, the features are handled by our specially designed graph encoder. We base this encoder on the Graph Attention Network (GAT) architecture [Veliˇckovi´c et al., 2018]. In GAT, the attention mechanism computes the relative weights of edges connecting node pairs, allowing the model to focus more on important structures within the molecule. Specifically, GAT updates the feature vector of each node by aggregating features from its neighbors, weighted by attention coefficients. However, the standard GAT only considers the neighbor node features, neglecting the edge feature and current node feature which are crucial for our task. To address this, we incorporate the edge feature and current node feature into the computation process. Mathematically, for a node i, the new node feature h′i and edge feature b′ij are computed as:

where N(i) denotes the neighbors of node i. αij is the attention coefficient between nodes i and j, computed using a softmax over a learned transformation of the node features:

W is a weight matrix applied to node features before computing attention coefficients. Note that we use different W for different features in our implementation. a is a learnable parameter vector used to compute the raw attention scores. ∥ denotes concatenation. The whole procedure can be summarized as follows:

After processing through the graph encoder, we use only the node feature ˜H, which now encapsulates all necessary information following multiple iterations of message passing. Initially, ˜H is input into two distinct MLP layers to compute the mean and variance values essential for the reparameterization technique described in Eq. 2. This step yields the latent embedding Z.

Subsequently, this latent embedding Z is fed into a vanilla transformer encoder [Vaswani et al., 2017], in conjunction with sine-cosine position-encoded tokenized SMILES strings S. This combination facilitates the final prediction Y:

3.3 MoE-based Latent Translation

The latent translation is a common technique in molecule optimization tasks [Du et al., ; Jin et al., 2020; Winter et al., 2019], where a single network is typically used to adjust a molecule's properties, such as converting a toxic molecule into a non-toxic one. However, unlike molecule optimization, the retrosynthesis task involves more complex chemical reactions and requires more precise predictions of reactants. To address this complexity, we propose a mixture-of-experts (MoE) network to better model the intricate chemical reactions involved in retrosynthesis.

As shown in Fig. 3, given a product molecule graph Gp and a reactant molecule graph Gr, we first input them into the previously mentioned graph encoder to obtain their respective latent embeddings, Zp for the product and Zr for the reactant. Subsequently, we employ a Mixture of Experts (MoE) network to translate Zp into Zr.

The MoE network comprises a gating network and several expert networks. All the gating networks and expert networks contain a series of MLP layers, ReLU layers, and LayerNorm layers. The gating network generates weights for each expert based on the input features. These weights are then used to combine the outputs from the various expert networks appropriately. By weighting and aggregating these expert outputs, we derive the translated latent embedding Zpr which is then optimized to closely match the reactant latent embedding Zr. To facilitate the accurate prediction of reactants, we still apply the transformer encoder to obtain the final output Ypr. The procedure can be formulated as follows:

3.4 Optimization

During training, our model undergoes a two-stage process. First, we train the non-symmetric VAE with the ELBO outlined in Eq. 1. Specifically, we apply a cross-entropy loss between tokenized SMILES string S and the prediction Y as the first reconstruction term. For the second term, we still use the KL divergence but with a regularization factor to encourage the more disentangled latent representation following [Higgins et al., 2017]. The loss function for this first stage is defined as:

where β represents the regularization factor. This stage includes all products and reactants from the training set to ensure an effective latent space for the subsequent translation task in the second stage.

Second, we train the MoE network and transformer decoder using a mean square error (MSE) loss between the translated latent embedding Zpr and the reactant latent embedding Zr. Additionally, we apply a cross-entropy loss to ensure the final output Yr accurately matches the tokenized reactant SMILES string Sr. This approach optimizes both the fidelity of the translated embeddings and the accuracy of the predicted sequences. The loss function for this second stage is defined as:

In the inference stage, given input product graph Gp, we use the trained model from the second stage to predict reactant SMILES string Yr.

4 Experiments

4.1 Datasets

In our study, we utilize two established retrosynthesis benchmark datasets: USPTO-50k [Schneider et al., 2016] and USPTO-MIT [Jin et al., 2017]. The USPTO-50k dataset comprises 50,016 atom-mapped reactions, categorized into 10 reaction classes, and is divided into training, validation, and test sets with 40,008, 5,001, and 5,007 reactions, respectively, following the partitioning used in previous work [Dai et al., 2019]. We assess model performance under two scenarios: with and without known reaction classes. The USPTO-MIT dataset contains approximately 479,000 atom-mapped reactions, with around 409,000 for training, 40,000 for validation, and 30,000 for testing. Unlike USPTO-50k, reaction classes are not used for benchmarking in the USPTO-MIT dataset.

4.2 Implementation Details

In our experiments on the USPTO-50K dataset, we use the Adam optimizer with an initial learning rate of 1.25e-4 for the first training stage and 1e-4 for the second stage. The training lasts for 45 epochs in the first stage and 170 epochs in the second. We apply an exponential scheduler for learning rate decay and set the β value in Loss1 to 0.001. Both the graph encoder and transformer decoder in each stage have a hidden size of 512, 8 layers for both encoder and decoder, and 8 attention heads. The MoE network consists of 3 gating layers, 8 expert layers, and 3 experts.

For the experiments on the USPTO-MIT dataset, we use the Adam optimizer with an initial learning rate of 1e-4 for the first phase and 5e-5 for the subsequent phase. The training duration is 85 epochs for the first stage and 300 epochs for the second. We continue using an exponential scheduler for learning rate decay, with the β value in Loss1 set to 0.001. In terms of model configuration, both stages feature a graph encoder and transformer decoder with a hidden size of 768, 8 encoder and decoder layers, and 12 attention heads. The MoE network includes 3 gating layers, 8 expert layers, and 3 experts. We conduct the experiments using the PyTorch framework on NVIDIA A5000 GPUs.

To maintain baseline performance, we use the SMILES alignment and data augmentation techniques similar to previous methods [Wan et al., 2022; Zeng et al., 2024].

4.3 Evaluation

During inference, we use beam search to generate the output SMILES with a beam size of 10. In all our experiments, we assess performance using the Top-k exact match accuracy for k = 1, 3, 5, and 10. This metric calculates the proportion of input products for which the tested method correctly predicts the entire set of reactants within its top-k predictions.

In the experiments on the USPTO-50k dataset, we also report the Top-k round-trip accuracy and coverage following [Igashov et al., 2023]. Specifically, we use the Molecular Transformer model [Schwaller et al., 2019] to predict forward reactions for the top-k samples of each input product. Roundtrip accuracy measures the proportion of correctly predicted reactants, considering predictions as correct if they either match the ground truth or lead back to the input product. Round-trip coverage evaluates whether at least one prediction within the top-k meets this correctness criterion. These metrics account for the possibility of multiple valid reactant sets for a single product.

4.4 Comparison with State-of-the-art Methods Results on USPTO-50K

The results on the USPTO-50K dataset are shown in Table. 1 and Table. 2. Note that a more thorough comparison can be found in the supplementary material. Additionally, we illustrate our method's ability to generate new molecules in the supplementary material.

Top-k Exact Match Accuracy. With known reaction classes, our model achieves 66.7% top-1, 86.5% top-3, 91.3% top-5, and 94.4% top-10 accuracy, setting a new benchmark for template-free methods and proving competitive with template-based and semi-template-based methods. It is noteworthy that our model matches the top-1 accuracy achieved by RetroKNN [Xie et al., 2023], marking a significant success in the template-free methods. Without reaction class information, the model reaches a 54.8% top-1, 76.7% top-3, 83.4% top-5 and 89.3% top-10 accuracy. While its performance trails behind RetroKNN [Xie et al., 2023], it surpasses most template-based and semi-template-based methods, showing our method's superiority.

Top-k Round-trip Coverage and Accuracy. Following RetroBridge [Igashov et al., 2023], we also report the top-k round-trip coverage and accuracy on the USPTO-50K dataset under the class-unknown setting. As shown in Table. 2, our method achieves the state-of-the-art performance on both round-trip coverage and accuracy. This demonstrates that our model is more effective at proposing valid and efficient synthetic routes for downstream applications.

The Ability to Generate NewMolecules. Our model can also successfully generate novel reactant molecules that can synthesize the same product. More details are introduced in the supplementary material.

Results on USPTO-MIT

The top-k exact match accuracy results on the USPTO-MIT dataset are shown in Table. 4. Our method outperforms other template-free methods and most template-based methods. Note that our model achieves competitive performance with the state-of-the-art template-based method which is under the class-unknown setting since the class information is not provided by the USPTO-MIT dataset. This shows the superiority of our method.

4.5 Ablation Analysis

Distance between the Latent Embeddings of Product and Reactant. Previous research [Du et al., ] has demonstrated that molecules with similar structures tend to cluster in the latent space. Despite this, we conduct further experiments to validate this phenomenon. Specifically, we calculated the average euclidean distance between the latent embeddings of product-reactant pairs within the USPTO-50K test set. For comparison, we also calculate the average euclidean distance for disparate product pairs and reactant pairs. The results, as presented in the Table. 3, show that the average distance between corresponding products and reactants is much lower than that between different products and reactants.

The Number of Experts. Here, we examine the influence of the number of experts in the MoE network on the performance of the USPTO-50K test dataset under the class-unknown setting. According to the results shown in Table 5, utilizing three experts yielded the best performance. Due to memory constraints, we limit our testing to configurations with 1, 2, 3, and 4 experts, noting that increasing the number of experts could potentially enhance performance further.

Different Loss Functions. Our model in the second stage can be trained using only the cross entropy loss. Thus, we investigate the importance of incorporating the mean square error loss. This loss measures the distance between the predicted reactant latent embedding and the ground truth latent embedding. As shown in Table. 6, employing mean square error loss, which effectively narrows the distance between the product and reactant latent embeddings, is crucial for enhancing performance.

The Effect of Non-symmetric VAE. In our work, we use a graph-transformer encoder-decoder design to model precise predictions of SMILES strings, creating an accurate latent space for the second stage. Here, we discuss the effect of this non-symmetric VAE. We design a graph-graph symmetric VAE, where the graph decoder mirrors the structure of the graph encoder introduced in the method section. After training this symmetric VAE, we apply the graph encoder to the second stage of our training, just as we do with the nonsymmetric VAE. As seen in Table 7, the graph-transformer non-symmetric structure is crucial for performance, indicating that the transformer decoder's ability to accurately predict SMILES strings helps build an effective latent space.

Effectiveness of the Latent Translation Scheme. To demonstrate the effectiveness of our latent translation scheme, we train the model using only the non-symmetric architecture without incorporating VAE or MoE components. Specifically, we removed the VAE modules (MLPs and Reparameterization) depicted in Fig. 2, using only the product graph as input and outputting the reactant SMILES string. As shown in Table 8, our two-stage latent translation process significantly outperforms the one-stage training approach.

5 Conclusion

In our work, we introduce RetroMoE, a novel generative model designed for the single-step retrosynthesis task. This model employs a non-symmetric variational autoencoder (VAE) that incorporates a graph encoder and a transformer decoder to effectively learn a molecular latent space. Furthermore, we apply a simple yet effective mixture-of-experts (MoE) network that adeptly translates the product latent embedding into the reactant latent embedding. As shown in the experiments on USPTO-50K and USPTO-MIT datasets, our approach not only surpasses other template-free and semitemplate-based methods, but also matches the performance of state-of-the-art template-based method, RetroKNN [Xie et al., 2023].