Reaction-Transformation-Aware Flow Matching for Generalizable Transition State Generation

Kaipeng Zeng; Wenxi Zhai; Shengrui Xu; Jie Zhao; Bowen Li; Shiyue Wang; Junchi Yan; Tong Zhu

Source-linked AI summary

Reaction-Transformation-Aware Flow Matching for Generalizable Transition State Generation

Kaipeng Zeng, Wenxi Zhai, Shengrui Xu, Jie Zhao, Bowen Li, Shiyue Wang, Junchi Yan, Tong Zhu

arXiv:2608.14076v1Aug 14, 2026physics.chem-phcs.AI

TL;DR

Transition-state generation remains difficult because conventional saddle-point searches require expensive quantum-mechanical calculations, while existing machine-learning methods underrepresent structural transformations. TransTS explicitly models atom-level transformations and improves TS initialization, especially under zero-shot distribution shift, where candidates more often recover intended reaction channels after refinement.

Problem
Existing machine-learning TS generators mainly learn endpoint–TS geometric correspondence, leaving elementary-reaction structural transformations implicitly represented.
Method
TransTS learns atom-level transformations from atom-mapped reactant–product pairs and integrates them with an atom-aligned geometric framework for reaction-aware TS generation.
Results
61.5% and 58.6% reaction correctness on GDB-10-rxn and GDB-17-rxn, respectively, were the highest rates after RGD1 augmentation.
Takeaways & Limitations
TransTS provides more reliable TS initial guesses under distribution shift, improving both pre-optimization geometry and recovery of intended reaction channels after refinement.
Takeaways & Limitations
Evaluation focuses on organic CHNO reactions, so broader chemistry involving charged species, radicals, transition metals, or explicit solvent and catalyst environments requires additional data and validation.

Abstract

from arXiv · show

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.

1 Introduction

Transition states govern reaction barriers and reveal elementary reaction mechanisms, but their direct characterization and computational discovery remain difficult and expensive. TransTS addresses this limitation by explicitly learning structural transformations from atom-mapped reactant–product pairs for reaction-aware TS generation.

Motivation: Transition states are first-order saddle points that define energetic bottlenecks, activation barriers, and structural changes along elementary reaction pathways.Their geometries help elucidate reaction mechanisms.
Challenges: Direct TS characterization is challenging because these structures are highly transient, while available experimental methods remain specialized rather than routine and general.Ultrafast electron diffraction has enabled observation of transient molecular structures but does not yet provide a general exploratory route.
Challenges: Automated PES exploration still requires expensive repeated energy and force evaluations and can encounter convergence difficulties on complex potential energy landscapes.Candidate TSs are refined using local saddle-point optimization methods, including Hessian-based approaches and modern saddle optimizers.
Research gap: Existing ML-based TS methods largely learn geometric correspondence between reaction endpoints and TS structures, leaving elementary-reaction transformations implicit.These transformations include bond rearrangements and changes in atomic environments between atom-corresponded reactants and products.
TransTS: TransTS explicitly learns atom-wise structural transformations between corresponding reactant and product environments and uses transformation-informed conditions for TS generation.The framework generates TSs from atom-mapped reactant-product pairs through a reaction-aware representation module.

2 Results

TransTS is evaluated through pre-optimization geometry and post-optimization saddle-point refinement, with its strongest advantages emerging under zero-shot distribution shift. Scaling reaction coverage and model capacity further improves geometry, refinement outcomes, and recovery of reaction-center structural evolution.

IID results: On Transition1x, TransTS achieves 94.1% optimization success and 71.8% reaction correctness, remaining competitive despite not attaining the lowest RMSD.Its median RMSD is 0.171 ˚A against original labels and 0.168 ˚A on the post-checked subset, versus 96.5% and 73.2% for React-OT in refinement.
OOD results: On zero-shot OOD benchmarks, TransTS has the lowest median RMSD, reaching 0.379 ˚A on GDB-10-rxn and 0.429 ˚A on GDB-17-rxn.The second-best values are 0.401 ˚A and 0.485 ˚A from React-OT.
OOD results: TransTS achieves 55.1% reaction correctness on GDB-10-rxn and 37.9% on GDB-17-rxn, exceeding the next-best model by 13.3 and 13.2 percentage points.Its optimization-success rate is best on GDB-17-rxn and competitive on GDB-10-rxn.
OOD results: With RGD1 augmentation, TransTS reaches median RMSDs of 0.310 ˚A on GDB-10-rxn and 0.362 ˚A on GDB-17-rxn, while reaction correctness rises to 61.5% and 58.6%.These correspond to absolute gains of 0.057 ˚A and 0.202 ˚A over MolGEN, the second-best coordinate-level baseline.
OOD results: TransTS’s OOD reaction-correctness advantage is broadly distributed across BnFm bond-rearrangement classes rather than driven by one dominant class.It exceeds React-OT in most retained classes, including higher-rearrangement classes.
Scaling and structural analysis: Scaling reaction diversity and model capacity improves TransTS across geometry and refinement, including optimization-success gains of 2.4, 7.2, and 20.1 percentage points.The gains occur on Transition1x, GDB-10-rxn, and GDB-17-rxn, respectively; qualitative examples also show recovery of coupled bond formation, cleavage, and bond-order rearrangements.

