Diffusion-based generative AI for exploring transition states from 2D molecular graphs

Transition states (TSs) govern the kinetics and mechanisms of chemical reactions, but are difficult to observe directly and expensive to compute reliably. Conventional TS search methods (single-ended and double-ended algorithms) and recent ML approaches typically require carefully prepared 3D geometries of reactants and products, with appropriate orientations, making them sensitive to input conformations and alignment. The research question addressed here is whether TS geometries can be accurately and efficiently predicted directly from 2D molecular reaction graphs, bypassing 3D input preparation, while also enabling exploration of multiple TS conformations corresponding to distinct reaction pathways. The study introduces TSDiff, a stochastic diffusion-based generative model conditioned on 2D reaction graphs, to reduce user effort, mitigate input sensitivity, and improve the discovery of favorable low-barrier pathways.

The paper reviews two principal families of TS optimization: single-ended methods (e.g., Berny optimization, AFIR, ADDF, single-ended GSM) that start from a single 3D structure, and double-ended methods (e.g., NEB, double-ended GSM) that rely on both reactant and product 3D geometries. While widely used, these methods can be computationally costly and sensitive to initial structures, with frequent convergence challenges. Recent ML approaches have targeted barrier height prediction and, more relevantly, direct TS geometry prediction. Prior TS geometry ML models generally require aligned 3D reactant and product conformations and have shown promising results on datasets like Grambow’s gas-phase reactions and on specific reaction classes. A concurrent diffusion model (OA-ReactDiff) predicts the highest-energy NEB image but still uses the double-ended 3D setup. Across domains, ML models using 3D geometries are known to be input-sensitive. This motivates learning from 2D reaction information to avoid conformer/orientation dependence and enable broader, more robust TS exploration.

TSDiff is a conditional generative diffusion model that learns the distribution of 3D TS geometries given 2D reaction information encoded as a condensed reaction graph built from SMARTS/SMILES. Input construction: molecular graphs for reactants (G_R) and products (G_P) are derived from atom and bond information; extended edges (within 3-hop graph distance) augment connectivity features. A condensed reaction graph G_rxn is formed by combining G_R and G_P via atom mapping to encode bond changes and atomic features (e.g., aromaticity, formal charge, hybridization, valency, chirality, ring membership). During training/inference, a geometric reaction graph is created by adding noisy 3D positions to nodes and connecting atom pairs within a radial cutoff, integrating bond, graph-distance, and spatial-distance features as edge inputs. Learning and architecture: The forward diffusion adds Gaussian noise to TS geometries across discrete timesteps with schedules α_t and β_t. The model learns the reverse process p_θ(C_{t-1}|C_t, G_rxn) as a Gaussian with predicted mean; training minimizes KL divergences equivalent to a score-matching loss. To ensure SE(3) invariance/equivariance, the score is predicted in pairwise distance space and mapped back to Cartesian coordinates via the chain rule. The denoiser is a GNN based on seven modified SchNet layers with message passing using radial basis kernels and edge-feature modulation. The model outputs denoised updates iteratively over 5000 timesteps during inference, starting from Gaussian noise to yield TS geometries. Training and data: The ωB97X-D3/def2-TZVP Grambow gas-phase organic reaction dataset (11,959 reactions after excluding two with non-reactive N2) was split 8:1:1; reverse reactions were included for data augmentation (19,132 training points). An ensemble of eight independently trained models was used; each trained ~22 hours on a single RTX 2080 Ti GPU. At inference, multiple samples can be drawn per reaction graph to explore TS conformational diversity. Evaluation and validation: Generative performance was assessed by coverage (COV) and matching (MAT) using the interatomic-distance MAE (D-MAE). Chemical validity was evaluated using DFT-based saddle point optimizations (Berny algorithm) and intrinsic reaction coordinate (IRC) calculations at ωB97X-D3/def2-TZVP with ORCA. Success criteria included convergence to a saddle point with a single imaginary frequency (< -100 cm^-1) and IRC connectivity to the intended reactants/products (checked by connectivity consistency with Open Babel). Additional analyses included clustering of generated conformations and identification of alternative pathways with different reaction coordinates.

