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

TL;DR: 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.

Abstract: The application of deep learning for partial differential equation (PDE)-constrained control is gaining increasing attention. However, existing methods rarely consider safety requirements crucial in real-world applications. To address this limitation, we propose Safe Diffusion Models for PDE Control (SafeDiffCon), which introduce the uncertainty quantile as model uncertainty quantification to achieve optimal control under safety constraints through both post-training and inference phases. Firstly, our approach post-trains a pre-trained diffusion model to generate control sequences that better satisfy safety constraints while achieving improved control objectives via a reweighted diffusion loss, which incorporates the uncertainty quantile estimated using conformal prediction. Secondly, during inference, the diffusion model dynamically adjusts both its generation process and parameters through iterative guidance and fine-tuning, conditioned on control targets while simultaneously integrating the estimated uncertainty quantile. We evaluate SafeDiffCon on three control tasks: 1D Burgers' equation, 2D incompressible fluid, and controlled nuclear fusion problem. Results demonstrate that SafeDiffCon is the only method that satisfies all safety constraints, whereas other classical and deep learning baselines fail. Furthermore, while adhering to safety constraints, SafeDiffCon achieves the best control performance. The code can be found at https://github.com/AI4Science-WestlakeU/safediffcon.

Lay Summary: Controlling physical systems with deep learning is promising but often unsafe, because most methods ignore real-world safety constraints. We introduce SafeDiffCon, a method that uses uncertainty estimates to ensure safe and effective control of physical systems. It improves a pre-trained model to follow safety constraints and dynamically adjusts its behavior during inference. In tests on fluid and fusion control tasks, SafeDiffCon is the only method that always stays safe while also achieving the best control performance.

Link To Code: https://github.com/AI4Science-WestlakeU/safediffcon

Primary Area: Applications->Chemistry, Physics, and Earth Sciences

Keywords: safe PDE control, PDE simulation, generative models, conformal prediction

Submission Number: 7353