Paper: RetroBridge: Modeling Retrosynthesis with Markov Bridges (ICLR-2024)

Reviewer comment (Novelty):
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.

Author response:
We respect the opinion of the reviewer on this very subjective matter, but we would like to highlight our contributions that, in our opinion, make our work clearly novel (beyond incremental improvements) and differentiate it from previous generative modeling frameworks. None of the existing approaches, including diffusion models (DM), conditional flow matching (CFM) and diffusion Schrödinger bridge (DSB) models, address the problem of mapping between two intractable, discrete distributions. DMs do not allow for an intractable prior. CFMs do not operate on discrete data and do not optimise p(y|x); DSBs either do not operate on coupled samples or do not work with discrete data. Therefore, to the best of our knowledge the proposed Markov Bridge Model (MBM) is the first method that fulfills all these requirements. We kindly ask the reviewer to provide specific pointers to other works if there is reason to believe our evaluation of the prior art is incomplete.

Apart from this general remark, we address each of your specific concerns below:
* We took the discrete data formalism from Austin et al., as other papers did too [1, 2], but we want to stress that **MBM is not a diffusion model.** The major differences are that neither a tractable prior nor a forward noising process are required.
* The loss function is a variational lower bound on the data log-likelihood which is common practice in the generative modelling community. Here, we derive the specific form of this lower bound for the proposed Markov bridge model and show empirical advantages over a previously used simplified cross-entropy loss.
* The network is a variable component of our framework that can depend on the task and data structure (e.g. if you apply MBMs to text, you can plug in an LSTM as a backbone model instead). As developing a new type of GNN is not the focus of this work, we took the same architecture as in DiGress which was shown to perform well on molecule-related tasks. We would add that it even makes our comparison with DiGress more fair and highlights our contribution on the methodological side w.r.t. the new MBM framework.
* Finally, we did not use major chemical inductive biases because our work is motivated by the statistics of the retrosynthesis data rather than chemical intuition. The main inductive bias we investigate in this paper is the one-to-many mapping between products and reactants that is naturally captured by our proposed framework. The benefits of this idea become apparent in the round-trip evaluation where RetroBridge shows significant empirical gains compared to baselines.

Reviewer comment (Evaluation):
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?

Author response:
We are happy to hear that you find value in our ablation studies. We believe that this kind of analysis should be independently conducted for each specific use case. For example, previous work on graph generation with diffusion models [1] found a simple cross-entropy loss to perform well, while we observed improved performance using the more principled VLB.

While we did not derive our method from MEGAN and cannot trace what exactly improves the performance, we argue that sampling proportional to p(reactants|product) is important to find a diverse set of valid reactants to provide more good options to a multi-step pipeline (as eventually supported by our top-k and round-trip results for k>1). Our framework is designed to approximate this probability while MEGAN optimizes the likelihood of a specific sequence of edits, imposing a more or less arbitrary order, even though different sequences can lead to the same outcome. We therefore expected our formalism to be better suited to the task, which might explain the superior performance.

We selected training and network hyperparameters among a small set of options (CE and VLB loss functions, context/no-context) and selected the best model based on its performance on the validation set before the final test set evaluation. We do not expect unusually large gains from network tuning alone. Besides, it is reasonable to believe that a similar amount of hyperparameter tuning has been performed for other retrosynthesis models.

Reviewer comment (Experiments):
There are many missing competing methods, such as RetroFormer and Graph2Smiles.

Author response:
We thank the reviewer for these comments. We added comparison with all three methods in our tables. Besides, we added even more methods in a (new) Table 4, where we provide more information on the technical aspects of the methods and evaluation methodology.

Reviewer comment (Novelty):
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.

Author response:
We thank the reviewer for these comments. We added comparison with all three methods in our tables. Besides, we added even more methods in a (new) Table 4, where we provide more information on the technical aspects of the methods and evaluation methodology.

Reviewer comment (Presentation):
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?

