RetroBridge: Modeling Retrosynthesis with Markov Bridges

Introduction

Computational and machine learning methods for de novo drug design show great promise as more cost-effective alternatives to experimental high-throughput screening approaches (Thomas et al. 2023) to propose molecules with desirable properties. While in silico results suggest high predicted target binding affinities and other favorable properties of the generated molecules, limited emphasis has so far been placed on their synthesizability (Stanley & Segler 2023). For laboratory testing, synthetic pathways need to be developed for the newly designed molecules, which is an extremely challenging and time-consuming task.

Retrosynthesis planning (Corey 1991; Strieth-Kalthoff et al. 2020; Tu et al. 2023) tools address this challenge by proposing reaction steps or entire pathways that can be validated and optimized in the lab. Single-step retrosynthesis models predict precursor molecules for a given target molecule (Segler & Waller 2017; Coley et al. 2017; Liu et al. 2017; Strieth-Kalthoff et al. 2020; Tu et al. 2023). Applying these methods recursively allows to decompose the initial molecule in progressively simpler intermediates and eventually reach available starting molecules (Segler et al. 2018).

While most works have used a discriminative formulation for retrosynthesis modeling (Strieth-Kalthoff et al. 2020; Tu et al. 2023; Jiang et al. 2022), we propose to view the task as a conditional distribution learning problem, as shown in Figure 1. This approach has several advantages, including the ability to model uncertainty and to generate new and diverse retrosynthetic pathways. Furthermore, and most importantly, the probabilistic formulation reflects the fact that the same product molecule can often be synthesized with different sets of reactants and reagents.

Diffusion models (Sohl-Dickstein et al. 2015; Ho et al. 2020) and other modern score-based and flow-based generative methods (Rezende & Mohamed 2015; Song et al. 2020; Lipman et al. 2022; Albergo & Vanden-Eijnden 2022; Albergo et al. 2023) may seem like good candidates for retrosynthesis modeling. However, as we show in this work, such models do not fit naturally to the formulation of the problem, as they are designed to approximate a single intractable data distribution. To do so, one typically samples initial noise from a simple prior distribution and then maps it to a data point that follows a complex target distribution. In contrast, we aim to learn the dependency between two intractable distributions rather than one intractable distribution itself. While this can be achieved by conditioning the sampling process on the relevant context and keep sampling from the prior noise, we show here that such use of the original unconditional generative idea is suboptimal for approximating the dependency between two discrete distributions.

In this work, we propose RetroBridge, a template-free probabilistic method for single-step retrosynthesis modeling. As shown in Figure 1, we model the dependency between the spaces of products and reactants as a stochastic process that is constrained to start and to end at specific data points. To this end, we introduce the Markov Bridge Model, a generative model that learns the dependency between two intractable discrete distributions through the finite sample of coupled data points. Taking a product molecule as input, our method models the trajectories of Markov bridges starting at the given product and ending at data points following the distribution of reactants. To score reactant graphs sampled in this way, we leverage the probabilistic nature of RetroBridge and measure its uncertainty at each sample. We demonstrate that RetroBridge achieves competitive results on standard retrosynthesis modeling benchmarks. Besides, we compare RetroBridge with the state-of-the-art graph diffusion model DiGress (Vignac et al. 2022), and demonstrate quantitatively and qualitatively that the proposed Markov Bridge Model is better suited to tasks where two intractable discrete distributions need to be mapped. Our source code is available at https://github.com/igashov/RetroBridge.

To summarise, the main contributions of this work are the following:

We introduce the Markov Bridge Model to approximate the probabilistic dependency between two intractable discrete distributions accessible via a finite sample of coupled data points.

We demonstrate the superiority of the proposed formulation over diffusion models in the context of learning the dependency between two intractable discrete distributions.

We propose RetroBridge, the first Markov Bridge Model for retrosynthesis modeling. RetroBridge is a template-free single-step retrosynthesis prediction method that achieves state-of-the-art results on standard benchmarks.

Related Work

Diffusion models (Sohl-Dickstein et al. 2015; Ho et al. 2020) form a class of powerful and effective score-based generative methods that have recently achieved promising results in many different domains including protein design (Watson et al. 2023), small molecule generation (Hoogeboom et al. 2022; Igashov et al. 2022; Schneuing et al. 2022), molecular docking (Corso et al. 2022), and sampling of transition state molecular structures (Duan et al. 2023; Kim et al. 2023). While most models are designed for the continuous data domain, a few methods were proposed to operate on discrete data (Hoogeboom et al. 2021; Johnson et al. 2021; Austin et al. 2021; Yang et al. 2023) and, in particular, on discrete graphs (Vignac et al. 2022). To the best of our knowledge, however, no diffusion models have been applied to modeling chemical reactions and recovering retrosynthetic pathways.

