Physics-Informed Distillation of Diffusion Models for PDE-Constrained Generation
Abstract
Diffusion models show growing promise for generative modeling of physical systems, but enforcing partial differential equation (PDE) constraints directly is infeasible during the stochastic denoising process. Current methods apply constraints to the expected clean sample, incurring a Jensen’s Gap that forces a trade-off between PDE satisfaction and generative accuracy. To bridge this gap, we propose Physics-Informed Distillation of Diffusion Models (PIDDM), a simple yet effective post-hoc distillation strategy that enforces PDE constraints after training. PIDDM enables fast single-step generation while improving both physical consistency and sample quality, supporting forward/inverse problems and reconstruction from partial observations. Extensive experiments across PDE benchmarks show PIDDM outperforms recent baselines, such as PIDM, DiffusionPDE, and ECI-sampling, in both accuracy and constraint satisfaction, with lower computation and minimal hyperparameter tuning, offering a more efficient pathway to physics-informed diffusion models.
Lay Summary
Many problems in science and engineering, such as simulating fluids, heat, or materials, are governed by physical equations. Machine learning models can make these simulations faster, but they must still obey the laws of physics. Diffusion models are powerful tools for generating realistic data, yet it is hard to force them to satisfy physical equations while they turn noise into a final sample. Existing methods often apply the physics rule to an average estimate of the final sample, which can differ from the actual sample produced. We propose Physics-Informed Distillation of Diffusion Models (PIDDM), a method that learns from a standard diffusion model and trains a faster student model whose final outputs are directly checked against the equations. This moves physics enforcement from many intermediate denoising steps to the final generated result, where the equations are meaningful. The resulting model can generate physically consistent samples in one step and can solve forward, inverse, and reconstruction problems. Across several physical-equation benchmarks, PIDDM improves both physical accuracy and sample quality while using less computation than recent alternatives. This could make generative simulation more practical for repeated scientific and engineering workflows, such as design optimization, data completion, and uncertainty analysis.