Accurate transition state generation with an object-aware equivariant elementary reaction diffusion model

Introduction

Breaking down complex chemical reactions into their constituent elementary reactions is key for understanding reaction mechanisms and designing processes that favor target reaction pathways. Due to the transient nature of the intermediate and transition state (TS) involved in these elementary reactions, it is difficult to isolate and characterize these structures experimentally. Instead, high throughput quantum chemistry computation, for example, with density functional theory (DFT) , provides valuable insights on potential reaction mechanisms by constructing comprehensive reaction networks. These networks are established by either iteratively enumerating potential elementary reactions on-the-fly given existing species or propagating biased ab initio molecular dynamics followed by elementary reaction refinement. Both approaches, however, require a tremendous number of quantum chemistry calculations due to the large number of species potentially involved in a chemical reaction.

Among all DFT energy evaluations, the overwhelming majority comes from locating an accurate TS structure solely based on reactant and product information. Nonetheless, obtaining these TS structures is vital for estimating reaction rates and determining dominant reaction pathways in a reaction network. Conventional TS search algorithms (for example, nudged elastic band, or NEB) are computationally intensive and notorious for their difficulty in convergence, yielding low success rates and wasting substantial computational resources. Recently, there has been growing interest in exploring the use of machine learning techniques for TS search. This includes ideas that formulate TS search as a 2D graph-to-structure conversion problem , a "shooting game" solved by reinforcement learning , generative tasks addressed alternately by graph neural networks (GNN) , a generative-adversarial network and a combination of gated recurrent neural network and transformer, and using an ML potential as a surrogate for DFT during TS optimizations . However, these approaches do not respect all the physical symmetries in describing an elementary reaction and require further reconstruction and optimization to obtain the final 3D TS structure. In addition, they are still far from reaching the high precision required (that is, 3 kcal/mol, corresponding to a change of one order of magnitude in reaction rate at 300 °C) for estimating a TS barrier height in lieu of DFT evaluation.

Diffusion models have recently been adapted in physical science problems, such as generating organic molecules and their conformations, protein-ligand docking , and structural-based drug design . There, an SE(3) equivariant GNN is used as the scoring function to preserve the required permutation, translation, and rotation symmetry for a 3D object (for example, molecule or protein) in the Euclidean space, which works ideally for systems that contain only one single object . However, there are many scenarios in chemistry and materials science where the desired system consists of multiple objects, for which the relative positioning does not influence the system itself. This includes the design of compounds with multiple building blocks (for example, metal organic frameworks ), pairs of molecules that have similar chemistry but demonstrate distinct properties (for example, the well-known activity cliff in protein binding ), and chemical processes that involve multiple distinct structures such as in chemical reactions . Existing diffusion models with SE(3) equivariant GNNs are problematic for modeling these systems as they do not respect all symmetries and constraints for describing these systems.

In this work, we developed a general procedure to adapt an SE(3) equivariant neural network to preserve all desired symmetries and constraints on systems that consist of multiple objects. These symmetries include permutation among atoms in a fragment of reactant or product, permutation among fragments in reactant or product, and rotation and translation for each fragment in reactant and product (Supplementary Text S1). With all these symmetries satisfied, we bypassed the need of atom order mapping and fragment alignment in TS search, generating TS structure with only the 3D geometry of fragments in reactant and product. We demonstrated this "object-aware" SE(3) GNN for generating sets of 3D molecules in elementary reactions under the diffusion model framework, which we refer to as OA-ReactDiff. In particular, we focused on TS search, an essential but computationally demanding step for estimating reaction barrier heights, rates, and exploring reaction networks. With OA-ReactDiff, the predicted TS structures are highly similar to the true TS structures with an average root mean square deviation (RMSD) of 0.18 Å within 6 seconds on a single GPU. We further built a recommender based on confidence ranking to select among samples generated by OA-ReactDiff, which reduced the average RMSD to 0.13 Å. Using the self-confidence score of OA-ReactDiff for uncertainty quantification, we obtain a mean absolute error (MAE) 2.6 kcal/mol on barrier height by only performing 14% of the DFT-based optimizations for the most challenging systems for the model, approaching the accuracy required (that is, 3 kcal/mol) for exploring reaction networks with unknown mechanisms. The high accuracy of generated 3D structures and reaction barrier estimate achieved by OA-ReactDiff provides the possibility of accelerating and even circumventing expensive quantum chemistry calculations normally required for TS search.

Results

