Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Ultra Fast PDE Solving via Physics Guided Few-step Diffusion

Cindy Xiangrui Kong, Yueqi Wang, Haoyang Zheng, Weijian Luo, Guang Lin

TL;DR

This work targets two bottlenecks in diffusion-based PDE solving: sampling efficiency and physical consistency. It introduces Phys-Instruct, a physics-guided distillation framework that compresses a multi-step diffusion teacher into a few-step generator while injecting PDE knowledge during training via a differentiable PDE residual, yielding fast, physics-faithful samples without test-time PDE guidance. The approach uses an IKL-based distillation objective augmented with a PDE residual loss, enabling a single trained model to support flexible step budgets and to serve as a reusable prior for downstream conditional PDE tasks. Empirically, Phys-Instruct delivers orders-of-magnitude faster inference and substantial PDE-error improvements across five PDE benchmarks, with strong performance on downstream forward/inverse problems while maintaining physical fidelity. This framework thus provides a practical, scalable path to ultra-fast, physics-aware PDE solving with deep generative models.

Abstract

Diffusion-based models have demonstrated impressive accuracy and generalization in solving partial differential equations (PDEs). However, they still face significant limitations, such as high sampling costs and insufficient physical consistency, stemming from their many-step iterative sampling mechanism and lack of explicit physics constraints. To address these issues, we propose Phys-Instruct, a novel physics-guided distillation framework which not only (1) compresses a pre-trained diffusion PDE solver into a few-step generator via matching generator and prior diffusion distributions to enable rapid sampling, but also (2) enhances the physics consistency by explicitly injecting PDE knowledge through a PDE distillation guidance. Physic-Instruct is built upon a solid theoretical foundation, leading to a practical physics-constrained training objective that admits tractable gradients. Across five PDE benchmarks, Phys-Instruct achieves orders-of-magnitude faster inference while reducing PDE error by more than 8 times compared to state-of-the-art diffusion baselines. Moreover, the resulting unconditional student model functions as a compact prior, enabling efficient and physically consistent inference for various downstream conditional tasks. Our results indicate that Phys-Instruct is a novel, effective, and efficient framework for ultra-fast PDE solving powered by deep generative models.

Ultra Fast PDE Solving via Physics Guided Few-step Diffusion

TL;DR

This work targets two bottlenecks in diffusion-based PDE solving: sampling efficiency and physical consistency. It introduces Phys-Instruct, a physics-guided distillation framework that compresses a multi-step diffusion teacher into a few-step generator while injecting PDE knowledge during training via a differentiable PDE residual, yielding fast, physics-faithful samples without test-time PDE guidance. The approach uses an IKL-based distillation objective augmented with a PDE residual loss, enabling a single trained model to support flexible step budgets and to serve as a reusable prior for downstream conditional PDE tasks. Empirically, Phys-Instruct delivers orders-of-magnitude faster inference and substantial PDE-error improvements across five PDE benchmarks, with strong performance on downstream forward/inverse problems while maintaining physical fidelity. This framework thus provides a practical, scalable path to ultra-fast, physics-aware PDE solving with deep generative models.

Abstract

Diffusion-based models have demonstrated impressive accuracy and generalization in solving partial differential equations (PDEs). However, they still face significant limitations, such as high sampling costs and insufficient physical consistency, stemming from their many-step iterative sampling mechanism and lack of explicit physics constraints. To address these issues, we propose Phys-Instruct, a novel physics-guided distillation framework which not only (1) compresses a pre-trained diffusion PDE solver into a few-step generator via matching generator and prior diffusion distributions to enable rapid sampling, but also (2) enhances the physics consistency by explicitly injecting PDE knowledge through a PDE distillation guidance. Physic-Instruct is built upon a solid theoretical foundation, leading to a practical physics-constrained training objective that admits tractable gradients. Across five PDE benchmarks, Phys-Instruct achieves orders-of-magnitude faster inference while reducing PDE error by more than 8 times compared to state-of-the-art diffusion baselines. Moreover, the resulting unconditional student model functions as a compact prior, enabling efficient and physically consistent inference for various downstream conditional tasks. Our results indicate that Phys-Instruct is a novel, effective, and efficient framework for ultra-fast PDE solving powered by deep generative models.
Paper Structure (48 sections, 1 theorem, 50 equations, 8 figures, 9 tables, 1 algorithm)

