[Paper Review] RetroBridge: Modeling Retrosynthesis with Markov Bridges

Ilia Igashov, Arne Schneuing+3 more|arXiv (Cornell University)|Aug 30, 2023

Machine Learning in Materials Science4 citations

TL;DR

RetroBridge introduces a novel Markov Bridge Model to address single-step retrosynthesis as a probabilistic distribution learning problem between intractable discrete distributions of products and reactants. By modeling stochastic trajectories pinned at known product and reactant molecules, it achieves state-of-the-art performance on USPTO-50k benchmarks, outperforming diffusion-based methods like DiGress in both accuracy and uncertainty estimation.

ABSTRACT

Retrosynthesis planning is a fundamental challenge in chemistry which aims at designing 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 planning 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.

Motivation & Objective

• Address the challenge of single-step retrosynthesis planning by modeling the probabilistic dependency between product and reactant molecule distributions.
• Overcome limitations of diffusion-based generative models that are not naturally suited for mapping between two intractable discrete distributions.
• Develop a template-free method that enables diverse, uncertainty-aware retrosynthetic pathway generation for drug discovery.
• Improve the reliability and interpretability of retrosynthesis predictions through confidence scoring and trajectory modeling.
• Demonstrate the superiority of the Markov Bridge framework over existing diffusion-based approaches in discrete, chemistry-relevant generation tasks.

Proposed method

• Introduce the Markov Bridge Model as a generative framework that learns the conditional distribution between two intractable discrete distributions using finite coupled data points.
• Model retrosynthesis as a stochastic process constrained to start at a given product molecule and end at a reactant molecule, with the process defined over molecular graph states.
• Use a conditional score-based approach to estimate the likelihood of sampled reactant molecules, enabling uncertainty quantification via confidence scores.
• Leverage the probabilistic nature of the model to score reactant candidates based on their likelihood under the learned Markov bridge process.
• Train the model using contrastive learning (RetroBridge-CE) and variational lower bound (RetroBridge-VLB) objectives to improve generation quality and uncertainty calibration.
• Incorporate the input product as context at each sampling step to guide the trajectory while preserving structural fidelity.

Experimental results

Research questions

• RQ1 Can a generative model based on Markov bridges effectively learn the dependency between two intractable discrete distributions in retrosynthesis?
• RQ2 How does the Markov Bridge framework compare to diffusion-based models in terms of accuracy and uncertainty estimation for retrosynthesis prediction?
• RQ3 To what extent does the proposed method improve zero-shot and few-shot generalization in retrosynthesis planning compared to existing template-free models?
• RQ4 Can the model generate diverse, chemically valid retrosynthetic pathways while maintaining high confidence in correct predictions?
• RQ5 Does the inclusion of context (product molecule) at each sampling step enhance the quality and reliability of generated reactants?

Key findings

• RetroBridge achieves a top-1 accuracy of 50.74% on the USPTO-50k validation set when using context, surpassing the state-of-the-art DiGress model (47.32%).
• With context, RetroBridge reaches 80.58% top-50 accuracy, outperforming DiGress (80.88%) in the same setting, demonstrating strong performance across all top-k metrics.
• Even without context, RetroBridge-CE achieves 48.71% top-1 accuracy, indicating strong structural retention and minimal reliance on external conditioning.
• The model exhibits reliable uncertainty quantification: correct reactants are assigned significantly higher confidence scores than incorrect ones, with a 0.66 vs. 0.12 gap in one example.
• RetroBridge-VLB with context achieves 86.31% top-50 accuracy, indicating strong performance in generating diverse and correct reactant candidates.
• Qualitative analysis shows that RetroBridge successfully recovers correct reactants in 100% of three randomly selected test cases, with high confidence in the top prediction.