3 Discussion

TransTS explicitly models atom-level reaction transformations and couples them with endpoint-aligned representations for transition-state generation. Its clearest benefits appear under distribution shift, while geometric fidelity depends on data scale and broader chemistry remains unevaluated.

Framework: TransTS compares corresponding reactant and product atoms to emphasize elementary-reaction transformations and align target transition states.The framework couples this reaction condition with an endpoint-symmetric canonical anchor.
Benchmark performance: TransTS shows its clearest advantage under distribution shift, achieving the strongest zero-shot OOD performance on both GDB-10-rxn and GDB-17-rxn.After RGD1 augmentation, it also achieves the best IID RMSD while remaining competitive after saddle-point refinement.
Limitations: Broader reaction coverage and increased model capacity improve TransTS geometric accuracy, indicating that data scale matters for high-fidelity TS coordinates.When trained only on Transition1x, TransTS is not the strongest RMSD model, although its refinement behavior remains competitive.
Limitations: Evaluation is currently limited to organic CHNO reaction space, leaving broader chemistry involving charged systems unevaluated.The supplied discussion identifies broader chemistry as a limitation of the current evaluation.

4 Methods

TransTS formulates transition-state generation as conditional, atom-aligned 3D structure generation from reactant–product pairs. It builds a symmetric unified reaction frame and combines transformation-aware endpoint encoding with SE(3)-equivariant flow-matching generation.

Problem formulation: TS generation is modeled as learning a conditional distribution over TS geometries given atom-mapped reactant and product structures.Reactants, products, and TSs share atom index set V and are represented by Cartesian coordinates x_R, x_P, and x_TS.
Unified reaction frame: Atom mapping defines the elementary reaction channel and enables reactant, product, and TS conformations to form one reaction-level geometric system.The unified representation uses object-aware transformation freedoms to preserve SE(3) equivariance across molecular roles.
Unified reaction frame: A Kabsch-aligned canonical anchor provides a deterministic reference frame that is symmetric to reversing the reactant and product endpoints.The TS is subsequently aligned to this inferred anchor, yielding a unique conformation for each aligned reactant–product pair.
Reaction encoder: The reaction encoder captures structural transformations through shared endpoint branches and atom-mapped Cross-State Fusion between corresponding atoms.This design exploits endpoint exchange symmetry while focusing representations on local environment changes between reactants and products.
TS generator and flow matching: The generator treats the TS interpolant as a dynamic state and predicts an SE(3)-equivariant velocity field using time embeddings and symmetric atom-mapped reaction-context injection.This allows the model to reuse unchanged scaffold information while emphasizing atoms and bonds undergoing structural changes.

Declarations

Supplementary materials document the endpoint-matching and RMSD-evaluation procedures, while data and implementation resources are publicly available through linked repositories and archives.

Supplementary materials: Supplementary information includes Python files documenting endpoint-matching and RMSD-evaluation procedures.The ancillary files are anc/xyz2smiles.py and anc/rmsd.py.
Data availability: GDB-10-rxn, GDB-17-rxn, Transition1x, RGD1, and processed data are available through Figshare, GitLab, GitHub, and the TransTS repository.The passage provides access links for each dataset or processed-data resource.
Code availability: The TransTS implementation code is available in the project’s GitHub repository.The repository is identified as github.com/zengkaipeng/TransTS.

Supplementary Information · 1 Data Information

The study constructs two OOD reaction datasets and validates transition-state pathways computationally, while post-checking the Transition1x test set to ensure reference TSs correspond to annotated elementary reactions. The resulting datasets and checks support evaluation of generalization and clean-label TS analysis.

1.1 Construction and Quantum-Chemical Validation of GDB-10-rxn and GDB-17-rxn: GDB-10-rxn contains 779 reactions with up to 10 heavy atoms and reaction types absent from training data.Its elemental composition and molecular size remain similar to the training distribution.
1.1 Construction and Quantum-Chemical Validation of GDB-10-rxn and GDB-17-rxn: Both OOD datasets began from 100 randomly selected seed molecules per database, followed by reaction-network enumeration with YARP.For retained reactions, TS geometries were optimized at ωB97X-D/def2-SVP and pathways were confirmed using IRC calculations.
1.1 Construction and Quantum-Chemical Validation of GDB-10-rxn and GDB-17-rxn: The released OOD datasets include IRC-path structures with energy and force labels for reactive machine-learning potential training.These labels accompany the computationally validated reaction pathways.
1.2 Post-Checked Transition1x Test Subset: The Transition1x post-check started from 287 test-split reactions and used each reference TS for saddle-point optimization followed by IRC analysis.Calculations used Gaussian 16 at the ωB97X/6-31G(d) level, matching the original Transition1x setup.
1.2 Post-Checked Transition1x Test Subset: Reference TSs were accepted only when both IRC endpoints matched the annotated reactant-product pair.Endpoint identity was assessed from optimized reactant, product, and IRC-endpoint geometries using mapped SMILES or RMSD matching; failed calculations were unmatched.
1.2 Post-Checked Transition1x Test Subset: 262 reactions passed the reference-TS-initialized post-check, and three additional reaction-consistent TSs were recovered using model-predicted TS guesses.Applying the same optimization, IRC, and endpoint-matching protocol yielded a final subset of 265 reactions out of 287.