This paper contains 48 sections, 1 theorem, 50 equations, 8 figures, 9 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Diffusion priors for PDE fields.
  4. Few-step diffusion and distillation.
  5. Background
  6. Problem Setup of PDE Solvers
  7. Diffusion Models
  8. Integral Kullback-Leibler Divergence
  9. Phys-Instruct
  10. Phys-Instruct Generator
  11. Unified $k$-Step Formulation.
  12. PDE Error and Physics Loss.
  13. Phys-Instruct Objective
  14. Training Procedure
  15. Experiments
  16. ...and 33 more sections

Key Result

Theorem 4.1

Consider the physics-constrained IKL matching problem eq:constrained_ikl. As a practical surrogate, we apply a Lagrangian relaxation of eq:constrained_ikl and define the generator objective, treating $\lambda_{\mathrm{phys}}>0$ as a trade-off hyperparameter: where $\mathcal{L}_{\mathrm{phys}}(\theta)$ is given by eq:physics_loss. Under mild regularity conditions stated in Appendix apx:proof, the

Figures (8)

  • Figure 1: Burgers' unconditional generation. Top: PDE Error (lower is better, indicated by bubble size) vs. per-sample sampling latency. Baselines use 4/10/100/200 steps; Phys-Instruct uses 1--4 steps. Bottom: Unconditional Burgers' samples (left to right): data from $p_{\text{data}}$, four-step Phys-Instruct, and four-step EDM teacher.
  • Figure 2: Overview of Phys-Instruct. We sample latent variable $\mathbf{z}\sim p_z$, step budget $k\sim \Delta_K$, and generate PDE state $\mathbf{x}_0=g_\theta(\mathbf{z};k)$ with the generator $g_\theta$. Then $\mathbf{x}_0$ is diffused to noisy state $\mathbf{x}_t$, which is evaluated by a frozen teacher diffusion model $s_\nu$ and an auxiliary diffusion model $s_\phi$. The auxiliary network is trained online via denoising score matching ($\nabla_\phi \mathcal{L}_{\mathrm{DSM}}$). The generator $g_\theta$ is updated by the distillation gradient $\nabla_\theta \mathcal{D}_{\mathrm{IKL}}(q_\theta\|p)$ plus a physics gradient $\nabla_\theta \mathcal{L}_{\mathrm{phys}}$ weighted by a coefficient $\lambda_{\mathrm{phys}}$.
  • Figure 3: Distributional discrepancy on the Burgers' equation benchmark measured by sliced Wasserstein distance (SWD, top) and Maximum Mean Discrepancy (MMD, bottom; reported as $\sqrt{\mathrm{MMD}^2}$), where lower is better.
  • Figure 4: Darcy-flow PDE error with one dtep Phys-Instruct during distillation under different physics-guidance schedules; stars indicate when guidance starts.
  • Figure 5: Sensitivity of PDE error to the physics guidance weight $\lambda_{\mathrm{phys}}$ on Darcy flow. (a) We plot PDE error (log scale) versus training progress for one-step Phys-Instruct with different $\lambda_{\mathrm{phys}}$ (all other settings match Default). Moderate $\lambda_{\mathrm{phys}}$ yields consistently lower PDE errors, while too small $\lambda_{\mathrm{phys}}$ leads to early instability and a higher error plateau. Diverged runs with overly large $\lambda_{\mathrm{phys}}$ are omitted. (b) Summary of the final PDE error across $\lambda_{\mathrm{phys}}$ with finer grids between $1\mathrm{e}{-3}$ and $1\mathrm{e}{-2}$. The red dot denotes the PDE error of default $\lambda_{\mathrm{phys}}$.
  • ...and 3 more figures

Theorems & Definitions (4)

  • Theorem 4.1: Score-Function Identity
  • proof
  • Remark 4.2: Amortized physics correction
  • proof : Proof of Theorem \ref{['thm:phys_impl_one']}