A Priori Sampling of Transition States with Guided Diffusion

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 from 2D potentials to biomolecular conformational changes and chemical reaction, ASTRA locates transition states with high precision and discovers multiple reaction pathways, enabling mechanistic studies of complex molecular systems.

💡 Research Summary

This paper introduces ASTRA (A Priori Sampling of TRAnsition States with Guided Diffusion), a novel framework that reframes the challenging problem of transition state (TS) search as an inference-time control task for generative models. Transition states, first-order saddle points on the potential energy surface (PES), are crucial for understanding reaction kinetics and mechanisms but are notoriously difficult to locate due to their fleeting nature and the topological complexity of reaction pathways. Conventional methods rely on heuristic assumptions about reaction coordinates or initial pathway guesses, limiting their generality and potentially overlooking alternative mechanisms.

ASTRA bypasses these limitations by leveraging score-based diffusion models. The workflow consists of three core stages. First, a conditional generative model is trained solely on configurations sampled from the known reactant (A) and product (B) metastable states, requiring no prior data from the transition region. Second, during inference, a principled guidance mechanism called Score-Based Interpolation (SBI) is applied. By composing the conditional scores from the two state-specific models ((S_A) and (S_B)), the reverse diffusion process is steered to sample configurations on the isodensity surface where the probability of belonging to either state is equal. This surface inherently approximates the transition region. Third, the generated samples are refined via a Score-Aligned Ascent (SAA) process. SAA uses the vector difference (S_A - S_B) as an approximation of the reaction coordinate. It then performs force-based updates: ascending the PES along this approximate reaction coordinate direction while descending along orthogonal directions using physical forces (negative gradients of the PES). This hybrid approach combines the generative model’s ability to navigate complex probability manifolds with the precision of physical energy minimization, converging configurations onto first-order saddle points without requiring the diffusion model to learn the PES perfectly.

The authors rigorously validate ASTRA across systems of increasing complexity. On 2D analytical potentials (Double Well, Müller-Brown, Double Path), ASTRA accurately samples points at the barrier tops and, crucially, discovers multiple distinct saddle points simultaneously where they exist, demonstrating inherent pathway exploration capability. For molecular systems, ASTRA successfully identifies the multiple known transition states for the conformational isomerization of alanine dipeptide. Samples generated by ASTRA exhibit committor values sharply peaked at 0.5, converge rapidly in single-ended (Dimer) TS optimizations, and are structurally nearly identical to reference TSs found by the Nudged Elastic Band (NEB) method. Furthermore, application to the coarse-grained model of the fast-folding protein Chignolin shows that ASTRA-generated samples localize in the transition region as predicted by a machine-learned committor model.

In summary, ASTRA presents a paradigm shift in transition state search. It operates a priori, requiring only data from stable states, eliminates the need for predefined reaction coordinates or iterative enhanced sampling, and naturally facilitates the discovery of multiple reaction pathways by sampling the TS ensemble. It is positioned as a universal, data-efficient method for generating high-quality initial guesses for subsequent TS optimization, potentially replacing conventional double-ended search protocols in hierarchical computational workflows.