# One Path to Model Them All: Learnable-Time Flow Matching for Crystal Structure and Energy Prediction [![Paper](https://img.shields.io/badge/Paper-ACM%20DL-B31B1B.svg)](https://dl.acm.org/doi/pdf/10.1145/3770855.3818902) [![Checkpoints](https://img.shields.io/badge/%F0%9F%A4%97%20Checkpoints-Wthinker%2FUniPath-FFD21E.svg)](https://huggingface.co/Wthinker/UniPath) [![Datasets](https://img.shields.io/badge/%F0%9F%A4%97%20Datasets-Wthinker%2FUniPath-FFD21E.svg)](https://huggingface.co/datasets/Wthinker/UniPath) UniPath KDD 2026 poster
- [Setup](#1-setup) · [Download](#2-download-checkpoints-and-data) · [Sampling](#3-run-sampling) · [Results](#4-results) · [Training](#5-training-from-scratch) · [Citation](#citation) · [Contact](#contact) --- ## 1. Setup ```bash git clone cd crystal-uniug conda env create -f environment.yml conda activate crystal_uniug pip install -e remote/cdvae pip install -e remote/DiffCSP-official pip install -e remote/riemannian-fm pip install -e . ``` Then set four environment variables. The easiest way is to copy them into a `.env` file at the project root and `source .env`: ```bash export PROJECT_ROOT="/path/to/crystal-uniug" export HYDRA_JOBS="/path/to/crystal-uniug" export WANDB_DIR="/path/to/crystal-uniug" export DATA_DIR="/path/to/data_download/dataset" # set in step 2 ``` `PROJECT_ROOT` is required: the training and sampling scripts `cd` to it and resolve checkpoint paths against it. Experiment logging via [SwanLab](https://swanlab.cn) is optional and off by default. To enable it, export your own key before running any script: ```bash export SWANLAB_API_KEY="" ``` --- ## 2. Download Checkpoints and Data Both are hosted on HuggingFace: - Checkpoints: [`Wthinker/UniPath`](https://huggingface.co/Wthinker/UniPath) (model repo, ~870 MB) - Datasets: [`Wthinker/UniPath`](https://huggingface.co/datasets/Wthinker/UniPath) (dataset repo, ~22 GB) If `huggingface.co` is unreachable, set `export HF_ENDPOINT=https://hf-mirror.com` first. ```bash # Checkpoints -> ckpt_download/ckpt/ hf download Wthinker/UniPath --repo-type model --local-dir ./ckpt_download # Datasets -> data_download/dataset/ hf download Wthinker/UniPath --repo-type dataset --local-dir ./data_download export DATA_DIR="$PWD/data_download/dataset" ``` ### Checkpoint paths | Name | Path under `ckpt_download/ckpt/` | |---|---| | `mpts52_PathRefine` | `mpts52_PathRefine/every_n_epochs/epoch=149-step=80976.ckpt` | | `mpts52_uns_PathRefine` | `mpts52_uns_PathRefine/every_n_epochs/epoch=149-step=36225.ckpt` | | `mp20_PathRefine` | `mp20_PathRefine/every_n_epochs/epoch=499-step=190587.ckpt` | | `mp20_uns_PathRefine` | `mp20_uns_PathRefine/every_n_epochs/epoch=299-step=87163.ckpt` | > ⚠️ **Do not move a `.ckpt` out of its directory.** Model loading resolves the architecture by > walking up the directory tree to find the sibling `.hydra/config.yaml`. Moving the file alone > will fail with `Failed to find hydra config!`. ### Expected data layout Each experiment needs one **CSP** dataset plus one **MD** dataset: ``` $DATA_DIR/ ├── mpts_52/ # CSP — both mpts52 checkpoints ├── mp_20/ # CSP — both mp20 checkpoints ├── mptsubmpts52/ # MD — mpts52_PathRefine ├── mptsubmpts52_uns_V0/ # MD — mpts52_uns_PathRefine ├── mptsubmp20_v0/ # MD — mp20_PathRefine └── mptsubmp20_uns_V0/ # MD — mp20_uns_PathRefine ``` CSP dirs contain `{train,val,test}.csv` + `{train,val,test}_ori.pt`; MD dirs contain `{train,val}.csv` + `{train,val}_ori.pt`. > Sampling only predicts on the test split, but the data module still loads the **train** split to > compute normalization scalers, and the MD splits because these are joint CSP+MD models. All of the > above must be present. --- ## 3. Run Sampling Sampling runs four steps in sequence. Set `CKPT` to any checkpoint from the table above: ```bash CKPT=./ckpt_download/ckpt/mp20_PathRefine/every_n_epochs/epoch=499-step=190587.ckpt SUBDIR=reproduce_1 python ./scripts_model/evaluate.py reconstruct ${CKPT} \ --subdir ${SUBDIR} --inference_anneal_slope 10.0 \ --stage test --num_evals 1 --batch_size 1024 --sample_way flow --single_gpu && \ python ./scripts_model/evaluate.py consolidate ${CKPT} --subdir ${SUBDIR} && \ python ./scripts_model/evaluate.py old_eval_metrics ${CKPT} --subdir ${SUBDIR} --stage test --do_not_log_wandb && \ python ./scripts_model/evaluate.py lattice_metrics ${CKPT} --subdir ${SUBDIR} --stage test ``` To do all four checkpoints, loop over them: ```bash for name in mpts52_PathRefine mpts52_uns_PathRefine mp20_PathRefine mp20_uns_PathRefine; do CKPT=$(ls ./ckpt_download/ckpt/${name}/every_n_epochs/*.ckpt) echo "=== ${name} ===" python ./scripts_model/evaluate.py reconstruct "${CKPT}" --subdir reproduce_1 \ --inference_anneal_slope 10.0 --stage test --num_evals 1 \ --batch_size 1024 --sample_way flow --single_gpu && \ python ./scripts_model/evaluate.py consolidate "${CKPT}" --subdir reproduce_1 && \ python ./scripts_model/evaluate.py old_eval_metrics "${CKPT}" --subdir reproduce_1 --stage test --do_not_log_wandb && \ python ./scripts_model/evaluate.py lattice_metrics "${CKPT}" --subdir reproduce_1 --stage test done ``` A ready-made template is also available at `UniPath/script/sample.sh` — edit `ckpt` and `target_dir_name`, then `bash UniPath/script/sample.sh --run_id 1`. ### Key flags | Flag | Meaning | |---|---| | `--inference_anneal_slope` | annealing slope for flow inference (`10.0` reproduces the paper) | | `--batch_size` | sampling batch size (`1024`) | | `--sample_way flow` | flow-based sampler | | `--num_evals` | number of sampling repeats per structure (`1`) | | `--single_gpu` | run on one GPU; **omit** for multi-GPU DDP | > **Runtime:** `reconstruct` dominates — roughly 1.5 h (mp_20) to 3.5 h (mpts_52) per checkpoint on a > single A100. The other three steps take a few minutes each. --- ## 4. Results Each run writes to `//`: | File | Contents | |---|---| | `old_eval_metrics_reconstruct_multi.json` | `match_rate`, `rms_dist` (multi-eval) | | `old_eval_metrics_reconstruct_single.json` | `match_rate`, `rms_dist` (single-eval) | | `consolidated_reconstruct.pt` | consolidated raw predictions | | `lengths.png` / `angles.png` | lattice parameter distributions | Compare the reported `match_rate` and `rms_dist` against the values in the paper. With `--num_evals 1` the multi- and single-eval JSONs are identical. --- ## 5. Training from Scratch Only needed if you want to retrain rather than use the released checkpoints. Budget several GPU-days per checkpoint. ### Script layout ``` UniPath/script/uni_train/main/ ├── mpts52/ -> mpts52_PathRefine (3 stages) ├── mpts52_uns/ -> mpts52_uns_PathRefine (3 stages) ├── mp20/ -> mp20_PathRefine (2 stages) └── mp20_uns/ -> mp20_uns_PathRefine (2 stages) ``` Outputs are written to `runs/trash/main//`. ### Launching Scripts launch `UniPath/run_uni.py` through `torchrun` and read the standard distributed variables (`WORLD_SIZE` = nodes, `NPROC_PER_NODE` = GPUs per node, `RANK`, `MASTER_ADDR`, `MASTER_PORT`). Each script declares a **total** batch size and divides it by the world size, so the per-device batch size adapts to your GPU count — but the total must be divisible by it. Single-node, 4 GPUs: ```bash export WORLD_SIZE=1 NPROC_PER_NODE=4 RANK=0 export MASTER_ADDR=127.0.0.1 MASTER_PORT=29500 ``` Each script `cd`s to the project root and sources `UniPath/running_setup.sh`, so it can be invoked from anywhere. ### Training order Each checkpoint is produced by a **multi-stage pipeline**: every stage after the first warm-starts from the previous stage's checkpoint (`partial_ckpt_load=True`, a partial weight load rather than a Lightning full-state resume). Stages must therefore run **in order**.
mpts52_PathRefine — 3 stages ```bash T=UniPath/script/uni_train/main/mpts52 # 1. Joint CSP+MD pretraining, from scratch (150 epochs) bash $T/uni_E_tclearn_freeze.sh # 2. Unsupervised-flow path refinement, warm-started from stage 1 @ epoch 109 bash $T/mpts52_mptsub_path_refine_wunstune.sh --w_uns 0.1 # 3. + consistency loss, warm-started from stage 2 @ epoch 129 -> RELEASED bash $T/mpts52_mptsub_path_refine_consist_wunstune.sh --w_uns 0.1 --w_consist 0.05 ```
mpts52_uns_PathRefine — 3 stages ```bash T=UniPath/script/uni_train/main/mpts52_uns bash $T/uni_E_tclearn_freeze.sh bash $T/mpts52_uns_path_refine_wunstune.sh --w_uns 0.02 bash $T/mpts52_uns_path_refine_consist_wunstune.sh --w_uns 0.02 --w_consist 0.02 # -> RELEASED ```
mp20_PathRefine — 2 stages ```bash T=UniPath/script/uni_train/main/mp20 bash $T/bz256_tclearn_freeze_epoch500.sh bash $T/mp20_wtune.sh --w_uns 0.001 --w_consist 0.01 # -> RELEASED ```
mp20_uns_PathRefine — 2 stages ```bash T=UniPath/script/uni_train/main/mp20_uns bash $T/uni_E_tclearn_freeze.sh bash $T/mp20_uns_wunte.sh --w_uns 0.02 --w_consist 0.5 # -> RELEASED ```
`--w_uns` sets the unsupervised-flow loss weight, `--w_consist` the consistency loss weight. The values above are the ones that produced the released checkpoints. ### Before you train - **`partial_ckpt_path` is built from `${PROJECT_ROOT}`** in each warm-starting script, so exporting `PROJECT_ROOT` is enough — no need to edit the paths by hand. - **Parent checkpoints are referenced by exact filename** (e.g. `epoch=109-step=424160.ckpt`). The step count depends on your world size and batch size, so a retrained parent produces a *different* filename — update the referencing script accordingly. - **`runs/` is gitignored**, so training outputs are never committed. ### Sampling your own checkpoint Use the Section 3 commands with `CKPT` pointing into your run directory, keeping the checkpoint beside its `.hydra/config.yaml`: ```bash CKPT=runs/trash/main/mpts52_PathRefine_wuns0.1_wc0.05_e109_e129/every_n_epochs/epoch=149-step=80976.ckpt ``` --- ## Citation If you find this work useful, please cite: ```bibtex @inproceedings{li2026one, title={One Path to Model Them All: Learnable-Time Flow Matching for Crystal Structure and Energy Prediction}, author={Li, Songyou and Li, Mingze and Liu, Qianpu and Cen, Jiacheng and Wu, Liming and Xu, Tingyang and Rong, Yu and Zhao, Deli and Huang, Wenbing}, booktitle={Proceedings of the 32nd ACM SIGKDD Conference on Knowledge Discovery and Data Mining V. 2}, pages={11351--11362}, year={2026} } ``` --- ## Contact Questions and issues are welcome — open a GitHub issue or reach out to [songyou_li@126.com](mailto:songyou_li@126.com).