Learning Condensed Graph via Differentiable Atom Mapping for Reaction Yield Prediction

We thank the reviewer for their comments, which we address as follows.

CGR meaningfully approximates TS?

During rebuttal, we performed quantum chemical validation analogous to that reported by Choi, in 'Prediction of transition state structures of gas-phase chemical reactions via machine learning' Nature Commun. 2023.

Here, we have used our CGR based TS as the initial guess for quantum chemical optimization for reactions in a GP dataset. We observed that 80% of these converged to TS. Subsequent QM calculations could produce correct reactant and product from the TS for 60% cases.

In addition, we also compare CGR-based TSs with true (QM-computed) TSs using two metrics:

(1) Accuracy: The fraction of total connections in the CGR-based TS which match with the QM computed TS.
(2) RMSD: Root mean squared deviation (RMSD -lower the better) between atom positions obtained from CGR-based TS and the QM computed TS. Results for the GP dataset for our method and other interpolation based methods as follows. We observe that our method performs better.

Our CGR serves as a surrogate for TS to enhance yield prediction, not to provide an accurate TS prediction. More accurate TS prediction requires TS supervision during training, which most yield datasets lack. Complex datasets like SC contain ~200 atoms per reaction. Here, QM calculations on TSs take a couple of days per reaction--- assuming a good initial TS guess, which is even more time-consuming given the complexity of the dataset. Hence, general evaluation for our CGR was infeasible.

However, as the reviewer suggested, we performed experiments on reaction instances in the GP dataset, which shows that despite absence of supervision, our CGR shows better proximity to the TS than its alternatives.

compare to quantum-ML models

We feed the QM-optimized-TS (TS*) in our model to compute yield prediction error $E^*(r) = |y^*(r)-y _{TS^*}(r)|$ for each reaction $r$ and then compute yield prediction error $E(r) = |y^*(r)-y _{CGR}(r)|$ provided by our CGR based TS. The average difference between these errors is $<$ 8% of QM-optimized-TS based yield prediction error.

generalize to different reaction classes

We compute MAE on test examples that include reactions from SC dataset, whereas the model is trained on DF dataset. We observe that our method shows better performance.

ignore reaction conditions?

We do incorporate both catalyst and solvent, when present, into the reactant set R. We will explicitly mention them in the revised paper. Temperature was excluded as it is invariant in the case of high-throughput experimentation datasets (e.g., DF and NS datasets), which maintain consistent reaction conditions throughout, or varies mildly for others (e.g., SC). Integrating temperature into the node/edge features might enhance CGR or yield prediction quality.

a catalyst as a part of the reagent

In i-th step of a multi-step reaction, the reactant set $R_i$ is a disjoint graph of all the participating species including catalyst, base, solvent, etc.

Failure mode

Upon close inspection, we observe that our method couldn't do well in reactions that are very unusual/rare in terms of bond formation and breaking, due to their under-representation in training samples. Poor predictions on such reactions are unlikely to make any adverse impact, since they are too rare. If the dataset labels (yield values) carry large measurement errors, then also our model (and the baselines) would perform poorly.

lack yield visualization

yhist.pdf in https://bit.ly/rbynet shows that our predicted yield distribution mimics true yield distribution, very closely.

yield for rearrangements, pericyclic processes

Our datasets include rearrangements (present in GP), pericyclic processes (present in GP), metal-catalyzed bond formations (all reactions in SC), etc. Following are MAE based comparisons.

experimental design has limitations, dataset, benchmarking

We evaluated over 10 SOTA datasets against six baselines across ten splits for all experiments, including ablation study, to maintain consistency. We believe that our work conducts experiments with more rigor and comprehensiveness than any previous yield prediction baselines.

