From Uncertain to Safe: Conformal Fine-Tuning of Diffusion Models for Safe PDE Control
This addresses safety-critical control problems in domains like fluid dynamics and nuclear fusion, offering a novel integration of uncertainty quantification for improved reliability.
The paper tackles the problem of ensuring safety in deep learning-based control for partial differential equations (PDEs) by proposing SafeDiffCon, a method that integrates uncertainty quantification via conformal prediction into diffusion models. Results show it satisfies all safety constraints across three control tasks, achieving the best control performance while other baselines fail.
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.