[논문리뷰] Physics Informed Diffusion Models (PIDM)

https://arxiv.org/abs/2403.14404

Physics-Informed Diffusion Models

아이디어: PINN objective를 Diffusion에 통합해서 생성된 데이터가 주어진 물리 (PDE) 제약조건을 따르도록 하자.

---

Introduction

Diffusion 모델은 이미지/텍스트/비디오/그래프 생성 등의 영역에서 데이터의 분포를 매우 잘 모델링하는 것으로 알려져 있다.

여기서 더 나아가서 Scientific machine learning의 관점에서 Diffusion을 이용하려는 시도들이 있어왔다 (새로운 분자/재료구조 생성 등).

이러한 문제들은 흔히 Inverse problem으로 모델링되는데, 이런 문제들은 유일한 해를 가지지 않을 가능성이 높다.

따라서 데이터 혹은 가능한 전체 solution의 분포를 모델링할 수 있는 Diffusion이 이런 문제들을 해결하기에 용이하다.

 

그러나 수치해석 등으로 생성한 물리기반 데이터를 학습하고자 할 때, 기존 Diffusion 모델은 데이터에만 의존하여 분포를 모델링하므로, Underlying governing equation (주로 PDE) 등에 대한 정보를 explicit하게 통합하여 생성한 데이터가 물리법칙을 따르게 만들도록 할 필요성이 있다.

 

따라서 저자들은 Neural network의 관점에서 PDE를 풀고자 하는 시도 중 가장 잘 알려져 있고 잘 정립된 PINN을 Diffusion model에 통합하여 물리 정보를 효과적으로 모델 학습 과정에 이용하도록 하는 Physics-Informed Diffusion Models (PIDM)을 제안한다.

 

Contributions

1. PINN의 objective를 Diffusion에 통합하여 PDE의 Residual에 대한 오차를 크게 줄일 수 있었다.

2. 추가된 PINN objective가 기존 Diffusion 과정을 방해하지 않으며 overfitting을 방지하기 위한 효과적인 regularization 항으로 작용함을 보였다.

3. 기존 Diffusion 구조에 통합하기 매우 용이한 구조를 가지고 있으며 inference 과정을 변화시키지 않는다.

4. 논문에서 제안하는 PDE 제약조건은 equality constraint의 특수한 경우로, 제약조건을 확장하여 미분가능한 어떤 inequality/equality constraint에 대해서도 Diffusion objective에 효과적으로 통합될 수 있음을 보였다.

Schematic Overview of PIDM

---

Background

Denoising Diffusion Models

본 연구에서는 DDPM loss를 쓰되 각 time step t에 대해 x0를 근사하는 x̂0을 찾는 것으로 모델링했다.

(time step별로 noise를 예측하는 것보다 성능이 좋다고 함 최신?테크닉)

Assembly of Governing Equations

물리 법칙은 아래와 같이 domain Ω ⊂ R^d에서 정의되는 PDE의 집합과 boundary condition으로 표현할 수 있다.

 

이 때 F는 differential operator, ξ는 Ω에서의 spatial coordinate, u는 우리가 모델링해야 할 solution function, B는 boundary condition operator이다.

F와 B가 모두 Residual의 형태로 정의되어 있는 것에 유의하자

 

이제 Diffusion model에서 샘플링한 x0 ~ pθ(x0)가 위의 제약조건을 만족하도록 하기 위해 Residual 벡터 R을 정의한다.

 

즉 우리의 목표는 R(x0)이 0에 최대한 가깝도록 만드는 것이다. 이를 계산하기 위해서는 PDE의 미분항을 직접 계산해야 하는데, 논문에서는 따로 numerical solver를 사용했다. (PINN에서는 automatic differentiation으로 구현함)

---

Physics-Informed Diffusion Models

Consideration of PDE Constraints

우리의 Diffusion 모델은 기존 Diffusion objective를 만족하면서도 물리 제약조건도 동시에 만족하여야 한다.

그러나 생성모델이 확률적인 과정을 기반으로 모델링되는 데에 반해 우리가 정의한 residual은 deterministic한 값이다.

따라서 확률 기반 objective에 좀더 효과적으로 PDE constraint를 통합하기 위해 virtual observation r̂ = 0을 정의한다.

 

