기술적 요약: 예측 - 투영 - 재노이즈 (PPR)

문제 정의

확산 모델은 고차원 데이터 분포 $p(x)$ 를 효과적으로 근사하는 최첨단 심층 생성 모델링 프레임워크로 자리 잡았습니다. 그러나 유체 역학, 의료 영상, 기후 모델링과 같은 과학 및 공학 분야에서는 단순히 분포 $p$ 를 근사하는 샘플 생성만으로는 부족합니다. 이러한 응용 분야는 강제 제약 조건 $x \in C$ (예: 정확한 보존 법칙, 경계 조건, 관측 일관성, 비압축성 유동의 발산 제로 등) 을 만족하는 샘플 $x$ 를 요구합니다.

제약 샘플링에 대한 기존 접근 방식은 다음과 같은 중대한 한계에 직면해 있습니다:

1. 학습 기반 방법은 종종 유연성이 부족하여 다른 제약 조건에 대해 재학습이 필요하며, 일반적으로 제약 조건을 약하게만 강제합니다.
2. 베이지안 분해 방법(제약 조건을 노이즈가 있는 관측으로 간주) 은 암시적 가능도가 디랙 델타에 접근하기 때문에 강제 제약 조건에는 적합하지 않습니다. 이로 인해 기울기가 불안정해지며, 실현 가능 집합의 즉각적인 근방에서만 유용한 정보를 제공합니다.
3. 투영 기반 방법은 명시적으로 실현 가능성을 강제할 수 있지만 이론적 및 실용적 함정에 시달립니다. 이러한 방법들은 종종 진정한 제약 분포를 샘플링하지 못해 세 가지 특정 실패를 초래합니다: (i) 제약 위반, (ii) 잘못된 제약 분포, (iii) 원래 데이터 분포의 지지대 $S$ 바깥에 위치한 샘플 (예: 연속 데이터에서 불연속성 도입). 또한, 노이즈가 있는 상태의 직접 투영은 종종 잘 정의되지 않았으며, 반복된 투영 단계는 샘플러를 체계적으로 편향시켜 모드 붕괴를 유발할 수 있습니다.

방법론: 예측 - 투영 - 재노이즈 (PPR)

저자들은 올바른 제약 샘플링이 제약 집합 $C$ 에서 시작하는 전방향 확산 확률 미분 방정식 (SDE) 에 의해 유도된 제약 주변 분포 $p^C_t$ 를 근사해야 한다고 주장합니다.

이론적 기초

이 논문은 초기 상태 $x_0$ 가 $p^C(x_0) \propto p(x_0)\mathbb{1}_C(x_0)$ 에서 추출되는 제약 전방향 과정을 정의합니다. 이는 일련의 제약 주변 분포 $\{p^C_t\}$ 를 유도합니다. 이에 상응하는 역방향 과정은 제약 스코어 함수를 포함합니다:
$$\nabla_{x_t} \log p^C_t(x_t) = \nabla_{x_t} \log p_t(x_t) + \nabla_{x_t} \log P(x_0 \in C | x_t)$$
두 번째 항은 로그 실현 가능성의 기울기를 나타내며, 깨끗한 샘플 $x_0$ 가 실현 가능할 가능성이 높은 영역을 가리킵니다.

PPR 알고리즘

재학습 없이 이러한 제약 주변 분포에서 샘플링하기 위해, 저자들은 고유한 정상 분포가 $p^C_t$ 인 정정자 커널 $K_t$ 를 식별합니다. 정확한 제약 사후 분포에서 샘플링하는 것은 계산적으로 불가능하므로, 저자들은 각 역방향 확산 단계에서 3 단계 반복 과정을 사용하여 이 커널을 근사합니다:

1. 예측 (Predict): 표준 예측기 (예: DDIM) 를 사용하여 $x_{t+1}$ 에서 $x_t$ 로 이동합니다.
2. 투영 (Project): 노이즈가 있는 상태 $x_t$ 또는 탈노이즈 추정치 $d_\theta(x_t)$ 를 직접 투영하는 대신, PPR은 탈노이즈기를 통해 투영합니다. 제약 조건을 만족하는 탈노이즈 출력 $d_\theta(x^*_t)$ 를 갖는 노이즈가 있는 상태 $x^*_t$ 를 찾기 위해 최적화 문제를 풉니다:
   $$x^*_t = \arg \min_{x} c(d_\theta(x))$$
   이는 투영된 깨끗한 추정치 $d_\theta(x^*_t)$ 가 데이터 매니폴드 $S$ 위에 유지되도록 하여 불연속성과 같은 아티팩트를 방지합니다.