A diffusion model contains two processes . In the forward (that is, diffusion) pass, Gaussian noise is continuously added to the original data distribution, which, over time, becomes an approximately normal distribution (Fig. 1a). In the reverse (that is, sampling) pass, a random sample is drawn from the normal distribution, after which a denoising neural network is iteratively applied to remove noise, recovering the original data distribution. This denoising network is trained to predict the noise added to the original data distribution (see Equivariant diffusion models.). Since a 3D molecule or macromolecule fulfills permutational, translational, and rotational symmetry, the denoising graph neural network (GNN) used in chemistry application requires SE(3) equivariance (Fig. 2a). For a molecule represented by atom types (that is, scalars) and their Cartesian coordinates (that is, vectors), as one applies an SE(3) transformation (for example, rotation), the predicted noise on atom types should be the same while that on coordinates should undergo the same transformation.

Despite the success of SE(3) GNN-based equivariant diffusion models (EDM) in many chemistry applications, they inherently lack the symmetries required for systems containing multiple objects (for example, molecules) whose interactions are independent of their coordinates in the 3D Euclidean space. An elementary reaction, as our system of interest, consists of three objects: reactant, TS, and product. If an SE(3) transformation (that is, rotation) is applied on one of these object (for example, reactant), the description of this elementary reaction should stay invariant/equivariant. . In addition, for a reactant or product that has multiple fragments, SE(3) transformations on individual fragments should also have no influence on the elementary reaction. A vanilla SE(3) GNN, however, would take these object-wise SE(3) transformations as if the entire system undergoes a non-SE(3) transformation and would yield non-equivariant results, breaking the symmetry required to predict the noise on atom types and Cartesian coordinates in EDMs (Fig. 2a).

There, we model an elementary reaction as a joint distribution of the 3D structures of the reactant, TS, and product (Fig. 1b). The diffusion process is essentially the same as the vanilla EDM, where independent Gaussian noise is added to reactant, TS, and product until they become independent normal distributions. In the denoising process, however, an object-aware SE(3) equivariant GNN is used to preserve correct physical symmetries and constraints in an elementary reaction (Supplementary Text S1). We consider two denoising schemes. One is unconditional generation where reactant, TS, and product are all sampled from the normal distribution, which can be used to generate new elementary reactions from scratch (Fig. 1b). In chemistry, however, many important applications are targeted for conditional generation, where some information of an elementary reaction of interest is known a priori. For example, in double-ended TS search, the 3D structure of both reactant and product is known, and the task is to find the unique corresponding TS structure. For these conditional generation tasks, we applied the inpainting scheme, which models the joint distribution of the reactant, TS, and product where the unknown objects are inpainted during the inference time. In TS search, specifically, it combines distributions from the diffused reactant and product (that is, known parts) and denoised TS structure (that is, unknown parts) at each step before proceeding to the following denoising step (Fig. 1b, see Inpainting for conditional generation.).

In an elementary reaction, any non-SE(3) transformation on a single object (for example, reactant) should simultaneously influence all three objects, while any object-based SE(3) transformation on reactant, TS, and product should not change a reaction (Fig. 2a). While a vanilla SE(3) GNN fulfills the former requirement, it violates the latter symmetry as it considers all atoms in a system as belonging to the same molecule (Supplementary Table S1). Here, we achieve all required physical symmetries in elementary reactions by developing a general procedure to adapt any SE(3) equivariant GNN as object-aware SE(3) equivariant with minimal effort (Fig. 2b and Supplementary Text S1). In this procedure, we build an object-aware SE(3) interaction layer from a regular SE(3) update layer, a series of non-parameterized operations (i.e, scalarization and concatenation), and a scalar-only message-passing update. In essence, the SE(3) update only operates on individual objects, where the relative positioning is not encoded, to learn a comprehensive representation for each molecule, while the scalar-only message passing layer learns interactions among atoms from different molecules (see Object-aware SE(3) implementation.). Similar to standard SE(3) GNNs, this object-aware SE(3) update repeats several times until a final SE(3) readout layer to pool out the final predicted noise on atom types and Cartesian coordinates. In this work, we choose LEFTNet, our recently-developed SE(3) GNN that reaches comparable state-of-the-art performance on QM9 and MD17 , as the vanilla SE(3) GNN for OA-ReactDiff (see LEFTNet.).

OA-ReactDiff training.

