Doubly Stochastic Graph-based Non-autoregressive Reaction Prediction

Doubly Stochastic Graph-based Non-autoregressive Reaction Prediction

Ziqiao Meng, Peilin Zhao, Yang Yu, Irwin King

Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence
Main Track. Pages 4064-4072. https://doi.org/10.24963/ijcai.2023/452

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.
Keywords:
Machine Learning: ML: Applications
Machine Learning: ML: Structured prediction
Machine Learning: ML: Attention models
Multidisciplinary Topics and Applications: MDA: Physical sciences