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

Technical Summary: Reaction-Transformation-Aware Flow Matching for Generalizable Transition State Generation

Problem Statement
Transition states (TSs) are first-order saddle points on the potential energy surface (PES) that define the energetic barriers and mechanistic pathways of elementary chemical reactions. Identifying these structures is computationally expensive, as conventional methods rely on repeated, costly quantum-mechanical calculations (e.g., Density Functional Theory) combined with saddle-point optimization algorithms. While recent machine learning (ML) approaches have accelerated TS generation by predicting structures from reaction endpoints (reactants and products), existing methods primarily learn the geometric correspondence between endpoints and the TS. They often treat the structural transformations underlying elementary reactions—such as bond rearrangements and changes in atomic environments—as implicitly represented. Furthermore, current geometric formulations frequently lack a unified, atom-aligned representation that facilitates the exchange of higher-order equivariant information across reactants, TSs, and products. This limitation hinders the generalizability of TS generators, particularly for out-of-distribution (OOD) reactions where the model must infer saddle-point geometries for unseen molecular sizes and stoichiometries.

Methodology: TransTS Framework
The authors introduce TransTS, a reaction-transformation-aware framework designed to generate generalizable TS geometries from atom-mapped reactant-product pairs. The methodology integrates three core components:

1. Reaction-Transformation-Aware Representation:
   Unlike previous models that treat endpoint geometry uniformly, TransTS explicitly learns atom-level structural transformations. It employs a shared-parameter dual encoder that processes atom-mapped reactant and product conformations. A key module, Cross-State Fusion, exchanges information between corresponding atoms of the reactant and product. This allows the network to distinguish between conserved scaffold regions (providing stable geometric context) and atoms undergoing local environmental changes (receiving reaction-specific conditioning). This design makes endpoint differences explicit, providing transformation-informed conditions for generation.

2. Unified Atom-Aligned Geometric Framework:
   To enable cross-state information exchange, TransTS represents reactants, TSs, and products within a single aligned frame. The framework maps coordinates to a center-of-mass-free (com-free) subspace and constructs a canonical anchor by rigidly aligning the product with the reactant using the Kabsch algorithm. This anchor is symmetric with respect to the reaction direction (forward or reverse). The TS generator operates in this unified frame, allowing it to reuse information from unchanged scaffold regions while focusing modeling capacity on atoms and bonds undergoing structural rearrangement.

3. Anchor-Aligned Flow Matching:
   The generation process utilizes Conditional Flow Matching (CFM). The model learns a time-conditioned equivariant velocity field that transports a noisy initial state to the target TS geometry. Crucially, the flow path is constructed using the inferred canonical anchor. The training trajectory aligns both the noise source and the target TS to this anchor, creating an exact straight-line interpolation path. This approach serves as a practical surrogate for optimal transport while ensuring consistency between training and inference, as the anchor is derived solely from observable reactant and product structures.

Key Contributions

• Explicit Transformation Learning: TransTS is the first framework to explicitly model atom-wise structural transformations between reaction endpoints as a conditioning signal for TS generation, moving beyond simple geometric correspondence.
• Unified Equivariant Representation: The paper proposes a unified, atom-aligned geometric framework that enables the direct exchange of higher-order equivariant information across reactants, products, and TSs.
• Robust OOD Generalization: By leveraging reaction-transformation awareness, the model achieves superior generalization to unseen reaction distributions (zero-shot OOD) compared to state-of-the-art baselines.
• Validation-Centric Evaluation: The authors introduce a rigorous evaluation protocol that assesses not only geometric similarity (RMSD) but also the utility of generated structures as initial guesses for quantum-chemical refinement, measuring optimization success and reaction correctness rates.

Results
The performance of TransTS was evaluated on the Transition1x dataset (In-Distribution, IID) and two challenging OOD benchmarks derived from the GDB chemical universe: GDB-10-rxn and GDB-17-rxn.

• IID Performance: When trained solely on Transition1x, TransTS did not achieve the lowest pre-optimization RMSD compared to baselines like React-OT. However, it remained competitive in post-optimization metrics (optimization success and reaction correctness). When the training set was augmented with the larger RGD1 dataset, TransTS achieved the lowest median RMSD on both original and post-checked subsets, demonstrating that scaling improves geometric fidelity.
• Zero-Shot OOD Performance: TransTS demonstrated significant advantages in zero-shot settings. On GDB-10-rxn and GDB-17-rxn, it achieved the lowest pre-optimization RMSD among all methods. More critically, it showed superior post-optimization performance. For instance, when trained only on Transition1x, TransTS recovered the intended reaction in 55.1% of GDB-10-rxn cases, outperforming the second-best method by 13.3 percentage points. On GDB-17-rxn, it achieved a 37.9% reaction correctness rate, exceeding the next-best model by 13.2 percentage points.
• Scaling Effects: Increasing reaction coverage (Transition1x + RGD1) and model capacity led to consistent improvements across all benchmarks. The gains were most pronounced in the most demanding OOD setting (GDB-17-rxn), where reaction correctness improved by 20.7 percentage points.
• Structural Evolution: Case studies on GDB-10-rxn showed that TransTS effectively recovers characteristic reaction-center structural evolutions, such as coupled bond formation and hydrogen migration, with lower mean absolute errors in reaction-core distances compared to baselines.

Significance and Claims
The paper posits that reaction-transformation-aware generation is a promising strategy for accelerating the computational exploration of chemical reaction spaces. The authors claim that by explicitly modeling the structural changes associated with elementary reactions, TransTS provides more reliable initial guesses for quantum-chemical refinement, particularly for reactions outside the narrow distribution of standard benchmarks. The framework's ability to generalize to unseen molecular sizes and complex stoichiometries suggests it can support the construction of broader reaction networks and the rational design of catalysts. The authors modestly note that the current evaluation is limited to organic CHNO reaction spaces and that accurate coordinate-level generation appears dependent on data scale, implying that further scaling is necessary for high-fidelity applications in more complex chemical environments (e.g., transition metals, radicals).