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

Motivation

• Known: Diffusion model은 물리 시스템의 forward/inverse/reconstruction 문제를 하나의 프레임워크로 다룰 수 있어 PINNs나 DeepONet, Fourier Neural Operator 같은 기존 operator-learning 방법보다 유연하지만, PDE와 같은 엄격한 low-level 제약을 만족시키기는 어렵다. PIDM, DiffusionPDE, ECI-sampling, CoCoGen 등 기존 방법들은 Tweedie's formula로 계산한 posterior mean E[x0|xt]에 PDE residual을 부과하는 방식을 사용한다.
• Gap: posterior mean에 제약을 부과하는 F[E[x0|xt]]는 실제로 원하는 E[F[x0]|xt]와 Jensen's inequality로 인해 다르며(Jensen's Gap), 이 불일치가 PDE 만족도와 생성 정확도 사이의 trade-off를 강제하여 물리적 충실도를 저하시킨다.
• Why: PDE 제약을 정확히 반영하면서도 생성 품질을 유지하는 diffusion 기반 물리 시뮬레이션은 과학 계산 분야에서 신뢰성과 실시간 적용 가능성을 동시에 요구하는 핵심 문제이며, 이 gap을 해소하면 forward/inverse/reconstruction을 아우르는 효율적이고 정확한 통합 solver를 얻을 수 있다.
• Approach: 기존처럼 diffusion 궤적 중간의 posterior mean에 제약을 거는 대신, teacher diffusion model로부터 얻은 noise-data pair를 이용해 student 모델을 distillation하고, student의 최종 출력 x0에 직접 PDE 제약을 부과하는 post-hoc distillation 전략(PIDDM)을 제안한다.

Achievement

Figure 2. Empirical illustration of the Jensen’s Gap in physics-constrained diffusion models. (a) Absolute velocity erro

1. Jensen's Gap의 실증적 규명: MoG(Mixture of Gaussians) 데이터셋과 Stokes 문제에서 DPS 샘플링과 실제 conditional ODE velocity 간의 오차, 그리고 학습 시 diffusion loss 비교를 통해 posterior-mean 기반 제약이 실제로 gap을 유발함을 정량적으로 보였다.
2. 최종 sample 단위 물리 제약 부과: PIDDM은 posterior-mean surrogate 없이 student의 최종 출력에 직접 PDE supervision을 적용하여 posterior-mean과 final-sample 간 mismatch를 student objective에서 제거했다.
3. 효율적이고 범용적인 추론: distillation된 student 모델은 teacher의 생성 능력(forward/inverse/reconstruction)을 유지하면서 one-step 생성을 가능케 하여, PIDM, DiffusionPDE, ECI-sampling 대비 정확도와 제약 만족도, 계산 비용, 하이퍼파라미터 튜닝 부담 모든 측면에서 우수함을 다양한 PDE 벤치마크에서 입증했다.

How

Figure 1. Illustration of physics-constrained diffusion generation and our proposed framework. (a) Existing methods (Hua

• 먼저 vanilla diffusion training/sampling으로 teacher model을 학습시키고, 이를 통해 noise-data pair 데이터셋을 생성한다(distillation을 위한 paired dataset 구축).
• teacher를 distillation하여 student 모델을 학습하되, 표준 distillation loss에 더해 최종 생성 sample x0에 대한 PDE residual loss ∥R(x0)∥²를 추가로 부과한다.
• 이를 통해 student는 diffusion 궤적 중간의 posterior-mean surrogate를 거치지 않고 one-step(또는 few-step) 생성 시 곧바로 물리 제약을 만족하는 clean sample을 출력하도록 학습된다.
• Darcy flow, Stokes, correlated MoG 등 다양한 PDE 벤치마크에서 forward/inverse/reconstruction 문제에 대해 PIDM, DiffusionPDE, ECI-sampling, CoCoGen과 비교 평가한다.

Originality

• posterior-mean 기반 PDE guidance에서 발생하는 Jensen's Gap을 이론적·실증적으로 명시적으로 규명한 최초의 분석
• 기존의 training-time 또는 inference-time guidance 방식과 달리, distillation을 통해 물리 제약 부과 시점을 최종 clean sample 단계로 옮기는 새로운 관점 제시
• 물리 제약을 만족시키면서도 single-step 생성이 가능하도록 distillation과 physics-informed loss를 결합한 점

Limitation & Further Study

• 논문에서 언급하듯 distributional fidelity(생성 분포의 정확성)에 대한 주장은 여전히 실증적(empirical)인 수준에 머물러 있으며 엄밀한 이론적 보장은 부족하다.
• distillation 과정 자체가 teacher model의 품질에 의존적이므로 teacher의 한계(예: 특정 PDE 유형이나 고차원 문제에서의 성능 저하)가 그대로 student에 전이될 가능성이 있다.
• 향후 연구로 더 복잡한 3차원 PDE, 시간에 따라 변화하는 시스템, 다중 물리 결합 문제 등으로의 확장, 그리고 Jensen's Gap에 대한 보다 엄밀한 이론적 상한 분석이 필요하다.

Evaluation

총평: Jensen's Gap이라는 기존에 간과되어온 문제를 명확히 규명하고, 이를 distillation 기반의 간단하면서도 효과적인 방법으로 해결하여 효율성과 물리적 정확성을 동시에 개선한 실용적이고 임팩트 있는 연구이다.