즉 x0에 대한 residual R(x0)를 평균으로 가지는 정규분포를 하나 정의하고, 여기서 r̂을 샘플링할 수 있다고 가정하는 것이다.

(쉽게 말해서 Residual을 확률변수처럼 다룬 것이 r̂ 임)

분산이 늘어날수록 분포가 R(x0)에 가까워져서 좀더 strict한 제약조건이 된다는 것을 염두에 두자

 

이제 우리가 근사하는 모델 pθ를 residual의 관점에서 maximum likelihood를 취해보자.

 

위에 언급했듯 우리는 r̂ = 0이 이미 관측되었다고 정의한다.

즉 pθ에서 샘플링한 x0가 Residual을 0으로 만들 확률을 최대화하는 것이다.

이는 PINN에서 PDE loss를 residual을 통해 계산하는 것과 본질적으로 같은 formulation이다.

 

Consideration of Observed Data

이제 Diffusion의 objective와 위에서 얻은 식을 합쳐보자.

 

이게 PIDM에서 제안하는 Objective이다.

 

이를 Score-based model의 관점에서 바라보면, PDE 제약조건의 추가는 기존 Diffusion의 objective를 방해하지 않고, Optimal한 score model은 두 항을 동시에 최적화한다고 한다.

 

첫번째 항은 Diffusion의 objective 그대로이다. 이건 그대로 ELBO를 통해 최적화할 수 있다.

그런데 두번째 항을 보면 우리가 정의했듯 pθ에서 샘플링한 x0가 필요하다.

즉 residual 계산을 위해 항상 t=T에서 0까지 샘플링을 해서 x0을 구해야 한다는 뜻인데 이는 매우 계산 비효율적이다.

따라서 x0를 바로 샘플링하지 않고 아래의 두가지 테크닉을 통해 근사한다.

 

Simplification of the Training Objective

앞서 설명했듯 x0를 pθ로부터 바로 구하는 것은 비효율적이기 때문에 구하기 쉬운 x*0을 계산해서 x0 대신 쓴다.

 

Mean Estimation

위에서 정의한 Diffusion model로 다시 돌아가보자.

즉 모델은 어떤 noisy input x_t가 주어졌을 때, 이를 바탕으로 x0을 근사하는 x̂0을 복원한다.

이는 x0과 x̂0 사이의 mean-squared distance를 구해서 최소화하는 것으로 모델링될 수 있다.

따라서 x̂0은 E[x0|x_t ]의 estimation이라고 할 수 있다.

 

우리가 구하고자 하는 것은 E[R(x0|x_t)]인데, 이 대신 Diffusion model에서 학습 도중 자연스럽게 계산되는 x̂0을 써서 R(E[x0|x_t])을 구하자는 것이다.

그러나 이 둘은 동등하지 않고 t가 0에 가까워질 때가 아니라면 diffusion objective와 conflict를 일으킬 수 있으므로, 추후 소개할 Variance scheduling을 통해 보정한다.

 

Sample estimation

샘플링을 pθ에서 하긴 하되 전체 trajectory를 따르지 말고 DDIM과 같은 가속화된 샘플링을 사용하자는 아이디어다.

 

 

이제 두가지 estimation 방법으로부터 근사한 x0를 x*0라 하고 우리의 objective에 다시 집어넣으면 아래와 같아진다.

 

 

x*0는 더이상 pθ에 의존하지 않고 대신 reverse process의 전체 trajectory pθ(x1:T)로부터 계산된다.

이제 q_R을 꺼내서 써보면

 

위 식과 같아진다.

추가적으로 앞서 설명했던 문제점을 해결하기 위해 Variance scheduling을 이용한다.

 

즉 PDE objective의 variance를 Diffusion의 reverse process에서 정의한 fixed variance Σt로 대체한다는 것이다.

(c는 상수, 하이퍼파라미터)

 

이렇게 하면 t가 0에 가까울수록 분산이 줄어들어서 PDE 제약조건을 더 strict하게 만족시키도록 하고,

t가 T에 가까울수록 PDE 제약조건을 덜 고려하도록 (데이터 분포 모델링에 집중하도록) 만들 수 있다.

 