Author response:
Equation 2 provides a required property of a Markov bridge. Figure 1 represents three different Markov bridges that start at the same point $\boldsymbol{x}$ and end at three different points $\boldsymbol{y}_1$, $\boldsymbol{y}_2$ and $\boldsymbol{y}_3$. Thus, for each of these bridges there is a specific requirement that it should end at a specific point. For example, for the first (orange) bridge the requirement from eq. 2 will be written as $p(\boldsymbol{z}_T=\boldsymbol{y}_1 | \boldsymbol{z}\_{T-1}, \boldsymbol{y}_1) =  1$.

Reviewer comment (Writing):
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?

Author response:
$t\sim U(0,\dots,T-1)$ means that the time step is sampled from the uniform distribution over a discrete set of $T$ values $0, 1, 2, \dots, T-1$. As shown in Appendix A.1, we replace the sum over all steps with the expectation over $U(0,\dots,T-1)$ as it is an unbiased estimator that is crucial for efficient simulation-free learning. In particular, we move from sum to expectation in transition (29) $\to$ (30) (Appendix A.1). It does not go to the end, $T$, because for a given $t$ we compute the KL divergence between distributions of $\boldsymbol{z}_{t+1}$. It means that for $t=T-1$ we will compute the KL divergence for distributions of $\boldsymbol{z}_T$.

Reviewer comment (Theory):
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?

Author response:
As we discuss in the answer to the previous question, for a given $t$ we compute the KL divergence between distributions of $\boldsymbol{z}\_{t+1}$. It means that for $t=T-1$ we will compute the KL divergence for distributions of the final state $\boldsymbol{z}\_T$. Therefore, for the last timestep $t=T-1$, we compute $D_{\text{KL}}(p(\boldsymbol{z}\_T|\boldsymbol{z}\_{T-1}, \boldsymbol{y})|q\_{\theta}(\boldsymbol{z}\_T|\boldsymbol{z}\_{T-1}))$ which is precisely the divergence between the ground-truth distribution of the final value $\boldsymbol{z}\_T$ given the previous value $\boldsymbol{z}\_{T-1}$ and target $\boldsymbol{y}$, and approximated (by the neural network) distribution of $\boldsymbol{z}\_T$ given $\boldsymbol{z}\_{T-1}$. No Dirac functions appear here.

Reviewer comment (Writing):
What is p(z|x,y)?

Author response:
$p(\boldsymbol{z}_t|\boldsymbol{x}, \boldsymbol{y})$ is the distribution of the intermediate state $\boldsymbol{z}_t$ at time step $t$ of the Markov bridge that starts at the point $\boldsymbol{x}$ and ends at the point $\boldsymbol{y}$.

Reviewer comment (Experiments):
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.

Author response:
We thank the reviewer for this useful remark and provide an ablation study in Appendix A.3.5. As shown in the new Figure 5A, the performance of our model does not significantly change with up to x2 speed improvement (i.e. when sampling with T=250 instead of T=500). The performance starts degrading notably only after a 5-fold reduction of the number of time steps.

Reviewer comment (Evaluation):
How much more time is needed during the training and inference phase compared with the baselines?

Author response:
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 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.

Regarding training times, we trained our model 72 hours on a single Tesla V100-PCIE-32GB GPU and selected the best model based on top-5 exact match accuracy on the validation set (see Section A.3.4). We did not use early stopping or perform ablations with earlier checkpoints since training is only a one-time cost and depends on many unrelated hyperparameters, e.g. the number of validation steps, which impedes comparisons with other methods.

Reviewer comment (Experiments):
What's the choice of T in the experiments, and how many samples are needed to obtain the results?

Author response:
We sample with T=500 for all our experiments, and generate 100 samples per input product. We additionally perform an ablation study on the number of samples (see Appendix A.3.5). As shown in the new Figure 5B, RetroBridge is very robust to the reduction of the number of samples.

Reviewer comment (Evaluation):
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.

Author response:
Since one product can be obtained through different reaction pathways and different starting sets of reactants, top-1 accuracy reports nothing but the ability of the model to guess what exact reaction for a given product was patented and therefore included in the dataset. We stress that **top-3 and top-5 are much more important metrics as they reflect the probabilistic aspect of the problem.** Besides, as it was discussed in [1], the typically expected or desired breadth of the multi-step retrosynthesis planning tree is 5 which makes top-5 accuracy the preferred metric. We additionally emphasize that, as shown in Table 3, **our model was optimized for top-5 accuracy.** We added these additional remarks in Section 4.3.