- Generative coverage and accuracy: With ensemble TSDiff, COV (δ=0.1 Å) and MAT improved with more samples per reaction: 1 sample: COV 49.1%, MAT 0.137 Å; 3: 67.3%, 0.096 Å; 5: 75.2%, 0.079 Å; 10: 84.0%, 0.063 Å; 100: 91.7%, 0.045 Å. Using δ=0.2 Å, COV rose to 73.9%, 87.5%, 92.6%, 95.6%, and 97.7% for 1, 3, 5, 10, and 100 samples, respectively. - Accuracy vs prior ML: Without conformer matching (single sample), TSDiff achieved D-MAE 0.137 Å, outperforming most prior models (Makoś 0.170, Jackson 0.244, Pattanaik 0.225) and approaching Choi (0.095). With conformer matching after saddle point optimization, TSDiff achieved D-MAE 0.063 Å (1 sampling; 53.2% coverage of test reactions) and 0.067 Å (8 samplings; 84.6% coverage), both better than all reported baselines. - Chemical validity and success rates: For 1197 test reactions, 8 generated geometries per reaction (9576 total) led to 9289 successful saddle points (97.0%). IRC validation on one random TS per reaction succeeded for 998 reactions (83.4%). After excluding 95 reactions whose reference TS failed IRC (refined set: 1102 reactions), saddle-point success was 97.4% (8588/8816), and single-sample IRC success was 90.6%. Over five sampling rounds, refined-set success rates reached 99.9% (saddle) and 98.5% (IRC); coverage of correct TSs was 98.6% (1087/1102 reactions) within five rounds. - Diversity of TS conformations: For a single reaction with 100 samples, all were optimized to saddle points, yielding nine distinct TS conformations; average D-MAE between generated and optimized structures was 0.045 Å. - Discovery of alternative and improved pathways: Across 1197 test reactions, optimization produced 3316 unique TS conformations; 2303 corresponded to different saddle points than the references. Among IRC-validated TSs, 309 had lower TS energies than the references by >0.1 kcal mol^-1. After further reactant optimization from IRC endpoints, 513 pathways exhibited lower barrier heights than the reference. A case study showed a 6.4 kcal mol^-1 lower barrier due to significant conformational differences (e.g., chair vs boat ring). - Multiple reaction coordinates: TSDiff identified distinct pathways with different bond formation/breaking sequences or different migrating atoms, despite sharing the same reaction graph, demonstrating capacity to sample diverse, mechanistically distinct TSs. - Efficiency: Although diffusion inference (5000 denoising steps) is iterative, per-reaction inference takes only seconds and is negligible compared to DFT optimization times.

The study demonstrates that learning TS geometries directly from 2D reaction graphs effectively addresses input sensitivity and preparation burdens associated with 3D conformations and alignment. TSDiff’s stochastic diffusion formulation captures the multimodal distribution of TS conformations, enabling sampling of diverse, chemically valid TSs that often include lower-barrier pathways than those in reference datasets. High saddle-point and IRC validation success rates show TSDiff’s reliability as an initial TS guesser, reducing quantum-chemical trial-and-error. Its superior accuracy compared to prior 3D-input ML models, combined with the ability to uncover different reaction coordinates and conformers, supports its relevance for kinetics modeling, mechanism elucidation, and reaction discovery. The generative diversity also facilitates downstream selection via clustering, focusing expensive quantum calculations on representative candidates.

TSDiff, a diffusion-based generative model conditioned on 2D reaction graphs, accurately and efficiently predicts TS geometries without requiring 3D reactant/product inputs or alignment. It outperforms prior ML approaches in D-MAE accuracy, achieves high DFT validation success, and systematically explores multiple TS conformations and distinct reaction coordinates, frequently identifying lower-barrier pathways than references. These capabilities promise substantial reductions in user effort and computational cost for TS exploration and mechanistic analysis. Future work includes expanding training to larger and more diverse reaction datasets, especially inorganic and transition-metal reactions, integrating more physics-based constraints, and further optimizing sampling/denoising for faster inference.

- Domain coverage: The current work focuses on organic gas-phase reactions (C, H, O, N) due to the availability of large, high-quality datasets. A lack of extensive inorganic/transition-metal TS datasets limits immediate generalization to those domains. - Reference data constraints: The benchmark dataset provides one reference TS per reaction; evaluating generative coverage for unseen conformers/pathways is constrained, and some reference TSs failed IRC checks, complicating comparisons. - Computational aspects: Diffusion inference requires many denoising steps (5000), which, although much cheaper than DFT, is higher than typical feed-forward ML predictors. - Implicit assumptions: Input limited to 2D reaction graphs may under-specify certain stereochemical or long-range effects; quantum validation remains necessary to confirm chemical validity.