ReCurveflow: A Flow Matching Framework that Learns Curved Reaction Trajectories to Predict Transition State Geometries

Seungheun Baek¹, Mogan Gim^2,†, Jaewoo Kang^1,†

¹Department of Computer Science, Korea University, Seoul, South Korea ²Department of Biomedical Engineering, Hankuk University of Foreign Studies, Yongin, South Korea ^†Corresponding authors

Overview

ReCurveflow predicts transition state (TS) geometries by learning a velocity field v_theta(x, t | R, P) whose flow follows the curved reaction trajectory from the reactant state to the product state. Instead of supervising on a straight linear interpolant, we fit a natural cubic spline through the full band of NEB-derived geometry images of each Transition1x reaction, and rescale flow time so that the TS always sits at t = 0.5. Integrating the ODE from R therefore yields both a predicted TS at the midpoint and a full generated reaction trajectory.

Supervising on a curved path invites exposure bias, so ReCurveflow adds dual off-path correction: the regression target carries a term that points from an off-path geometry state back to the reference path, applied both to perturbed states around the path (perturbation-based) and to the states the model reaches during its own stop-gradient rollout (rollout-based). The field therefore learns to correct itself when the inference rollout drifts.

Installation

conda create -n recurveflow python=3.12 -y
conda activate recurveflow
# install torch matching your CUDA version first, e.g.
pip install torch==2.5.1 --index-url https://download.pytorch.org/whl/cu121
# torch_cluster provides radius_graph and needs a matching torch/CUDA wheel
pip install torch_cluster -f https://data.pyg.org/whl/torch-2.5.1+cu121.html
pip install -r requirements.txt

Data

We use Transition1x (Schreiner et al., 2022), which its authors distribute as a single HDF5 file (~6.6 GB):

# either the official downloader
git clone https://gitlab.com/matschreiner/Transition1x
python Transition1x/download_t1x.py .          # -> ./transition1x.h5

# or a direct download of the same file
wget -O transition1x.h5 https://figshare.com/ndownloader/files/36035789

Then build the three data splits of the paper:

python scripts/prepare_data.py --h5 transition1x.h5

Alternatively, download the already-built splits (see Pretrained weights and data below) and skip the 6.6 GB download.

The exact partitions are fixed by the reaction-id lists in data/splits/*.json. This writes (see data/README.md for the record layout):

data/t1x_native/         # Native: the original split, partitioned by molecular formula
data/t1x_reaction_core/  # Reaction-Core: partitioned by atom-mapped reaction cores
data/t1x_barrier/        # Barrier: tails of the activation-energy distribution held out
    train_t1x.pkl  val_t1x.pkl  test_t1x.pkl

Each pickle provides the RS, TS and PS geometries plus the ordered sequence of On-Trajectory States for every reaction. The curved reference paths are built from them when the dataset is constructed, so no further preprocessing is needed.

Training

python main.py                                        # Native split, seed 0 (configs/train.yaml)
python main.py --config-name reaction_core            # Reaction-Core split
python main.py --config-name barrier                  # Barrier split
python main.py seed=3 run_name=recurveflow-native-s3  # any Hydra override

The default config is the reported setting: 1500 epochs, 4 GPUs, batch 128 per GPU (effective batch 512), Adam with linear warmup and cosine decay, gradient clipping, and an EMA of the weights for evaluation. With fewer GPUs, raise trainer.accumulate_grad_batches to keep the effective batch at 512:

python main.py 'trainer.devices=[0]' trainer.accumulate_grad_batches=4

run_test_after_fit=true (the default) evaluates the best checkpoint on the test split when training finishes. To reproduce all reported runs (3 splits × 5 seeds):

bash scripts/run_seeds.sh                 # logs/<run_name>.log

Pretrained weights and data

Large files are published separately from the code:

Asset	Size	Contents
recurveflow-{native,reaction-core,barrier}-s0.ckpt	44 MiB each	trained weights, seed 0 of each data split
t1x_{native,reaction_core,barrier}.tar.gz	~30 MiB each	the three data splits, already built

bash scripts/download_assets.sh all      # or: weights | data

The checkpoints are weight-only exports of the training checkpoints (EMA weights, optimizer state stripped): they evaluate identically but cannot resume training. See checkpoint/README.md for the metrics each one reproduces. The data archives are exactly what scripts/prepare_data.py writes, so either route gives the same files.

Evaluation

python main.py --config-name native mode=test \
    ckpt_path=checkpoint/recurveflow-native-s0.ckpt 'trainer.devices=[0]'
# or
bash scripts/test_checkpoint.sh native checkpoint/recurveflow-native-s0.ckpt 0

This reports test/rmsd (RMSD), test/rmsd_chiral (RMSD°), test/d_mae, test/angle_mae, test/dihedral_mae and test/steric_clashes, together with trajectory-level errors (test/traj_rmsd, test/end_rs_rmsd, test/end_ps_rmsd). Evaluate on a single GPU: with several GPUs Lightning's DistributedSampler repeats samples to pad the last batch, which shifts the averages slightly.

Citation

@article{baek2026recurveflow,
  title         = {ReCurveflow: A Flow Matching Framework that Learns Curved Reaction
                   Trajectories to Predict Transition State Geometries},
  author        = {Baek, Seungheun and Gim, Mogan and Kang, Jaewoo},
  year          = {2026},
  eprint        = {TODO-ARXIV-ID},
  archivePrefix = {arXiv},
  primaryClass  = {TODO-ARXIV-CATEGORY}
}

This code builds on the Transition1x dataset (Schreiner et al., 2022); please cite it as well if you use this repository.

License

Released under the MIT License.