From Uncertain to Safe: Conformal Fine-Tuning of Diffusion Models for Safe PDE Control

TL;DR

This paper introduces SafeDiffCon, a novel method that enhances diffusion models for PDE control by incorporating uncertainty quantification to ensure safety constraints are met during both training and inference, outperforming existing approaches.

Contribution

The paper proposes SafeDiffCon, a new framework that integrates conformal prediction-based uncertainty quantification into diffusion models for safe PDE control, with dynamic adjustment during inference.

Findings

SafeDiffCon satisfies all safety constraints in tested control tasks.

It achieves superior control performance while maintaining safety.

Outperforms classical and deep learning baselines in safety and effectiveness.

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.