🤖 machine learning

From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control

This paper introduces SafeDiffCon, a novel framework that enhances pre-trained diffusion models for PDE-constrained control by integrating conformal prediction-based uncertainty quantification during both post-training and inference to guarantee safety constraints while achieving superior control performance.

Original authors: Peiyan Hu, Xiaowei Qian, Wenhao Deng, Rui Wang, Haodong Feng, Ruiqi Feng, Tao Zhang, Long Wei, Yue Wang, Zhi-Ming Ma, Tailin Wu

Published 2026-06-25

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Peiyan Hu, Xiaowei Qian, Wenhao Deng, Rui Wang, Haodong Feng, Ruiqi Feng, Tao Zhang, Long Wei, Yue Wang, Zhi-Ming Ma, Tailin Wu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a very talented but slightly reckless chef (the Diffusion Model) how to cook a perfect meal. The chef is great at guessing the right ingredients and steps to make a delicious dish (solving complex physics problems), but they have a bad habit: sometimes, in their excitement to make the food taste amazing, they accidentally add too much salt or set the stove on fire. In the real world, this is like a control system for a nuclear reactor or a fluid flow that gets too hot or turbulent, causing a disaster.

The paper introduces a new method called SafeDiffCon to fix this. It's like giving the chef a "safety supervisor" and a "reality check" so they can cook safely without losing their flair.

Here is how it works, broken down into simple steps:

1. The Problem: The "Reckless Chef"

Existing AI methods for controlling physical systems (like fluids or nuclear fusion) are great at getting the job done efficiently, but they often ignore safety rules. They might find a solution that works perfectly on paper but would cause a real-world explosion or damage. It's like a self-driving car that drives the fastest route possible but ignores stop signs because it's never been told to care about them.

2. The Solution: The "Safety Supervisor" (Conformal Adaptation)

The authors realized that the AI is often uncertain about whether its plan is safe. It might think it's safe, but it could be wrong.

To fix this, they use a statistical trick called Conformal Prediction. Think of this as the chef keeping a "safety logbook."

The Logbook: Before the chef starts cooking the final meal, they look at a list of past cooking attempts (data). They check: "How far off was my guess on the salt level compared to what actually happened?"
The Safety Margin: Instead of just guessing a single number, the chef calculates a "safety buffer." If the logbook says, "I was usually off by 2 grams," the chef adds a 2-gram buffer to their plan. This ensures that even if the chef is slightly wrong, they stay within the safe zone.

3. Step One: The "Safety Training" (Post-Training)

First, the AI is given a special training session.

The Old Way: The AI was trained to just minimize the "badness" of the result (like making the soup taste bad).
The New Way (SafeDiffCon): The AI is retrained with a Reweighted Loss. Imagine the teacher telling the chef: "If you are uncertain about whether a step is safe, we are going to punish you extra hard for that step."
The Result: The AI learns to avoid risky moves. It stops trying to be the "fastest" chef and starts being the "safest" chef who still makes a great meal. It learns to steer its "guesses" toward areas where it is confident the safety rules won't be broken.

4. Step Two: The "Live Correction" (Inference-Time Fine-Tuning)

Even after training, the AI might face a specific, tricky situation (like a sudden wind gust in a fluid experiment).

The Adjustment: Before the AI makes its final move, it does a quick, real-time check. It looks at the specific goal, checks its "safety logbook" again, and tweaks its plan one last time.
The Analogy: It's like a pilot doing a final pre-flight check. They look at the weather, check the fuel, and make tiny adjustments to the controls to ensure they don't hit a storm, even if the flight plan looked perfect on the ground.

5. The Results: The "Perfect Meal"

The authors tested this "Safe Chef" on three very difficult scenarios:

1D Burgers' Equation: Like managing a wave in a pipe.
2D Incompressible Fluid: Like guiding smoke through a maze without letting it touch a fire.
Tokamak Fusion: Like controlling a super-hot plasma ball for nuclear energy without letting it melt the machine.

The Outcome:

Other Methods: Most other AI methods (like standard deep learning or older control methods) either failed to keep the system safe (the smoke hit the fire, or the plasma became unstable) or they were so cautious they failed to do the job at all.
SafeDiffCon: It was the only method that successfully kept the system safe and achieved the best performance. It didn't just avoid the fire; it guided the smoke exactly where it needed to go.

Summary

The paper proposes a way to take powerful, creative AI models and teach them to respect safety boundaries. By using a "safety logbook" to understand their own uncertainty and adjusting their training and real-time decisions accordingly, they can control complex physical systems without causing disasters. It turns an "uncertain" AI into a "safe" one, without sacrificing its ability to get the job done.

Technical Summary: From Uncertain to Safe: Conformal Adaptation of Diffusion Models for Safe PDE Control

Problem Statement

The control of physical systems governed by Partial Differential Equations (PDEs)—such as fluid dynamics, incompressible flow, and controlled nuclear fusion—is critical in engineering and science. While deep learning, particularly diffusion models, has shown promise in modeling high-dimensional, nonlinear data for control tasks, existing methods largely overlook safety constraints. In real-world applications, violating safety constraints (e.g., causing turbulence, structural damage, or plasma disruptions) can lead to catastrophic failures.

