RetroBridge: Modeling Retrosynthesis with Markov Bridges
Authors: Ilia Igashov
Arne Schneuing
Marwin Segler
Michael M. Bronstein
Bruno Correia
Venue: ICLR 2024 spotlight (2024)
Keywords: Retrosynthesis
Reactions
Chemistry
Drug Discovery
Markov Bridge
Primary area: applications to physical sciences (physics, chemistry, biology, etc.)

Abstract: Retrosynthesis planning is a fundamental challenge in chemistry which aims at designing multi-step reaction pathways from commercially available starting materials to a target molecule. Each step in multi-step retrosynthesis planning requires accurate prediction of possible precursor molecules given the target molecule and confidence estimates to guide heuristic search algorithms. We model single-step retrosynthesis as a distribution learning problem in a discrete state space. First, we introduce the Markov Bridge Model, a generative framework aimed to approximate the dependency between two intractable discrete distributions accessible via a finite sample of coupled data points. Our framework is based on the concept of a Markov bridge, a Markov process pinned at its endpoints. Unlike diffusion-based methods, our Markov Bridge Model does not need a tractable noise distribution as a sampling proxy and directly operates on the input product molecules as samples from the intractable prior distribution. We then address the retrosynthesis planning problem with our novel framework and introduce RetroBridge, a template-free retrosynthesis modeling approach that achieves state-of-the-art results on standard evaluation benchmarks.

Published as a conference paper at ICLR 2024
RETROBRIDGE: MODELING RETROSYNTHESIS WITH MARKOV BRIDGES
Ilia Igashor Arne Schneuing' Marwin Segler", Michael Bronstein & Bruno Correia* Ecooe Polytechnique Fédérale de Lausanne, "Microsoft Research, University of Oxford {ilia. www.shgggsaaanaoooo arne. achneuing, bruno. cow.......eeaaaaeee....ccoo
ABSTRACT
Retrosynthesis planning is a fundamental challenge in chemistry which aims at de- signing multi-step reaction pathways from commercially available starting mate- rials target molecule. Each step in multi-step retrosynthesis planning requires accurate prediction of possible precursor molecules given the target molecule and confidence estimates to guide heuristic search algorithms. We model single-step retrosynthesis as a distribution learning problem in a discrete state space. First, we introduce the Markov Bridge Model, generative framework aimed approxi- mate the dependency between two intractable discrete distributions accessible via finite sample of coupled data points. Our framework is based on the concept of a Markov bridge, Markov process pinned at its endpoints. Unlike diffusion- methods, our Markov Bridge Model does not need a tractable noise distribution as sampling proxy and directly operates on the input product molecules samples from the intractable prior distribution. We then address the retrosynthesis planning problem with our novel framework and introduce RetroBridge, template- -ree ret- rosynthesis modeling approach that achieves state-of of-the-an art results on standard evaluation benchmarks.
INTRODUCTION
Computational and machine learning methods for de novo drug design show great promise 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 larget 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-Kalthen et al. 2020 Tu et al. 2023) tools ad-
dress this challenge by proposing reaction steps or entire pathways that can be validated and op- timized in the lab. Single-step retrosynthesis models predict precursor molecules for given target molecule (Segler Waller 2017; Coley et al., 20177 Liu et al. 2017 StrithhKallthn 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 al. 2022 we propose to view the task conditional distribution learning problem, as shown in Figure 1 This approach has several advantages, including the ability model uncertainty and to generate new and diverse retrosynthetic pathways. Further- more, 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- Docl--Dickttein et al. 2015 Ho et al. 2020 and other modern score- score-based and flow-based generative methods (Rezende & Mohamed 2015 Song et al. 2020 Lipman al. 2022
"These authors contributed equally

---

Published as a conference paper at ICLR 2024
Products
Reactants
Px
1
t= T
Py(y)
Y
Y2
y
p(z Z1-1, y)
Figure 1: Markov bridges between the distribution of products and distribution of reactants.
Albergo Vanden-Eijnden 2022 Albergo et al., 2023) may seem like good candidates for ret- rosynthesis 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 single intractable data distribu- tion. To do so, one typically samples initial noise from 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, show here that such of the original unconditional generative idea is suboptimal for approximating the dependency between two discrete distributions.
In this work, propose RetroBridge, a template-free probabilistic method for single-step retrosyn- thesis modeling. As shown Figure [ we model the dependency between the spaces of prod- uets and reactants as a stochastic process that is constrained start and to end 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 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 probabiliisic nature of RetroBridge and measure uncertainty at each sample. We demonstrate that RetroBridge achieves compet- itive results standard retrosynthesis modeling benchmarks. Besides, 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 better suited tasks where two intractable discrete distributions need to be mapped. Our source code is available at https: /ttpttttttt/bt/// com/ igasshv /eeeeeetooooo RetroBiidge
To summarise, the main contributions of this work are the following:
We introduce the Markov Bridge Model approximate the probabilistic dependency be- tween two intractable discrete distributions accessible via finite sample of coupled data points.
We demonstrate the superiority of the proposed formulation 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 template-free single-step retrosynthesis prediction method that achieves state-of-ther the--artttt-tttt results on standard benchmarks.
RELATED WORK
Diffusion Models Diffusson 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
2

---

Published as a conference paper at ICLR 2024
results in many different domains including protein design (Watson et al., 2023 small molecule generation (Hoogeboom et al. 20222 Igashov et al. 2022 Schneuing et al. 2022) molecular dock-
ing (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 discrete data (Hoogehoom et al.] [021] Johnson al.] [02T]] Austin
iet al.) 2021] Yang et al. (202)) 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.
Schrödinger Bridge Problem 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-Leible divergence on path spaces. While most recent methods employ the SB formalism in the context of unconditional generative modeling Varrgs et al. 2021 Wang et al., 2021] De Bortoli et al., 2021 Chen et al. [202]] Bunne al.] 2023 Liu et al. 2022). a few works aimed approximate the reference stochastic process through training on coupled samples from two continuous distributions (Holdijk et al., [2022 Som- nath 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.
Retrosynthesis Modeling Recent retrosynthesis prediction methods can be divided into two main groups: remplate-based and template-free methods (Jiang et al.) 2022]. While template-basee meth- ods depend on predefined sets of specific reaction templates or leaving groups, template-fres meth- ods 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. variety of language models have been recently proposed (Liu et al.
2017 Zheng et al. 2019 Sun et al. 2021 Tetko et al.) 2020) operate SMILES, some of which (Jiang et al. 2023 Zhong et al. 2022) were additionally pre-trained on larger datasets. Due the nature of the
translation problem, all these methods are template-free.
Among the existing graph-based methods (Segler & Waller 2017). the most recent template-based
are GLN (Dai et al.] 2019), GraphRetro (Somnath al. T202T) and LocalRetro (Chen & Jung. 2021]. and template-free free approaches are G2G (Shi et al. (2020) and MEGAN (Sacha et al. 2021]
Template-free methods GTA (Seo al.] 2021), Graph2SMILES (Tu & Coley (022) and Retro-
former (Wan et al. 2022) leverage both graph and SMILES representations. In this work, we pro-
pose a novel template-free free graph-bace -bssed method.
3 RETROBRIDGE
We frame the retrosynthesis prediction task as generative problem of modeling a stochastic pro- cess between two discrete-valuet distributions of products px and reactants py. These distributions are intractable and are represented by a finite collection of D coupled samples {<x, ) where 2; ~ (x.) is a product molecule and yı ~ py (y:) is corresponding set of reactant molecules. While products and reactants follow distributions px and py respectively, there is a dependency between these variables that can be expressed in the form of the joint distribution Px y such that (x, y)dz = (3) and J (x, yd = (x)x)) The joint distribution px,y is also in- tractable and accessible only through the discrete sample of coupled data points { y:) P
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 simple but rather effective way of scoring RetroBridge samples based on the statistical uncertainty of the model.
3.1 MARKOV BRIDGE MODEL
We model the dependency between two discrete spaces X and y by a Markov bridge Fitzsimmons et al. 1992, Çetin& & Danilova, 2016), which is Markov process pinned to specific data points in the beginning and in the end. For pair of samples (z, y) ~ (x, y) and sequence of time steps = 0, 1 . T, define the corresponding Markov bridge sequence of random variables
3