Given two distributions and a reference stochastic process between them, solving the Schrödinger bridge (SB) problem (Schrödinger 1932; Léonard 2013) amounts to finding a process closest to the reference in terms of Kullback-Leibler divergence on path spaces. While most recent methods employ the SB formalism in the context of unconditional generative modeling (Vargas et al. 2021; Wang et al. 2021; De Bortoli et al. 2021; Chen et al. 2021; Bunne et al. 2023; Liu et al. 2022), a few works aimed to approximate the reference stochastic process through training on coupled samples from two continuous distributions (Holdijk et al. 2022; Somnath et al. 2023). To the best of our knowledge, there are no methods operating on categorical distributions, which is the subject of the present work.

Recent retrosynthesis prediction methods can be divided into two main groups: template-based and template-free methods (Jiang et al. 2022). While template-based methods depend on predefined sets of specific reaction templates or leaving groups, template-free methods are less restricted and therefore are able to explore new reaction pathways. Two common data representations used for retrosynthesis prediction are symbolic representations SMILES (Weininger 1988) and molecular graphs. A variety of language models have been recently proposed (Liu et al. 2017; Zheng et al. 2019; Sun et al. 2021; Tetko et al. 2020) to operate on SMILES, some of which (Jiang et al. 2023; Zhong et al. 2022) were additionally pre-trained on larger datasets. Due to the nature of the sequence-to-sequence translation problem, all these methods are template-free. Among the existing graph-based methods (Segler & Waller 2017), the most recent template-based ones are GLN (Dai et al. 2019), GraphRetro (Somnath et al. 2021) and LocalRetro (Chen & Jung 2021), and template-free approaches are G2G (Shi et al. 2020) and MEGAN (Sacha et al. 2021). Template-free methods GTA (Seo et al. 2021), Graph2SMILES (Tu & Coley 2022) and Retroformer (Wan et al. 2022) leverage both graph and SMILES representations. In this work, we propose a novel template-free graph-based method.

RetroBridge

We frame the retrosynthesis prediction task as a generative problem of modeling a stochastic process between two discrete-valued distributions of products $p_{\mathcal{X}}$ and reactants $p_{\mathcal{Y}}$. These distributions are intractable and are represented by a finite collection of $D$ coupled samples $\{(\bm{x}_{i},\bm{y}_{i})\}_{i=1}^{D}$, where $\bm{x}_{i}\sim p_{\mathcal{X}}(\bm{x}_{i})$ is a product molecule and $\bm{y}_{i}\sim p_{\mathcal{Y}}(\bm{y}_{i})$ is a corresponding set of reactant molecules. While products and reactants follow distributions $p_{\mathcal{X}}$ and $p_{\mathcal{Y}}$ respectively, there is a dependency between these variables that can be expressed in the form of the joint distribution $p_{\mathcal{X},\mathcal{Y}}$ such that $\int p_{\mathcal{X},\mathcal{Y}}(\bm{x},\bm{y})d\bm{x}=p_{\mathcal{Y}}(\bm{y})$ and $\int p_{\mathcal{X},\mathcal{Y}}(\bm{x},\bm{y})d\bm{y}=p_{\mathcal{X}}(\bm{x})$. The joint distribution $p_{\mathcal{X},\mathcal{Y}}$ is also intractable and accessible only through the discrete sample of coupled data points $\{(\bm{x}_{i},\bm{y}_{i})\}_{i=1}^{D}$.

First, we introduce the Markov Bridge Model, a general framework for learning the dependency between two intractable discrete-valued distributions. Next, we discuss a special case where random variables are molecular graphs. Upon this formulation, we introduce RetroBride, a Markov Bridge Model for single-step retrosynthesis modeling. Finally, we explain a simple but rather effective way of scoring RetroBridge samples based on the statistical uncertainty of the model.

