Doubly Stochastic Graph-based Non-autoregressive Reaction Prediction

1 Introduction

Reaction prediction is a crucial task in computational chemistry. With a reliable prediction model, chemists can verify the accuracy of retrosynthetic routes and potentially uncover new chemical reactions, which could significantly benefit the drug discovery industry. However, given a set of reactants, the number of potential products can be combinatorially large, necessitating automatic tools to narrow down the search space. Consequently, the use of deep learning techniques to automatically infer possible reaction products has become prevalent and essential.

∗This work is done when Ziqiao Meng worked as an intern in Tencent AI Lab.

†Corresponding authors: Peilin Zhao, and Irwin King.

Translation-based autoregressive models [Schwaller et al., 2019; Tetko et al., 2020; Tu and Coley, 2021; Lu and Zhang, 2022; Irwin et al., 2022; Zhao et al., 2022] have dominated the design of template-free end-to-end reaction models in recent years. These models represent molecules as sequences of characters, called SMILES [Weininger, 1988], and formulate the reaction prediction problem as a neural machine translation problem [Sutskever et al., 2014; Vaswani et al., 2017; Devlin et al., 2019]. This class of models achieves great predictive performance and does not rely on pre-computed atom-mapping information. However, autoregressive models have two major shortcomings. First, autoregressive sampling is very inefficient because the model has to generate the final predictions step-by-step from intermediate parts. Second, training autoregressive models requires pre-defined generation orders, while the absolute generation order of molecules is ambiguous.

To avoid these issues, a non-autoregressive model called NERF [Bi et al., 2021] has been proposed. It predicts the redistribution of electrons, which is an important observation because chemical reactions can be reflected by bond reconnection, since the elements of the involved atoms remain unchanged. Bond reconnection inherently involves electron redistribution. For example, forming a single bond between a carbon atom and a nitrogen atom involves sharing a pair of electrons between them. Therefore, predicting electron redistribution ∆E is sufficient for inferring possible complete products. With a novel design for the electron redistribution decoder, NERF has achieved state-of-the-art top-1 accuracy and much faster sampling speed with parallel inference.

However, an important issue that has been neglected is that ∆ˆE must simultaneously follow the electron counting rule and the symmetry rule to better approximate ∆E. The chemical reaction that adheres to both of these rules is illustrated in Figure 1. We can clearly see that ∆E is row-wise conservative (conforming to the electron counting rule) and symmetric. Currently, the NERF decoder neglects to symmetrize ∆ˆE, only ensuring that it follows the electron counting rule through combinations of BondBreaking and BondFormation self-attention matrices. This leads to inexact modeling of ∆E and impairs model performance. The predicted product adjacency matrix ˆEP is forced to be symmetric, but this operation implicitly corrupts the row-wise conservatism of ∆ ˆE.

In this work, we propose a novel framework called ReactionSink to refine the decoder. The framework applies Sinkhorn's algorithm to iteratively update the BondBreaking and BondFormation self-attention matrices. With a sufficient number of iterations, each self-attention matrix converges to a doubly stochastic matrix. To generalize this solution to multi-head attention mechanisms, we impose augmented constraints on the Sinkhorn normalization operation. We theoretically show that the NERF decoder leads to an implicit contradiction in preserving both physical constraints. In contrast, doubly stochastic self-attention mappings lead to valid ∆ ˆE, which preserve both rules simultaneously and are not corrupted by the symmetrization on ˆEP . We further discuss the implicit connections between the reaction prediction problem and the discrete optimal transport problem to validate the rationality of imposing the doubly stochasticity constraint on the decoder as an information prior. Comprehensive empirical results on the open benchmark dataset USPTOMIT demonstrate that our approach consistently outperforms baseline non-autoregressive reaction prediction models. Our contributions can be summarized as follows:

• We introduce two chemical rules that the predicted electron redistribution ∆ˆE should follow and theoretically show that the current non-autoregressive decoder inevitably violates either of these chemical rules;
• We propose a novel decoder design that conforms to both chemical rules through leveraging combinations of doubly stochastic self-attention mappings and we generalize our approach to multi-head attention mechanism;
• We provide some new interpretations on electron redistribution modeling, which facilitates further research on potential connections between the reaction prediction problem and the optimal transport problem.

2 RelatedWork

Reaction Prediction. In the early years, template-based methods were widely adopted for reaction prediction. Experts deduced possible products with the help of concluded reaction templates. Although these methods are reliable, they suffer from poor generalization. Quantum computation methods [Dral, 2020; Lam et al., 2020; Raucci et al., 2022] have also been adopted for reaction predictions, but they suffer from low computation speeds, despite being templatefree. With the rise of deep learning, various deep learning architectures have been utilized in reaction modeling. Models that learn neural matching between reactions and templates have been proposed [Segler and Waller, 2017; Segler and Waller, 2016a]. WLDN [Jin et al., 2017] proposes a template-free model that splits reaction prediction into a reaction center identification stage and candidate ranking stage. MEGAN and GTPN [Sacha et al., 2020; Do et al., 2019] model reactions as a sequence of graph edits, while Electro [Bradshaw et al., 2019] regards reactions as a sequence of electron transfers. A major class of models are translationbased reaction prediction models [Schwaller et al., 2019; Tetko et al., 2020; Tu and Coley, 2021; Lu and Zhang, 2022; Irwin et al., 2022; Zhao et al., 2022], which mainly transfer techniques from language modeling to reaction modeling. NERF [Bi et al., 2021] refines the decoder to model electron redistribution and achieves non-autoregressive modeling with state-of-the-art top-1 accuracy and parallel inference.

Sinkhorn and Attention. The Sinkhorn algorithm [Sinkhorn, 1966; Cuturi, 2013; Peyr´e and Cuturi, 2019] is a well-studied tool for approximating solutions to the optimal transport problem. For instance, a new set pooling algorithm is presented in [Mialon et al., 2021], which embeds the optimal transport plan between input sets and reference sets, and then uses the Sinkhorn algorithm as a fast solver. Sinkhorn normalization is presented in [Marshall and Coifman, 2019; Wormell and Reich, 2021], and its convergence properties to a doubly stochastic matrix have been theoretically validated. SparseMAP [Niculae et al., 2018] uses doubly stochastic attention matrices in LSTM-based encoder-decoder networks. Sinkhorn updates are also applied for differentiable ranking over internal representations in [Adams and Zemel, 2011]. Gumbel-Sinkhorn [Mena et al., 2018] applies the Sinkhorn algorithm to learn stochastic maximization over permutations. The Sinkhorn algorithm is used to iteratively generate doubly stochastic matrices for approximating smooth filters in [Milanfar, 2013]. It has also been applied to refine self-attention mechanisms in [Vaswani et al., 2017]. The Sparse Sinkhorn Transformer [Tay et al., 2020] learns sparse self-attention by introducing a sorting network that generates a doubly stochastic matrix to permute input sequence elements. Sinkformer [Sander et al., 2022] applies the Sinkhorn algorithm to derive doubly stochastic transformers.

3 ProposedMethods

This section introduces the learning framework of our new non-autoregressive reaction prediction model, ReactionSink. We first describe its problem formulation and general training process. Then we introduce its encoder architecture, which obtains the conditional reaction latent embedding. Most importantly, we present our novel decoder mechanism that leverages doubly stochastic self-attention matrices to obtain ∆E. In the theoretical analysis section, we mathematically compare our decoder with the NERF [Bi et al., 2021] decoder to demonstrate the advantages of our approach. Finally, we briefly discuss the new insights brought by ReactionSink.

