Physics-Informed Distillation of Diffusion Models for PDE-Constrained Generation

TL;DR

PIDDM tackles the challenge of enforcing PDE constraints in diffusion models by decoupling physics from the diffusion trajectory. It demonstrates Jensen’s Gap when constraints are applied to the posterior mean and employs a post-hoc distillation to enforce constraints directly on the final sample ${\mathbf{x}}_0$, enabling one-step generation and shared forward/inverse PDE solving. Empirically, PIDDM outperforms baseline constraint methods in PDE satisfaction while maintaining competitive generative fidelity, and its distilled student enables fast, physics-consistent generation with optional refinement and robust downstream tasks. This approach offers a practical, efficient pathway for integrating strict physical laws into diffusion-based generative modeling for scientific computing.

Abstract

Modeling physical systems in a generative manner offers several advantages, including the ability to handle partial observations, generate diverse solutions, and address both forward and inverse problems. Recently, diffusion models have gained increasing attention in the modeling of physical systems, particularly those governed by partial differential equations (PDEs). However, diffusion models only access noisy data $\boldsymbol{x}_t$ at intermediate steps, making it infeasible to directly enforce constraints on the clean sample $\boldsymbol{x}_0$ at each noisy level. As a workaround, constraints are typically applied to the expectation of clean samples $\mathbb{E}[\boldsymbol{x}_0|\boldsymbol{x}_t]$, which is estimated using the learned score network. However, imposing PDE constraints on the expectation does not strictly represent the one on the true clean data, known as Jensen's Gap. This gap creates a trade-off: enforcing PDE constraints may come at the cost of reduced accuracy in generative modeling. To address this, we propose a simple yet effective post-hoc distillation approach, where PDE constraints are not injected directly into the diffusion process, but instead enforced during a post-hoc distillation stage. We term our method as Physics-Informed Distillation of Diffusion Models (PIDDM). This distillation not only facilitates single-step generation with improved PDE satisfaction, but also support both forward and inverse problem solving and reconstruction from randomly partial observation. Extensive experiments across various PDE benchmarks demonstrate that PIDDM significantly improves PDE satisfaction over several recent and competitive baselines, such as PIDM, DiffusionPDE, and ECI-sampling, with less computation overhead. Our approach can shed light on more efficient and effective strategies for incorporating physical constraints into diffusion models.

Figures (5)

• Figure 1: Illustration of physics-constrained diffusion generation and our proposed framework.(a) Existing methods huang2024diffusionpdecheng2025ecibastek2025pdimjacobsen2024cocogen impose PDE losses or guidance on the posterior mean $\mathbb{E}[x_0|x_t]$ in diffusion training and sampling, introducing Jensen’s Gap. (b) We propose to train and sample diffusion model using vanilla methods to generate a noise-image data paired dataset for distillation. (c) Our proposed framework distills the teacher diffusion model and directly enforces physical constraints on the final generated sample $x_0$, avoiding Jensen’s Gap .
• Figure 2: Empirical illustration of the Jensen’s Gap in physics-constrained diffusion models.(a) Absolute velocity error and angular discrepancy ($1 - \cos(\theta)$) between Diffusion Posterior Sampling (DPS) and the ground-truth conditional ODE velocity on the MoG dataset. (b) and (c) Histograms comparing the first (unconstrained) and second (hard-constrained) dimensions of DPS-sampled MoG data against the ground truth MoG. (d) Training-time manifestation: diffusion loss comparison between vanilla training and PIDM on a Stokes Problem dataset.
• Figure 3: Ablation studies on the effect of several factors on the performance of PIDDM on Darcy dataset. (a), (b) and (c) refer to the effect of the $N_s$, $\lambda_{\text{train}}$ and $\lambda_{\text{infer}}$, respectively.
• Figure 4: Qualitative comparison on the Darcy forward problem. Each column shows (left) the predicted solution field, (middle) point-wise data error, and (right) PDE residual error. Our PIDDM (bottom row) delivers visibly lower data and PDE errors than other baselines while maintaining sharp solution details.
• Figure 5: Constraint satisfaction on correlated MoG. Comparison of generated samples using DPS, ECI, and PIDDM. PIDDM closely matches the target distribution while satisfying constraints.