{"type":"ScholarlyArticle","title":"Doubly Stochastic Graph-based Non-autoregressive Reaction Prediction","authors":[{"name":"Ziqiao Meng","slug":"ziqiao-meng","given":"Ziqiao","family":"Meng","affiliation":"The Chinese University of Hong Kong"},{"name":"Peilin Zhao","slug":"peilin-zhao","given":"Peilin","family":"Zhao","affiliation":"Tencent AI Lab"},{"name":"Yang Yu","slug":"yang-yu","given":"Yang","family":"Yu","affiliation":"Tencent AI Lab"},{"name":"Irwin King","slug":"irwin-king","given":"Irwin","family":"King","affiliation":"The Chinese University of Hong Kong"}],"abstract":"Organic reaction prediction is a critical task in drug discovery. Recently, researchers have achieved non-autoregressive reaction prediction by modeling the redistribution of electrons, resulting in state-of-the-art top-1 accuracy, and enabling parallel sampling. However, the current non-autoregressive decoder does not satisfy two essential rules of electron redistribution modeling simultaneously: the electron-counting rule and the symmetry rule. This violation of the physical constraints of chemical reactions impairs model performance. In this work, we propose a new framework called ReactionSink that combines two doubly stochastic self-attention mappings to obtain electron redistribution predictions that follow both constraints. We further extend our solution to a general multi-head attention mechanism with augmented constraints. To achieve this, we apply Sinkhorn's algorithm to iteratively update self-attention mappings, which imposes doubly conservative constraints as additional informative priors on electron redistribution modeling. We theoretically demonstrate that our ReactionSink can simultaneously satisfy both rules, which the current decoder mechanism cannot do. Empirical results show that our approach consistently improves the predictive performance of non-autoregressive models and does not bring an unbearable additional computational cost.","doi":"10.24963/ijcai.2023/452","url":"https://ijcai.org/proceedings/2023/452","pdf":"https://ijcai.org/proceedings/2023/0452.pdf","markdown":"https://ijcai.org/proceedings/2023/452.md","bibtex":"https://ijcai.org/proceedings/2023/bibtex/452","volume":{"slug":"2023","title":"Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence","short_title":"IJCAI 2023","url":"https://ijcai.org/proceedings/2023/","year":2023,"location":"Macao, SAR","dates":"19-25 August 2023","editors":"Edith Elkind","isbn":"978-1-956792-03-4","publisher":"International Joint Conferences on Artificial Intelligence Organization","doi_prefix":"10.24963/ijcai.2023"},"section":"Machine Learning","track":"Main Track","number":452,"pages":{"first":4064,"last":4072},"keywords":["Multidisciplinary Topics and Applications: MDA: Physical sciences","Machine Learning: ML: Applications","Machine Learning: ML: Attention models","Machine Learning: ML: Structured prediction"],"videos":[],"references":[{"position":1,"raw":"[Adams and Zemel, 2011] Ryan Prescott Adams and Richard S. Zemel. Ranking via sinkhorn propagation. CoRR, abs/1106.1925, 2011.","title":"Ranking via sinkhorn propagation","authors":"Ryan Prescott, Adams, Richard S Zemel","year":"2011","venue":null,"doi":null,"url":null,"internal_url":null},{"position":2,"raw":"[Bi et al., 2021] Hangrui Bi, Hengyi Wang, Chence Shi, Connor W. Coley, Jian Tang, and Hongyu Guo. Non- autoregressive electron redistribution modeling for reac- tion prediction. In Proceedings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Ma- chine Learning Research, pages 904-913. PMLR, 2021.","title":"Nonautoregressive electron redistribution modeling for reaction prediction","authors":"Bi","year":"2021","venue":"Proceedings of the 38th International Conference on Machine Learning, ICML 2021","doi":null,"url":null,"internal_url":null},{"position":3,"raw":"[Bradshaw et al., 2019] John Bradshaw, Matt J. Kusner, Brooks Paige, Marwin H. S. Segler, and José Miguel Hernández-Lobato. A generative model for electron paths. In 7th International Conference on Learning Representa- tions, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019. OpenReview.net, 2019.","title":"A generative model for electron paths","authors":"John Bradshaw, Matt J Kusner, Brooks Paige, Marwin H S Segler, José Miguel Hernández-Lobato","year":"2019","venue":"7th International Conference on Learning Representations, ICLR 2019","doi":null,"url":null,"internal_url":null},{"position":4,"raw":"[Cuturi, 2013] Marco Cuturi. Sinkhorn distances: Light- speed computation of optimal transport. In Advances in Neural Information Processing Systems 26: 27th An- nual Conference on Neural Information Processing Sys- tems 2013. Proceedings of a meeting held December 5- 8, 2013, Lake Tahoe, Nevada, United States, pages 2292- 2300, 2013.","title":"Sinkhorn distances: Lightspeed computation of optimal transport","authors":"Marco Cuturi","year":"2013","venue":"Advances in Neural Information Processing Systems 26: 27th Annual Conference on Neural Information Processing Systems 2013. Proceedings of a meeting","doi":null,"url":null,"internal_url":null},{"position":5,"raw":"[Devlin et al., 2019] Jacob Devlin, Ming-Wei Chang, Ken- ton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understand- ing. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume 1 (Long and Short Papers), pages 4171-4186. Association for Computational Linguistics, 2019.","title":"BERT: pre-training of deep bidirectional transformers for language understanding","authors":"Jacob Devlin, Ming-Wei Chang, Kenton Lee, Kristina Toutanova","year":"2019","venue":"Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, NAACL-HLT 2019","doi":null,"url":null,"internal_url":null},{"position":6,"raw":"[Do et al., 2019] Kien Do, Truyen Tran, and Svetha Venkatesh. Graph transformation policy network for chemical reaction prediction. In Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, KDD 2019, Anchorage, AK, USA, August 4-8, 2019, pages 750-760. ACM, 2019.","title":"Graph transformation policy network for chemical reaction prediction","authors":"Kien Do, Truyen Tran, Svetha Venkatesh","year":"2019","venue":"Proceedings of the 25th ACM SIGKDD International Conference on Knowledge Discovery & Data Mining, KDD 2019","doi":null,"url":null,"internal_url":null},{"position":7,"raw":"[Dral, 2020] Pavlo Dral. Quantum chemistry in the age of machine learning. Journal of Physical Chemistry Letters, 11:2336-2347, 03 2020.","title":"Pavlo Dral. Quantum chemistry in the age of machine learning","authors":"Dral","year":"2020","venue":"Journal of Physical Chemistry Letters","doi":null,"url":null,"internal_url":null},{"position":8,"raw":"[Irwin et al., 2022] Ross Irwin, Spyridon Dimitriadis, Ji- azhen He, and Esben Jannik Bjerrum. Chemformer: a pre-trained transformer for computational chemistry. Ma- chine Learning: Science and Technology, 3(1):015022, jan 2022.","title":"Chemformer: a pre-trained transformer for computational chemistry","authors":"Ross Irwin, Spyridon Dimitriadis, Jiazhen He, Esben Jannik Bjerrum","year":"2022","venue":"Machine Learning: Science and Technology","doi":null,"url":null,"internal_url":null},{"position":9,"raw":"[Jin et al., 2017] Wengong Jin, Connor W. Coley, Regina Barzilay, and Tommi S. Jaakkola. Predicting organic re- action outcomes with weisfeiler-lehman network. In Ad- vances in Neural Information Processing Systems 30: An- nual Conference on Neural Information Processing Sys- tems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 2607-2616, 2017.","title":"Predicting organic reaction outcomes with weisfeiler-lehman network","authors":"Jin","year":"2017","venue":"Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017","doi":null,"url":null,"internal_url":null},{"position":10,"raw":"[Kingma and Ba, 2015] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. In 3rd Inter- national Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015.","title":"Adam: A method for stochastic optimization","authors":"P Diederik, Jimmy Kingma, Ba","year":"2015","venue":"3rd International Conference on Learning Representations, ICLR 2015","doi":null,"url":null,"internal_url":null},{"position":11,"raw":"[Kovacs et al., 2021] David Kovacs, William McCorkindale, and Alpha Lee. Quantitative interpretation explains ma- chine learning models for chemical reaction prediction and uncovers bias. Nature Communications, 12, 03 2021.","title":"Quantitative interpretation explains machine learning models for chemical reaction prediction and uncovers bias","authors":"David Kovacs, William Mccorkindale, Alpha Lee","year":"2021","venue":"Nature Communications","doi":null,"url":null,"internal_url":null},{"position":12,"raw":"[Lam et al., 2020] Yu-hong Lam, Yuriy Abramov, Ravi S. Ananthula, Jennifer M. Elward, Lori R. Hilden, Sten O. Nilsson Lill, Per-Ola Norrby, Antonio Ramirez, Edward C. Sherer, Jason Mustakis, and Gerald J. Tanoury. Applica- tions of quantum chemistry in pharmaceutical process de- velopment: Current state and opportunities. Organic Pro- cess Research & Development, 24(8):1496-1507, 2020.","title":"Applications of quantum chemistry in pharmaceutical process development: Current state and opportunities","authors":"Yu-Hong Lam, Yuriy Abramov, Ravi S Ananthula, Jennifer M Elward, Lori R Hilden, O Sten, Per-Ola Nilsson Lill, Antonio Norrby, Edward C Ramirez, Jason Sherer, Gerald J Mustakis, Tanoury","year":"2020","venue":"Organic Process Research & Development","doi":null,"url":null,"internal_url":null},{"position":13,"raw":"[Lowe, 2012] Daniel Mark Lowe. Extraction of chemical structures and reactions from the literature. Apollo -Uni- versity of Cambridge Repository, 2012.","title":"Extraction of chemical structures and reactions from the literature","authors":"Daniel Mark, Lowe","year":"2012","venue":null,"doi":null,"url":null,"internal_url":null},{"position":14,"raw":"[Lu and Zhang, 2022] Jieyu Lu and Yingkai Zhang. Uni- fied deep learning model for multitask reaction predictions with explanation. J. Chem. Inf. Model., 62(6):1376-1387, 2022.","title":"Jieyu Lu and Yingkai Zhang. Unified deep learning model for multitask reaction predictions with explanation","authors":"Zhang Lu","year":"2022","venue":"J. Chem. Inf. Model","doi":null,"url":null,"internal_url":null},{"position":15,"raw":"[Marshall and Coifman, 2019] Nicholas F Marshall and Ronald R Coifman. Manifold learning with bi-stochastic kernels. IMA Journal of Applied Mathematics, 84(3):455- 482, 2019.","title":"Manifold learning with bi-stochastic kernels","authors":"F Nicholas, Ronald R Marshall, Coifman","year":"2019","venue":"IMA Journal of Applied Mathematics","doi":null,"url":null,"internal_url":null},{"position":16,"raw":"[Mena et al., 2018] Gonzalo E. Mena, David Belanger, Scott W. Linderman, and Jasper Snoek. Learning la- tent permutations with gumbel-sinkhorn networks. In 6th International Conference on Learning Representa- tions, ICLR 2018, Vancouver, BC, Canada, April 30 -May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018.","title":"Learning latent permutations with gumbel-sinkhorn networks","authors":"Gonzalo E Mena, David Belanger, Scott W Linderman, Jasper Snoek","year":"2018","venue":"6th International Conference on Learning Representations, ICLR 2018","doi":null,"url":null,"internal_url":null},{"position":17,"raw":"[Mialon et al., 2021] Grégoire Mialon, Dexiong Chen, Alexandre d'Aspremont, and Julien Mairal. A trainable optimal transport embedding for feature aggregation and its relationship to attention. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021.","title":"A trainable optimal transport embedding for feature aggregation and its relationship to attention","authors":"Grégoire Mialon, Dexiong Chen, Alexandre Aspremont, Julien Mairal","year":"2021","venue":"9th International Conference on Learning Representations, ICLR 2021, Virtual Event","doi":null,"url":null,"internal_url":null},{"position":18,"raw":"[Milanfar, 2013] Peyman Milanfar. Symmetrizing smooth- ing filters. SIAM J. Imaging Sci., 6(1):263-284, 2013.","title":"Symmetrizing smoothing filters","authors":"Peyman Milanfar","year":"2013","venue":"SIAM J. Imaging Sci","doi":null,"url":null,"internal_url":null},{"position":19,"raw":"[Niculae et al., 2018] Vlad Niculae, André F. T. Martins, Mathieu Blondel, and Claire Cardie. Sparsemap: Dif- ferentiable sparse structured inference. In Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan, Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learn- ing Research, pages 3796-3805. PMLR, 2018.","title":"Sparsemap: Differentiable sparse structured inference","authors":"Vlad Niculae, F T André, Mathieu Martins, Claire Blondel, Cardie","year":"2018","venue":"Proceedings of the 35th International Conference on Machine Learning, ICML 2018, Stockholmsmässan","doi":null,"url":null,"internal_url":null},{"position":20,"raw":"[Peyré and Cuturi, 2019] Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data science. Foundations and Trends® in Machine Learn- ing, 11(5-6):355-607, 2019.","title":"Gabriel Peyré and Marco Cuturi. Computational optimal transport: With applications to data science","authors":"Cuturi Peyré","year":"2019","venue":"Foundations and Trends® in Machine Learning","doi":null,"url":null,"internal_url":null},{"position":21,"raw":"[Qian et al., 2020] Wesley Wei Qian, Nathan T. Russell, Claire L. W. Simons, Yunan Luo, Martin D. Burke, and Jian Peng. Integrating deep neural networks and sym- bolic inference for organic reactivity prediction. Chem- Rxiv, 2020.","title":"Integrating deep neural networks and symbolic inference for organic reactivity prediction","authors":"Wesley Wei Qian, Nathan T Russell, Claire L W Simons, Yunan Luo, Martin D Burke, Jian Peng","year":"2020","venue":"Chem-Rxiv","doi":null,"url":null,"internal_url":null},{"position":22,"raw":"[Raucci et al., 2022] Umberto Raucci, Valerio Rizzi, and Michele Parrinello. Discover, sample, and refine: Explor- ing chemistry with enhanced sampling techniques. The Journal of Physical Chemistry Letters, 13:1424-1430, 02 2022.","title":"Discover, sample, and refine: Exploring chemistry with enhanced sampling techniques","authors":"Umberto Raucci, Michele Valerio Rizzi, Parrinello","year":"2022","venue":"The Journal of Physical Chemistry Letters","doi":null,"url":null,"internal_url":null},{"position":23,"raw":"[Sacha et al., 2020] Mikolaj Sacha, Mikolaj Blaz, Piotr Byrski, Pawel Wlodarczyk-Pruszynski, and Stanislaw Jas- trzebski. Molecule edit graph attention network: Model- ing chemical reactions as sequences of graph edits. CoRR, abs/2006.15426, 2020.","title":"Molecule edit graph attention network: Modeling chemical reactions as sequences of graph edits","authors":"Mikolaj Sacha, Mikolaj Blaz, Piotr Byrski, Pawel Wlodarczyk-Pruszynski, Stanislaw Jastrzebski","year":"2020","venue":null,"doi":null,"url":null,"internal_url":null},{"position":24,"raw":"[Sander et al., 2022] Michael E. Sander, Pierre Ablin, Math- ieu Blondel, and Gabriel Peyré. Sinkformers: Transform- ers with doubly stochastic attention. In International Con- ference on Artificial Intelligence and Statistics, AISTATS 2022, 28-30 March 2022, Virtual Event, volume 151 of Proceedings of Machine Learning Research, pages 3515- 3530. PMLR, 2022.","title":"Sinkformers: Transformers with doubly stochastic attention","authors":"Michael E Sander, Pierre Ablin, Mathieu Blondel, Gabriel Peyré","year":"2022","venue":"International Conference on Artificial Intelligence and Statistics, AISTATS 2022","doi":null,"url":null,"internal_url":null},{"position":25,"raw":"[Schwaller et al., 2019] Philippe Schwaller, Teodoro Laino, Théophile Gaudin, Peter Bolgar, Christopher A. Hunter, Costas Bekas, and Alpha A. Lee. Molecular transformer: A model for uncertainty-calibrated chemical reaction pre- diction. ACS Central Science, 5(9), 2019.","title":"Molecular transformer: A model for uncertainty-calibrated chemical reaction prediction","authors":"Philippe Schwaller, Teodoro Laino, Théophile Gaudin, Peter Bolgar, Christopher A Hunter, Costas Bekas, Alpha A Lee","year":"2019","venue":"ACS Central Science","doi":null,"url":null,"internal_url":null},{"position":26,"raw":"[Segler and Waller, 2016a] Marwin H. S. Segler and Mark P. Waller. Modelling chemical reasoning to predict reac- tions. Chemistry-A European Journal, 23(25):6118-6128, 2016a.","title":"Modelling chemical reasoning to predict reactions","authors":"Segler, H S Waller ; Marwin, Mark P Segler, Waller","year":"2016","venue":"Chemistry-A European Journal","doi":null,"url":null,"internal_url":null},{"position":27,"raw":"[Segler and Waller, 2017] Marwin H. S. Segler and Mark P. Waller. Neural-symbolic machine learning for retrosyn- thesis and reaction prediction. Chemistry -A European Journal, 23(25):5966-5971, 2017.","title":"Neural-symbolic machine learning for retrosynthesis and reaction prediction","authors":"H S Marwin, Mark P Segler, Waller","year":"2017","venue":"Chemistry -A European Journal","doi":null,"url":null,"internal_url":null},{"position":28,"raw":"[Sinkhorn, 1966] Richard Sinkhorn. A relationship between arbitrary positive matrices and stochastic matrices. Cana- dian Journal of Mathematics, 18:303-306, 1966.","title":"A relationship between arbitrary positive matrices and stochastic matrices","authors":"Richard Sinkhorn","year":"1966","venue":"Canadian Journal of Mathematics","doi":null,"url":null,"internal_url":null},{"position":29,"raw":"[Sohn et al., 2015] Kihyuk Sohn, Honglak Lee, and Xinchen Yan. Learning structured output representation using deep conditional generative models. In Advances in Neural In- formation Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Montreal, Quebec, Canada, pages 3483-3491, 2015.","title":"Learning structured output representation using deep conditional generative models","authors":"Sohn","year":"2015","venue":"Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015","doi":null,"url":null,"internal_url":null},{"position":30,"raw":"[Sutskever et al., 2014] Ilya Sutskever, Oriol Vinyals, and Quoc V. Le. Sequence to sequence learning with neural networks. In Advances in Neural Information Process- ing Systems 27: Annual Conference on Neural Information Processing Systems 2014, December 8-13 2014, Montreal, Quebec, Canada, pages 3104-3112, 2014.","title":"Le. Sequence to sequence learning with neural networks","authors":"Sutskever","year":"2014","venue":"Advances in Neural Information Processing Systems 27: Annual Conference on Neural Information Processing Systems 2014","doi":null,"url":null,"internal_url":null},{"position":31,"raw":"[Tay et al., 2020] Yi Tay, Dara Bahri, Liu Yang, Donald Met- zler, and Da-Cheng Juan. Sparse sinkhorn attention. In Proceedings of the 37th International Conference on Ma- chine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pages 9438-9447. PMLR, 2020.","title":"Sparse sinkhorn attention","authors":"Yi Tay, Dara Bahri, Liu Yang, Donald Metzler, Da-Cheng Juan","year":"2020","venue":"Proceedings of the 37th International Conference on Machine Learning, ICML 2020","doi":null,"url":null,"internal_url":null},{"position":32,"raw":"[Tetko et al., 2020] Igor Tetko, Pavel Karpov, Ruud Deursen, and Guillaume Godin. State-of-the-art aug- mented nlp transformer models for direct and single-step retrosynthesis. Nature Communications, 11, 11 2020.","title":"State-of-the-art augmented nlp transformer models for direct and single-step retrosynthesis","authors":"Igor Tetko, Pavel Karpov, Ruud Deursen, Guillaume Godin","year":"2020","venue":"Nature Communications","doi":null,"url":null,"internal_url":null},{"position":33,"raw":"[Tu and Coley, 2021] Zhengkai Tu and Connor W. Co- ley. Permutation invariant graph-to-sequence model for template-free retrosynthesis and reaction prediction. CoRR, abs/2110.09681, 2021.","title":"Zhengkai Tu and Connor W. Coley. Permutation invariant graph-to-sequence model for template-free retrosynthesis and reaction prediction","authors":"Coley Tu","year":"2021","venue":null,"doi":null,"url":null,"internal_url":null},{"position":34,"raw":"[Vaswani et al., 2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In Advances in Neural Information Processing Sys- tems 30: Annual Conference on Neural Information Pro- cessing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, pages 5998-6008, 2017.","title":"Attention is all you need","authors":"Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, Illia Polosukhin","year":"2017","venue":"Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017","doi":null,"url":null,"internal_url":null},{"position":35,"raw":"[Weininger, 1988] David Weininger. Smiles, a chemical lan- guage and information system. 1. introduction to method- ology and encoding rules. J. Chem. Inf. Comput. Sci., 28(1):31-36, 1988.","title":"Smiles, a chemical language and information system. 1. introduction to methodology and encoding rules","authors":"David Weininger","year":"1988","venue":"J. Chem. Inf. Comput. Sci","doi":null,"url":null,"internal_url":null},{"position":36,"raw":"[Wormell and Reich, 2021] Caroline L. Wormell and Sebas- tian Reich. Spectral convergence of diffusion maps: Im- proved error bounds and an alternative normalization. SIAM J. Numer. Anal., 59(3):1687-1734, 2021.","title":"Spectral convergence of diffusion maps: Improved error bounds and an alternative normalization","authors":"Caroline L Wormell, Sebastian Reich","year":"2021","venue":"SIAM J. Numer. Anal","doi":null,"url":null,"internal_url":null},{"position":37,"raw":"[Zhao et al., 2022] Ming Zhao, Lei Fang, Li Tan, Yves Lep- age, and Jian-Guang Lou. Leveraging reaction-aware substructures for retrosynthesis and reaction prediction. CoRR, abs/2204.05919, 2022.","title":"Leveraging reaction-aware substructures for retrosynthesis and reaction prediction","authors":"Ming Zhao, Lei Fang, Li Tan, Yves Lepage, Jian-Guang Lou","year":"2022","venue":null,"doi":null,"url":null,"internal_url":null}],"license_note":"Copyright International Joint Conferences on Artificial Intelligence."}