[1] Maziarz, Krzysztof, et al. "Re-evaluating Retrosynthesis Algorithms with Syntheseus." arXiv preprint arXiv:2310.19796 (2023).

Reviewer comment (Experiments):
Some baseline methods are missing, which may perform much better than RetroBridge, like DualTF, RSMILES, PMSR, etc.

Author response:
We thank the reviewer for these suggestions. We significantly extended Tables 1 and 2 and, in particular, added DualTF to Table 1. We would like to mention, however, that DualTF is an ensemble of 3 models, one of which is a ranking model in the forward direction, which all the other baseline models do not use. We would also like to note that the comparison with PMSR and RSMILES is not fair. While all the methods in Table 1 (including ours) are trained on the same subset of USPTO-50k, PMSR is additionally pre-trained on 3.7M reactions from Pistachio. Likewise, RSMILES is pre-trained on USPTO-FULL and their test set is x20 augmented. We therefore add these methods only for a full picture in the extended Table 4 in the Appendix, where we provide more information on the technical aspects of the methods.

Reviewer comment (Experiments):
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.

Author response:
While RetroBridge uses atom mapping for training, this is not conceptually required. It is true that a training set with all correctly mapped molecules significantly simplifies the training process. However, although we did not perform such experiments, we speculate that the probabilistic nature of our model allows for minor inconsistencies in the atom mapping. With the modern high-performant tools for atom mapping such as RXNMapper [1], we believe that the level of possible mapping inconsistencies is negligible for the performance of RetroBridge. Also, there have been to the best of our knowledge no studies so far in the literature which empirically demonstrate that errors in the atom mapping lead to significant practical performance loss for any kind of retrosynthesis model.

[1] Schwaller P. et al. “Extraction of organic chemistry grammar from unsupervised learning of chemical reactions”

Reviewer comment (Writing):
This paper provides a fresh perspective on AI for retrosynthesis; however, it could benefit from further elaboration on the chemical aspects involved.

Author response:
Indeed, some prior works are inspired by the arrow pushing formalism from organic chemistry using discrete graph edits. Here, we instead build on statistical intuition and the fact that each product molecule can be obtained starting from several sets of reactants. We argue that sampling proportional to p(reactants|product) is important to find a diverse set of valid reactants to provide more good options to a multi-step pipeline. Our framework is designed to approximate this probability while autoregressive models typically optimize the likelihood of specific sequences of edits, imposing a more or less arbitrary order, even though different sequences can lead to the same outcome. We therefore expected our formalism to be better suited to the task, which might explain the superior performance on the round-trip metrics.

Additionally, MEGAN also makes heavy use of chemical input features (atomic number, formal charge, chiral tag, and number of explicit hydrogen atoms is aromatic, bond type, bond stereometry) which might provide additional cues to solve the task. We believe it could be an interesting future direction to investigate if endowing RetroBridge with more chemically-relevant input features can further improve its empirical performance.

Reviewer comment (Theory):
Some previous works, like GraphRetro and MEGAN somehow, mimic the process of chemical reactions. So, how does Markov bridge formulate chemical reactions?

Author response:
Indeed, some prior works are inspired by the arrow pushing formalism from organic chemistry using discrete graph edits. Here, we instead build on statistical intuition and the fact that each product molecule can be obtained starting from several sets of reactants. We argue that sampling proportional to p(reactants|product) is important to find a diverse set of valid reactants to provide more good options to a multi-step pipeline. Our framework is designed to approximate this probability while autoregressive models typically optimize the likelihood of specific sequences of edits, imposing a more or less arbitrary order, even though different sequences can lead to the same outcome. We therefore expected our formalism to be better suited to the task, which might explain the superior performance on the round-trip metrics.

Additionally, MEGAN also makes heavy use of chemical input features (atomic number, formal charge, chiral tag, and number of explicit hydrogen atoms is aromatic, bond type, bond stereometry) which might provide additional cues to solve the task. We believe it could be an interesting future direction to investigate if endowing RetroBridge with more chemically-relevant input features can further improve its empirical performance.