---

Published as a conference paper at ICLR 2024
Product a
2(2 F..1.
Cat (z))
Reactant y
t= 0
Figure 2: The process of changing atom types along the trajectory of the Markov bridge. The trajectory starts at time step = O with the product molecule and several disconnected "dummy" atoms that will be included in the final reactant molecule. The probability of sampling the target atom type increases as grows. Five circles filled with different colors represent these probabilities. To make the illustration less bulky, we omitted part of the product molecule and of two reactant molecules.
(z that starts at æ, i.e., 20 æ, and satisfies the Markov property, p(z, ko, ......, 21-1, y) 2 3). To pin the process at the data point we introduce an additional requirement,
(1)
p(zr 2 1.
Assuming that both distributions Px and Py are categorical with finite sample space {1..... K}, we can represent data points as K-dimensional one-hot vectors: æ, y. E RRK To model a Markov bridge defined by equations (1 2 similar to Austin al. (202T), we introduce sequence of transi-
tion matrices Qo Q1..... Qr-1 E x defined
Q. = Qx = a, I + (1 aayyl
where I is a K K identity matrix, 1 is K -ddimensional all-one vector, and ar is schedule parameter transitioning from ag = to ar- 0. Transition probabilities can be written as follows,
p Lc. 8) = Cat (
(4)
where Cat(- :p) is a categorical distribution with probabilitics given by p. We note that setting QT-1 = 0 ensures the requirement (
Using the finite set of coupled samples {(x,, v:) ~ px.y, our goal is to learn a Markov bridge (12) to be able to sample y when only is available. To do this, replace y with approximation computed with a neural network (og:
Ú yo (2(
(5)
and define an approximated transition kernel,
(0(2++12z) Cat (z++1: (g))).))
(6)
We train So by maximizing lower bound of log-likelihood log go (g(la))) As shown in Appendix A.1 it has the following closed-form expression,
log qo (z) -T. E E, Ex p(ss....)))
DRi. (p(z44 I y) (2
For any x E X,y E y, and = 1,...,,, sampling of 2, can be effectively performed using a cumulative product matrix Q; = Q, can be written closed form,
As shown in Appendix A.2 the cumulative matrix
Q, =
T = + (1- (1-
where a. -T'= ag.. Therefore, p(z |20, x can be written follows,
(9) To sample a data point y = 2y starting from a given 2n = x ~ px(x), one iteratively predicts = <po(2,, t) and then derives z++1 (0(2++1l21) Cat (z141 Q(())z) for t 0..... 1. Training and sampling procedures the Markov Bridge Model are provided in Algorithms and respectively.
p(z |2, 2r) Cat (z++1: Q,20)
4

---

Published as a conference paper at ICLR 2024
Algorithm Training of the Markov Bridge Model
Input: coupled sample (æ, y) ~ px,y, neural network 4p t t~ -O.........
1), 24 Cat (z)) Q.-1
> Sample time step and intermediate state Output of ypg is vector of probabilities
ÿ <- (0((1)))
p(2++1/21, V) Cat (2(+1: (v))) (2(++1| + Cat (Z:41; Q,())z))
b Reference transition distribution Approximated transition distribution
Minimize DLL (p(=++1|2) y) g0 (20 (z)+1
Algorithm Sampling
Input: starting point a px. neural network y0
20 +
for l in O..., T -1:
ÿ t p0(2,,t)
Output of (p is vector of probabilities Approximated transition distribution
(0 (2)) t Cat (z(+1: ())))z)
2,41 2o (z (z,)
Return 2T
3.2 RETROBRIDGE: MARKOV BRIDGE MODEL FOR RETROSYNTHESIS PLANNING
In setup, each data point is molecular graph with nodes representing atoms and edges responding to covalent bonds. We represent the molccular graph with a matrix of node features H R x.K, which are, for instance, one-hot encoded atom types, and tensor of edge features E E R No.KK, which can be one-hot encoded bond types.
In the scope of our probabilistic framework, we consider such a molecular graph representation collection of independent categorical random variables. More formally, denote product and reactants data points x and as tuples of the corresponding node and edge feature tensors: H Ez] and y HHy Ey)- For such complex data points, modify the definitions of transition matrices and probabilities accordingly:
( = JK, + - a1) H [Q]1ia Ik, + (1 ox) 1 Eylj,k
P HH,1 H,)) = Cat (H++1; H.Q.")
(10)
E Er, E,y) Cat (Ee+1: EE++))
(11)
where [H] E IR IxI, is the j-th row of the feature matrix H (i.e. transposed feature vector of the j-th node), and E E glxk, is the transposed feature vector of the edge between j-th and k-th nodes.
Because atoms present in the reactant molecules 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 transformed into atoms of reactant molecules. In experiments, always add 10 "dummy" nodes to the initial product graphs.
3.3 CONFIDENCE AND SCORING
It is important to have reliable scoring method that selects the most relevant sets of reactants 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
scoring function. We
estimate the confidence of the model by computing the likelihood qe (y)z) of a set of reactants
for input product molecule æ. For of M samples (=(9)).) LM generated by RetroBridge for an input product x, we compute the likelihood- confidence score for the set of reactants as follows.
1 M
go (y)2x) = Eyy'ro 1xy = y) ~
Σ
1[2 = =y}. t=1
(12)
M
5

