IMPORTANT — Documentation disclaimer: Supplementary documentation in this repository (including files under docs/ and related auxiliary guides) was produced with AI assistance and is still being checked for accuracy and completeness. The reaction-center predictor and its isolated TorchDrug environment are now included under rc_prediction/; see rc_prediction/README.md before generating a new prediction file.

---

RetroDiT: Order Matters in Retrosynthesis

Official codebase for the ICML 2026 paper: Order Matters in Retrosynthesis: Structure-aware Generation via Reaction-Center-Guided Discrete Flow Matching.

Overview

RetroDiT is a structure-aware template-free framework that encodes the two-stage nature of chemical reactions as a positional inductive bias. By strategically placing reaction center atoms at the sequence head, we transform implicit chemical knowledge into explicit positional patterns.

Coupled with Discrete Flow Matching, this approach decouples training from sampling, enabling highly efficient and accurate retrosynthesis.

Key Highlights

• State-of-the-Art Performance: Achieves 61.2% Top-1 accuracy on USPTO-50k and 51.3% on the large-scale USPTO-Full using predicted reaction centers.
• Exceptional Upper Bound: With oracle centers, performance reaches 71.1% (USPTO-50k) and 63.4% (USPTO-Full), surpassing foundation models trained on 10 billion reactions with a fraction of the data.
• Lightning Fast Sampling: Enables generation in just 20–50 steps, compared to the 500 steps required by prior diffusion methods.
• Novel Architecture: Introduces RetroDiT, a graph transformer backbone equipped with Rotary Position Embeddings (RoPE) to effectively capture structural priors.

Repository Structure

.
├── src/          # Core training, sampling, and evaluation entrypoints
├── configs/      # Hydra configurations for experiments, data, and models
├── rc_prediction/ # Isolated TorchDrug reaction-center predictor
├── scripts/      # Ready-to-use launch scripts for common experiments
└── docs/         # Detailed documentation (setup, data, reproducibility)

Quick Start

1. Environment Setup

We recommend conda for dependency consistency. Create and activate the environment:

bash create_pytorch_env.sh retrodit
conda activate retrodit

If you use a different environment name, activate it explicitly. See docs/installation.md for details.

Reaction-center prediction uses a separate Python 3.10 / TorchDrug environment and cannot be run in the RetroDiT training environment. See rc_prediction/README.md before producing a new RC JSON file.

2. Data Preparation

Download and prepare datasets under ${PROJECT_ROOT}/data, following the expected layout and filenames described in docs/data_preparation.md.

3. Training

Set the variables expected by the launch scripts, then run the script for your dataset:

export PROJECT_ROOT=/path/to/retro_icml2026_submission
export CONDA_PATH=/path/to/miniconda3
export NODE_COUNT=1
export NODE_RANK=0
export PROC_PER_NODE=1
export MASTER_ADDR=127.0.0.1
export MASTER_PORT=6000

bash scripts/grid_dit_model_scale_rc_all.sh

scripts/grid_dit_model_scale_rc_all.sh runs the USPTO-50k RC-all grid over model sizes. For USPTO-Full RC-all training, use scripts/dit_uspto_full_rc_all.sh with the same environment variables.

More options and Hydra overrides: docs/training_and_sampling.md.

4. Sampling and Prediction

Use the same distributed-related variables as in training, plus CKPT_PATH:

export PROJECT_ROOT=/path/to/retro_icml2026_submission
export CONDA_PATH=/path/to/miniconda3
export NODE_COUNT=1
export NODE_RANK=0
export PROC_PER_NODE=1
export MASTER_ADDR=127.0.0.1
export MASTER_PORT=6000
export CKPT_PATH=/path/to/model.ckpt

bash scripts/dit_uspto_50k_rc_all_sampling_pred.sh

Pre-trained checkpoints (when released): docs/checkpoints_and_models.md.

Three-seed benchmark

The reproducibility benchmark uses seeds 0, 1, and 42. RC prediction runs on one NVIDIA H20 per seed; RetroDiT training and sampling run on eight H20 GPUs per seed. The fixed RC export policy is strict root top-7, temperature 1.25, and split-edge aggregation. USPTO-50K preprocessing, training, and sampling all explicitly use n_dummy=10. Scripts are under scripts/experiments/seed_benchmark/.

Deduplicated top-k exact-match accuracy on the 4,944-example USPTO-50K test set is reported in percent. Mean and sample standard deviation use the three seeds (n=3, standard-deviation denominator n-1).

Top-k	Seed 0	Seed 1	Seed 42	Mean ± std
1	59.69	60.58	61.35	60.54 ± 0.83
3	81.33	82.56	82.48	82.13 ± 0.69
5	87.10	87.54	87.48	87.37 ± 0.24
10	90.94	91.16	91.02	91.04 ± 0.11
20	92.23	92.62	92.35	92.40 ± 0.20
50	92.52	92.96	92.76	92.75 ± 0.22
100	92.54	92.96	92.76	92.75 ± 0.21

Reproducibility

See docs/reproducibility.md for the mapping from paper experiments to scripts and suggested hyperparameter overrides.

License

This project is licensed under the MIT License; see the LICENSE file in the repository root.