2 Implementation Details

TransTS variants were trained with Adam under a shared 2000-epoch schedule and selected using validation Kabsch RMSD. Inference and evaluation used fixed ODE integration, standardized geometric metrics, and Gaussian 16 saddle-point refinement with frequency analysis.

Training: Both variants used Adam, a 2 × 10−4 warm-up peak learning rate, exponential decay, 2000 epochs, effective batch size 256, and validation every 5 epochs.Validation generated samples with a midpoint ODE solver using 10 fixed steps.
Training: Checkpoint selection minimized direct validation Kabsch RMSD before atom reindexing.This criterion was applied during validation of the trained variants.
Inference: Inference integrated the learned ordinary differential equation from an anchor-aligned Gaussian-noise state to t = 1 using torchdiffeq’s midpoint solver with 15 ODE steps.All models, including TransTS, used one sampled trajectory per reaction for fair comparison.
Evaluation: Geometric accuracy used reindexed, chirality-insensitive, Kabsch-aligned RMSD, with atom matching performed by pymatgen v2026.5.4 and pymatgen-core v2026.5.18.The evaluator implementation was provided as anc/rmsd.py.
Evaluation: Saddle-point refinement used Gaussian 16 with ωB97X/6-31G(d) for Transition1x and ωB97X-D/def2-SVP for the OOD test sets, followed by frequency analysis.Each generated transition-state guess was refined with Opt=(TS, CalcFC, NoEigenTest) Freq.

3 Supplementary Experimental Results

Supplementary results report pre-optimization RMSD distributions, post-optimization outcome rates, and an OOD bond-rearrangement stratification for GDB-10-rxn and GDB-17-rxn. TransTS shows lower reported RMSD medians than the listed alternatives in one reported setting, while the stratification retains dominant reaction classes covering most datasets.

Pre-optimization RMSD: TransTS reports RMSD medians of 0.069 Å, 0.067 Å, 0.310 Å, and 0.362 Å on T1x-original, T1x-clean, GDB-10-rxn, and GDB-17-rxn, respectively.Values are reported as median [Q1, Q3] in Å, with lower median RMSD preferred.
Post-optimization outcomes: Post-optimization evaluation defines invalid, success, and correctness rates over all reactions, with lower invalid and higher success and correctness rates preferred.The supplied table excerpt provides the metric headers but not the corresponding numerical outcome rates.
OOD bond-rearrangement analysis: OOD bond-rearrangement classes are defined as BnFm according to the numbers of broken and formed bonds between atom-mapped reactant and product graphs, including hydrogen-involving bonds.Graphs are obtained from optimized endpoint geometries using anc/xyz2smiles.py, the same ancillary file used for IRC endpoint matching.
OOD bond-rearrangement analysis: Retained BnFm classes cover 86.78% of GDB-10-rxn and 87.36% of GDB-17-rxn, while remaining long-tail chemical classes are grouped into Other and NO SMILES remains standalone.Because some exact classes are small, stratified post-optimization comparisons are reported as absolute counts rather than within-class rates.

4 Justification for Anchor Alignment

The section justifies anchor-based alignment as a causality-preserving surrogate for optimal transport in equivariant 3D flow matching. It combines efficient aligned transport with identical, endpoint-observable initialization during training and inference.

The Optimal Transport Ideal: Optimal alignment reduces transport-cost minimization to an Orthogonal Procrustes problem with an efficient closed-form SVD solution.The optimal rotation is obtained from the covariance matrix with determinant correction enforcing membership in SO(3).
The Optimal Transport Ideal: After alignment, Wasserstein-optimal displacement interpolation yields a time-independent constant-velocity field that reduces learned-flow complexity and structural collapse.The aligned path is a Euclidean straight line between coupled noise and target states.
The Observability Constraint: Target-aware alignment is unusable at inference because the ground-truth TS is unobservable, creating train-test distribution shift and preventing computation of the optimal rotation.The problematic training boundary condition depends directly on the target TS.
Equivariance and the Jointly Moving System: Whole-reaction SE(3) equivariance makes global orientation arbitrary, allowing a deterministic canonical frame constructed solely from observable reactant and product geometries.The model treats reactant, product, and TS as a jointly moving system while preserving relative geometries.
The Surrogate Anchor Alignment: The endpoint-symmetric anchor aligns noise and TS structures in an inferable canonical frame, preserving causality and train-test consistency while reducing rotational variance.The anchor satisfies Fanchor(xR, xP) = Fanchor(xP, xR), so it is not tied to reaction direction.