We trained OA-ReactDiff on Transition1x, a dataset that contains climbing-image NEB calculated reactant, TS structure, and product at the $\omega$B97x/6-31G(d) level of theory on 10,073 organic reactions of various types originated from a quite exhaustive enumeration of product-reactant pairs based on the GDB7 dataset. Each reaction consists of up to seven heavy atoms including C, N, and O, with the largest system consisting of 23 total atoms. The use of climbing-image NEB ensures a relatively accurate TS structure, making each elementary reaction in Transition1x a unique set of reactant, TS, and product, which guarantees the necessary condition for training OA-ReactDiff. We trained OA-ReactDiff on 9,000 elementary reactions randomly partitioned from Transition1x, leaving 1,073 unseen reactions as the test set. Despite the potential overlap of certain chemical species in the training and test set, there are always at least two species (reactant and TS or TS and product) that are distinct in any test reaction compared to all training data (see Details for model training.).

In OA-ReactDiff, a molecule is represented by atom types with one-hot encoding and nuclear charges and Cartesian coordinates of its constituent atoms. It is common to consider all components of the atom representation in the diffusion and denoising process as none of them is considered known a priori. In chemical reactions, however, it is reasonable to assume that we know the atom types due to the conservation of atoms . Therefore, we only diffused and denoised the Cartesian coordinates of the reactant, TS structure, and product in OA-ReactDiff (Fig. 1b). Since OA-ReactDiff satisfies all the symmetries and constraints for describing an elementary reaction, it does not require any pre-processing of reaction data, such as atom order matching for different species and careful alignment of reactants and products, which sometimes can be infeasible to obtain (see Equivariant diffusion models. and Object-aware SE(3) implementation.). Transition1x dataset has atom mapping and aligned product molecules as constructed. To showcase of the capability of OA-ReactDiff not replying on the atom mapping and fragment alignment, we intentionally swapped the atom ordering and broke the alignment in Transition1x beforehand and verified that OA-ReactDiff functions well without these pre-processing requirement (Supplementary Figure S1). Due to the use of an object-aware SE(3) GNN, OA-ReactDiff breaks the reflection symmetry and thus can distinguish chiral molecules (see LEFTNet.). OA-ReactDiff also bypasses the need for data augmentation in the case of reversing reaction direction by enforcing the same graph embedding layer for reactant and product, and thus guarantees the outputs are invariant to the order of reactant and product as inputs. Lastly, there are no post-processing steps (for example, reconstructing the 3D structure from a distance matrix through optimizations) required as OA-ReactDiff directly yields the Cartesian coordinates of reactant, TS structure, and product . These outstanding features make OA-ReactDiff an end-to-end model for elementary reaction generation and TS search.

Overcoming the stochastic nature of diffusion models with confidence ranking.

OA-ReactDiff models the joint distribution of a set of reactant, TS, and product, and thus can generate new elementary reactions without any conditions, including for those which the chemical composition is unseen during the model training (Supplementary Figure S2). Evaluating the accuracy and building a reaction network from these generated elementary reactions, however, require substantial computational resources for running DFT optimizations and may be subject to selection bias on which chemical compositions are included during evaluation. Therefore, we focus on evaluating OA-ReactDiff under the scheme of conditional generation, specifically for TS search where the task is to identify the 3D TS structure provided a pair of reactant and product.

We first consider three example reactions in Transition1x that break and form a varied number of bonds, representing different levels of complexity (Fig. 3a). Due to the stochastic nature of diffusion models, sampled TS structures from OA-ReactDiff will not be unique with a fixed reactant-product pair. For each of the three reactions, we ran the OA-ReactDiff under the inpainting scheme 128 times, generating 128 distinct samples. We then computed the Hessian for these 128 generated TS structures with $\omega$B97x/6-31G(d) and found many only contained one imaginary frequency and thus are good candidates for single-ended TS optimizations (Fig. 3b). To evaluate the differences among the structures with one imaginary frequency, we performed TS optimization on these structures with $\omega$B97x/6-31G(d). We then computed the pairwise RMSD and identified multiple distinct TS structures discovered by OA-ReactDiff. (Fig. 3c). Further internal reaction coordinate calculations with $\omega$B97x/6-31G(d) confirmed that these TS structures lead to different reactant and product conformation or connectivity, and thus form distinct elementary reactions compared to the intended one (Supplementary Data). These generated TS structures sometimes show relatively large structure deviations (that is, > 0.2 Å) compared to the optimized TS structures, which is expected since the input reactant and product do not actually correspond to the generated TS structure (Fig. 3d). However, these different TS structures can still be exploited during the reaction network exploration for elementary reactions that would otherwise be neglected.

