Alignment is Key for Applying Diffusion Models to Retrosynthesis (2405.17656v1)
Abstract: Retrosynthesis, the task of identifying precursors for a given molecule, can be naturally framed as a conditional graph generation task. Diffusion models are a particularly promising modelling approach, enabling post-hoc conditioning and trading off quality for speed during generation. We show mathematically that permutation equivariant denoisers severely limit the expressiveness of graph diffusion models and thus their adaptation to retrosynthesis. To address this limitation, we relax the equivariance requirement such that it only applies to aligned permutations of the conditioning and the generated graphs obtained through atom mapping. Our new denoiser achieves the highest top-$1$ accuracy ($54.7$\%) across template-free and template-based methods on USPTO-50k. We also demonstrate the ability for flexible post-training conditioning and good sample quality with small diffusion step counts, highlighting the potential for interactive applications and additional controls for multi-step planning.
- Structured denoising diffusion models in discrete state-spaces. In Advances in Neural Information Processing Systems, volume 34, pages 17981–17993. Curran Associates, Inc., 2021.
- Tweedie moment projected diffusions for inverse problems. arXiv preprint arXiv:2310.06721, 2023.
- Maskgit: Masked generative image transformer. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11315–11325, 2022.
- S. Chen and Y. Jung. Deep retrosynthetic reaction prediction using local reactivity and global attention. JACS Au, 1(10):1612–1620, 2021.
- Diffusion posterior sampling for general noisy inverse problems. In International Conference on Learning Representations, 2023.
- Prediction of organic reaction outcomes using machine learning. ACS Central Science, 3(5):434–443, 2017a.
- Computer-assisted retrosynthesis based on molecular similarity. ACS Central Science, 3(12):1237–1245, 2017b.
- E. Corey and X.-M. Cheng. The logic of chemical synthesis. Journal of the American Chemical Society, 118(43):10678–10678, 1996.
- Computer-assisted design of complex organic syntheses: Pathways for molecular synthesis can be devised with a computer and equipment for graphical communication. Science, 166(3902):178–192, 1969.
- Retrosynthesis prediction with conditional graph logic network. In Advances in Neural Information Processing Systems, volume 32, pages 8872–8882. Curran Associates, Inc., 2019.
- Z. Dou and Y. Song. Diffusion posterior sampling for linear inverse problem solving: A filtering perspective. In International Conference on Learning Representations, 2024.
- V. P. Dwivedi and X. Bresson. A generalization of transformer networks to graphs. In AAAI Workshop on Deep Learning on Graphs, 2021.
- P. Ertl and A. Schuffenhauer. Estimation of synthetic accessibility score of drug-like molecules based on molecular complexity and fragment contributions. Journal of Cheminformatics, 1:1–11, 2009.
- User-defined event sampling and uncertainty quantification in diffusion models for physical dynamical systems. In International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 10136–10152. PMLR, 2023.
- Compositional sculpting of iterative generative processes. In Advances in Neural Information Processing Systems, volume 36, pages 12665–12702. Curran Associates, Inc., 2023.
- Denoising diffusion probabilistic models. In Advances in Neural Information Processing Systems, volume 33, pages 6840–6851. Curran Associates, Inc., 2020.
- Argmax flows and multinomial diffusion: Learning categorical distributions. In Advances in Neural Information Processing Systems, volume 34, pages 12454–12465. Curran Associates, Inc., 2021.
- Autoregressive diffusion models. In International Conference on Learning Representations, 2022a.
- Equivariant diffusion for molecule generation in 3D. In International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 8867–8887. PMLR, 2022b.
- Graphgdp: Generative diffusion processes for permutation invariant graph generation. In 2022 IEEE International Conference on Data Mining (ICDM), pages 201–210. IEEE Computer Society, 2022.
- RetroBridge: Modeling retrosynthesis with Markov bridges. In Conference on Learning Representations, 2024a.
- Equivariant 3d-conditional diffusion model for molecular linker design. Nature Machine Intelligence, pages 1–11, 2024b.
- Categorical reparameterization with gumbel-softmax. In International Conference on Learning Representations, 2016.
- Learning chemical rules of retrosynthesis with pre-training. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 5113–5121, 2023.
- Predicting organic reaction outcomes with Weisfeiler-Lehman network. In Advances in Neural Information Processing Systems, volume 30, pages 2607–2616. Curran Associates, Inc., 2017.
- Elucidating the design space of diffusion-based generative models. In Advances in Neural Information Processing Systems, volume 35, pages 26565–26577. Curran Associates, Inc., 2022.
- Valid, plausible, and diverse retrosynthesis using tied two-way transformers with latent variables. Journal of Chemical Information and Modeling, 61(1):123–133, 2021.
- Diffusion-lm improves controllable text generation. In Advances in Neural Information Processing Systems, volume 35, pages 4328–4343. Curran Associates, Inc., 2022.
- FusionRetro: Molecule representation fusion via in-context learning for retrosynthetic planning. In International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 22028–22041. PMLR, 2023.
- T.-Y. Liu et al. Learning to rank for information retrieval. Foundations and Trends® in Information Retrieval, 3(3):225–331, 2009.
- D. M. Lowe. Extraction of Chemical Structures and Reactions from the Literature. PhD thesis, University of Cambridge, 2012.
- Dpm-solver: A fast ode solver for diffusion probabilistic model sampling in around 10 steps. In Advances in Neural Information Processing Systems, volume 35, pages 5775–5787. Curran Associates, Inc., 2022.
- Repaint: Inpainting using denoising diffusion probabilistic models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11461–11471, 2022.
- Re-evaluating retrosynthesis algorithms with syntheseus. In NeurIPS 2023 AI for Science Workshop, 2023.
- Permutation invariant graph generation via score-based generative modeling. In Proceedings of the Twenty Third International Conference on Artificial Intelligence and Statistics, volume 108 of Proceedings of Machine Learning Research, pages 4474–4484. PMLR, 2020.
- Improving diffusion models for inverse problems using optimal posterior covariance. arXiv preprint arXiv:2402.02149, 2024.
- Molecule edit graph attention network: modeling chemical reactions as sequences of graph edits. Journal of Chemical Information and Modeling, 61(7):3273–3284, 2021.
- Adversarial diffusion distillation. arXiv preprint arXiv:2311.17042, 2023.
- What’s what: The (nearly) definitive guide to reaction role assignment. Journal of Chemical Information and Modeling, 56(12):2336–2346, 2016.
- “Found in Translation”: predicting outcomes of complex organic chemistry reactions using neural sequence-to-sequence models. Chemical Science, 9(28):6091–6098, 2018.
- Molecular transformer: A model for uncertainty-calibrated chemical reaction prediction. ACS Central Science, 5(9):1572–1583, 2019.
- Extraction of organic chemistry grammar from unsupervised learning of chemical reactions. Science Advances, 7(15), 2021.
- Neural-symbolic machine learning for retrosynthesis and reaction prediction. Chemistry–A European Journal, 23(25):5966–5971, 2017.
- GTA: Graph truncated attention for retrosynthesis. Proceedings of the AAAI Conference on Artificial Intelligence, 35(1):531–539, 2021.
- A graph to graphs framework for retrosynthesis prediction. In International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, pages 8818–8827. PMLR, 2020.
- Parallel sampling of diffusion models. In Advances in Neural Information Processing Systems, pages 4263–4276. Curran Associates, Inc., 2023.
- Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning, volume 37 of Proceedings of Machine Learning Research, pages 2256–2265. PMLR, 2015.
- Learning graph models for retrosynthesis prediction. In Advances in Neural Information Processing Systems, volume 34, pages 9405–9415. Curran Associates, Inc., 2021.
- Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations, 2021.
- Consistency models. In International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pages 32211–32252. PMLR, 2023.
- Towards understanding retrosynthesis by energy-based models. In Advances in Neural Information Processing Systems, volume 34, pages 10186–10194. Curran Associates, Inc., 2021.
- State-of-the-art augmented NLP transformer models for direct and single-step retrosynthesis. Nature Communications, 11(1):5575, 2020.
- B. M. Trost. The atom economy—a search for synthetic efficiency. Science, 254(5037):1471–1477, 1991.
- Z. Tu and C. W. Coley. Permutation invariant graph-to-sequence model for template-free retrosynthesis and reaction prediction. Journal of Chemical Information and Modeling, 62(15):3503–3513, 2022.
- Dual use of artificial-intelligence-powered drug discovery. Nature Machine Intelligence, 4(3):189–191, 2022.
- DiGress: Discrete denoising diffusion for graph generation. In International Conference on Learning Representations, 2023.
- Retroformer: Pushing the limits of end-to-end retrosynthesis transformer. In International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 22475–22490. PMLR, 2022.
- RetroDiff: Retrosynthesis as multi-stage distribution interpolation. arXiv preprint arXiv:2311.14077, 2023.
- Practical and asymptotically exact conditional sampling in diffusion models. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 31372–31403. Curran Associates, Inc., 2023.
- A comprehensive survey on graph neural networks. IEEE Transactions on Neural Networks and Learning Systems, 32(1):4–24, 2021.
- Retrosynthesis prediction with local template retrieval. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 37, pages 5330–5338, 2023.
- RetroXpert: Decompose retrosynthesis prediction like a chemist. In Advances in Neural Information Processing Systems, volume 33, pages 11248–11258. Curran Associates, Inc., 2020.
- SwinGNN: Rethinking permutation invariance in diffusion models for graph generation. arXiv preprint arXiv:2307.01646, 2023.
- Predicting retrosynthetic reactions using self-corrected transformer neural networks. Journal of Chemical Information and Modeling, 60(1):47–55, 2019.
- Root-aligned SMILES: a tight representation for chemical reaction prediction. Chemical Science, 13(31):9023–9034, 2022.