---

Published as a conference paper at ICLR 2024
4 RESULTS
4.1 EXPERIMENTAL SETUP
Dataset 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/validationntest splits provided by Dai et al. 2019). Sommath et al. (202T) report that the dataset contains a shortcut in that the product atom with atom-mapping 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 al. (202T). Besides, randomly permute graph nodes once SMILES are read and converted to graphs.
Baselines We compare RetroBridge with template-bas based methods GLN (Dai et al. 2019), Local-
Retro (Chen & Jung
2021), and GraphRetro (Somnath et al. 2021), and template frree methods
MEGAN (Sacha et al.
2021). G2G (Shi et al. 2020), Augmented Transformer Tetko et al. (2020),
2019h, Tied Transformer (Tetko al. 2020), GTA (Seo et al.] 202 1) Du- alTF (Sun et al. 2021). J. Graph2SMILES (Tu & Coley][202) and Retroformer (Wan et al.) 2022).
SCROP (Zheng et al.
We obtained GLN predictions using the publicly available code and model weight: and used the latest LecalRetro 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 slate-of-the-ar 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, 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 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 imput product. Round-trip coverage, on the other hand, measures there is at least one correct prediction in the top-k according 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] [
Evaluation
4.2 NEURAL NETWORK
We use a graph transformer network (Dwivedi & Bresson 2020 Vignac et al. 2022) to approxi- mate the final state of the Markov bridge process. We represent molecules as culnecco 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). Sim- ilarly 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
4.3 RETROSYNTHESS 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 provides exact match accuracy results, and Table 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-basse methods. However, template-free modeling is more challenging task, and RetroBridge is template-free, therefore we primarily focus on the latter group of methods. As shown in Table | RetroBridge performs on with the state of the art in top-1 1 accuracy and, most im- portantly, outperforms other template-free methods in all other top-k accuracy metrics, in particular
Note that the top-k exact match accuracies differ from the one originally reported values because deduplicate outputs for our evaluation.

---