We model the dependency between two discrete spaces $\mathcal{X}$ and $\mathcal{Y}$ by a Markov bridge (Fitzsimmons et al. 1992; Çetin & Danilova 2016), which is a Markov process pinned to specific data points in the beginning and in the end. For a pair of samples $(\bm{x},\bm{y})\sim p_{\mathcal{X},\mathcal{Y}}(\bm{x},\bm{y})$ and a sequence of time steps $t=0,1,\dots,T$, we define the corresponding Markov bridge as a sequence of random variables $(\bm{z}_{t})_{t=0}^{T}$, that starts at $\bm{x}$, i.e., $\bm{z}_{0}=\bm{x}$, and satisfies the Markov property,

To pin the process at the data point $\bm{y}$, we introduce an additional requirement,

where $\bm{I}_{K}$ is a $K\times K$ identity matrix, $\bm{1}_{K}$ is a $K$-dimensional all-one vector, and $\alpha_{t}$ is a schedule parameter transitioning from $\alpha_{0}=1$ to $\alpha_{T-1}=0$. Transition probabilities (1) can be written as follows,

where $\text{Cat}(\cdot\ ;\bm{p})$ is a categorical distribution with probabilities given by $\bm{p}$. We note that setting $\alpha_{T-1}=0$ ensures the requirement (2).

Using the finite set of coupled samples $\{(\bm{x}_{i},\bm{y}_{i})\}_{i=1}^{D}\sim p_{\mathcal{X},\mathcal{Y}}$, our goal is to learn a Markov bridge (1-2) to be able to sample $\bm{y}$ when only $\bm{x}$ is available. To do this, we replace $\bm{y}$ with an approximation $\hat{\bm{y}}$ computed with a neural network $\varphi_{\theta}$:

and define an approximated transition kernel,

We train $\varphi_{\theta}$ by maximizing a lower bound of log-likelihood $\log q_{\theta}(\bm{y}|\bm{x})$. As shown in Appendix A.1, it has the following closed-form expression,

For any $\bm{x}\in\mathcal{X},\bm{y}\in\mathcal{Y}$, and $t=1,\dots,T$, sampling of $\bm{z}_{t}$ can be effectively performed using a cumulative product matrix $\overline{\bm{Q}}_{t}=\bm{Q}_{t}\bm{Q}_{t-1}...\bm{Q}_{0}$. As shown in Appendix A.2, the cumulative matrix $\overline{\bm{Q}}_{t}$ can be written in closed form,

where $\overline{\alpha}_{t}=\prod_{s=0}^{t}\alpha_{s}$. Therefore, $p(\bm{z}_{t+1}|\bm{z}_{0},\bm{z}_{T})$ can be written as follows,

To sample a data point $\bm{y}\equiv\bm{z}_{T}$ starting from a given $\bm{z}_{0}\equiv\bm{x}\sim p_{\mathcal{X}}(\bm{x})$, one iteratively predicts $\hat{\bm{y}}=\varphi_{\theta}(\bm{z}_{t},t)$ and then derives $\bm{z}_{t+1}\sim q_{\theta}(\bm{z}_{t+1}|\bm{z}_{t})=\text{Cat}\left(\bm{z}_{t+1};\bm{Q}_{t}(\hat{\bm{y}})\bm{z}_{t}\right)$ for $t=0,\dots,T-1$. Training and sampling procedures of the Markov Bridge Model are provided in Algorithms 1 and 2 respectively.

2 RetroBridge: Markov Bridge Model for Retrosynthesis Planning

In the scope of our probabilistic framework, we consider such a molecular graph representation as a collection of independent categorical random variables. More formally, we denote product and reactants data points $\bm{x}$ and $\bm{y}$ as tuples of the corresponding node and edge feature tensors: $\bm{x}=[\bm{H}_{x},\bm{E}_{x}]$ and $\bm{y}=[\bm{H}_{y},\bm{E}_{y}]$. For such complex data points, we modify the definitions of transition matrices and probabilities accordingly:

Because some atoms present in the reactant molecules can be absent in the corresponding product molecule, we add “dummy” nodes to the initial graph of the product. As shown in Figure 2, some “dummy” nodes are transformed into atoms of reactant molecules. In our experiments, we always add 10 “dummy” nodes to the initial product graphs.

3 Confidence and Scoring

It is important to have a reliable scoring method that selects the most relevant sets of reactants out of all generated samples. In order to rank RetroBridge samples, we benefit from the probabilistic nature of the model and utilize its confidence in the generated samples as a scoring function. We estimate the confidence of the model by computing the likelihood $q_{\theta}(\bm{y}|\bm{x})$ of a set of reactants $\bm{y}$ for an input product molecule $\bm{x}$. For a set of $M$ samples $\{\bm{z}_{T}^{(i)}\}_{i=1}^{M}$ generated by RetroBridge for an input product $\bm{x}$, we compute the likelihood-based confidence score for the set of reactants $\bm{y}$ as follows,