Provided the 3D conformations of reactant and product, there is only one unique TS structure . Yet the stochastic nature of OA-ReactDiff generates TS structures in a non-deterministic manner. To address this challenge, we further trained an object-aware SE(3) LEFTNet as a confidence model, which also satisfies all the symmetries and constraints for elementary reactions. There, provided a set of input reactant, TS, and product, the confidence model predicts its probability of being a true elementary reaction. During its training, we provided elementary reactions sampled by OA-ReactDiff, which are labeled as good (i.e, 1) if the RMSD between sampled and true TS structure is < 0.2 Å and bad (that is, 0) otherwise (see Details for model training.). Once trained, the confidence model successfully distinguishes different 3D TS structures generated by OA-ReactDiff, assigning them a distinct probability score (Fig. 3d). Moreover, the confidence model always give the highest probability score to generated structures with among the lowest RMSD with respect to the true TS structure for all three example reactions. Without the confidence model, random selection from samples generated by OA-ReactDiff may result in a TS structure of completely different conformation or connectivity compared to the reactant and product. It would in turn yield a large (> 10 kcal/mol) energy difference for the predicted and true TS structure, which would lead to orders of magnitude differences in predicted reaction rates. Even though it is not guaranteed that the confidence model always selects the sample with the lowest RMSD compared to the true TS structure, the confidence model will likely avoid choosing samples that have incorrect connectivity or geometries with large deviations, especially in reactions that have multiple bonds breaking and forming.

High quality TS structures from OA-ReactDiff.

We next systematically evaluated the structural similarity between the OA-ReactDiff and true TS structures for 1,073 set-aside unseen reactions in Transition1x, as judged by RMSD. Notably, in contrast to an average runtime of 12 hours using climbing image NEB with DFT, it only takes on average within 6 seconds to generate a TS structure with OA-ReactDiff with proper batching on a V100 GPU (single sample generation without batching takes 17 seconds, Supplementary Table S2) . Compared to bond lengths, angles, and dihedrals that mostly compare local geometry for a subset of atoms, the RMSD should provide a more accurate assessment on overall structural agreement, which is the ultimate goal of TS search. For each reaction, we ran OA-ReactDiff 40 times, generating 40 independent guess TS structures. For a random selection of 40 samples, OA-ReactDiff has already reached an average RMSD of 0.183 Å with a median being 0.076 Å for the 1073 test elementary reactions (Fig. 4a). More than half (two thirds) of the TS structures have an RMSD < 0.1 (0.2) Å compared to their corresponding true TS structures identified by climbing image NEB. We observed a near power-law dependence of RMSD for OA-ReactDiff generated samples with the number of training data (Supplementary Table S3 and Fig. S3). The performance of OA-ReactDiff shows no sign of depletion when we reach the maximum of our training data available, suggesting room for further improvement through training OA-ReactDiff on larger datasets. More interestingly, we find that, given a similar amount of training data, an OA-ReactDiff model trained only on system within 15 atoms performs similarly well to that trained on randomly sampled reactions (Supplementary Table S3). Meanwhile, the OA-ReactDiff model yields similar RMSD distribution regardless of system size (Supplementary Figure S4). This observation showcases great size extensiveness of OA-ReactDiff, which is encouraging for its application in practical reaction exploration for large molecules.

With the confidence model, we can further improve the procedure of sample selection using a recommender approach. Together with the true reactant and product, these guess TS structures are fed into the confidence model to get their probability score. The sample with highest probability score (that is, top-1 confidence) is chosen as the final predicted TS structure from OA-ReactDiff. With this recommender approach, the quality of selected TS structures is greatly improved, most likely due to the removal of TS structures with incorrect connectivity and geometries with large deviations (Fig. 3). Moreover, the recommended structures mostly reside in the low RMSD and high confidence region, which demonstrates the effectiveness of our combined OA-ReactDiff and confidence recommender approach (Fig. 4b). The average and median of the error in RMSD become 0.129 and 0.058 Å, respectively, with approximately two thirds of the recommended TS structures having RMSD < 0.1 Å (Fig 4 a). We observed a systematically-improving performance for the OA-ReactDiff + recommender approach as the number of total independent runs increases (Supplementary Figure S5). Here, we took 40 runs for each sample for a balance between total run time (4 minutes in total) and sampling accuracy. Despite the fact that the recommender is still far from perfect for distinguishing structures with low RMSD (i.e, < 0.2 Å) structures, it helps avoiding TS samples that are very different from the true TS (that is, > 0.45 Å, Fig. 4c). We also trained a regressor that predicts the RMSD between generated and true TS given a set of reactant, product, and generated TS structure as the confidence model, where the same quantitative behavior is observed (Supplementary Figure S6). This recommender feature is particularly useful in end-to-end applications for ML models.

