FragmentFlow: Scalable Transition State Generation for Large Molecules

TL;DR

FragmentFlow introduces a fragmentation-based approach to transition-state generation that isolates the reactive core of a reaction, trains a generative model on core fragments, and re-attaches substituents to recover full TS geometries. By factorizing the TS as $p(x)=p(x_c)p(x|x_c)$, it mitigates size-induced distribution shifts and enables scalable TS generation for molecules up to $33$ heavy atoms on the LargeT1x benchmark, achieving about $90\%$ of TSs within $1~\text{kcal mol}^{-1}$ of reference while using $30\%$ fewer Sella optimization steps. The method combines Partial ReactOT (flow-matching on cores), IDPP-based attachment, and UMA-driven energy evaluation, and validates the reactive-core hypothesis that core quality strongly governs full TS accuracy. This work advances high-throughput reactivity studies by delivering faster, scalable TS predictions for large molecules and provides avenues for further improvements in attachment strategies and ensemble TS generation.

Abstract

Transition states (TSs) are central to understanding and quantitatively predicting chemical reactivity and reaction mechanisms. Although traditional TS generation methods are computationally expensive, recent generative modeling approaches have enabled chemically meaningful TS prediction for relatively small molecules. However, these methods fail to generalize to practically relevant reaction substrates because of distribution shifts induced by increasing molecular sizes. Furthermore, TS geometries for larger molecules are not available at scale, making it infeasible to train generative models from scratch on such molecules. To address these challenges, we introduce FragmentFlow: a divide-and-conquer approach that trains a generative model to predict TS geometries for the reactive core atoms, which define the reaction mechanism. The full TS structure is then reconstructed by re-attaching substituent fragments to the predicted core. By operating on reactive cores, whose size and composition remain relatively invariant across molecular contexts, FragmentFlow mitigates distribution shifts in generative modeling. Evaluated on a new curated dataset of reactions involving reactants with up to 33 heavy atoms, FragmentFlow correctly identifies 90% of TSs while requiring 30% fewer saddle-point optimization steps than classical initialization schemes. These results point toward scalable TS generation for high-throughput reactivity studies.

Figures (24)

• Figure 1: FragmentFlow: A Fragmentation Based Approach for Transition State (TS) Generation. (a) Given the reactants and products of chemical reactions, we use a flow matching model to generate the TS geometry of the reactive core only, which can be completed to the full structure upon the attachment of the substituents. While the generation of the full TS geometry is prone to errors that stem from a distribution shift for large molecules, FragmentFlow introduces an alternative that keeps the generative task in the training distribution. (b) FragmentFlow outperforms direct TS generation on full molecular structures, which suffers from a distribution shift. The success rate is defined as the fraction of molecules that are within $1~\text{kcal mol}^{-1}$ of reference structures after saddle-point optimization (see Figure \ref{['fig:evals']}(a)). (c) FragmentFlow achieves a lower average wall-clock time required for Sella TS optimization per structure on LargeT1x (see Section \ref{['sec:dataset_construction']}), which is the most time consuming step in our TS generation procedure.
• Figure 2: Dataset Curation Workflow. The different stages of the dataset curation consist of: reactive core identification, TS generation with TS-tools, Sella optimization using UMA, and two-stage validation (IRC + vibrational frequency analysis).
• Figure 3: Ablation Studies for Partial ReactOT: Examining the Average Number of Sella Optimization Steps.(a) Adding additional random fragments on top of the reactive core improves the quality of our TS guesses. The tested models were trained until convergence. (b) For the dataset with six additional random fragments, the quality of the generated TS guesses improves with the number of epochs.
• Figure 4: TS Generation Quality and Efficiency.(a) Empirical distributions of generated structures by RMSD and absolute energy difference to the reference TS structures. (b) FragmentFlow is more efficient than the IDPP method, despite having a comparable quality of structures. (c) The scaling laws of the average number of Sella optimization steps as a function of the number of heavy atoms in the molecules suggest that the FragmentFlow gets increasingly more efficient for larger molecules.
• Figure 5: Examining the Core Hypothesis.(a) The empirical distribution of the RMSD between the generated reactive cores and the reference reactive cores. (b) Structures with reactive cores that are closer to the reference require less Sella optimization steps, which boosts the efficiency of FragmentFlow. FragmentFlow (Partial ReactOT) achieves a Pearson correlation of 0.3 and a Spearman correlation of 0.5. IDPP achieves a Pearson correlation of 0.14 and a Pearson correlation of 0.43.