Results

For all the experiments we use the USPTO-50k dataset (Schneider et al. 2016) which includes 50k reactions found in the US patent literature. We use standard train/validation/test splits provided by Dai et al. 2019. Somnath et al. 2021 report that the dataset contains a shortcut in that the product atom with atom-mapping 1 is part of the edit in almost 75% of the cases. Even though our model does not depend on the order of graph nodes, we utilize the dataset version with canonical SMILES provided by Somnath et al. 2021. Besides, we randomly permute graph nodes once SMILES are read and converted to graphs.

We compare RetroBridge with template-based methods GLN (Dai et al. 2019), LocalRetro (Chen & Jung 2021), and GraphRetro (Somnath et al. 2021), and template-free methods MEGAN (Sacha et al. 2021), G2G (Shi et al. 2020), Augmented Transformer Tetko et al. 2020, SCROP (Zheng et al. 2019), Tied Transformer (Tetko et al. 2020), GTA (Seo et al. 2021), DualTF (Sun et al. 2021), Graph2SMILES (Tu & Coley 2022) and Retroformer (Wan et al. 2022). We obtained GLN predictions using the publicly available code and model weights Note that the top-$k$ exact match accuracies differ from the one originally reported values because we deduplicate outputs for our evaluation. and used the latest LocalRetro predictions provided by its authors. Additionally, MEGAN was originally trained and evaluated on random data splits, so we retrained and evaluated it ourselves. Finally, we compare RetroBridge with the state-of-the-art discrete graph diffusion model DiGress Vignac et al. 2022 and a template-free baseline based on a graph transformer architecture (Dwivedi & Bresson 2020; Vignac et al. 2022).

For each input product, we sample 100 reactant sets and report top-$k$ exact match accuracy ($k=1,3,5,10$) which is measured as the proportion of input products for which the method managed to produce the correct set of reactants in its top-$k$ samples. Subsequently, for top-$k$ samples produced for every input product, we run the forward reaction prediction model Molecular Transformer (Schwaller et al. 2019) and report round-trip accuracy and coverage (Schwaller et al. 2020). Round-trip accuracy is the percentage of correctly predicted reactants among all predictions, where reactants are considered correct either if they match the ground truth or if they lead back to the input product. Round-trip coverage, on the other hand, measures if there is at least one correct prediction in the top-$k$ according to the definition above. These metrics reflect the fact that one product can be mapped to multiple different valid sets of reactants, as shown in Figure 1.

2 Neural Network

We use a graph transformer network (Dwivedi & Bresson 2020; Vignac et al. 2022) to approximate the final state of the Markov bridge process. We represent molecules as fully-connected graphs where node features are one-hot encoded atom types (sixteen atom types and additional “dummy” type) and edge features are covalent bond types (three bond types and additional “none” type). Similarly to Vignac et al. 2022 we compute several graph-related node and global features that include number of cycles and spectral graph features. Details on the network architecture, hyperparameters and training process are provided in Appendix A.3.

3 Retrosynthesis Modeling

Here we report top-$k$ and round-trip accuracy for RetroBridge and other state-of-the-art methods on the USPTO-50k test set. Table 1 provides exact match accuracy results, and Table 2 reports round-trip accuracy computed using Molecular Transformer (Schwaller et al. 2019).

For completeness, we compared exact match accuracy results of RetroBridge with both template-free and template-based methods. However, template-free modeling is a more challenging task, and RetroBridge is template-free, therefore we primarily focus on the latter group of methods. As shown in Table 1, RetroBridge performs on par with the state of the art in top-1 accuracy and, most importantly, outperforms other template-free methods in all other top-$k$ accuracy metrics, in particular highly optimized transformer-based models. We emphasize that performance in top-5 accuracy is the most relevant metric as this number is the closest to the typically expected or desired breadth of the multi-step retrosynthesis planning tree (Maziarz et al. 2023). Hence, as shown in Table 3, our model was optimized for top-5 accuracy. Additional baseline methods and more information on their technical aspects are provided in Table 4.