We found that the pre-final layer embedding $z_r$ (Eq. 19) correlates with activation energies (corr.pdf in https://bit.ly/rbynet). Extracting mechanistic signals from yield alone is challenging and no prior yield prediction work attempts it-- they only focus on black box yield predictors. Hence, our work makes notable progress, paving the way for future research.

We thank the reviewer for their comments, which we address below.

results for atom mapping

Appendix E.3 contains the table with +- std error. Paired t-test in majority of cases revealed that the performance gain achieved by our method is statistically significant with $p=0.01$

Approach is novel but memory intensive

Our method utilizes the most compute on computing the permutation matrix using Sinkhorn iterations, which has $O(N^2)$ complexity if $N$ is the total number of nodes. However, this is not a hard bottleneck and can be easily overcome using low rank OT.

Low rank OT provides highly efficient Sinkhorn iterations (see Scetborn et al. Low-Rank Sinkhorn Factorization, ICML 2021). We first apporximate $M = \sum _{i=1} ^d \max (H _R[u,i], H _I [u',i])\approx AB^T$ where $A$ and $B$ are $N \times d$ and $N \times d$ low rank matrices. Similarly, we approximate the transport matrix $P$ as $P \approx Q D(1/g) R^T$ where $Q, R$ are $N \times r$ matrices, $g$ is $r \times 1$, and $D(1/g)$ is a diagonal matrix with entries $1/g_i$. To ensure $P$ is doubly stochastic, we enforce:

$$Q\mathbf{1} = \mathbf{1}, \quad \mathbf{1}^T Q = g^T$$ $$R\mathbf{1} = \mathbf{1}, \quad \mathbf{1}^T R = g^T$$ $$g^T \mathbf{1} = 1, \quad Q, R, g > 0$$

Instead of minimizing $Tr(P^T M)$, we minimize: $Tr((Q D(1/g) R^T)^T AB^T)$ using alterting minimization wrt $Q, R, g$. To optimize $Q$, we rewrite:

$$Tr((Q D(1/g) R^T)^T AB^T) = \langle Q, AB^T R D(1/g) \rangle$$

We apply Sinkhorn iterations to: $X = AB^T R D(1/g)$, computed in 3 steps. Note that complexity of multiplying $m$ x $n$ and $n$ x $p$ matrices is $O(mnp)$

1. Compute $R D(1/g)$, which is $O(Nr)$ since $R$ is $N \times r$, $g$ is $r \times 1$.
2. Compute $B^T R D(1/g)$ which is $O(dNr)$ since $B^T: d \times N$ and $R D(1/g): N\times r$
3. Compute $X = AB^T R D(1/g)$ which is $O(Ndr)$ since $A: N \times d$ and $B^T R D(1/g): d\times r$

Thus, complexity reduces from $O(N^2)$ to $O(Ndr)$. The same efficiency gain applies when optimizing $R$ and $g$.

Moreover, in terms of time efficiency, our model is comparable with the some of the GNN based models (HGT and TAG) but not as high as YieldBERT.

USPTO ...not in the main Table 1?

We agree with the reviewer. We would revise Table 1 to include the USPTO results in the main text.

....activation energy directly impacts yield; what is the precise relationship

For a single-step reaction assuming first order kinetics, the relationship is as follows. Suppose, $\Delta G^\ddagger$ represents activation energy in Jmol$^{-1}$. $k_B, h, R$ are the Boltzmann constant ($1.38\times 10^{-23}$ JK$^{-1}$), Planck constant ($6.626\times 10^{-34}$Js), and universal gas constant ($8.314$ JKmol$^{-1}$) respectively. Here, $t$ is the reaction time in sec. If "Temp" is the reaction temperature, then $k$, the first order rate constant is: $k = \dfrac{k_B \text{Temp}}{h} \exp\left(\dfrac{-\Delta G^{\ddagger}}{R\cdot\text{Temp}}\right),$ Yield is given by: $yield = (1 - \exp(-k t)) \times 100$

conditions such as person executing the experiment, temperature, etc. .... I guess it's fair to distinguish between theoretical idealized yield and practical yield.

There are several factors that could influence yields. We do incorporate both catalyst and solvent, when present, into the reactant set R. Temperature was excluded as it is invariant in the case of high-throughput experimentation datasets (e.g., DF and NS datasets), which maintain consistent reaction conditions throughout, or varies mildly for others (e.g., SC). Other variabilities (e.g., purification procedure) are lacking in the state-of-the-art datasets, which is why we were unable not consider them. We discussed these factors in the conclusion. We will elaborate them further, if our paper gets accepted.

why not incorporate the whole reaction and perform attention over that

Yes, $H_{CGR}$ captures the signals from atom mapping. We experimented with the transformer model where we incorporate whole reaction and perform attention over all reaction components, without explicit CGR modeling. Following MAE values show that our method performs significantly better.

Despite the increased model size, whole-reaction-attention performs worse. We believe this is due to the fact that attention is non-injective in nature, i.e., multiple atoms from reactants can be mapped one single atom in product. In contrast, permutation is injective-- they provide one-to-one mapping between atoms. One way to mitigate the problem of transformer is to design permutation induced transformer, which we believe has a strong potential in this direction.

We thank the reviewer for their suggestions, which we address as follows.

theoretical analysis

We provide the following theoretical underpinning. We start with the QAP in Eq 8. If $\mathcal{P}$ is the set of permutations, then QAP is

$$\underbrace{\sum _{u,v} \max (Adj _R,P Adj _I P^{\top}) [u,v]} _{:=c(R,I;P)} --- (QAP)$$

The standard way to minimize it is to use Gromov Wasserstein (GW) based projection [A1], where P is updated as

$P \leftarrow \text{argmin} _{P} \textrm{Trace}\left(P^T\nabla _{P}\ c(R,I;P) \right) -- (grad)$ In our work, we provide a neural approximation of the gradient $\nabla _{P}\ c(R,I;P) \approx [\sum _{i=1} ^d \max (H _R[u,i], H _I [u',i])] _{u,u'}$. We use max to keep aligned with max() in Eq (QAP). This approximates Eq. (grad) as the following linear OT (Eq 9):

$\min _{P\in \mathcal{P}} \underbrace{ \sum _{u,u'} \sum _{i=1} ^d \max(H _R[u,i], H _I [u',i]) P _{uu'}} _{F(P)} -- (a)$

Eq (a) is optimization over permutations, which may appear to be computationally hard. However, we can write it as an instance of the following optimization over the space of doubly stochastic matrices $B = (P: P \ge 0, P\mathbf{1} = P^T\mathbf{1} = 1)$.

$\min _{P}F(P), \text{ s.t. } P\in B --- (b)$

Eq (b) is equivalent to Eq (a) since it is an linear program.The optimal solution of a linear program always coincides with some vertex of the space induced by the set of linear constraints. Thus, the optimal solution of Eq (b) is a permutation matrix $P^*$, as the set of permutation matrices are the vertices of B, the space of doubly stochastic matrices.

However, $\arg\min _{P}F(P)$ in Eq (b) is non-differentiable. To enable differentiation, we approximate it using Sinkhorn iterations, analogous to how argmax is approximated using softmax. Given n numbers $G_1,..G_n, \ \ \mathbb{1}[i=\arg \min _{j} G_j] \approx p_i=\frac{e^{-G_i/\lambda}}{\sum _{j}e^{-G _{j}/\lambda}}$, where $\mathbb{1}$ is indicator function. One can show that $p_i: i\in [n]$ minimizes the following entropy regularized optimization problem:

$\min _{p} \sum _{j} p _j G _j +\lambda. \sum _{j} p _j \log(p _j),\text{ such that } 0 \le p _j \le 1, \sum _{j} p _j =1$

Using the same technique, we can show that Sinkhorn mechanism minimizes the entropy regularized linear OT problem $\min _P F(P) - \lambda \text{Entropy}(P)$ [A2].

[A1] Xu et al. Gromov-wasserstein learning for graph matching and node embedding, ICML 2019

[A2] Mena et al. Learning Latent Permutations with Gumbel-Sinkhorn Networks. ICLR 2018.

Computational overhead

Following table shows the computational overhead in terms of inference time. It shows: our model is comparable with the some of the GNN based models (HGT, TAG) and not as high as YieldBERT.

Moreover, the training time per epoch for our model is 8 seconds (s), which is comparable with baselines: 6s for HGT, 10s for YieldBERT (with same batch size)

Efficiency enhancement: Low rank OT provides highly efficient Sinkhorn iterations. See Scetborn et al. Low-Rank Sinkhorn Factorization, ICML 2021. One first apporximates $M = \sum _{i=1} ^d \max (H _R[u,i], H _I [u',i])\approx AB^T$ where $A$ and $B$ are $N \times d$ and $N \times d$ low rank matrices. Similarly, we approximate the transport matrix $P$ as $P \approx Q D(1/g) R^T$ where $Q, R$ are $N \times r$ matrices, $g$ is $r \times 1$, and $D(1/g)$ is a diagonal matrix with entries $1/g_i$. To ensure $P$ is doubly stochastic, we enforce:

$Q\mathbf{1} = \mathbf{1}, \quad \mathbf{1}^T Q = g^T$

$R\mathbf{1} = \mathbf{1}, \quad \mathbf{1}^T R = g^T$

$g^T \mathbf{1} = 1, \quad Q, R, g > 0$

Instead of minimizing $Tr(P^T M)$, we minimize: $Tr((Q D(1/g) R^T)^T AB^T)$ using alternating minimization wrt $Q, R, g$. To optimize $Q$, we rewrite:

$$Tr((Q D(1/g) R^T)^T AB^T) = \langle Q, AB^T R D(1/g) \rangle$$

We apply Sinkhorn iterations to: $X = AB^T R D(1/g)$, computed in 3 steps. Note that complexity of multiplying $m \times n$ and $n \times p$ matrices is $O(mnp)$

1. Compute $R D(1/g)$, which is $O(Nr)$ since $R$ is $N \times r$, $g$ is $r \times 1$.
2. Compute $B^T R D(1/g)$ which is $O(dNr)$ since $B^T: d \times N$ and $R D(1/g): N\times r$
3. Compute $X = AB^T R D(1/g)$ which is $O(Ndr)$ since $A: N \times d$ and $B^T R D(1/g): d\times r$

Thus, complexity reduces from $O(N^2)$ to $O(Ndr)$. The same efficiency gain applies when optimizing $R$ and $g$.

lighter model

We replaced some complex components with simpler components in our model. MAE numbers are as follows.

We observe that there is some drop in performance. Depending on available resources, one can decrease the complexity to build a lighter model.