A Priori Sampling of Transition States with Guided Diffusion
Authors: Hyukjun Lim, Soojung Yang, Lucas Pinède, Miguel Steiner, Yuanqi Du, Rafael Gómez-Bombarelli
Organizations: Department of Materials Science and Engineering, Seoul National University, 1 Gwanak-ro, Gwanak-gu, Seoul, 08826, Republic of Korea. · Computational and Systems Biology, MIT, 77 Massachusetts Avenue, Cambridge, 02139, Massachusetts, USA. · Chimie ParisTech, PSL University, 11 rue Pierre et Marie Curie, Paris, 75005, France. · Department of Materials Science and Engineering, MIT, 77 Massachusetts Avenue, Cambridge, 02139, Massachusetts, USA. · Department of Computer Science, Cornell University, 127 Hoy Road, Ithaca, 14850, NY, USA.
Transition states, the first-order saddle points on the potential energy surfaces, govern the kinetics and mechanisms of chemical reactions and conformational changes. Locating them is challenging because transition pathways are topologically complex and can proceed via an ensemble of diverse routes. Existing methods address these challenges by introducing heuristic assumptions about the pathway or reaction coordinates, which limits their applicability when a good initial guess is unavailable or when the guess precludes alternative, potentially relevant pathways. We propose to bypass such heuristic limitations by introducing ASTRA, A Priori Sampling of TRAnsition States with Guided Diffusion, which reframes the transition state search as an inference-time scaling problem for generative models. ASTRA trains a score-based diffusion model on configurations from known metastable states. Then, ASTRA guides inference toward the isodensity surface separating the basins of metastable states via a principled composition of conditional scores. A Score-Aligned Ascent (SAA) process then approximates a reaction coordinate from the difference between conditioned scores and combines it with physical forces to drive convergence onto first-order transition states. Validated on benchmarks ranging from 2D potentials to biomolecular conformational changes and a chemical reaction, ASTRA locates transition states with high precision and discovers multiple reaction pathways, enabling mechanistic studies of complex molecular systems.
Figures & tables
Figure 1: Overview of ASTRA . The method consists of three stages: (1) training a conditional generative model, (2) sampling from an isodensity surface, and (3) inference-time guidance that combines a reaction coordinate approximated from score differences of the two conditional models θA,θB ( R=SθA−SθB ) with physical forces to rapidly sample transition states defined as first-order saddle points.
Figure 2: Our method discovers transition state regions with high precision. It finds a transition state for the (a) double well potential, and multiple transition states for the (b) Müller-Brown potential and (c) double path potential.
Figure 3: Application of our method to chemical systems. (a) For the alanine dipeptide, the generated samples (circles) localize the transition region between two states. (b) For chignolin, the free energy landscape is projected onto the two slowest collective modes from time-lagged independent component analysis fitted from a converged simulation Lindorff-Larsen et al. (2011) . The background densities in both plots correspond to the training data distribution where distinct colors indicate the defined State A and State B. The color scale for the alanine dipeptide plot represents the MD-based committor value while the scale for chignolin indicates the machine-learned committor value Kang et al. (2024) .
Figure 4: Detailed performance of ASTRA on alanine dipeptide. a) The ASTRA samples covering all three TSs are compared to the Nudged Elastic Band reference method. One ASTRA sample is selected randomly for each TS for readability. The background shows the training data distribution and colors represent the two defined states for the classifier free guidance training. The three-dimensional structures are overlayed for each identified TS. b) shows the distribution of committor value for the stable ASTRA samples peaked around 0.5. c-e) characterize how close ASTRA samples are from their Dimer-optimized counterpart. c) shows the distribution of number of iterations necessary to converge Dimer from the stable ASTRA samples. d) displays the Root Mean Square Deviation (RMSD) of ASTRA samples after Dimer optimization compared to before, while e) shows that difference in terms of the energy of the structures. We indicate the number of stable ASTRA samples, compared to the total number of structures drawn from our algorithm for that analysis, alongside the Dimer convergence rate from ASTRA stable samples. Stable ASTRA samples correspond to structures sampled by ASTRA for which running 2ps MD from is stable.
Figure 5: Visualization of chignolin transition state mechanisms. Representative structures from ASTRA-sampled TS down and TS up ( left ) ensembles are overlaid with transparent tubes on the reference conformations from Ref. 32 ( right ). The structural agreement validates our method’s ability to resolve distinct folding pathways.
Figure 6: Lewis structure representation of the chemical reaction and three dimensional reactant structure with the two coordinates highlighted that span the plot. The two chosen degrees of freedom separate the reactant and product state well, but were only used for visualization and not included in any part of our algorithm. The MD simulation data points are plotted as densities. The structures generated by the ASTRA algorithm and the structures optimized with the Dimer algorithm and converged to a true transition state are shown as differently shaped points that are colored based on the committor value predicted by the pretrained model.
Appendix figures & tables24 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 7: Potential Energy Surface (PES) of alanine dipeptide and chignolin.
Figure 8: Potential Energy Surface (PES) of alanine dipeptide sample from the trained diffusion model.
Figure 9: Two dimensional representation of the reactive potential energy surface (PES) spanned by two important internuclear distances. The MD simulation data points are plotted as densities. The various points represent structures optimized with the Dimer algorithm where the initial structure was generated with different TS guess structure generation algorithms. The various baseline methods are detailed in section 11.2.2 . The zoomed-in version of the PES features representative structures for each cluster of points that are numbered from 1 to 5.
Figure 10: Validation of Score Difference Approximation. The score difference vector (red) closely aligns with the true reaction coordinate (blue) on the (a) Müller-Brown potential and (b) double path potential.
Figure 11: Impact of Score-Based Interpolation and Score-Aligned Ascent on double well potential.
Figure 12: Impact of Score-Based Interpolation and Score-Aligned Ascent on Müller-Brown potential.
Figure 13: Impact of Score-Based Interpolation and Score-Aligned Ascent on double path potential.
Potential
II
II+SAA
SA
SA+SAA
Double Well
0.1174 ± 0.0101
0.0002 ± 0.0000
0.5870 ± 0.5000
0.0002 ± 0.0000
Müller-Brown
7.3632 ± 1.0707
0.0376 ± 0.0044
6.8723 ± 3.1354
0.0449 ± 0.0107
Double Path
0.6810 ± 0.1008
0.0511 ± 0.0377
0.6998 ± 0.1511
0.0384 ± 0.0358
Appendix
Table 1: L2 Distance statistics for analytical potentials from the closer transition state. The best results are bolded , and the second-best results are underlined .
Figure 14: Impact of Score-Based Interpolation and Score-Aligned Ascent on alanine dipeptide.
Figure 15: Impact of Score-Based Interpolation and Score-Aligned Ascent on Chignolin.
Figure 16: Impact of Score-Based Interpolation and Score-Aligned Ascent on DASA reaction.
Figure 17: Histogram of calculated committors on alanine dipeptide.
Figure 18: Histogram of ML predicted committors on Chignolin.
Figure 19: Histogram of ML predicted committors on DASA reaction.
Method
GS
PR
OS
SS
Conv. Rate ↑
Conv. Step ↓
Pos. RMSD ↓
Energy Diff. ↑
(Å)
(kcal ⋅ mol -1 )
Main
1.0
0.05
1000
0.0100
0.84
14
0.010
-0.50
Ablation on the Guidance Scale (GS)
GS = 2.0
2.0
0.05
1000
0.0100
0.86
34
0.011
-0.51
GS = 3.0
3.0
0.05
1000
0.0100
0.87
11
0.011
-0.46
Ablation on the Pause Ratio (PR)
Appendix
Table 2: Hyperparameter Ablation Study evaluated by running the Dimer method on ASTRA generated samples. The main result configuration is shown first. Each subsequent row ablates a single hyperparameter, with the modified value in bold . No value means that the SAA was unstable and the force field eventually returns NaN. We report median values for all metrics. Higher energy difference is better as all differences are negative.
Method
GS
PR
OS
SS
% of Comm. ↑
Num. stable samples
in [0.4,0.6]
(out of 100)
Main
1.0
0.05
1000
0.0100
85
76
Ablation on the Guidance Scale (GS)
GS = 2.0
2.0
0.05
1000
0.0100
81
88
GS = 3.0
3.0
0.05
1000
0.0100
80
78
Ablation on the Pause Ratio (PR)
Appendix
Table 3: Hyperparameter Ablation Study evaluated by computing the MD committor. The main result configuration is shown first. Each subsequent row ablates a single hyperparameter, with the modified value in bold . No value means that the SAA was unstable and the force field eventually returns NaN. The number of stable samples corresponds to the number of structure for which running MD for 2 ps starting, from these configurations, is stable. The metrics were evaluated by sampling 100 configurations from ASTRA.
Figure 20: Comparison of linear and geodesic interpolation for alanine dipeptide. The top row visualizes the interpolated pathways on the potential energy surface (PES), while the bottom row shows representative intermediate structures.
Figure 21: Comparison of linear and geodesic interpolation for the folding of chignolin. The top row displays the pathways on the PES, and the bottom row shows intermediate molecular structures.
Figure 22: Committor histograms comparison of a latent-level optimization and a clean structure-level optimization. This is for PR=0.05.
Figure 23: NEB transition state guesses and minimum energy paths computed from linear and geodesic interpolation guess paths. The yellow empty circles correspond to the initial guesses computed by interpolating between the start (red) and end (blue) configurations using the given interpolation technique. The optimized paths are colored with the viridis color scheme.
Figure 24: Interpolation paths using geodesic interpolation with the start and end structures used to evaluate SAA performance and the NEB baseline. We highlight in orange the 5 middle points for each path that were taken as input for SAA. The two graphs (a) and (b) correspond to whether the initial configurations are optimized with energy minimization before computing the interpolated path. As required by NEB, the minimized configuration (b) where used for NEB calculations.
Method
% of Comm. ↑
Num. stable samples
Conv. Rate ↑
Conv. Step ↓
Pos. RMSD ↓
Energy Diff. ↑
[0.4,0.6]
(out of 56)
(Å)
(kcal ⋅ mol -1 )
Geodesic Interpolation (minimized)
SAA
0.83
54
0.93
20
0.013
-0.82
NEB (baseline)
56
0.98
0.07
-0.009
Geodesic Interpolation (non-minimized)
SAA
0.9
40
0.85
23
0.012
-0.69
Appendix
Table 4: Evaluation of TS samples generated by running SAA or NEB after initializing guess structures from classical interpolation algorithms (linear and geodesic interpolation). We report metrics associated with running the Dimer method from these structures. The "minimized" or "non-minimized" initialization corresponds to minimizing the energy of the structure before interpolating the initial guess path. Higher energy difference is better as all differences are negative.
Figure 25: Effect of the guidance scale when sampling 100 structures using Isodensity Interpolation to guide the reverse diffusion process.
Figure 26: Effect of the guidance scale when sampling 100 structures using Simple Averaging to guide the reverse diffusion process.
Transition-state (TS) structures define the energetic barriers and mechanistic pathways of elementary chemical reactions, yet their identification remains computationally demanding because conventional saddle-point searches require expensive quantum-mechanical calculations. Recent machine-learning approaches have accelerated TS generation by predicting structures from reaction endpoint information, but they primarily learn geometric correspondence between endpoints and TSs, leaving the structural transformations underlying elementary reactions implicitly represented. To address this limitation, we introduce TransTS, a reaction-transformation-aware framework for generalizable TS generation from atom-mapped reactant-product pairs. TransTS explicitly learns atom-level structural transformations between reaction endpoints and integrates them with a unified atom-aligned geometric representation of reactants, TSs and products, enabling reaction-aware equivariant generation of TS geometries. TransTS is designed to provide reliable TS initial guesses for subsequent quantum-chemical refinement, where generated structures are evaluated not only by geometric similarity but also by their ability to converge to validated saddle points and recover the intended reaction pathways. Across IID and zero-shot OOD benchmarks, TransTS demonstrates improved TS initialization quality, with particularly strong generalization to unseen reaction distributions. On the challenging GDB-10-rxn and GDB-17-rxn OOD benchmarks, TransTS generates TS candidates that more frequently converge to validated saddle points and recover the intended elementary reactions after refinement than existing approaches under the same training regime. Scaling reaction coverage and model capacity further improves both geometric fidelity and refinement outcomes.
Kaipeng Zeng, Wenxi Zhai, Shengrui Xu +5
School of Computer Science & School of Artificial Intelligence, Shanghai Jiao Tong University, Shanghai, 200240, Shanghai, China. · Shanghai Innovation Institute, Shanghai, 200231, Shanghai, China. · School of Artificial Intelligence and Data Science, University of Science and Technology of China, Hefei, 230026, Anhui, China. +3
Physics-guided sampling with diffusion model priors has shown promise for solving partial differential equation (PDE) governed problems, but applications to chemically meaningful reaction-transport systems remain limited. We apply diffusion-based guided sampling to gas-phase chemical reactions by training on solutions of the advection-reaction-diffusion (ARD) equation across varying parameters. The method generates physically consistent concentration fields and accurately predicts outlet concentrations, including at unseen parameter values, demonstrating the potential of diffusion models for inference in reactive transport.
Andrew Millard, Henrik Pedersen
Department of Computer and Information Science, Linköping University, Linköping, Sweden · Department of Physics, Chemistry and Biology, Linköping University, Linköping, Sweden
Predicting transition states (TS) in chemical reactions is crucial, as they provide insights into reaction mechanisms. Recent work on TS prediction have focused on flow matching supervised on straight linear paths that do not align with actual reaction trajectories. We propose a novel flow matching-based framework ReCurveflow that learns to predict TS geometries supervised on continuously curved reference paths interpolated from a full NEB-derived band of molecular geometries. We also introduce off-path correction, which grants ReCurveflow with the ability to produce corrective velocity fields when engaged off-path geometry states during inference rollout, leading to better resistance against exposure bias and accuracy in TS prediction. Across three data splits and six evaluation metrics, ReCurveflow achieves the best result on the majority of split-metric combinations against seven baselines. Qualitative analyses further show that ReCurveflow generates reaction trajectories with energy profiles that closely track the reference NEB path, provides initializations that ease the NEB optimization bottleneck, and exhibits the intended corrective behavior in its learned velocity fields. The ReCurveflow codebase is publicly available at https://github.com/dmis-lab/ReCurveflow.
Seungheun Baek, Mogan Gim, Jaewoo Kang
Department of Computer Science Korea University Seoul, South Korea · Department of Biomedical Engineering Hankuk University of Foreign Studies Yongin, South Korea