Because exact match accuracy cannot reflect the complete picture of dependencies between spaces of products and reactants, we additionally measure round-trip results using an orthogonal method for forward reaction prediction. We evaluate predictions of the template-based methods GLN (Dai et al. 2019) and LocalRetro (Chen & Jung 2021), and template-free methods Graph2SMILES (Tu & Coley 2022) and Retroformer (Wan et al. 2022) as well as the retrained version of template-free MEGAN (Sacha et al. 2021) ourselves for comparison. The results are reported in Table 2. RetroBridge clearly outperforms the template-free baseline and, in spite of a much higher difficulty of the template-free setup, even achieves higher round-trip coverage and accuracy values than state-of-the-art template-based methods. These results support our hypothesis that retrosynthesis should be modeled in a probabilistic framework considering the absence of a unique set of reactants for a given initial product molecule.

4 Additional Experiments

In this section, we compare RetroBridge with the naïve adaptation of DiGress (Vignac et al. 2022) for retrosynthesis prediction, which can be considered the most comparable diffusion-based method for this task, and with a graph transformer network that predicts reactants in a one-shot fashion. Furthermore, we study the effect of adding an input product molecule as context to the neural network $\varphi_{\theta}$ at each sampling step. More precisely, for models with no context we use the formulation (5), while models with context compute predictions as $\hat{\bm{y}}=\varphi_{\theta}(\bm{z}_{t},\bm{x},t)$. Finally, we try a simpler cross-entropy (CE) loss function, as used in DiGress, that directly compares approximated reactants with the ground-truth,

In all experiments we use the same neural network architectures and sets of hyperparameters. We perform our evaluation on the USPTO-50k validation set and report top-$k$ accuracy of the generated samples in Table 3.

First of all, we observe that iterative sampling as performed by diffusion models or the Markov Bridge Model is essential for solving the problem of mapping between two graph distributions. Our one-shot graph transformer model trained for a comparable amount of time does not manage to recover any of the reactants. Indeed, it is an extremely challenging task as even a single incorrectly predicted bond or atom type is detrimental for the accuracy metric. To the best of our knowledge, all one-shot graph-based models proposed for retrosynthesis prediction fall into the category of template-based methods. In this case, the networks do not predict the entire set of reactants right away, but instead aim to solve much simpler tasks such as prediction of graph edits. To obtain the final set of reactants, graph edit predictions should be further processed by an additional block that typically relies on the predefined reaction templates or leaving group dictionaries.

Next, we compare RetroBridge with DiGress to demonstrate that the Markov bridge formulation is more suitable than diffusion models when two discrete graph distributions are to be mapped. As shown in Table 3, RetroBridge outperforms DiGress with context in all metrics. We note that if we do not pass the context, DiGress predictably does not manage to recover any reactants. This result illustrates that the Markov bridge framework captures the underlying structure of the task much more naturally than diffusion models. A diffusion model maps sampled noise to the reactants having access to the input product molecule only through the additional context while the Markov bridge model starts each sampling trajectory with a product molecule from the intractable distribution $p_{\mathcal{X}}(\bm{x})=\int p_{\mathcal{X},\mathcal{Y}}(\bm{x},\bm{y})d\bm{y}$.

Finally, we demonstrate that the variational lower bound loss (7) works better than a simplified cross-entropy loss (13) proposed by Vignac et al. 2022. Ultimately, we find it beneficial to include the input product as context at each sampling step. Unlike the diffusion approach, however, RetroBridge achieves reasonable accuracy values even without additional context as large parts of the product structure are retained throughout the sampling trajectory.

5 Examples

Figure 3 shows three examples of reactions we randomly selected from the USPTO-50k test set. For each of these examples, we provide top-3 RetroBridge samples and the corresponding confidence scores. In all three cases our model managed to recover the correct set of reactants. In the first case, RetroBridge predicts the correct reactants with a high confidence. In this case, the gap between the scores of the first and the second prediction is remarkably high (0.66 vs 0.12). In both other cases, the model is not as confident in the answer. This uncertainty is reflected in the scores: top-1 samples (which are not correct) have scores 0.18 and 0.38 respectively (cf. 0.17 and 0.2 for the correct ones). More examples are provided in Figure 6.

Conclusion

In this work, we introduce the Markov Bridge Model, a new generative framework for tasks that involve learning dependencies between two intractable discrete-valued distributions. We furthermore apply the new methodology to the retrosynthesis prediction problem, which is an important challenge in medicinal chemistry and drug discovery. Our template-free method, RetroBridge, achieves state-of-the-art results on common evaluation benchmarks. Importantly, our experiments show that choosing a suitable probabilistic modeling framework positively affects the performance on this task compared to the straightforward adaptation of diffusion models.

