Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

From Uncertain to Safe: Conformal Fine-Tuning 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

This work addresses safety in PDE-constrained control by introducing SafeDiffCon, a diffusion-model approach that quantifies safety uncertainty with conformal prediction. It integrates an uncertainty quantile into both a post-training reweighting scheme and inference-time fine-tuning to enforce safety constraints while maintaining control performance. The method yields conformal intervals that bound the safety score with high probability under distribution shift, and demonstrates safety-compliant and high-quality control across Burgers' equation, 2D Navier–Stokes flow, and tokamak fusion scenarios. The results show SafeDiffCon uniquely satisfies all safety constraints among baselines and delivers superior objective values in safe regimes, highlighting its potential for deploying ML-based controllers in high-stakes PDE-domain applications.

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.

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

TL;DR

This work addresses safety in PDE-constrained control by introducing SafeDiffCon, a diffusion-model approach that quantifies safety uncertainty with conformal prediction. It integrates an uncertainty quantile into both a post-training reweighting scheme and inference-time fine-tuning to enforce safety constraints while maintaining control performance. The method yields conformal intervals that bound the safety score with high probability under distribution shift, and demonstrates safety-compliant and high-quality control across Burgers' equation, 2D Navier–Stokes flow, and tokamak fusion scenarios. The results show SafeDiffCon uniquely satisfies all safety constraints among baselines and delivers superior objective values in safe regimes, highlighting its potential for deploying ML-based controllers in high-stakes PDE-domain applications.

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.

Paper Structure

This paper contains 50 sections, 6 theorems, 39 equations, 5 figures, 20 tables, 3 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Control of PDE Systems
  4. Conformal Prediction
  5. Preliminary
  6. Problem Setup
  7. Diffusion Models and Diffusion Control
  8. Method
  9. Uncertainty Quantification of Diffusion Models
  10. Post-training with Reweighted Loss
  11. Inference-time Fine-tuning
  12. Experiment
  13. 1D Burgers' Equation
  14. 2D Incompressible Fluid
  15. Tokamak Fusion Reactor
  16. ...and 35 more sections

Key Result

Theorem 4.1

Assume that samples in the calibration set $D_\textrm{cal}\sim p$ are independent, and the test set $(\mathbf{u},\mathbf{w})\sim \Tilde{p}$ is also independent with the calibration set. Assume $p$ is absolutely continuous with respect to $\Tilde{p}$, then (See Appendix app:theory for proof.)

Figures (5)

  • Figure 1: Overview of SafeDiffCon. First, we pre-train a diffusion model $p_{\theta}$ on the training data. Then, combined with the uncertainty quantile, we post-train the model to steer its distribution to safer regions with better control objectives. Finally, to improve performance and safety for specific control tasks, we conduct inference-time fine-tuning, again incorporating the uncertainty quantile into the process.
  • Figure 2: Visualizations of the 1D Burgers' equation. The top row shows the original trajectory corresponding to the control target, and the bottom row is the trajectory controlled by SafeDiffCon.
  • Figure 3: Visualization of the 2D incompressible fluid control by our SafeDiffCon. By controlling the fluid on the outside margin, the yellow smoke is successfully maneuvered to the center top exit while avoiding the red unsafe region.
  • Figure 4: Visualization of the 1D Burgers’ equation.
  • Figure 5: Visualizations of the 2D incompressible fluid control problem..

Theorems & Definitions (6)

  • Theorem 4.1
  • Theorem 4.2
  • Theorem 4.3
  • Theorem 1.1
  • Theorem 1.2
  • Theorem 1.3