Problem Formulation. A chemical reaction is represented as a pair (GR, GP ), where GR = (V, ER, f R) denotes the set of reactants and GP = (V, EP , f P ) denotes the set of products, V is the set of |V | atoms, ER and EP denote the adjacency matrix of GR and GP respectively, f R and f P denote the atom feature matrix of GR and GP respectively. Following the atom-mapping principle, reactants and products share the same set of atoms V . Unlike most of the papers representing E as a 3D tensor, where E ∈ R|V |×|V | is a 2D matrix and Eij denotes the number of shared electrons between atom i and atom j (i.e. single bond is 1, double bond is 2). For special aromatic bond, we denote it as Eij = 1.5. The goal of reaction prediction is to predict the set of products GP given a set of reactants GR.

Following NERF [Bi et al., 2021], we also adopt the conditional variational autoencoder [Sohn et al., 2015] (CVAE) architecture to approximate p(GP |GR) by introducing a latent variable z. Instead of directly maximizing the loglikelihood log p(GP |GR), CVAE is to maximize its evidencelower bound (ELBO):

where q(z|GP ,GR) is the reaction encoder with reaction (GR,GP ) as input and low-dimensional embedding hz as output, p(GP |GR, z) is product decoder with reactants GR and latent embedding hz as input, p(z|GR) denotes the prior distribution of latent variable z. KL term is minimizing the gap between q(z|GP ,GR) and p(z|GR). In real implementations, the backbone network architectures of the encoder q(z|GP ,GR) are graph neural networks (GNNs) and transformers. GNNs capture local interactions within each molecule and transformers capture global interactions across different molecules. With this architecture, GR and GP will be projected to reactants embedding hR and products embedding hP respectively. The cross attention layer is used for mapping hR to latent hz with hP as teacher forcing during training. Finally, the latent conditional reaction embedding is derived such that ˆhz = hR + hz . Note that similar reaction embedding networks are widely adopted in many related works as long as transformer-based reaction encoder architectures are used. Please refer the detailed architectures in Appendix 1 and the general architectures of ReactionSink is illustrated in Figure 2.

3.1 Decoder with Physical Constraints

The core mechanism of the decoder p(GP |GR, z) is first decoding electron redistribution matrix ∆ˆE from conditional latent ˆhz , and then add it to ER to predict ˆEP , i.e., ˆEP = ER+∆ ˆE. Before presenting ReactionSink decoder in details, we first introduce two important physical rules, the electroncounting rule and the symmetry rule, which should be followed by the predicted ∆ˆE. However, these rules are not explicitly stated in the previous literature. We will mathematically characterize these two rules as matrix properties.

Electron Counting Rule. This is an important chemical rule of thumb: the octet rule. It states that main-group elements (carbon, nitrogen, oxygen, halogens) tend to bond in a way that each atom has eight electrons in its valence shell. Following this rule, electron redistribution of each atom is conservative, allowing it to maintain eight electrons in its valence shell. For example, when a carbon atom forms a new bond with one atom, it tends to break an existing bond with another atom in order to maintain eight electrons in its valence shell. This empirical rule leads to an important mathematical property of ∆E, which is row-wise conservative (i.e., ∑j ∆Eij = 0 for each atom i). Note that this rule is more informative than the charge conservation rule, which can be mathematically expressed as ∑i,j ∆Eij = 0.

Symmetry Rule. The ground-truth ∆E should clearly be symmetric. For example, if the edge number between atom i and atom j increases by 1, then the edge number between atom j and atom i should also increase by 1. To better approximate ∆E, ∆ ˆE should also be symmetric.

We have shown that ∆ˆE should preserve two mathematical properties which are ∑j ∆ˆEij = 0 for any j and ∆ˆEij = ∆ˆEji for any i, j. This can easily come up with a new conclusion that ∆ ˆE should also be column-wise conservative such that ∑i ∆ ˆEij = 0 for any j. Therefore, ∆ ˆE should be both row-wise conservative and column-wise conservative in the strict sense, which is doubly conservative. Inspired by NERF decoder that leverages the combinations of row-wise stochastic self-attention mappings to achieve rowconservative ∆ˆE, we achieve doubly conservative ∆ˆE by using combinations of two doubly stochastic self-attention mappings. Specifically, we decode two self-attention mappings BondFormation W + and BondBreaking W - from the conditional latent ˆhz . BondFormationW + is to predict bond addition and BondBreaking W - is to predict bond removal. The detailed self-attention decoding mechanism is shown in Appendix 2. Then the combinations of W + and W - should satisfy the following equations:

To satisfy the above equations, we force W + and W - to be doubly stochastic such that:

Note that forcing W + and W - to be doubly stochastic is probably not the only solution to Eq 2. But stochasticity (sum to 1) is the most straightforward solution and more compatible with normalization operations. Therefore, if W + and W - can be doubly stochastic matrices as stated in Eq 3, their combinations can lead to a doubly conservative ∆ ˆE. Then we can apply symmetrization operation on ∆ ˆE such that ∆ ˆE+(∆ ˆE)T 2 , which is averaging predictions of ∆ˆEij and ∆ˆEji for each i, j. Interestingly, symmetrization in this case would not corrupt the conservative of ∆ˆE. We will theoretically validate these claims in the next section.

Unfortunately, decoding doubly stochastic self-attention mappings from conditional latent variables is non-trivial. Recall that the vanilla self-attention mapping is obtained as SoftMax(exp(QKT )), where Q and K are queries and keys linearly projected from latent variable ˆhz . Note that the SoftMax operator is basically a row-wise normalization technique that forces self-attention mappings to be row-wise stochastic, while it is not necessarily column-wise stochastic. Therefore, a specific operator is required to extend Bondformation W + and Bondbreaking W - to doubly stochastic matrices. Furthermore, this operator should be differentiable so that it can be naturally merged into decoder neural networks.

3.2 Sinkhorn's Algorithm

An iterative algorithm to convert positive matrix to doubly stochastic matrix is Sinkhorn's algorithm. The main idea of Sinkhorn's algorithm is alternatively applying row-wise normalization and column-wise normalization on the given positive matrix. Given a positive square matrix X ∈ RN×N , we define the Sinkhorn operator as follows:

where Tr denotes the row-wise normalization and Tc denotes the column-wise normalization. They are formally described as follows:

where ⊘ denotes the element-wise division operator and 1 denotes the all one vector. We only need to apply the above normalization operations alternatively with enough number of iterations. In real implementation, to stabilize the updating process, we put the computation of Eq 5 in log domain. Details of this implementation are shown in the Appendix 6.

Sinkhorn [Sinkhorn, 1966] proves that S∞(X) belongs to the Birkhoff polytope, the set of doubly stochastic matrices, such that S∞(X)1N = 1N and S∞(X)T 1N = 1N. Hence, applying this Sinkhorn operator on W + and W - separately will make them converge to doubly stochastic matrices. This Sinkhorn operator can be easily coupled with the self-attentionmechanismas an additional layer. In addition, it can be perfectly suited for backpropagation updates of neural networks. Specifically, this Sinkhorn operator can be stacked after the vanilla self-attention, since the vanilla self-attention mapping can be regarded as Tr(exp(QKT )), which is equivalent to the initialization step plus the row-wise normalization step in Eq 4. Therefore, in real implementations, we start with the column-wise normalization Tc after the SoftMax layer in the vanilla self-attention.

Currently, each entry of ∆ˆE is in range [-1, 1] while edge number changes more than 1 in many cases. Therefore, we apply multi-head attention mechanism such that ∆ˆE = ∑4d=1 W +d - ∑4d=1 W -d with d = 4 attention heads. In this way, ∆ˆEij is in range [-4, 4], which could cover enough support of edge number changes of chemical reactions. We only need to slightly change Eq 3 by increasing conservation constraints to 4 and impose the augmented doubly conservative constraint on the sum of attention heads ∑4d=1 W +d and ∑4d=1 W -d respectively as follows:

For applying the Sinkhorn normalization operator to the multi-head attention mechanism, we only need to impose the augmented constraints to both row-wise normalization and column-wise normalization such that:

To conclude, with Sinkhorn's algorithm, we can efficiently convert W + and W - to doubly stochastic matrices, to finally achieve symmetric doubly conservative ∆ ˆE.

4 Theoretical Analysis

In this section, we first theoretically show that the NERF decoder cannot output valid ∆ˆE that follows the previously mentioned two physical constraints. In comparison, we mathematically validate that our ReactionSink decoder can generate valid ∆ˆE and the symmetry is not contradictory to the conservative.

In NERF decoder, it generates two vanilla self-attention mappings W + and W - based on conditional latent for BondFormation and BondBreaking respectively and approximates ∆ ˆE through subtraction with multi-head attention such that ∆ˆE = ∑4d=1 W +d - ∑4d=1 W -d . Then it applies symmetrization on the final predicted ˆEP = ER +∆ ˆE such that EP +(EP )T 2 . Therefore, we begin to verify that ∆ˆE has also been implicitly symmetrized by NERF decoder, which is the following lemma:

Lemma 4.1 Symmetrization on the final predicted product adjacency matrix ˆEP is equivalent to implicit symmetrization on the predicted electron redistribution ∆ ˆE.

The proof procedures are shown in Appendix 2.

Then we argue that NERF decoder design cannot satisfy the two proposed physical constraints simultaneously. Specifically, the symmetrization operation on ˆEP would result in that ∆ˆE violates the first rule. Since we have already shown that the symmetrization of ˆEP is equivalent to the implicit symmetrization of ∆ ˆE in lemma 4.1, we only need to further prove the following lemma:

Lemma 4.2 Implicit symmetrization on ∆ˆE would result in that ∆ˆE violates the electron counting rule such that ∑j ∆ ˆEij ̸= 0 for some i.

The proof procedures are shown in Appendix 3.

To summarize, electron redistribution predictions generated by NERF cannot satisfy both constraints simultaneously. Implicit symmetrization operation applied on ∆ˆE makes it violate the first rule. Otherwise it will violate the second rule if no symmetrization operation is adopted. In short, NERF decoder actually leads to a contradiction between two physical constraints.

In comparison, with doubly stochastic self-attention mapping W + and W -, ReactionSink can conform to both physical rules concurrently. We can prove the following theorem:

Theorem 4.3 If the self-attention mappings for bondformation W + and bondbreaking W - are doubly stochastic, then ∆ˆE would be guaranteed to be doubly conservative and its doubly conservation would not be corrupted by the symmetrization of ˆEP .

The proof procedures are shown in Appendix 4.

To extend the above conclusion to multi-head attention mechanism. The following corollary can be easily proved:

Corollary 4.3.1 If the multi-head self-attention mechanisms for bondformation ∑Dd=1 W +d and bondbreaking ∑Dd=1 W -d are doubly conservative with D-sum constraints, then ∆ ˆE would be guaranteed to be doubly conservative and its doubly conservation would not be corrupted by the symmetrization of ˆEP .

The proof procedures are shown in Appendix 5.

5 Further Discussions

Our novel framework, ReactionSink, provides new insights for the non-autoregressive reaction prediction problem. Doubly stochastic matrices usually have special meaning in mathematics, with the most well-known being the permutation matrix. This suggests that electron redistribution can be formulated as a learned permutation of electrons. Note that permuting a matrix is a multiplication operation, which is different from the current formulation. The permutation formulation is intuitively closer to the discrete nature of electron redistribution, but the detailed formulation method requires further research. More generally, reaction prediction can be connected with the optimal transport (OT) problem. The permutation problem is the most elementary discrete OT problem. Specifically, electron redistribution can also be regarded as allocating electrons to lower the energy of the entire reaction system, which is implicitly connected to the discrete OT problem. Our work presents these new insights to inspire future research, although they are still in their nascent stages and have not been revealed by existing research.

6 Experiments