In order to make RetroBridge more applicable in practice, our future work aims to address several limitations present in RetroBridge and also common to other similar retrosynthesis modeling methods. First, despite its constraints, template-based planning might still be preferred by many chemists as it is more likely to adhere to established sets of reactants, reagents and reaction types. Therefore, guiding the generation process towards a specific set of compounds is a potential improvement that will make our method more useful in practice. Second, chemical reactions are heavily dependent on experimental conditions and additional reagents. Like most related methods, RetroBridge does not directly predict such conditions, nor does it provide explicit information about the required reaction type, but could be adapted towards this setup. While suggesting entire lab protocols seems out-of-reach with current methods, conditioning the generation process on this additional context could already make our method more applicable to real world scenarios. Finally, most pharmaceutically relevant reaction pathways consist of several steps. Therefore, future work will also assess RetroBridge’s performance in a multi-step setting by combining it with existing multi-step planning algorithms.

While this work is focused on the retrosynthesis modeling task, we note that application of Markov Bridge Models is not limited to this problem. The proposed framework can be used in many other settings where two discrete distributions accessible via a finite sample of coupled data points need to be mapped. Such applications include but are not limited to image-to-image translation, inpainting, text translation and design of protein binders. We leave exploration of Markov Bridge Models in the scope of these and other possible challenges for future work.

We thank Max Welling, Philippe Schwaller, Rebecca Neeser, Clément Vignac and Anar Rzayev for helpful feedback and insightful discussions. Ilia Igashov has received funding from the European Union’s Horizon 2020 research and innovation programme under the Marie Skłodowska-Curie grant agreement No 945363. Arne Schneuing is supported by Microsoft Research AI4Science.

References

Appendix A Appendix

The log-likelihood of reactants $\bm{y}\equiv\bm{z}_{T}$ given product $\bm{x}\equiv\bm{z}_{0}$ can be written as follows,

Using Jensen’s inequality (JI) and the fact that $p(\bm{z}_{1:T}|\bm{z}_{0},\bm{z}_{T})=p(\bm{z}_{1:T-1}|\bm{z}_{0},\bm{z}_{T})$ ($*$) we can derive a lower bound of this log-likelihood,

Here, the first term $\mathcal{L}_{1}$ can be written as follows,

Using Bayes’ rule (BR) and the Markov propery (MP) of the Markov bridge $p$, we can derive a similar expression for all intermediate terms $\mathcal{L}_{t}$,

We combine $\mathcal{L}_{1}$ and $\mathcal{L}_{t}$ to obtain the final expression for the variational lower bound of the log-likelihood:

Finally, we obtain the form provided in (7) by replacing the sum over all $T$ terms with its unbiased estimator,

A.2 Cumulative Transition Matrix 𝑸¯t\bar{\bm{Q}}_{t}

A proof by induction. For $t=0$, by definition, $\overline{\bm{Q}}_{0}=\bm{Q}_{0}$ and $\overline{\alpha}_{0}=\alpha_{0}$.

Assume that for $t>0$ Equation 8 holds, i.e. $\overline{\bm{Q}}_{t}=\overline{\alpha}_{t}\bm{I}_{K}+(1-\overline{\alpha}_{t})\bm{y}\bm{1}_{K}^{\top}$ and $\overline{\alpha}_{t}=\prod_{s=0}^{t}\alpha_{s}$.

Note that $\bm{y}\bm{1}_{K}^{\top}\bm{y}\bm{1}_{K}^{\top}=\bm{y}\bm{1}_{K}^{\top}$ and $\alpha_{t+1}\overline{\alpha}_{t}=\overline{\alpha}_{t+1}$. Therefore, we get

A.3 Implementation Details

In all experiments, we use the cosine schedule (Nichol & Dhariwal 2021)

with $s=0.008$ and number of time steps $T=500$.

A.3.2 Additional Features

We represent molecules as fully-connected graphs where node features are one-hot encoded atom types (sixteen atom types and additional “dummy” type) and edge features are covalent bond types (three bond types and additional “none” type). Besides, we use a global graph feature $\bm{y}$ which includes the normalized time step: $\bm{y}=t/T$. Similarly to Vignac et al. 2022 we compute additional node features. For completeness, we provide the description of these features as in (Vignac et al. 2022) below.