Published as a conference paper at ICLR 2024
Table 1: Top-k accuracy (exact match) on the USPTO-50K test dataset. The best performing template-based (TB) and template-free (TF) methods are highlighted in bold.
Model
k
k = 3
k = 5
k = 10
GLN (Dai et al.
2019)
52.5
74.7
81.2
87.9
GraphRetro Somnath al. 2021
53.7
68.3
72.2
75.5
LocalRetro (Chen & Jung 2021
52.6
76.0
84.4
90.6
SCROP (Zheng al. 2019)
43.7
60.0
65.2
68.7
G2G (Shi et al. 2020)
48.9
67.6
72.5
75.5
Aug. Transtormer TTtko al.
2020
48.3
73.4
77.4
DualTF aug (Sun et al. 2021)
53.6
70.7
74.6
77.0
.MEAAN (Sacha al. 2021
48.0
70.9
78.1
85,4
Tied Transtormer Kim et a.
2021
47.1
67.1
73.1
76.3
GTA (Seo al.
2021
51.1
67.6
74.8
81.6
Graph2SMILESS [T & Coley
2022
52.9
66.5
70.0
72.9
Retroformer (Wan et al. 2022)
52.9
68.2
72.5
76.4
RetroBridge (ours)
50.8
74.1
80.6
85.6
Table Top-k round-trip coverage and accuracy on the USPTO-50k test dataset. The best perform- ing template-based (TB) and template-free (TF) methods highlighted in bold.
Coverage
Accuracy
Model
k
1
k
h
k
k 5
2019 LocalRetro (Chen & Jung 2021)
GLN (Dai
82.5
92.0
94.0
82.5
71.0
66.2
TB
82.1
92.3
94.7
82.1
71.0
66.7
MEGAN (Sacha al. 202
78.1
88.6
91.3
78.1
67.3
61.7
E
Graph2SMILES Iu Coley 2022)
76.7
56.0
46.4
Retroformer.ng (Wan et al. 2022)
78.6
71.8
67.1
RetroBridge (ours)
85.1
95.7
97.1
85.1
73.6
67.8
highly optimized transformer-Be based models. We emphasize that performance in top-5 accuracy is the most relevant metric this number is the closest to the typically expected or desired breadth of the multi-step retrosynthesis planning tree (Maziarz et al. (203).. Hence, as shown in Table 3 model was optimized for top-5 accuracy. Additional baseline methods and more information their technical aspects are provided in Table
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-bass
methods GLN DDai
et al. 2019) and LocalRetro (Chen & Jung. 2021) , and template-free frree methods Graph2SMILES, (Tu & Coley| (2022) and Retroformer (Wan et 2022) as well as the retrained version of template- free MEGAN (Sacha et al. (021) 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 free setup, even achieves higher round-trip coverage and accuracy values than state- of-the-an art template-bas based methods. These results support our hypothesis that retrosynthesis should be modeled in probabilistic framework considering the absence of unique set of reactants for given initial product molecule.
4.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 considered the most comparable diffusion bused method for this task, and with 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 net- work (o each sampling step. More precisely, for models with no context the formulation

---

Published as a conference paper at ICLR 2024
Table 3: Additional experiments: top-k accuracy on the USPTO-50k validation set.
Model
k ==1
== 3
k ==
k 10
k: 50
DiGress (context)
47.32
68.56
73.93
78.45
80.88
RetroBridge-CC CE (no context)
48.71
66.84
72.33
76.08
79.38
RetroBridge- CE (context)
50.74
71.50
76.58
79.50
80.58
RetroBridge-VLL VLB (no context)
47.42
69.46
75.21
79.40
83.82
RetroBridge- VLB (context)
48.92
73.04
79.44
83.74
86.31
5). while models with context compute predictions as ÿ (0 2,t). Finally, we try a simpler cross-entropy (CE) loss function, as used in DiGress, that directly compares approximated reactants with the ground-truth,
CCr (0) =T T- Crosstttop (zr, (z),
(13)
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
First of all, we observe that iterative sampling 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 comparable amount of time does not manage to recover any of the reactants. Indeed, it is an extremely challenging task as even single incorrectly predicted bond atom type is detrimental for the accuracy metric. To the best of our knowledge, all one-shot graph-base 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 additional block that typically relies on the predefined reaction templates or leaving group dictionaries.
Next, 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 all metrics. We note that 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 product molecule from the intractable distribu- tion (x) = Px,y(x, y)dy.
Finally, we demonstrate that the variational lower bound loss ( works better than simplified cross- entropy loss (3) proposed by Vignac et al. 2022).. Ultimately. we find it beneficial to include the input product 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 retained throughout the sampling trajectory.
4.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. the first case, RetroBridge predicts the correct reactants with 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).. both other cases, the model is not as confident in the answer. This uncertainty is reflected in the scores: top-1 1 samples (which are not correct) have scores 0.18 and 0.38 respectively (cf. 17 and 0.2 for the correct ones). More examples are provided in Figure
8

---

Published as a conference paper at ICLR 2024
VV
Hone VV
Figure 3: Examples of modeled reactants. We selected three random inputs from the test set and for each them we provide the top-3 RetroBridge predictions along with their confidence scores. Two check marks indicate that sampled reactants are the same as the ground truth, and one check mark means that reactants are different, but Molecular Transformer (Schwaller et al. 2019) predicts the product molecule used as input.
CONCLUSION
In this work, we introduce the Markov Bridge Model, a new generative framework for tasks that in- volve learning dependencies between two intractable discrete-valued distributions. We furthermore apply the new methodology to the retrosynthesis prediction problem, which is an important chal- lenge in medicinal chemistry and drug discovery. Our template-fiee frree method, RetroBridge, achieves state-of-the-ar results on common evaluation benchmarks. Importantly, our experiments show that choosing suitable probabilistic modeling framework positively affects the performance this task compared to the straightforward adaptation of diffusion models.
In order to make RetroBridge more applicable in practice, our future work aims address several limitations present in RetroBridge and also common to other similar retrosynthesis modeling meth- ods. 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 specific 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 re- action type, but could be adapted towards this setup. While suggesting entire lab protocols seems out-of-reach with current methods, conditioning the generation process this additional context could already make our method more applicable to real world scenarios. Finally, most pharmaceu- tically relevant reaction pathways consist of several steps. Therefore, future work will also assess RetroBridge's performance 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 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.
ACKNOWLEDGMENTY
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 's Horizon 2020 research and innovation programme under the Marie Sklodowska- grant agreement No 945363. Ame Schneuing is supported by Microsoft Research AI4Science.

---

Published as a conference paper at ICLR 2024
REFERENCES
Michael S Albergo and Eric Vanden-Fijnden. Building normalizing flows with stochastic inter- polants. arXiv preprint arXiv:220.//557/ 2022.
Michael S Albergo, Nicholas M Boffi, and Eric Vanden-Eijndedn Stochastic interpolants: A unifying framework for flows and diffusions. arXiv preprint 230.::3733008777, 2023.
Jacob Austin, Daniel D Johnson, Jomathan Ho, Daniel Tarlow, and Rianne van den Berg. Structured denoising diffusion models in discrete state-spaces. Advances in Neural Information Processing Systems, 34:1 14:17991-17993, 2021.
Charlotte Bunne, Ya-Ping Hsich, Marco Cuturi, and Andreas Krause. The schrödinger bridge be- tween gaussian measures has closed form. In International Conference Artificial Intelligence and Statistics, pp. 5802-5833. PMLR, 2023.
Unut Çetin and Albina Danilova. Markov bridges: Sde representation. Stochastic Processes and their Applications, 126(3):651-679) 2016.
Shuan Chen and Yousung Jung. Deep retrosynthetic reaction prediction using local reactivity and global attention. JACS Au, 1(10):6112-120, 2021.
Tianrong Chen, Guan-Horng Liu, and Evangelos A Theodorou. Likelihood training of schr\ odinger bridge using forward-backward sdes theory. arXiv preprint arXiv:2110. /29/, 2021.
Zhengdao Chen, Lei Chen, Soledad Villar, and Joan Bruna. Can graph neural networks count substructures? Advances in neural information processing systems, 33:10383-10399 10395, 2020.
Connor W Coley, Luke Rogers, William H Green, and Klavs F Jensen. Computer-assister retrosyn- thesis based on molecular similarity. ACS central science, 3(12):1237-1245, 2017.
Elias James Corey. The logic of chemical synthesis. 1991.
Gabriele Corso, Hannes Stärk, Bowen Jing, Regina Barzilay, and Tommi Jaakkola. Diffdock: Dif- fusion steps, twists, and turns for molecular docking, arXiv preprint arXlv:2210.01776, 2022.
Hanjun Dai, Chengtao Li, Connor Coley, Bo Dai, and Le Song. Retrosynthesis prediction with conditional graph logic network. Advances in Neural Information Processing Systems, 32, 2019.
Valentin De Bortoli, James Thornton, Jeremy Heng, and Arnaud Doucet. Diffusion schrödinger bridge with applications to score-based generative modeling. Advances in Neural Information Processing Systems, 14955-9-177 2021.
Chenru Duan, Yuanqi Du Haojun Jia, and Heather J Kulik. Accurate transition state gener- ation with an object-aware equivariant elementary reaction diffusion model. arXiv preprint arXiv:2304.06114, 2023.
Vijay Prakash Dwivedi and Xavier Bresson. generalization of transformer networks graphs. arXiv preprint arXi::2012.09699 2020.
Pat Fitzsimmons, Jim Pitman, and Mare Yor. Markovian bridges: construction, palm interpretation, and splicing. Seminar on Stochastic Processes, 1992, pp. 101-134. Springer, 1992.
Jonathan Ho, Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840 6851, 2020.
Lars Holdijk, Yuanqi Du, Ferry Hooft, Priyank Jaini, Bernd Ensing, and Max Welling, Path integral stochastic optimal control for sampling transition paths. arXiv preprint arXiv:2207.02149) 2022.
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, and Max Welling, Argmax flows and multinomial diffusion: Learning categorical distributions. Advances in Neural Information Processing Systems, 14:12254-124666 2021.
Emiel Hoogeboom, Victor Garcia Satorras, Clément Vignac, and Max Welling. Equivariant diffu- sion for molecule generation in 3d. In International conference on machine learning, pp. 8867- 8887. PMLR, 2022.
10

---

Published as a conference paper at ICLR 2024
Ilia Igashov, Hannes Stärk, Clément Vignac, Victor Garcia Satorras, Pascal Frossard, Max Welling, Michael Bronstein, and Bruno Correia. Equivariant 3d-conditional diffusion models for molecular linker design. arXiv preprint arXiv:2210.05274. 2022.
Yinjie Jiang, Yemin Yu, Ming Kong, Yu Mei, Luotian Yuan, Zhengxing Huang, Kun Kuang, Zhi- hua Wang, Huaxiu Yao, James Zou, et al. Artificial intelligence for retrosynthesis prediction. Engineering, 2022.
Yinjie Jiang, WEI Ying, Fei Wu, Zhengxing Huang, Kun Kuang, and Zhihua Wang. Learning chemical rules of retrosynthesis with pre-training. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pp. 5113-5121, 2023.
Daniel D Johnson, Jacob Austin, Rianne van den Berg, and Daniel Tarlow. Beyond in-place comrp-- tion: Insertion and deletion in denoising probabilistic models. arXiv preprint arXiv:2107.0767. 2021.
Eunji Kim, Dongseon Lee, Youngchun Kwon, Min Sik Park, and Youn-Suk Choi. Valid, plausi- ble, and diverse retrosynthesis using tied two-way transformers with latent variables. Journal of Chemical Information and Modeling, 61(1):123-133, 2021.
Seonghwan Kim, Jeheon Woo, and Woo Youn Kim. Diffusion-basse generative ai for exploring transition states from 2d molecular graphs. 2023.
Christian Léonard. survey of the schr\ " odinger problem and some its connections with optimal transport. arXiv preprint
2013.
Yaron Lipman, Ricky TQ Chen, Heli Ben-Hamu, Maximilian Nickel, and Matt Le. Flow matching for generative modeling. arXiv preprint arXiv:2210.02747, 2022.
Bowen Liu, Bharath Ramsundar, Prasad Kawthekar, Jade Shi, Joseph Gomes, Quang Luu Nguyen, Stephen Ho, Jack Sloane, Paul Wender, and Vijay Pande. Retrosynthetic reaction prediction using neural loqquence-to-sequeee models. ACS central science, ((110):100-111)) 2017.
Guan-Homg Liu, Tianrong Chen, Oswin So. and Evangelos Theodorou. Deep generalized schrödinger bridge. Advances in Neural Information Processing Systems, 35:9374-9388, 2022.
Ilya Loshchilov and Frank Hutter. Decoupled weight decay regularization. arXiv preprint arXiv:1711.05101 2017.
Krzysztof Maziarz, Austin Tripp. Guoging Liu, Megan Stanley, Shufang Xie, Piotr Gaiński, Philipp Seidl, and Marwin Segler. Re-evaluating retrosynthesis algorithms with syntheseus. arXiv preprint arXiv:2310, 19796, 2023.
Alexander Quinn Nichol and Prafulla Dhariwal. Improved denoising diffusion probabilistic models. In International Conference on Machine Learning, pp.
PMLR, 2021.
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. Film: Visual reasoning with general conditioning layer. In Proceedings of the AAAI conference on artificial intelligence, volume 32, 2018
Danilo Rezende and Shakir Mohamed. Variational inference with normalizing flows. In Interna- tional conference on machine learning, pp. 1530-1538. PMLR, 2015.
Mikolaj Sacha, Mikolaj Blaz, Piotr Byrski, Pawet BabrowskiiTTumanse Mikolaj Chrominski, Rafal Loska, Pawet Wlodarczyk-Pruszznnce and Stanistaw Jastrzebski. Molecule edit graph attention network: modeling chemical reactions as sequences of graph edits. Journal of Chemical Infor- mation and Modeling, 6177)3323--2284, 2021.
Nadine Schneider, Nikolaus Stiefl, and Gregory A Landrum. What's what: The (nearly) definitive guide to reaction role assignment. Journal of chemical information and modeling, 56(1 56(12):2336- 2346, 2016.
Arne Schneuing, Yuanqi Du, Charles Harris, Arian Jamasb, Ilia Igashov, Weitao Du, Tom Blun- dell. Pietro Lió, Carla Gomes, Max Welling, et al. Structure- drug design with equivariant diffusion models. arXiv preprint arXiv:22 10. 13695, 2022.

---

Published as a conference paper at ICLR 2024
Erwin Schrödinger. Sur la théorie relativiste de l'électron et l'interprétation de la mécanique quan- tique. In Annales de l'institut Henri Poincaré, volume 2, pp. 269 -100, 1932.
Philippe Schwaller, Teodoro Laino, Théophile Gaudin, Peter Bolgar, Christopher A Hunter, Costas Bekas, and Alpha A Lee. Molecular transformer: model for envertaintt-calibbatte chemical reaction prediction. ACS central science, 5(9):1572-1583) 2019.
Philippe Schwaller, Riccardo Petraglia, Valerio Zullo, Vishnu H Nair, Rico Andreas Haeuselmann, Riccardo Pisoni, Costas Bekas, Anna Iuliano, and Teodoro Laino. Predicting retrosynthetic path- ways using transformer-basee models and hyper-graph exploration strategy. Chemical science, ((112):331-33222 2020.
Marwin HS Segler and Mark P Waller. Neural-symbolic machine learning for retrosynthesis and reaction prediction. Chemistry-A European Journal, 23(25):5966 5971, 2017.
Marwin HS Segler, Mike Preuss, and Mark P Waller. Planning chemical syntheses with deep neural networks and symbolic ai. Nature, 555(7698):604 610, 2018
Seung-Wov Woo Seo, You Young Song, June Yong Yang, Seohui Bae, Hankook Lee, Jinwoo Shin, Sung Ju Hwang, and Eunho Yang. Gta: Graph truncated attention for retrosynthesis. In Pro- ceedings of the AAAI Conference on Artificial Intelligence, volume 35, pp. 531-539, 2021.
Chence Shi, Minkai Xu, Hongyu Guo, Ming Zhang, and Jian Tang. A graph to graphs framework for retrosynthesis prediction. In International conference on machine learning, pp. 8818-8827. PMLR, 2020.
Jascha Dohh-Diccstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli. Deep unsupervised learning using nonequilibrium thermodynamics. In International conference on machine learn- ing. pp. 2256-2263 PMLR, 2015.
Vignesh Ram Somnath, Charlotte Bunne, Connor Coley, Andreas Krause, and Regina Barzilay. Learning graph models for retrosynthesis prediction. Advances in Neural Information Processing Systems, 34:9405-9411 2021.
Vignesh Ram Somnath, Matteo Pariset, Ya-Ping Hsieh, Maria Rodriguez Martinez, Andreas Krause, and Charlotte Bunne. Aligned diffusion schr\." "ddnnge bridges. arXiv preprint arXiv:2302 1419, 2023.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint arXiv:2011.13456. 2020.
Megan Stanley and Marwin Segler. Fake until you make it? generative de novo design and virtual screening of synthesizable molecules. Current Opinion in Structural Biology, 82:102651 2023,
Felix
Frederik Sandfort, Marwin HS Segler, and Frank Glorius. Machine leam- ing the ropes: principles, applications and directions in synthetic chemistry. Chemical Society Reviews, 49(17):6154 6168, 2020.
Ruoxi Sun, Hanjun Dai, Li Li, Steven Kearnes, and Bo Dai. Towards understanding retrosynthesis by energy-based models. Advances in Neural Information Processing Systems, 14:10166-10194, 2021.
Igor V Tetko, Pavel Karpov, Ruud Van Deursen, and Guillaume Godin. State-of-the-art augmented nlp transformer models for direct and single-step retrosynthesis. Nature communications, 11(1): 5575, 2020.
Morgan Thomas. Andreas Bender, and Chris de Graaf. Integrating structure-based approaches in generative molecular design. Current Opinion in Structural Biology, 79:102550 2023.
Zhengkai Tu and Connor W Coley. Permutation invariant graph-to- traph-to-sequence model for template- free retrosynthesis and reaction prediction. Journal of chemical information and modeling, 62 (15):3503-3513, 2022.
12

---

Published as a conference paper at ICLR 2024
Zhengkai Tu, Thijs Stuyver, and Connor W Coley. Predictive chemistry: machine learning for reaction deployment, reaction development, and reaction discovery. Chemical Science, 2023.
Francisco Vargas, Pierre Thodoroff, Austen Lamacraft, and Neil Lawrence. Solving schrödinger bridges via maximum likelihood. Entropy, 23(9): 1344, 2021.
Clement Vignae, Igor Krawezuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pas- cal Frossard. Digress: Discrete denoising diffusion for graph generation. arXiv preprint arXiv:2209. 14734, 2022.
Yue Wan, Chang-Ya Yu Hsich, Ben Liao, and Shengyu Zhang. Retroformer: Pushing the limits of end-to-end retrosynthesis transformer. In International Conference Machine Learning, pp- 12475-22490. PMLR, 2022.
Gefei Wang, Yuling Jiao, Qian Xu, Yang Wang, and Can Yang. Deep generative learning via schrödinger bridge. In International Conference on Machine Learning, pp. 10794-10804, PMLR, 2021.
Joseph I. Watson, David Juergens, Nathaniel R Bennett, Brian I. Trippe, Jason Yim, Helen E Eise- nach, Woody Ahern, Andrew Borst, Robert J Ragotte, Lukas F Milles, al. De novo design of protein structure and function with rfdiffusion. Nature, pp. 1-3, 2023.
David Weininger. Smiles, chemical language and information system. 1. introduction method- ology and encoding rules. Journal of chemical information and computer sciences, 28(1):31-36, 1988.
Dongchao Yang. Jianwei Yu, Helin Wang, Wen Wang, Chao Weng, Yuexian Zou, and Dong Yu. Diffsound: Discrete diffusion model for text-to-sound generation. IEEE/ACM Transactions on Audio, Speech, and Language Processing, 2023.
Shuangjia Zheng, Jiahua Rao, Zhongyue Zhang, Jun Xu, and Yuedong Yang. Predicting retrosyn- thetic reactions using self-corrected transformer neural networks, Journal of chemical information and modeling, 60(1):47-555 2019.
Zipeng Zhong, Jie Song, Zunlei Feng, Tiantao Liu, Lingxiang Jia, Shaolun Yao, Min Wu, Tingjun Hou, and Mingli Song. Root-aligned smiles: tight representation for chemical reaction predic- tion. Chemical Science, 13(31):9023-9033 2022.
13

---

Published as a conference paper at ICLR 2024
A APPENDIX
A.1 VARIATIONAL. LOWER BOUND
The log- logliielihood of reactants y = zr given product 20 can be written as follows,
log (0())z) log go (20)
(14)
log
dzi:1 1 190 (11:1 (zu)
(15)
log / dz
90 (Ze (x1-1)
(16)
1=1
/ d
p(Z1:T]
20: 2r)
-log|
2o 2 (2-1)
(17)
p((21:1 Zo: Z1)
|z0, ØT) / dEp-T-I((1111
T
90 (Z+ z
log
II
(18)
p(zzl 2 ZT)
t=1
Using Jensen's 's inequality (JI) and the fact that p( (20 a lower bound of this log-likelihood,
p(21:T -1| |20 zT) (*) we can derive
log go (y/a)) 2 f d dz1:1-1((21:1 (20 24) log
n
2o (z) (zt-1)
(19)
p(zz # (r)
1=1
/
T
90 (2 (2/2--1)
dde:r-ip(eirr-lla 27) log
(20)
p( z Z1
[=1
Qa ga(zz|ze Zt-1)
d21:T-1P(EIIT-11 2r) log
(21)
p(z: 21-1, 2T)
T
Lz 2T)
+
Σ Ce(zu, 2r).
(22)
Here, the first term L1 can be written as follows,
Go (z( (2n) (1(20, = / dzıp(zzlzo, xT) log p(2 ZT)
-DKL (p(z1|z0, 2T) (100(11
(23)
Using Bayes' rule (BR) and the Markov propery (MP) of the Markov bridge p. we can derive similar expression for all intermediate terms Lt,
g0(z1|21-)
(z 2T) =
dz dz 2/ zr) log
(24)
p( p(z) 11 ZT)
yo (Z
BR
dz, 12x-1dzp((1-1120000 2 2T) log
(25)
p( z 12 Z1
MP
/ dz 12 ( log
9o (2+ (2-11)
(26)
p(z) z zT)
/ d21 1120 zT)
(p(z, 2 2T) (
(z. (z)
(27)
We combine Ci and L, to obtain the final expression for the variational lower bound of the log- likelihood:
T
log ())()) 2
E 1<p(z; y) DKL
(28)
T-1
Σ
E Bz u) DKL (20
(29)
14

---

Published as a conference paper at ICLR 2024
Finally, obtain the form provided in 7 by replacing the sum all T terms with its unbiased estimator,
log (3([)) -T T-- DKL (ww (p( (0110(
(30)
A.2 CUMULATIVE TRANSITION MATRIX Q:
proof by induction. For 0, by definition, Qo Qo and [ Qg-
Assume that for > 0 Equation 8 holds, i.e. Q. = I + +1 - o. and ā. III 0s.
Then for t + 1 we have
Q = Qi+1Q;
(31)
[actIIK + (1 o141) T + (1-0)) y
(32)
a1410 k + (0++1 a+ + a at 10t
(33)
+ (1 0++1 - at 1
(34)
Note that y k = =ylk and 0+41 Je+1- Therefore, get
Q+11 =T:+1 Ix + (as+1 ae+1 + a++1 + - art1 a. + ac+1)11k = -e+t Ik + (1 - -1+1)y1k
(35)
(36)
A.3 IMPLEMENTATION DETAILS
A.3.1 NOISE SCHEDULE
In experiments, we use the cosine schedule (Nichol Dhariwal 2021
t/T +
0.5m
(37)
1+
with 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 covalent bond types (three bond types and additional "none" type). Besides, we use global graph feature which includes the normalized time step: y L/T.. Similarly to Vignac al. (2022) compute additional node features. For completeness, we provide the description of these features as in (Vignac et al. 2022 below.
Cycles Rings different sizes are crucial features many bioactive molecules but graph neural networks are unable to detect them (Chen et al., 2020). We therefore add both global cycle counts yk. that capture the overall number of cycles in the graph, as well as local cycle counts Hy, that measure how many eycles each node belongs to. These quantities computed for k-eveles to
size l = 5 and = for local and global counts, respectively. As proposed by Vignac et al. 2022
We corrected equation for H3.
15

---

Published as a conference paper at ICLR 2024
use the following equations that can be efficiently computed on GPUs
H3 diag(A*
H1 (diag( (A*) @ (d-1,,)) Ad)
H, (diag( (A)) - -TTd - 12d dddgia( diag(A²) ddiagAA)
13
H3 1n
94 H 1n
y5 H₁ 1n
+ 9 @ Al) - 6 6diag (di) diag(A A²) y6 = Tr (T - 3 Tr (A @ +
+ ('T - - (T((A')) 4'TT & A²) + 3 | 12 Tr @ A²) + (Tr
where A E = is the adjacency matrix and d E R" the node degree vector. Matrix T A ④ A indicates how many triangles each nodes shares with every other node.
Spectral features Again following (Vignac al. 2022), we include spectral features based the eigenvalues and eigenvectors of the graph Laplacian. We use the multiplicity of eigenvalue 0 and the first five nonzero eigenvalues graph-level features, and indicator of the largest connected component (approximated based on the eigenvectors corresponding to zero eigenvalues) as well as two eigenvectors corresponding the first nonzero cigenvalues as node-level features.
Molecular 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 I and 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 molecular features.
A.3.3 NEURAL NETWORK ARCHITECTURE
We use graph transformer network (Dwivedi & Bresson 2020 Vignac et al., (022) to approximate final state of the Markov bridge process. The architecture the network is provided in Figure AA. First, node, edge and global features are passed through 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 attentin and combines its output with edge features via FiLM (Perez et al. 2018):
X X²) X, W, (X, X₂ W2.
(38)
where W, and W₂ 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:
PNA(X) concat(max(X) min(minnn (X), mean X), (tdd))))s st))))
(39)
where W is learnable parameter.
Finally, 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 single GPU Tesla V100-PCIE-3200 using Adam W 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 subset of the USPTO-50k validation set).
16

---

Published as a conference paper at ICLR 2024
A
Neural Network Architecture
H
Graph Transformer Layer
Graph Transformer Layer
H
E
Encoder
Decoder
E
B
Graph Transformer Layer
E
#heads
Flatten heads
FILM
Linear
HH'
Linear
FILM
Outer product
Softmax
staling
H
Σ
Faatton heads
FILM
Linear
E
Linear
H
y
E
PNA
Σ
Linear
y
y
Linear
Figure 4: Architecture of the network that approximates the final state of the Markov bridge pro- cess (A) and scheme of the Graph Transformer Layer (B).
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 (50 per input) using different numbers of steps T' 2200//20/10/ 100/50 (with the same model trained for T 500). As shown in Figure 5A, the results do not significantly change with up speed improvement. The performance starts de- grading dramatically after only 5-fold reduction of the number of time steps.
We additionally analysed the dependency of the performance the number of samples. The statis- tical 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, 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 [ 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 second for sampling set of reactants for given product is 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 multistep retrosynthesis planning pipeline.
17

---

Published as a conference paper at ICLR 2024
A
Top-k accuracy different sampling steps
B
Top-k accuracy different samples
0.85
0.7
0.00
0.75
03
top-1
top-3
0.85
top-5
0.3
0.80
top-1
0.55
top-3
lop-5
0.50
top-10
0.0
Number staps
Number of samples N
Figure 5: (A) Dependency of the RetroBridge performance on the number of sampling steps. Exper- iment performed on the RetroBrdige with CE loss the validation set with 50 samples per product. (B) Dependency of the RetroBridge performance on the number of samples. Experiment performed on the RetroBrdige with VLB loss on the test with 500 sampling steps.
Table 4: Top-k accuracy (exact match) on the USPTO-50k test dataset and classification of meth- ods. TB: template-ber based; TF: template-free; G: graph-based S: SMILES- FP: molecular fingerprints; PT: pre- pre-trained on larger datasets. Horizontal lines separate methods that had access to significantly more training data (pre-trained) or operate on a limited vocabulary (template-hasedl The best performing methods in each category are highlighted in bold.
Top-k accuracy
Classification
Model
k:
1
3
k
h 10
Templ.
Data
PT
RSMILES Zhong et al. 2022
56.3
79.2
86.2
91.0
TF
S
PMSR (Jiang et
62.0
78.4
82.9
86.8
TF
S
/
Retrosym Coley al. (2017
37.3
54.7
63.3
74.1
TB
FP
X
GLN (Dai el al. 2019
52.5
74.7
81.2
87.9
TB
G
X
GraphRetro Sommath et al..
2021
53.7
68.3
72.2
75.5
TB
G
X
Leto Chen & Jung
2021
52.6
76.0
84.4
90.6
TB
G
X
SCROP (Zheng et al.
2019
43.7
60.0
65.2
68.7
TF
G2G (hhi et al. 2020
48.9
67.6
72.5
75.5
TF
G
Aug. Transformer TTtko et al.
2020
48.3
73.4
77.4
TF
X
DualTF aug (Sun et al.
202T)
53.6
70.7
74.6
77.0
TF
S
MEGAN (Sacha et al.
2021
48.0
70.9
78.1
85.4
X
Tied Transtormer
Kim et al.
2021
47.1
67.1
73.1
76.3
TF
S
X
GTAaue (Sco
2021
51.1
67.6
74.8
81.6
TF
GIS
X
X
Graph2SMILES Tudd & Coley
2022
52.9
66.5
70.0
72.9
TF
GIS
X
Retroformer (Wan al. 2022
52.9
68.2
72.5
76.4
TF
G/S
RetroBridge (ours)
50.8
74.1
80.6
85.6
TF
G
X
A.4 EXTENDED RESULTS
In Table we provide an extended version of the exact-match results from Table tional models and methodological details.
TileeIInche includding addi-
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 in other experiments. As shown in Table 5 our models (ForwardBridge) demonstrate comparable performance with other state-of-the-ar methods. However, stress that the probabilistic formulation is less applicable
18

---

Published as a conference paper at ICLR 2024
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.
Table 5: Top-k accuracy for forward reaction prediction.
Dataset
Model
k =
k
3
5
USPTO-50K
ForwardBridge (ours)
89.9
93.9
94.0
ForwardBridge (ours)
81.6
88.5
89.8
USPTO-MIT
MEGAN
86.3
92.4
94.0
Mol. Transformer
88.7
93.1
94.2
Input product
Sample
Sample
Sample 3
True reactants
Q - th
Rexx + +6
Sott+
www VV
inmeeoo of
XX ty
XX
SET VV
12=0)~
100% A30 VV
10m MI
AV
100
Box #bgg2
x
Bop 0+ + Dzg
x
12 VV
V
1236 ON
JOW ST
Figure 6: Examples of modeled reactants. We selected 10 random inputs from the USPTO-50k test set and for each of them we provide the top-3 RetroBridge predictions along with their confidence scores. Two check marks indicate that sampled reactants are the same the ground truth, and one check mark means that reactants are different, but Molecular Transformer (Schwaller et al. 2019) predicts the product molecule used as input.
19


Reviews:
[{"review_id": "aFJhFu2Z18", "reviewer": "Reviewer_vNHE", "summary": "This paper proposes a novel single-step retrosynthesis method called RetroBridge, utilizing the Markov Bridge Model. RetroBridge formulates the reactants and products in single-step retrosynthesis as two discrete distributions and learns the dependencies between them. By performing conditional sampling, the generated reactions demonstrate high accuracy.", "questions": "This paper provides a fresh perspective on AI for retrosynthesis; however, it could benefit from further elaboration on the chemical aspects involved. Some previous works, like GraphRetro and MEGAN somehow, mimic the process of chemical reactions. So, how Markov bridge formulates chemical reactions? I would consider raising my rating if the authors addressed my concerns.", "limitations": "", "rating": 6, "confidence": 5, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. To the best of my knowledge, this paper is the first to introduce the concept of the single-step retrosynthesis problem by framing it as a distribution fitting task, akin to a diffusion model.\n2. The proposed model achieves the SOTA Top k accuracy (k > 3) on USPTO-50K.\n3. The comparison with diffusion model makes a perfect sense.", "weaknesses": "1. The top 1 accuracy achieved on the USPTO-50K dataset is relatively low, and a similar trend is observed in the forward prediction task. It would be helpful if the authors could provide an explanation for these results.\n2. Some baseline methods are missing, which may perform much better than RetroBridge, like DualTF, RSMILES, PMSR, etc.\n3. In comparison to other TF methods, RetroBridge requires atom-to-atom mappings, which could potentially pose limitations on its feasibility. This constraint could become a hindrance when applying RetroBridge to larger datasets such as USPTO-Full.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VIvkStb5d5", "reviewer": "Reviewer_4mon", "summary": "The authors introduce a novel approach to retrosynthesis planning by modeling it as a distribution learning problem in a discrete state space. The core of this approach is the Markov Bridge Model, which is a generative framework designed to approximate the relationship between two discrete distributions that are intractable. By optimizing a variational lower bound, similar to that in diffusion models, the method learns a Markov process that transitions between the product and reactant distributions. Empirical results suggest that the proposed method surpasses existing ones, achieving state-of-the-art results on standard benchmarks.", "questions": "- How much more time is needed during the training and inference phase compared with the baselines?\n- What's the choice of T in the experiments, and how many samples are needed to obtain the results?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "- The paper offers a novel perspective on retrosynthesis planning, framing it as a distribution learning problem on a discrete distribution instead of a discriminative learning problem. The introduction of the Markov Bridge Model in this context and its comparison with the diffusion model is novel and interesting.\n- The presentation of the methodology is clearly written and easy to follow.\n- The derivation of the variational objective, which involves maximizing a lower bound of log-likelihood, is theoretically sound.\n- The empirical result is competitive against existing one-step retrosynthesis models.\n- The proposed method is general and can potentially be extended to other application domains involving the learning of transitions between discrete distributions.", "weaknesses": "- The paper does not discuss the time required during training and inference stage of the proposed method, including the time step T, the number of samples needed to converge on the top-k prediction, etc.\n- From the atom/bond level parametrization in Section 3.2, it seems that the required T could be quite large. An ablation study to report the top-k performance vs. different Ts could be very useful.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "thelhCkcZL", "reviewer": "Reviewer_CFGE", "summary": "The paper presents a retrosynthesis model where the reactant is diffused to product, inducing a Markov Bridge. The method is simple and gets state-of-the-art performance.", "questions": "- The paper misses quite a lot of related works. Retroformer, Graph2Smiles, GTA, Tied transformer all seem to be missing. The paper needs to cite methods from 2022 and 2023, and include USPTO benchmark comparisons.\n- If you pin the outcome to y in eq 2, why does fig1 show that we start from x, and diverge into a bunch of y’s? Shouldn’t fig1 show starting from single x, fanning out, and reducing all mask down to single y?\n- What does t ~ U(0,…,T-1) mean in eq 7? Seems very odd… Why doesn’t it go to the end T? Why do we have multiple values?\n- What happens to the last timestep in eq 7? This seems to be a Dirac, which probably poses problems. Currently it seems that the final prediction is never evaluated, and we only predict until z_{T-1} instead of matching the final z_T to y. How do you evaluate the final prediction z_T?\n- What is p(z|x,y)?\n- What is p(z_t+1 | z_0, z_T)? Where does this go? Why do we talk about this? Does this have something do with p(zt | z0, y)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 2, "strengths": "- The presentation of the paper is clean, and the paper is easy to read and understand.\n- The method is principled and simple, and seems to work well. The idea of flowing product to reactant is very sensible.", "weaknesses": "- The method is incremental: the diffusion is directly from Austin, and I think the loss function is also only a variant of diffusion losses. The network is directly from literature. There is no chemical inductive biases that I can identify.\n- The performance is opened up somewhat, but not sufficiently. I like the ablations that show the context and VLB help, but these are generally known from diffusion literature already. It would have been good to try to expose why the performance is better than MEGAN. For instance, how much of the results are from network tuning or training tricks?\n- There are many missing competing methods, such as RetroFormer and Graph2Smiles.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "bbS9BxlDhu", "reviewer": "Reviewer_g93h", "summary": "This paper addresses single-step retrosynthesis.\nThe single-step retrosynthesis is formalized as the mapping problem from product molecule graphs to reactant molecule graphs.\nThe authors propose the Markov Bridge Model, which learns the dependency between two discrete variables, and the model is applied to the graph mapping problem.\nThe proposed method outperforms state-of-the-art template-free methods on the USPTO-50K benchmark.", "questions": "No question.", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "* The Markov Bridge Model is a novel and interesting approach for discrete variables and can be applied to other fields.\n* The evaluation results using top-k round-trip coverage and accuracy support the potential of the proposed method.\n* The future work described in Section 5 clearly describes the limitations of this work.", "weaknesses": "* The reason for selecting graph-based representation needs to be explained.  It is not clear to me why molecules are represented as a graph instead of the 3D coordinates of atoms since 3D coordinates are frequently used in the unconditional generation of molecules  (Hoogeboom et al., 2022).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)