Dataset. Following previous work, we evaluate our approach on the open public benchmark dataset USPTO-MIT [Jin et al., 2017], which contains 479K reactions filtered by removing duplicates and erroneous reactions from Lowe's original data [Lowe, 2012]. An exisitng work discovers that about 0.3% of reactions in USPTO-479K do not satisfy the non-autoregressive learning settings for various reasons [Bi et al., 2021]. Following this work [Bi et al., 2021], we also filter 0.3% reactions from dataset and deduct 0.3% from the top-k accuracies of our model correspondingly.

Experimental Setting. We conduct experiments on three different splits of reaction prediction, which are random split, tanimoto-0.4 split and tanimoto-0.6 split. Random split is adopted by most of the previous work. Scaffold splits, tanimoto-0.4 split and tanimoto-0.6 split, are adopted by [Kovacs et al., 2021] to test the generalization of reaction model with larger distribution shift between training and testing. Tanimoto similarity is measuring whether two reactions are structurally similar. Higher tanimoto index indicates that two reactions are more similar and otherwise two reactions are dissimilar with each other (e.g. Tanimoto-0.4 has larger distribution gap between training and testing). For original random split, the training set, validation set and testing set follows split ratio 409K:30K:40K. For scaffold split, the split ratio is 392K:30K:50K.

Comparison Baselines. We compare ReactionSink with the following baseline models: WLDN [Jin et al., 2017] is a two-stage model, which firstly predicts reaction centers and then ranks enumerated products; GTPN [Do et al., 2019] models reaction predictions as a series of graph transformations and use policy networks to learn the transformations; MT-base [Schwaller et al., 2019] is a transformerbased autoregressive modeling with both input and output in SMILES sequence format; MEGAN [Sacha et al., 2020] models reaction predictions as series of graph edit operations and generates edit sequences in an autoregressive manner; MT [Schwaller et al., 2019] is MT-base models with data augmentation techniques applied to SMILES sequence input; Symbolic [Qian et al., 2020] introduces the chemical rules to reaction modeling using symbolic inference; Chemformer [Irwin et al., 2022] leverages molecular SMILES encoder pretrained on 100M molecular datasets with three selfsupervised tasks; Graph2Smiles [Tu and Coley, 2021] leverages the similar backbone network of NERF [Bi et al., 2021]; AT×100 [Tetko et al., 2020] leverages molecular transformer with ×100 SMILES augmentations; Sub-reaction [Zhao et al., 2022] leverages motif tree to achieve substructure-aware reaction prediction; NERF [Bi et al., 2021] uses CVAE to achieve non-autoregressive modeling;

Model Configuration and Reproducibility Setting. The major encoder and decoder architectures are following NERF [Bi et al., 2021]. To ensure fair comparison with major baseline models Molecular Transformer and NERF, we set the number of transformer encoder layers and transformer decoder layers (cross-attention layer) to be 4, the same as previous work. And we set the dimension of latent embedding to be 256. For multi-head attention decoder, BondFormation and BondBreaking both have 4 attention heads. The model is optimized using Adam optimizer [Kingma and Ba, 2015] at learning rate 10-4 with linear warm-up and linear learning rate decay. The number of iterations l of Sinkhorn normalization is also an hyperparameter to fine-tune. Larger l indicates more rounds of Sinkhorn normalization. Finally, we train our ReactionSink for 100 epochs with a batch size of 128 using 8 Nvidia V100 GPUs in this work.

EvaluationMetrics. Following the convention of previous work, we adopt the top-k accuracies to evaluate the performance of all the compared algorithms. The top-k accuracy is the percentage of reactions that have the ground-truth product in the set of the top-k predicted molecules. As long as the set of predicted products contains the ground-truth main product, then the prediction would be counted as a correct one. In this work, the value of k is set as 5 different values: {1, 2, 3, 5, 10}.