Rings of different sizes are crucial features of many bioactive molecules but graph neural networks are unable to detect them (Chen et al. 2020). We therefore add both global cycle counts $\bm{y}_{k}$, that capture the overall number of cycles in the graph, as well as local cycle counts $\bm{H}_{k}$, that measure how many cycles each node belongs to. These quantities are computed for $k$-cycles up to size $k=5$ and $k=6$ for local and global counts, respectively. As proposed by Vignac et al. 2022 We use a corrected equation for $\bm{H}_{5}$., we use the following equations that can be efficiently computed on GPUs

Again following (Vignac et al. 2022), we include spectral features based on the eigenvalues and eigenvectors of the graph Laplacian. We use the multiplicity of eigenvalue 0 and the first five nonzero eigenvalues as graph-level features, and an indicator of the largest connected component (approximated based on the eigenvectors corresponding to zero eigenvalues) as well as two eigenvectors corresponding to the first nonzero eigenvalues as node-level features.

We also tried adding the molecular weight as a graph-level feature and each atom’s valency as node-level features, following (Vignac et al. 2022). Models trained with these features were used everywhere except experiments reported in Tables 1 and 2. Later on, we found out that removing these features improves the performance on the validation set. Therefore, the final model from Tables 1 and 2 does not use molecular features.

A.3.3 Neural Network Architecture

We use a graph transformer network (Dwivedi & Bresson 2020; Vignac et al. 2022) to approximate the final state of the Markov bridge process. The architecture of the network is provided in Figure 4A. First, node, edge and global features are passed through an encoder which is implemented as an MLP. Next, the encoded features are processed by a sequence of Graph Transformer Layers (GTL). As shown in Figure 4B, GTL first updates node features using self-attention and combines its output with edge features via FiLM (Perez et al. 2018):

where $\bm{W}_{1}$ and $\bm{W}_{2}$ are learnable parameters. Then, edge features are updated using attention scores and global features. To update the global features, GTL combines encoded global features and node and edge features aggregated with PNA:

Finally, to obtain the final graph representation, updated node and edge features are passed through the decoder which is implemented as an MLP.

A.3.4 Training

We train our models on a single GPU Tesla V100-PCIE-32GB using AdamW optimizer (Loshchilov & Hutter 2017) with learning rate $0.0002$ and batch size $64$. We trained models for up to 800 epochs (which takes 72 hours) and then selected the best checkpoints based on top-5 accuracy (that was computed on a subset of the USPTO-50k validation set).

A.3.5 Sampling

To establish the dependency between the model performance and the number of sampling steps $T$, we sampled reactants for the validation set (50 per input) using different numbers of steps ${T=500/250/100/50}$ (with the same model trained for $T=500$). As shown in Figure 5A, the results do not significantly change with up to x2 speed improvement. The performance starts degrading dramatically after only 5-fold reduction of the number of time steps.

We additionally analysed the dependency of the performance on the number of samples. The statistical nature of the scoring method suggests that the performance should strongly correlate with the number of samples. However, as shown in Figure 5B, even for 5-fold reduction of the number of samples RetroBridge achieves competitive performance.

Despite the huge success and the wide applicability of diffusion models and similar score-based generative methods, a known bottleneck of such algorithms is sampling. We also cannot avoid iterative sampling and hence the inference procedure requires hundreds of forward passes. Our model reported in Table 1 takes about 1.3 seconds to sample reactants. For comparison, it takes 0.0064 and 0.0065 seconds for MEGAN and GLN respectively. We however note that 1 second for sampling a set of reactants for a given product is a completely feasible time for applications of models like ours. In particular, this speed does not prevent the use of our method as a component of a multistep retrosynthesis planning pipeline.

A.4 Extended Results

In Table 4, we provide an extended version of the exact-match results from Table 1 including additional models and methodological details.

A.5 Forward Reaction Prediction

We additionally trained and evaluated two models for forward reaction prediction using USPTO-50k and USPTO-MIT datasets. We used the same hyperparameters as in other experiments. As shown in Table 5, our models (ForwardBridge) demonstrate comparable performance with other state-of-the-art methods. However, we stress that the probabilistic formulation is less applicable to the forward reaction prediction task, and under certain assumptions this problem can be considered as completely deterministic. Therefore, we leave the study of capabilities of the Markov Bridge Model in the context of forward reaction prediction out of the scope of this work.