Current deep learning approaches for PDE control face two primary challenges:

Distribution Shift: Offline training data is often suboptimal and contains unsafe samples. This creates a gap between the distribution learned by the model and the desired optimal, safe distribution.
Safety-Performance Trade-off: Optimizing control performance (minimizing error) often conflicts with satisfying safety constraints. Aggressive policies that reduce control error may violate safety thresholds.
Lack of Formal Guarantees: Existing methods rarely provide formal guarantees that predicted control sequences will satisfy safety bounds, especially when model uncertainty is high.

Methodology: SafeDiffCon

The authors propose SafeDiffCon (Safe Diffusion Models for PDE Control), a framework that adapts diffusion models to varying safety constraints by quantifying model uncertainty using conformal prediction. The method operates in three stages:

1. Uncertainty Quantification via Conformal Prediction

To address the distribution shift between training/calibration data and the generated control sequences, the authors introduce an uncertainty quantile.

Calibration: A portion of the data is set aside as a calibration set ( $D_{cal}$ ).
Shifted Score Set: The authors define a score set based on the prediction errors of safety scores. To account for distribution shifts (covariate shift), they apply weighting to these scores, creating a "shifted score set" ( $\tilde{S}$ ).
Uncertainty Quantile ( $Q$ ): The uncertainty quantile is calculated as a specific quantile of the shifted score set. This value defines a conformal interval around the predicted safety score. With a guaranteed probability ( $1-\alpha$ ), the true safety score of the system interacting with the environment lies within this interval.

2. Post-Training with Reweighted Loss

Following initial pre-training, the model undergoes a post-training phase to steer its output distribution toward safer regions.

Reweighted Loss: The authors propose a reweighted diffusion loss ( $L_{post-train}$ ). The weighting factor is derived from the uncertainty quantile and the safety constraint. Specifically, the weight is $e^{-W(u,w)}$ , where $W(u,w)$ includes a term that penalizes the model if the predicted safety score plus the uncertainty quantile exceeds the safety bound ( $s_0$ ).
Objective: This loss function effectively reweights the training samples, assigning higher weights to those that satisfy safety constraints when accounting for uncertainty, thereby shifting the model's distribution $p(u,w)$ to a safer distribution $p^*(u,w)$ .

3. Inference-Time Fine-Tuning

To further ensure safety for specific control targets, the model undergoes iterative fine-tuning during inference.

Guidance: The generation process is guided by a function $G$ that incorporates the control objective and the safety constraints (including the uncertainty quantile).
Iterative Adjustment: The model parameters are fine-tuned using a loss function ( $L_{fine-tune}$ ) based on the guidance and the uncertainty quantile. This process iteratively adjusts the model to improve safety and performance within a few iterations without requiring full retraining.

Key Contributions

Integration of Safety into PDE Control: The paper introduces safety constraints into deep learning-based PDE control, proposing SafeDiffCon as a novel solution.
Conformal Adaptation: The authors design a conformal adaptation process that quantifies safety uncertainty under distribution shifts. This involves:
- Defining an uncertainty quantile to create valid conformal intervals for safety scores.
- Implementing post-training with a reweighted loss to align the model's distribution with safe regions.
- Introducing inference-time fine-tuning to dynamically adjust the model for specific tasks.
Comprehensive Evaluation: The method is evaluated on three distinct control tasks:
- 1D Burgers' Equation: A fundamental fluid dynamics equation.
- 2D Incompressible Fluid: Controlling smoke flow in a complex environment with obstacles and hazardous regions.
- Controlled Nuclear Fusion (Tokamak): Managing plasma parameters to prevent disruptions.

Experimental Results

The authors compare SafeDiffCon against several baselines, including Behavior Cloning (BC), Safe BC (BC-Safe), Constrained Decision Transformer (CDT), diffusion-based TREBI, and classical methods like PID and Model Predictive Control (MPC) combined with Lagrangian approaches.

Safety Compliance: SafeDiffCon is the only method that successfully satisfies all safety constraints across all three tasks. Other methods, including CDT and TREBI, fail to meet safety requirements in one or more scenarios, often resulting in unsafe trajectories.
Control Performance: While adhering to strict safety constraints, SafeDiffCon achieves the best control performance (lowest control objective $J$ ) among all methods that satisfy safety. For instance, in the 2D fluid task, it achieves a safety violation magnitude of 0.0000 while maintaining competitive control objectives.
Ablation Studies: Experiments removing post-training, fine-tuning, or the uncertainty quantile ( $Q$ ) result in the failure to meet safety constraints, demonstrating that all components are essential for the method's success.

Significance and Claims

The paper claims that SafeDiffCon addresses a critical bottleneck in applying deep learning to high-stakes scientific and engineering problems: the lack of formal safety guarantees. By leveraging conformal prediction to quantify uncertainty and integrating it into both the training and inference phases, the method ensures that control sequences remain within safe boundaries even when faced with distribution shifts and model uncertainty.

The authors position this work as a step toward the safer deployment of machine learning in real-world applications where safety is paramount, such as nuclear fusion and fluid dynamics. They note that while the current work focuses on diffusion models, the approach could potentially be extended to other generative methods requiring constraints. The code is made publicly available to facilitate further research in safe PDE control.