마지막으로 학습 과정에서는 고정된 q(x1:T)가 이미 주어지므로, reverse process의 trajectory에 의존하던 x*0의 계산을 q(x1:T)에 의존하도록 바꿀 수 있다. 최종 objective를 loss 형태로 나타내면 다음과 같다.

 

이 때 Σ̄t = Σt /c 이다.

 

학습 과정 정리

 

여기서 중요한 점은 우리가 정의한 PDE constraint가 다른 어떠한 미분 가능한 constraint로도 확장될 수 있다는 것이다.

즉 constraint가 f(x)=c 꼴을 만족하기만 하면 Residual 꼴로 정의할 수 있고 이를 그대로 objective에 집어넣어도 formulation을 해치지 않는다는 것이다.

여기서 PDE는 equality constraint로 정의되었지만 저자들은 여기서 확장해서 inequality constraint와 auxiliary optimization objective도 정의할 수 있음을 확인하였다.

 

이를 이용해서 surrogate model Rθ(x0)을 정의하고 학습 가능한 제약 조건 아래에서 조건을 만족하도록 하는 등 다양한 시도를 해볼 수 있다.

---

Experiments

2D Darcy Flow

아래와 같이 permeability K가 주어졌을 때, velocity u와 pressure distribution p를 구하는 문제다.

 

논문에서는 pressure distribution을 구하는 것을 목표로 하였다. (Forward problem이라고 볼 수 있을듯)

 

 

Baseline 방법론들과 비교해 Mean estimation / Sample estimation을 적용한 PIDM의 residual error와 test data loss를 비교했을 때, Residual이 특히 크게 줄어드는 것을 확인할 수 있다.

또한 multi-objective loss function을 사용했을 때 일반적으로는 data loss와 residual loss 사이의 tradeoff가 발생하지만, 재미있게도 test data loss가 vanilla diffusion model과 비슷하게 / 혹은 더 안정적으로 감소하며 overfitting에 강한 것을 확인할 수 있다. 이는 특히 Mean estimation을 사용할 경우에 두드러진다.

 

 

생성 예시이다. 마찬가지로 Vanilla diffusion과 비교할 때 residual이 훨씬 작게 나타나는 것을 확인할 수 있다.

 

Topology optimization

주어진 조건 (equilibrium) 아래 최적의 구조 특성 (주로 optimal stiffness)을 찾기 위한 문제이다. Inverse problem에 해당한다.

예를 들어 Linear elasticity를 가지는 물질에 대해 (Stress-strain 사이 관계가 Hooke's Law를 따름) 최적의 material distribution을 찾는 문제는 아래와 같이 formulation할 수 있다.

 

이런 최적화 문제는 보통 iterative하게 해를 구해야 하기 때문에 계산비용이 매우 크다. (논문에서는 SIMP와 비교)

 

 

마찬가지로 Baseline model들과 비교하여 sample estimation을 적용한 PIDM의 결과를 seen/unseen boundary condition에 대해 각각 살펴보면 Residual의 MAE가 가장 낮게 나타나는 것을 확인할 수 있다. 

 

 

마찬가지로 Vanilla Diffusion (a)과 PIDM (b)를 통해 생성한 샘플 예시. PIDM이 특별히 좋아보이지는 않는다...

---

Discussion

논문을 읽으면서 생각할 수 있는 현재 프레임워크의 매우 치명적인 단점은 PDE의 시간 항을 못 푼다는 것이다.

대부분의 PDE는 time dependent하므로 아직까지는 적용분야가 많지 않을 것 같다...

또한 Residual을 numerical solver (FDM과 FEM 둘다 이용했음)로 계산하고 있는데 domain 크기가 커지면 (2d에서 3d로만 늘어나도) 매우 비효율적일 것이므로 Autograd를 이용하는 것이 적절하지 않은가 하는 생각이 든다.

하지만 다양한 condition을 통합할 수 있다는 점에서 좀더 개선된다면 좋은 물리모델로 기능할 수 있지 않을까 싶다

 

물리기반 네트워크 여러가지 공부하면서 드는 생각은 모델이

이럴 바에 솔버 여러번 돌리면 되는 거 아님? <- 하는 질문에 어떻게 대답할 수 있느냐가 매우 중요한 것 같다는것 이다