Sampling Top-k Predictions. Since ReactionSink follows CVAE architecture, multi-modal outputs are generated through sampling k different latent vectors ˆhz by increasing temperature values. To sample top-k predictions, we multiply a scalar temperature parameter t to the variance of standard Gaussian distribution, such that hz s are sampled from different N(0, tI). Increasing the temperature value t would make the model predict different products. Specifically, we sample the first k predictions as our top-k predictions. Lower temperatures are set to output predictions with higher rank (e.g. prediction using t = 1 is treated as the top-1 prediction).

Main Experiments with Random Split and Tanimoto Splits. We conduct experiments under randomsplit and tanimoto splits, which can be seen as a cross-validation process. Table 1 is the major benchmark adopted by previous work under random split conducted by WLDN [Jin et al., 2017]. From this table, we can see that our method reaches the stateof-the-art top-1 accuracy over all baseline methods and consistently outperforms the non-autoregressive model NERF with top-k accuracies. Note that these results are evaluated on models without knowing reagents, which means these models do not know which reactant is reagent during inference stage. We can see that currently non-autoregressive models are performing worse than autoregressive models when k is higher. We conjecture that this may be caused by the simple sampling method for CVAE, or simple latent distribution assumption on uncertainty modeling (Gaussian distribution). In the future, it is worth to find better sampling method for CVAE, and explore stronger generative model than CVAE for reaction modeling. From Table 2 and Table 3, we can observe that non-autoregressive models have stronger generalization than autoregressive models. We conjecture that electron redistribution modeling proposed by NERF [Bi et al., 2021] is closer to the nature of this problem. Under tanimoto splits, ReactionSink consistently improves the performance of nonautoregressive models in terms of top-k accuracies.

Ablation Studies on Sinkhorn Iteration. Since the Sinkhorn algorithm is also an iterative normalization process, we conduct ablation studies on the number of iterations l to check its effect on model performance. From table 5, we can see that the model performance is consistently improved when increasing l from 1 to 5. However, when l increases to 10, it does not bring further performance gains compared to l = 3. This demonstrates that the Bondformation and BondBreaking matrices are converted to doubly stochastic matrices with few iterations of the Sinkhorn normalization.

Computational Efficiency. The main concern about ReactionSink is its computational complexity, which is caused by the additional layers of Sinkhorn normalization. Fortunately, with a few iterations of Sinkhorn normalization, the selfattention matrices will converge to doubly stochastic matrices with negligible error. Therefore, the computational burden would not be significantly increased due to additional normalization. Following NERF, we report the wall-time and latency of the inference model. Wall-time is the total time cost for inferring all testing samples, and latency is the inference time cost for a single testing sample. The detailed wall-time and latency computation standards are stated in NERF [Bi et al., 2021]. From Table 4, we can see that the proposed method only adds a few additional seconds to Wall-time and a few more milliseconds to latency. This demonstrates that the additional normalization layers do not create a trade-off between accuracy and speed. Regarding the training process, additional Sinkhorn normalization only increases the training time by a few seconds for each epoch, which is almost negligible compared to the total training time. Overall, ReactionSink is more suitable for high-throughput reaction predictions than auto-regressive methods.

7 Conclusion

In this work, we first introduce two chemical rules that govern electron distribution and characterize their corresponding mathematical properties. We then identify that the current electron redistribution decoder does not maintain these two physical constraints simultaneously, and we theoretically prove our claims. To address this issue, we propose the ReactionSink architecture to extend the current self-attention mapping to doubly stochastic matrices. We also prove that the predicted electron distribution generated by our proposed methods better adheres to the physical constraints. Additionally, we establish the connections between electron redistribution and the learned optimal transport problem. This new perspective has the potential to lead to a new problem formulation for reaction prediction. Experimental results demonstrate that ReactionSink consistently outperforms the stateof-the-art non-autoregressive reaction prediction model and does not require expensive computational costs.

Acknowledgments

The work described here was partially supported by grants from the National Key Research and Development Program of China (No. 2018AAA0100204) and from the Research Grants Council of the Hong Kong Special Administrative Region, China (CUHK 14222922, RGC GRF, No. 2151185)