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

作者：Difan Zou 等发布时间：2025-05-28（arXiv 预印本）arXiv 号：2505.22391PDF 与代码：https://arxiv.org/abs/2505.22391

一、研究背景与痛点

1. 扩散模型（Diffusion Models, DM）在物理系统生成中表现优异，可一次性求解正/逆问题、给出不确定性估计。
2. 现有“物理-约束”做法把 PDE 残差加在“后验均值” 𝔼[x₀|xₜ] 上，而非真正的干净样本 x₀，导致 Jensen Gap——物理约束与生成保真度之间的固有 trade-off。
3. 多步采样计算量大，难以实时应用。

二、核心贡献提出 PIDDM（Physics-Informed Distillation of Diffusion Models）

一种“后处理蒸馏”框架：

• 先按标准扩散流程训练教师模型（可多步）； - 再蒸馏出单步学生网络，并在蒸馏阶段直接把 PDE 残差施加于最终样本 x₀，彻底绕开 Jensen Gap；

• 学生模型一次前向即可生成满足物理方程的样本，兼顾“物理严格”与“速度提升”。

三、方法流程

1. 教师训练

完全沿用 vanilla diffusion：线性加噪、score-matching 损失，不加任何物理约束，保证生成质量。

2. 数据对采集

用教师模型 + 确定性采样器（Euler 或 Rectified Flow）生成 {ε → x₀} 配对数据；线性扩散路径保证 ε 与 x₀ 欧氏距离小，利于后续回归学习。

3. 学生蒸馏

学生网络 sθ′ 直接学习“噪声→数据”映射： ℒdistill = 𝔼‖sθ′(ε) − x₀‖² 同时在 x₀ = sθ′(ε) 上施加 PDE 残差： ℒpde = ‖ℛ(x₀)‖² 总损失：ℒ = ℒdistill + λ ℒpde

4. 单步生成

推理阶段仅需一次前向：x₀ = sθ′(ε)，ε∼𝒩(0,I)，即可得到既服从数据分布又近似满足 PDE 的样本。

四、理论/实证分析

• 给出 Jensen Gap 的定量测量：

在 Darcy 流与 Burgers 方程上，传统“后验均值”约束法的 PDE 残差比 PIDDM 高 1.7×–3.2×。

• 单步学生 vs 教师 1000 步：

PIDDM 的 PDE 残差反而更低，同时 FID 保持同一水平，实现“既快又准”。

五、实验结果

1. 2-D Darcy 流 – PDE 残差 MSE：

PIDDM 1.8×10⁻³，最佳基线 DiffusionPDE 5.9×10⁻³；
– 推理时间：PIDDM 0.02 s，基线 4.8 s（1000 步）。

2. 1-D Burgers 激波

– 激波宽度误差降低 35%； – 逆问题（源项识别）相对误差 4.1%，比 PIDM 降低 42%。

3. 2-D Navier–Stokes 涡街

– 单步生成涡量场与参考解相关系数 0.97； – 在 256×256 网格上，GPU 内存占用仅为教师多步采样的 1/15。

4. 部分观测重构

仅给定 10% 网格点的压力值，PIDDM 重构整场压力，PSNR 32.4 dB，优于 ECI-sampling 3.2 dB。

六、消融与鲁棒性

• λ 增大 → PDE 误差单调降，但 FID 先降后升，λ=10 为最佳平衡点；

• 采用 Rectified Flow 做教师时，蒸馏成功率 98%，DDPM 教师仅 74%，说明“直线化”轨迹对蒸馏至关重要；

• 对教师模型进行 50% 参数量剪枝后，PIDDM 仍能保持 PDE 残差低于原始教师，显示框架对教师质量变化具有鲁棒性。

七、代码与复现作者已开源 PyTorch 实现，包含 9 个 PDE 案例（Poisson、Heat、Stefan、Helmholtz 等）。

训练脚本： bashpython train_teacher.py --pde darcy # 1. 训练教师python distill.py --pde darcy --lambda 10 # 2. 蒸馏python sample.py --steps 1 --n 100 # 3. 单步生成单卡 RTX 4090 上，2 h 完成教师训练，0.5 h 完成蒸馏。

八、局限与未来

1. 依赖“预训练好 + 采样路径直线化”的教师，若教师严重偏离数据分布，学生误差会被放大；

2. 目前仅针对时间无关或抛物型 PDE，对双曲型守恒律（需熵条件）的扩展仍待研究；

3. 单步生成虽快，但难以像多步 DM 那样进行“渐进修正”，在需要超高精度场景可能受限。

九、一句话总结

PIDDM 通过“后处理蒸馏”把物理约束直接施加在最终样本，绕过 Jensen Gap，首次实现扩散模型单步、高保真、PDE 一致的物理生成，为实时物理模拟与逆问题求解提供了新范式。