Approaching the energetic accuracy needed in TS search.

With OA-ReactDiff and the recommender, we reach an average of 4.4 kcal/mol and median of 1.6 kcal/mol for the absolute energy difference between the generated and true TS structure, with 71% of TS barrier errors < 3 kcal/mol (Fig. 5b). OA-ReactDiff + recommender far outperforms semi-empirical methods such as density functional tight binding, which was reported to have an MAE of 16.1 kcal/mol and average runtime of 82 seconds on Transition1x dataset. Interestingly, the MAE would only improve marginally to 4.0 kcal/mol if we were able to select the OA-ReactDiff sample with the lowest RMSD compared to the true TS, indicating the power of the recommender for selecting TS structures with low energy deviations (Table 1). The performance of OA-ReactDiff with the recommender on elementary reactions with multiple reactants and/or products is comparable to the performance for rearrangement reactions that only contain one single reactant and product (Supplementary Figure S9) . The slight deterioration of the performance is likely due to the imbalance of reaction types included in Transition1x, where only one fourth of the elementary reactions contain multiple reactants or products. It is known that the TS structure is less sensitive to the choice of DFT functional, especially or small organic molecules with only CNOH. Here, we also evaluated the error for barrier height estimation of OA-ReactDiff with two other functionals and found quantitatively similar performance (Supplementary Table S4).

We also compare OA-ReactDiff + recommender with two pioneering works where non-diffusion-based approaches were developed for generating 3D TS structures on large diverse organic reaction datasets such as Transition1x or its predecessor . Choi developed a "PSI-based" model combining transformer and bidirectional gated recurrent unit that generates TS structures from refining the linear interpolation of reactant and product, which, however, requires the prior knowledge of atom mapping and careful alignment among fragments in reactant and product for multi-molecular reactions. Schreiner et al. trained a machine learning potential on 10M structures (with both energy and forces) collected during the generation of Transition1x, and, for the first time, applied the trained potential to the TS search problem in place of DFT. There, similar to the problem of DFT-based TS search (for example, NEB), an attempt may still encounter convergence issues during the saddle point optimization, leading to a null prediction for the final TS structure. We find OA-ReactDiff + recommender systematically outperforms the prior approaches on both the RMSD and barrier height estimate in terms of both the mean and median of the error distribution (Table 1). This superior performance is attributed the fact that OA-ReactDiff manages to respect all physical symmetries and constraints for describing an elementary reaction, without the need for atom order mapping, reactants or products alignment, reconstruction of 3D geometry from distance matrix, and data augmentation of any kind. We ascribe the slightly lower Pearson’s r between RMSD and barrier height error of OA-ReactDiff to the presence of many high RMSD but low barrier height TS samples, which correspond to cases where OA-ReactDiff can generate good 3D structures for nearly non-interacting individual fragments but not their alignment (Fig. 5a and Supplementary Figure S10). In addition, a more gradual increase of absolute energy difference with respect to RMSD was identified in OA-ReactDiff + recommender compared with the two other approaches, suggesting a more accurate barrier estimate can be obtained by OA-ReactDiff at the same level of structural similarity between generated and true TS (Table 1 and Supplementary Figure S10).

We would ideally aim to select one single TS structure sampled by OA-ReactDiff with the recommender. The recommended sample, however, may not be confident if all 40 samples generated by OA-ReactDiff suffer from a low confidence score (i.e, $p$ < 0.5) due to the limited amount of training data. Further removal of these reactions (153, or 14%) from the test set leads to a substantially improved energy difference with a mean of 3.1 kcal/mol and median of 1.4 kcal/mol (Table 1). Moreover, we observe a monotonic behavior between the MAE for barrier height estimates and the confidence threshold imposed for TS structure generation that we consider as valid (Fig. 5c). This desired monotonic behavior suggests that we can use the confidence score for uncertainty quantification to balance the accuracy and number of DFT calculations required in a practical workflow that combines OA-ReactDiff, recommender, and DFT-based TS search. For a set of TS structures generated by OA-ReactDiff and their corresponding confidence score evaluated by the confidence model and recommender, we can decide whether we would accept the recommended TS structure depending on its confidence score or would rather launch a DFT-based NEB. With a confidence threshold of 0.5, we would only perform NEB with DFT on 14% of reactions while directly accepting TS structures from OA-ReactDiff + recommender for the remaining 86% reactions, leading to an overall accuracy of 2.6 kcal/mol. This strategy showcases the power of combining OA-ReactDiff, recommender, and DFT-based NEB for efficient generation of TS structures given a target accuracy level.

