# Doubly Stochastic Graph-based Non-autoregressive Reaction Prediction

arXiv: https://arxiv.org/abs/2306.06119 · HTML: https://mlsanity.com/paper/2306.06119/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 19 of its 44 references that are in the map.
- First posted on arXiv 5 Jun 2023.

## Builds on (most cited first, 19 of 19)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Sequence to Sequence Learning with Neural Networks](https://mlsanity.com/paper/1409.3215/) (2014)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Sparse Sinkhorn Attention](https://mlsanity.com/paper/2002.11296/) (2020)
- [Molecular Transformer - A Model for Uncertainty-Calibrated Chemical Reaction Prediction](https://mlsanity.com/paper/1811.02633/) (2018)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [Sinkformers: Transformers with Doubly Stochastic Attention](https://mlsanity.com/paper/2110.11773/) (2021)
- [Predicting Organic Reaction Outcomes with Weisfeiler-Lehman Network](https://mlsanity.com/paper/1709.04555/) (2017)
- [Molecule Edit Graph Attention Network: Modeling Chemical Reactions as Sequences of Graph Edits](https://mlsanity.com/paper/2006.15426/) (2020)
- [Permutation invariant graph-to-sequence model for template-free retrosynthesis and reaction prediction](https://mlsanity.com/paper/2110.09681/) (2021)
- [SparseMAP: Differentiable Sparse Structured Inference](https://mlsanity.com/paper/1802.04223/) (2018)
- [Graph Transformation Policy Network for Chemical Reaction Prediction](https://mlsanity.com/paper/1812.09441/) (2018)
- [A Generative Model For Electron Paths](https://mlsanity.com/paper/1805.10970/) (2018)
- [Spectral convergence of diffusion maps: improved error bounds and an alternative normalisation](https://mlsanity.com/paper/2006.02037/) (2020)
- [Manifold learning with bi-stochastic kernels](https://mlsanity.com/paper/1711.06711/) (2017)
- [Non-Autoregressive Electron Redistribution Modeling for Reaction Prediction](https://mlsanity.com/paper/2106.07801/) (2021)
- [Modelling Chemical Reasoning to Predict Reactions](https://mlsanity.com/paper/1608.07117/) (2016)

## Built on by (most cited first, 6 of 6)

- [Challenging reaction prediction models to generalize to novel chemistry](https://mlsanity.com/paper/2501.06669/) (2025)
- [A Unified View of Deep Learning for Reaction and Retrosynthesis Prediction: Current Status and Future Challenges](https://mlsanity.com/paper/2306.15890/) (2023)
- [Doubly Stochastic Adaptive Neighbors Clustering via the Marcus Mapping](https://mlsanity.com/paper/2408.02932/) (2024)
- [ReactionTeam: Teaming Experts for Divergent Thinking Beyond Typical Reaction Patterns](https://mlsanity.com/paper/2310.04674/) (2023)
- [ReactAIvate: A Deep Learning Approach to Predicting Reaction Mechanisms and Unmasking Reactivity Hotspots](https://mlsanity.com/paper/2407.10090/) (2024)
- [Driving Reaction Trajectories via Latent Flow Matching](https://mlsanity.com/paper/2602.10476/) (2026)

## Abstract (arXiv)

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 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 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.