Riemannian Diffusion Models on General Manifolds via Physics-Informed Neural Networks

기술 요약: 물리 정보 신경망(PINN)을 통한 일반 다양체 상의 리만 확산 모델

1. 문제 정의

리만 확산 모델은 리만 브라운 운동(Riemannian Brownian motion)을 통해 전방 확산 과정을 정의함으로써, 매끄러운 다양체(smooth manifolds)에 존재하는 데이터로 스코어 기반 생성 모델링을 일반화합니다. 이러한 모델을 학습하고 샘플링하는 데 있어 핵심적인 병목 현상은 다양체의 열 커널(heat kernel, 브라운 운동의 전이 밀도)과 그 기울기(gradient)를 구해야 한다는 점입니다. 구체적으로, 학습 목적 함수인 디노이징 스코어 매칭(Denoising Score Matching)은 조건부 스코어 $\nabla_x \log p_{t|0}(x|x_0)$를 평가해야 하며, 전방 샘플링을 위해서는 열 커널으로부터 표본을 추출해야 합니다.

열 커널은 고도로 대칭적인 다양체(예: 유클리드 공간, 구, 토러스)에 대해서는 폐쇄형(closed form)으로 알려져 있지만, 일반적인 리만 다양체에 대해서는 이를 명시적으로 구할 수 없는 경우가 대부분입니다. 바라단 공식(Varadhan's formula)과 같은 단시간 점근 전개나 스펙트럼 확장(spectral expansion) 같은 기존의 근사법들은 특정 영역(작은 시간 $t$)에 국한되거나 강한 구조적 가정(예: 컴팩트 릭 군)을 필요로 하므로, 다양한 기하학적 설정에 대한 실질적인 적용 가능성을 제한합니다.

2. 방법론

저자들은 **물리 정보 신경망(Physics-Informed Neural Networks, PINNs)**을 사용하여 다양체의 열 방정식을 직접 해결함으로써 열 커널을 근사하는 일반적인 프레임워크를 제안합니다. 이 방법은 다음의 네 가지 주요 단계로 진행됩니다.

A. 좌표계 선택 및 임베딩
국소 좌표 차트(local coordinate charts)를 패치하는 대신, 이 방법은 대상 다양체 $M$을 더 큰 주변 다양체 $\tilde{M}$(일반적으로 $\mathbb{R}^N$) 내에 부분 다양체(submanifold) 또는 몫 다양체(quotient manifold)로 임베딩합니다.

• 부분 다양체: 제약 조건 $f_i(x)=0$에 의해 정의됩니다. 메트릭(metric)과 미분 연산자는 주변 공간을 $M$의 접공간(tangent space)으로 매핑하는 투영 행렬 $P(x)$를 사용하여 유도됩니다.
• 몫 다양체: 군 작용(group action)에 의해 점들이 동일하게 식별되는 경우입니다(예: 치환 대칭성). 이 방법은 궤도 점(orbit points)을 동일한 상태로 취급하며, 종종 정준 대표원(canonical representatives)을 사용합니다.

B. 열 방정식 유도
주변 좌표와 투영 행렬을 사용하여 저자들은 라플라스-벨트라미 연산자 $\Delta_M$과 그에 대응하는 포커-플랑크(Fokker-Planck, 열) 방정식의 명시적 형태를 유도합니다:
$$\partial_t p_{t|0}(x|x_0) = \Delta_M p_{t|0}(x|x_0)$$
이 방정식은 주변 메트릭과 투영 행렬만을 사용하여 표현되므로, 내재적 좌표 차트를 요구하지 않고도 좌표 기반 계산이 가능합니다.

C. 초기 조건 근사
디락 델타 초기 조건 $p_{0|0} = \delta_{x_0}$는 수치적으로 부과할 수 없으므로, 저자들은 원칙적인 단시간 점근 근사를 사용하여 작은 시간 $t_0 > 0$에서 PINN을 초기화합니다. 이들은 파라메트릭스 방법(Minakshisundaram–Pleijel 재귀)을 활용하여 다음과 같은 고차 확장을 얻습니다:
$$p_{t_0}(x|x_0) \approx G_t^{(n)}(x, x_0) = G_t(x, x_0) \left( u_0 + u_1 t + \dots + u_{n-1} t^{n-1} \right)$$
여기서 $G_t$는 가우시안 유사 항이며, $u_k$는 기하학적 특성(예: 부피 왜곡 및 곡률)에서 유도된 계수입니다.

D. PINN 학습
신경망 $\phi_\theta(t, x)$는 로그 열 커널 $\log p_{t|0}$를 근사하도록 학습됩니다. 학습 목적 함수는 다음의 조합을 최소화합니다:

1. 초기/경계 손실(Initial/Boundary Loss): 네트워크 출력과 $t=t_0$에서의 점근적 초기화 값 사이의 차이.
2. PDE 잔차 손실(PDE Residual Loss): 로그 밀도에 대해 재작성된 열 방정식 $\partial_t \phi - \Delta_M \phi - \|\nabla \phi\|^2 = 0$의 잔차.

저자들은 안정성을 보장하기 위해 적응형 손실 가중치(adaptive loss weighting)와 커리큘럼 학습(시간 간격을 점진적으로 확장)을 채택합니다.

E. 확산 모델로의 통합
학습된 PINN은 열 커널의 대리물(surrogate) 역할을 하여 두 가지 핵심 연산을 가능하게 합니다:

1. 전방 샘플링: PINN이 평가한 밀도를 사용하여 마르코프 체인 몬테카를로(MCMC) 방법(예: 측지선 랜덤 워크)을 통해 표본을 생성합니다.
2. 스코어 평가: PINN을 미분하여 조건부 스코어 $\nabla_x \log p_{t|0}$를 계산하며, 이는 디노이징 스코어 매칭(DSM)을 통해 역방향 시간 확산 모델을 학습하는 데 사용됩니다.

3. 주요 기여

• 일반 PDE 기반 대리물: 저자들은 열 방정식을 수치적으로 해결함으로써 임의의 다양체 상에서 열 커널을 근사하는 방법을 도입하였으며, 이를 통해 폐쇄형 해나 제한적인 스펙트럼 가정이 필요하지 않음을 보여주었습니다.
• 좌표 불변 유도: 이 프레임워크는 유클리드 공간의 부분 다양체 또는 몫 다양체로 정의된 다양체에 대한 열 방정식을 체계적으로 유도하는 방법을 제공하며, 이를 통해 $S^2$, $SO(3)$, $SPD(n)$, 그리고 치환 몫(permutation-quotient)을 가진 점 구름(point clouds)과 같은 복잡한 기하학에 적용할 수 있게 합니다.
• 안정성 분석: 저자들은 근사된 초기 조건으로부터 열 방정식을 푸는 것이 모든 $t \ge t_0$에 대해 유계된 오차를 갖는 해를 생성한다는 이론적 정당성(정리 4.1)을 제공하여, PINN 대리물의 신뢰성을 보장합니다.
• 치환 불변 확산(Permutation-Invariant Diffusion): 이 방법은 몫 다양체 $\mathbb{R}^{k \times n}/S_n$에 적용되어, 분자 생성 시 원자 치환을 명시적으로 제거하는 확산 과정을 가능하게 합니다. 이는 표준 유클리드 확산이 구조적 대칭성을 존중하지 못하는 설정에서 매우 중요합니다.

4. 실험 결과

이 방법은 다양한 다양체와 데이터셋에서 평가되었습니다:

• 구형 데이터 ($S^2$): 지질학적 사건(화산, 지진, 홍수, 산불)의 밀도 추정에서, 제안된 방법은 바라단 근사와 스펙트럼 방법에 의존하는 리만 스코어 기반 생성 모델(RSGM)과 비교하여 경쟁력 있거나 더 우수한 테스트 로그-가능도(log-likelihood)를 달erm했습니다.
• 회전 데이터 ($SO(3)$): 합성 혼합 모델링에서 RSGM 베이스라인보다 개선된 로그-가능도를 보여주었습니다.
• 양의 정부호 행렬 ($SPD(n)$):
  • 교통 분석 ($SPD(10)$): 뉴욕 택시 데이터로부터 교통 흐름 행렬의 조건부 생성을 수행했습니다. 이 방법은 SPD-DDPM에 비해 실제 값과의 프로베니우스 거리(Frobenius distance)를 줄였습니다.
  • 뇌 연결성 ($SPD(n)$): EEG 연결 행렬의 조건부 생성을 수행했습니다. 이 방법은 높은 정밀도(precision)와 F1 점수를 달성하며, 디페오모픽 매핑(diffeomorphic mapping)을 사용하는 DiffeoCFM 및 LowTriDDPM과 경쟁하거나 이를 능가하는 성능을 보였습니다.
• 분자 생성 ($\mathbb{R}^{k \times n}/S_n$): 치환 몫 공간에서의 분자 생성 실험을 수행했습니다. 이 방법은 분자의 유효성(validity)과 안정성(stability) 측면에서 VE EDM 베이스라인보다 높은 수치를 기록했으나, 알려진 유효성-참신성 트레이드오프(validity-novelty trade-off)를 반영하여 참신성(novelty)은 다소 낮게 나타났습니다.

검증: 저자들은 학습된 PINN 대리물을 학습된 로그 밀도 및 스코어를 참조 근사치와 비교하고, 정규화된 PDE 잔차 및 경계 오차를 보고함으로써, 대리물이 열 방정식을 만족함을 확인하여 검증하였습니다.

5. 의의 및 주장

본 논문은 PDE 기반 열 커널 근사가 리만 확산 모델을 위한 실용적이고 일반적인 대안을 제공한다고 주장합니다. 열 커널을 알려진 PDE의 해로 취급함으로써, 이 접근법은 해석적 커널을 구할 수 없는 다양체로 드리프트 없는(drift-free) 리만 확산을 확장합니다.

저자들은 이 연구를 "열 방정식 풀기"라는 관점으로 위치시키며, 이를 통해 단순한 기하학을 넘어 다양체 확산 모델의 적용 범위를 넓혔다고 주장합니다. 그러나 다음과 같은 한계점도 겸허히 인정합니다:

• 계산 비용: 고차 자동 미분(autodifferentiation)과 콜로케이션(collocation) 평가로 인해 PINN 학습에는 상당한 자원이 필요합니다.
• 안정성: 고차원이나 복잡한 기하학적 구조에서는 학습이 불안정할 수 있으며, 더 강력한 귀납적 편향(inductive bias)이 필요할 수 있습니다.
• 복잡성: 이 파이프라인은 열 커널 추정을 아예 배제하는 Flow Matching보다 더 복잡합니다.

이 연구는 Flow Matching을 대체하려는 것이 아니라, 측지선 보간(geodesic interpolation)이나 지수 사상(exponential map)에 대한 의존도가 높거나 특정 열 커널 구조가 반드시 필요한 다양체 상의 스코어 기반 모델을 위한 필수적인 백엔드를 제공하는 것을 목표로 합니다.