Discussion

Elucidating TS structures is essential for uncovering the underlying microscopic mechanisms of chemical reactions and estimating reaction barriers for building large reaction networks. In this work, we extended SE(3) equivariant diffusion models to respect the object-wise symmetries, leading to OA-ReactDiff, an object-aware SE(3) equivariant diffusion model that first fulfills all the symmetries and constraints for generating elementary reactions. In addition, we built a confidence-model-enabled recommender to overcome the stochastic nature of diffusion model to select from sampled generated by OA-ReactDiff in multiple runs. OA-ReactDiff + recommender gives an RMSD of 0.129 Å and MAE of 4.4 kcal/mol compared to the true TS structure obtained by computationally demanding climbing image NEB calculations. By further using the confidence score for uncertainty quantification, we can selectively perform climbing image NEB only for 14% of elementary reactions that OA-ReactDiff is most uncertain about, leading to a reduced MAE of 2.6 kcal/mol.

The current OA-ReactDiff approach has two major limitations. First, we describe an elementary reaction as a set of 3D structures (say $N$ atoms for reactant, TS, and product), which leads to a system that is 3x larger (that is, $3N$ atoms). Although the most expensive equivariant update is still object-wise (that is, scales with $N$), the scalar message-passing update requires building a fully-connected graph for the $3N$ atoms, which will be the bottleneck for applying OA-ReactDiff on chemical systems > 100 atoms on a single GPU. Second, despite the workaround of using a confidence model and recommender to select a unique sample generated by OA-ReactDiff, the stochastic nature of diffusion model cannot be avoided. This leads to uncertainty for the sample quality of generated TS structure and accumulated runtime for running OA-ReactDiff repeatedly. These limitations are inherent for diffusion models, which can potentially be addressed by reformulating elementary reaction generation as a transport problem, where optimal transport via flow matching or Schrödinger bridge can be applied. The authors are actively exploring along this direction as a future work. During the revision of this manuscript, we found a concurrent work from Kim et al. named TSDiff , which showcases the use of diffusion model for TS generation in another perspective. There, the TS structure is generated from a 2D graph composed of molecular connectivity of reactant and product via a diffusion model. They circumvent the need to handle object-wise SE(3) symmetry by removing 3D conformations of reactant and product as inputs. This simplification comes with the price of not being able to select “the best” TS structure out of all generated samples . However, it would be of interest to compare OA-ReactDiff plus 2D-to-3D conformation sampling and TSDiff in reaction exploration in future work.

Together with uncertainty quantification, OA-ReactDiff + recommender reached both the structural and energetic accuracy required in TS search, which can be readily integrated in current high throughput computation workflows for reaction network exploration. In this work, we focus on the relatively well defined TS search problem such that we can evaluate our newly-developed OA-ReactDiff more easily and demonstrate the promise of this new model. OA-ReactDiff, however, models the joint distribution of structures in elementary reactions and thus is not limited to double-ended TS search problem and can be applied in single-ended (i.e, only the reactant is provided) or zero-ended (that is, only the chemical composition of a system is provided) scenarios. Although we perform DFT calculations for evaluating the barrier height throughout this work, OA-ReactDiff framework can be readily adopted to predict the reaction barrier by swapping the model output layer. Very recently, a more diverse elementary reaction dataset 17 times larger than Transition1x, named as RGD1, has been established . Provided that the quality of diffusion model is highly dependent on the size of training data, RGD1 has the potential of unleashing the power of OA-ReactDiff for establishing large reaction networks and exploring chemical reactions with unknown mechanisms with a greatly reduced number of DFT calculations. Lastly, despite solely focusing on chemical reactions, the object-aware SE(3) equivariant diffusion model developed in this work can be applied to diverse chemical problems where the system of interest consists of multiple 3D objects, in which their interactions do not depend on their locations in Euclidean space.

Methods

Equivariance.– A function $f$ is said to be equivariant to a group of actions $G$ if $g\circ f(x)=f(g\circ x)$ for any $g\in G$ acting on $x$ . In this paper, we specifically consider the Special Euclidean group in 3D space (SE(3)) which includes permutation, translation and rotation transformations. We intentionally break the reflection symmetry so that our model can describe molecules with chirality.

