SafeDiffCon (From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control) (ICML 2025)

arXiv | Paper

Official repo for the paper From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control
Peiyan Hu*, Xiaowei Qian*, Wenhao Deng, Rui Wang, Haodong Feng, Ruiqi Feng, Tao Zhang, Long Wei, Yue Wang, Zhi-Ming Ma, Tailin Wu
ICML 2025.

We propose safe diffusion models for PDE Control, which introduces the uncertainty quantile as model uncertainty quantification to achieve optimal control under safety constraints through both post-training and inference phases.

Environment

Run the following commands to install dependencies. In particular, when running the 2D control task, the Python version must be 3.8 due to the requirement of the Phiflow software.

bash env.sh

Dataset

The dataset files can be downloaded via this link.

Please place the 1D Burgers' dataset in the 1D/datasets/free_u_f_1e5/ folder, place the 2D Smoke dataset in the 2d/data folder, and place the Tokamak dataset on the tokamak/tokamak_dataset folder.

Checkpoints

The checkpoints can be downloaded via this link.

Please place the 1D Burgers' checkpoint in the 1D/experiments/checkpoints/turbo-1 folder, place the 2D Smoke checkpoint in the 2d/results folder, and place the Tokamak checkpoint on the tokamak/experiments folder.

Experiments

1D Burgers' Equation

If the checkpoint is downloaded, reproduce the result with

bash /1D/reproduce_InfFT.sh

Or pretrain the model by yourself with

bash /1D/pretrain_eval.sh

2D Smoke

If the checkpoint is downloaded, reproduce the result with

bash /2d/scripts/finetune.sh

Or pretrain and posttrain the model by yourself with

bash /2d/scripts/train.sh && bash /2d/scripts/posttrain.sh 

Tokamak

If the checkpoint is downloaded, reproduce the result with

bash /tokamak/scripts/finetune.sh

Or pretrain and posttrain the model by yourself with

python /tokamak/pretrain.py && bash /tokamak/scripts/posttrain.sh && bash /tokamak/scripts/finetune.sh

Related Projects

• WDNO (ICLR 2025): We introduce Wavelet Diffusion Neural Operator (WDNO), a novel method for generative PDE simulation and control, to address diffusion models' challenges of modeling system states with abrupt changes and generalizing to higher resolutions.

• CL-DiffPhyCon (ICLR 2025): We introduce an improved, closed-loop version of DiffPhyCon. It has an asynchronous denoising schedule for physical systems control tasks and achieves closed-loop control with significant speedup of sampling efficiency.

• DiffPhyCon (NeurIPS 2024): We introduce DiffPhyCon which uses diffusion generative models to jointly model control and simulation of complex physical systems as a single task.

Citation

If you find our work and/or our code useful, please cite us via:

@inproceedings{
  hu2025from,
  title={From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe {PDE} Control},
  author={Peiyan Hu and Xiaowei Qian and Wenhao Deng and Rui Wang and Haodong Feng and Ruiqi Feng and Tao Zhang and Long Wei and Yue Wang and Zhi-Ming Ma and Tailin Wu},
  booktitle={Forty-second International Conference on Machine Learning},
  year={2025},
  url={https://openreview.net/forum?id=XGJ33p4qwt}
}