3. 재노이즈 (Renoise): 탈노이즈되고 실현 가능한 추정치 $d_\theta(x^*_t)$ 를 전방향 커널을 통해 재노이즈하여 현재 노이즈 수준 $t$ 에서 새로운 샘플 $x'_t$ 를 얻습니다.

이 예측 - 투영 - 재노이즈 시퀀스는 각 확산 단계에서 $M$ 번 반복됩니다. 재노이즈 단계는 올바른 노이즈 수준을 복원하고 혼합에 필요한 확률성을 주입하여 샘플러가 고정점으로 붕괴되는 것을 방지하므로 중요합니다.

주요 기여

1. 제약 전방향/역방향 공식화: 논문은 제약 전방향 과정을 정의하고 관련 역방향 과정 및 스코어 함수를 유도하며, 제약 주변 분포를 고유한 정상 분포로 갖는 정정자 커널을 식별합니다.
2. PPR 알고리즘: 탈노이즈기를 통한 반복적 투영과 재노이즈를 통해 정정자 커널을 근사하는 학습 없는 샘플링 알고리즘입니다.
3. 구성 요소의 필요성: 저자들은 세 가지 구성 요소 (예측, 탈노이즈기를 통한 투영, 재노이즈) 가 모두 필요함을 입증합니다. 어떤 구성 요소라도 제거하면 성능이 극적으로 저하되거나 제약 위반이 증가하거나 데이터 지지대 밖의 샘플이 생성됩니다.
4. 이론적 오차 한계: 논문은 근사 오차 분석을 제공하여, 오차가 혼합 항 (반복 횟수 $M$ 에 따라 지수적으로 감소) 과 점 추정 근사로 인한 불가피한 편향으로 분해됨을 보여줍니다.

실험 결과

이 방법은 세 가지 다른 설정에서 평가되었습니다:

1. DATA2D: 무작위 제약 조건을 가진 2D 분포 (체크보드 및 가우시안 혼합).
2. Kuramoto-Sivashinsky (KS): 난류를 모델링하는 1 차원 편미분 방정식으로, 128 차원 궤적에서 테스트됨.
3. WEATHER: Appa 잠재 기상 모델을 사용한 전 세계 기상 예보 (108 차원 대기 상태).

성능 지표:

• 실현 가능성: PPR 은 KS 데이터셋에서 82.7%, DATA2D 에서 79.0% 의 실현 가능성률 ($c(x) \le \epsilon$ 을 만족하는 샘플) 을 달성했습니다. 반면, 다음으로 좋은 방법들은 KS 에서 약 1%, DATA2D 에서 31% 의 실현 가능성률만 달성했습니다.
• 제약 위반: PPR 은 모든 데이터셋에서 일관되게 가장 낮은 제약 위반을 보여주었습니다.
• 분포 충실도: PPR 은 데이터 지지대 $S$ 내에 머무르는 샘플을 생성하여 연속성과 규칙성을 유지한 반면, 많은 기준 방법들 (예: CPS, PDM, HDC) 은 데이터 매니폴드 밖으로 떨어진 불연속 샘플을 생성했습니다.
• 정확도: KS 및 Weather 작업에서 PPR 은 DPS, MMPS, Trust Sampling (TS) 과 같은 기준 방법들에 비해 더 낮은 RMSE, 더 나은 Skill 점수, 그리고 더 낮은 CRPS(연속 순위 확률 점수) 를 달성했습니다.

중요성 및 주장

이 논문은 PPR 이 재학습 없이 강제 제약 조건 하의 확산 모델에서 샘플링하는 오랜 과제를 성공적으로 해결했다고 주장합니다. 그 주요 중요성은 낮은 제약 위반과 높은 분포 충실도를 동시에 달성하는 능력에 있으며, 이는 기존 방법들이 제공하지 못하는 조합입니다.

저자들은 PPR 이 계산 비용과 샘플 품질 사이의 의미 있는 절충안을 제공한다고 강조합니다. 커널 적용 횟수 ($M$) 를 늘리면 제약 분포 근사가 일관되게 개선됩니다. 또한, 탈노이즈기를 통해 투영함으로써 이 방법은 샘플이 학습된 데이터 매니폴드와 일관되게 유지되도록 보장하여 $x_0$ 공간 투영 방법에서 흔히 발생하는 아티팩트 (예: 불연속성) 를 방지합니다.

이 연구는 PPR 이 물리 법칙이나 관측 데이터에 대한 정확한 준수가 필요한 과학적 응용 분야를 위한 견고한 학습 없는 솔루션이라고 결론지으며, 비미분 가능 또는 이산 제약 조건에 관한 한계점을 지적합니다.