Diffusion models.– Diffusion models are originally inspired from non-equilibrium thermodynamics . A diffusion model has two processes, the forward (diffusing) process and the reverse (denoising) process. The noise process gradually adds noise into the data until it becomes a prior (Gaussian) distribution:

where $\alpha_{t}$ controls the signal retained and $\sigma_{t}$ controls the noise added. A signal-to-noise ratio is defined as $\text{SNR}(t)=\frac{\alpha_{t}^{2}}{\sigma_{t}^{2}}$. We set $\alpha_{t}=\sqrt{1-\sigma_{t}^{2}}$ following the variance preserving process in .

The true denoising process can be written in a closed form due to the property of Gaussian noise:

where $s$ < $t$ refer to two different timesteps along the diffusion/denoising process ranging from 0 to T, $\alpha_{t|s}=\frac{\alpha_{t}}{\alpha_{s}}$, $\sigma^{2}_{t|s}=\sigma^{2}_{t}-\alpha^{2}_{t|s}\sigma^{2}_{s}$. However, this true denoising process is dependent on $x_{0}$ which is the data distribution and not accessible. Therefore, diffusion learns the denoising process by replacing $x_{0}$ with $\hat{x}=\epsilon_{\theta}(x_{t},t)$ predicted by a denoising network $\epsilon_{\theta}$. The training objective is to maximize the variational lower bound (VLB) on the likelihood of the training data:

Empirically, a simplified objective has been found to be efficient to optimize :

Equivariant diffusion models.– To build an SE(3)-equivariant diffusion model, it has been proven that we need an SE(3)-invariant prior and an SE(3)-equivariant transition kernel . To guarantee equivariance on permutation, rotation, and translation, a necessary condition is to use an SE(3)-equivariant transition kernel (that is denoising network), as we will explain in details at a later section. (see LEFTNet.). There are additional requirements for rotation and translation. For rotations, the isotropic Gaussian prior has the nice property to transform equivariantly. For translations, we need to limit the distribution on the linear subspace where the center of mass is the origin .

Inpainting for conditional generation.

Inpainting is a flexible technique to formulate the conditional generation problem for diffusion models. Instead of modeling the conditional distribution, inpainting models the joint distribution during training. During inference, inpainting methods combine the conditional input as part of the context through the noising process of the diffusion model before denoising both the conditional input and the inpainting region together. The resampling technique has demonstrated excellent empirical performance in harmonizing the context of the denoising process as there is sometimes mismatch between the noised conditional input and the denoised inpainting region. Specifically, resampling increases the total number of sampling steps in each denoising step by sampling the inpainting region back and forth together with the conditional input. Despite resampling increases the number of total denoising steps, this can be compensated by decreasing the number of total denoising steps accordingly by striding the sampling schedule without significantly sacrificing the model performance.

LEFTNet.

We build our denoising network on top of a recently proposed SE(3)-equivariant GNN, LEFTNet . The main idea of LEFTNet relies on building local frames to scalarize the vector (for example position, velocity) and higher order tensor (for example stress) which becomes invariant to SE(3) transformations. Tensorization can be applied to invert the scalar back to vector and higher order tensor without information loss in each layer to update these quantities. The benefit of scalarization is demonstrated by the flexibility of neural network parameterizations without breaking the symmetry and further proved by the universal approximation theorem such that the resulting neural network has the universality in the space of continuous SE(3) and permutation equivariant functions.

Scalarization and tensorization.– Scalarization and tensorization are two operations in differential geometry to convert geometric quantities. Specifically, scalarization transforms geometric tensors into scalars while tensorization is the inverse of scalarization, transforming scalars back to geometric tensors. In this case, scalarization is used to transform equivariant quantities by three equivariant orthonormal frames:

For simplicity, we use vector as an example. The geometric tensors are scalarized by the inner product between the frames $\mathcal{F}$ and the input vector $x$ as follows:

On the contrary, tensorization reverses the process by:

where $(\mathbf{x_{1}},\mathbf{x_{2}},\mathbf{x_{3}})$ is the input scalar tuple and $x$ is the converted tensor.

where $M$ is the message function, $\mathcal{N}(i)$ denotes neighbors of node $i$ and $m_{i}$ is the message. Then the message is used to update the node feature as:

where $U$ is the update function. After a number of layers, the global embedding is calculated by:

where $s_{ij}$ is obtained by scalarizing the input coordinate over the frames $\mathcal{F}_{ij}$ between each pair of nodes $i$ and $j$:

Building towards LEFTNet.– Motivated by distinguishing local 3D geometric isomorphisms, LEFTNet introduced a local structure encoding module to encode the local atomic environment of each atom in the scalarization operation. In addition, LEFTNet designed another frame transition encoding block to consider the transition between two frames (the atom and its neighbor atom) when calculating the message between them.

Object-aware SE(3) implementation.

In general, a system consists of multiple objects (molecules or proteins) that do not have interactions through the 3D Euclidean space can be described by a set of independent graphs, $\{\mathcal{G}_{i}=(\mathcal{V}_{i},\mathcal{E}_{i})\}$. In elementary reaction, for example, we have three objects which can index, reactant ($i=0$), TS ($i=1$), and product ($i=2$). One important addition symmetry for these systems is object-wise SE(3) equivariance, meaning any SE(3) transformation on each individual object in a system should not influence its description:

$g_{i}$ represents an SE(3) transformation on $i^{th}$ object, which is not necessarily the same for all objects. On the other hand, any non-SE(3) transition on any objects in a system should influence its description:

where $q_{i}$ represents a non-SE(3) transformation on $i^{th}$ object. An SE(3) GNN would hold for the latter but violate the former symmetry.

Details for model training.

Dataset and train/test partitioning.– Built on top a large chemically diverse dataset by Grambow et al., Transition1x dataset consists of 10,073 elementary reactions optimized by climbing image NEB. We partitioned Transition1x randomly, with 9,000 reactions used in training and validation and the remaining 1,073 reactions as set-aside test set. It is not guaranteed that all species in test reactions are unseen by a trained model due to the overlapping structures in different reactions. However, due to the uniqueness of elementary reaction, there is, at most, one chemical species (specifically, reactant or product) that may overlap in multiple reactions. We think this partition is reasonable because all TS structures in the test set are completely new and unseen from model training. In addition, having a certain degree of overlap in reactants/products for the training and set-aside test set is useful to judge whether the diffusion model only memorizes training samples rather than learning to generate new samples. Throughout this manuscript, we do not generate new reaction data through DFT-based TS optimizations, except for ones used as examples in Fig. 3. Therefore, all DFT results discussed here are computed with $\omega$B97x/6-31G(d).

Confidence model training.– The confidence model shares exactly the same set of hyperparameters as the scoring network, with the only change being the use of sigmoid function at the final output layer. To get data for training the confidence model, we ran OA-ReactDiff on the 9,000 training reactions for 40 runs, generating 360,000 synthetic reactions. We labeled a reaction as "good" (that is, 1) if the generated TS structure has a RMSD < 0.2 Å compared to the true TS, and labeled it as "bad" (that is, 0) otherwise. Lastly, we train the confidence model as a binary classifier, where the predicted probability is used as the confidence score to estimate the quality of generated TS structure. Note that we used the same partition for both scoring network and confidence model, which ensures the 1,073 reactions in the set-aside test set are unseen to both models during evaluation.

Code availability

Code for OA-ReactDiff is available as a open source repository on github, https://github.com/chenruduan/OAReactDiff

Acknowledgements

This work was supported by the U.S. Office of Naval Research under grant no. N00014-20-1-2150 (C.D. and H.J.K.) and National Science Foundation grant CBET-1846426 (H.J. and H.J.K.). C.D. thanks the Molecular Sciences Software Institute for the fellowship support under NSF grant OAC-1547580. C.D. thanks Q. Zhao and M. Monkey for discussions about elementary reactions. C.D. thanks A. Nandy and W. Du for discussions about equivariant graph neural networks. C.D. and Y.D. thank G.-H. Liu and T. Chen for discussions about diffusion model and Schrödinger bridge. C.D. and H.J. thank Y. Zhao for his help on preparing a demo jupyter notebook for this work. The authors thank S. Choi and M. Schreiner for communications and providing their raw data that makes the comparison in Table 1 possible.

Author contributions

C.D.: conceptualization, methodology, software, validation, investigation, data curation, writing of original draft, review and editing, and visualization. Y.D.: methodology, software, writing of original draft, and review and editing. H.J.: data curation, review and editing. H.J.K.: writing of original draft, review and editing.

Competing interests

The authors declare no competing financial interest at this moment.

References

Abbreviation

The following is the list of abbreviation utilized in the main paper.

OA-ReactDiff: Object-aware SE(3) GNN for generating sets of 3D molecules in elementary reactions under the diffusion model

SE(3): Special Euclidean group in 3D space.

Appendix S1 Physical symmetries and constraints in an elementary reaction.