기술적 요약: 빠른 확률적 추론을 위한 확률적 전이 지도 증류

문제 제기

확산 모델은 GAN 및 VAE 와 같은 이전 모델들에 비해 생성 품질, 다양성, 분포 커버리지 측면에서 최첨단 성능을 입증했습니다. 그러나 이러한 성능은 종종 많은 단계에 걸쳐 연속 시간 상미분 방정식 (ODE) 또는 확률 미분 방정식 (SDE) 을 적분해야 한다는 필요성으로 인해 비싼 추론 비용과 동반됩니다.

증류 (예: 일관성 모델) 및 최적 수송 기반 접근법과 같은 기존 가속화 기술이 존재하지만, 이들은 종종 생성 품질/다양성과 추론 복잡성 사이의 트레이드오프에 직면합니다. 현재 증류 방법의 중요한 한계는 결정론적 추론 (예: 확률 흐름 ODE 를 통한) 에 의존한다는 점입니다. 고정된 잡음 샘플이 주어지면 출력은 유일합니다. 이는 확산 사후 샘플링을 통한 역문제, 에너지 기반 미세 조정, 또는 사용자 쿼리와의 정렬 개선과 같이 본질적으로 확률적 추론이 필요한 하위 작업에서 그 활용도를 제한합니다. 또한, 많은 기존 증류 프레임워크는 사전 훈련된 교사 모델, 이중 수준 최적화, 또는 궤적의 시뮬레이션 및 캐싱을 요구하여 훈련과 확장성을 복잡하게 만듭니다.

방법론: 확률적 전이 지도 증류 (STMD)

저자들은 **확률적 전이 지도 증류 (STMD)**를 제안합니다. 이는 기본 과정의 확률적 전이 구조를 유지하면서 확산 모델 추론을 가속화하도록 설계된 교사 없는 프레임워크입니다.

핵심 개념

일반적으로 사후 분포의 평균을 모델링하는 점수 기반 확산 모델과 달리, STMD 는 샘플링 SDE 와 관련된 전체 전이 지도를 증류합니다. 이 방법은 역 SDE 의 비용이 많이 드는 수치적 적분을 피하면서, 단일 또는 소수 단계에서 $t' < t$에 대한 전이 확률 $p(x_{t'} | x_t)$를 예측하는 조건부 생성기를 학습합니다.

기술적 공식화

1. 조건부 평균 흐름 (MF): 이 프레임워크는 SDE 전이를 조건부 평균 흐름 모델을 사용하여 매개변수화합니다. 결정론적 흐름 지도를 학습하는 대신, 확산 과정의 궤적을 따라 평균 속도를 학습합니다.
2. $x_0$-매개변수화: 전이 $p(x_{t'} | x_t)$를 모델링하기 위해 저자들은 $x_0$-매개변수화를 활용합니다. 그들은 전이를 사후 분포 $p(x_0 | x_t)$에 대한 적분과 알려진 브리지 측도 $p(x_{t'} | x_0, x_t)$로 표현합니다. 분산 보존 (VP) SDE 의 경우 브리지 측도가 폐형식으로 제공되므로, 문제는 조건부 모델 $p_\theta(x_0 | x_t)$를 학습하는 것으로 축소됩니다.
3. 훈련 목적 함수: 모델 $u_\theta$는 조건부 평균 속도를 근사하도록 훈련됩니다. 저자들은 조건부 평균 흐름 (CMF) 손실 함수를 정의합니다:
   $$L_{dCMF} = \mathbb{E} \left[ \left\| u_\theta(z_s, r, s, x_t, t) - \text{sg}\left( (z_1 - x_0) - (s-r) \frac{d}{ds}u_\theta(z_s, r, s, x_t, t) \right) \right\|^2 \right]$$
   여기서 $z_s$는 확산 상태 $x_t$와 시간 $t$에 조건부인 데이터 $x_0$와 사전 샘플 $z_1$ 사이의 직선 보간입니다. $\text{sg}$ 연산자는 그라디언트 정지를 나타냅니다.
4. 교사 없는 훈련: 훈련 과정은 사전 훈련된 교사 모델을 요구하지 않습니다. 이는 $x_0$를 데이터 분포에서 샘플링하고 폐형식으로 알려진 전방 SDE 커널을 적용하여 생성할 수 있는 결합 분포 $p(x_0, x_t)$의 샘플에 의존합니다.

추론

추론은 학습된 조건부 평균 흐름을 시간 역방향으로 적분하여 수행됩니다. 이 과정은 다음을 포함합니다:

1. 초기 잡음 $x_1$(또는 $z_1$) 을 샘플링합니다.
2. 학습된 속도장 $u_\theta$를 반복적으로 적용하여 시간 $s$에서 $s-\Delta s$로 단계별로 이동합니다.
3. 각 확산 시간 단계 $t$에서 브리지 분포 $p(x_{t-\Delta t} | x_0, x_t)$에서 샘플링하여 상태를 업데이트함으로써, 큰 단계로 확률적 전이를 효과적으로 시뮬레이션합니다.

주요 기여

1. 새로운 프레임워크: 사전 훈련된 교사 없이 회귀 목적 함수를 통해 단일 단계 또는 소수 단계 생성기를 사용하여 고차원 확산 SDE 의 전이 지도를 학습할 수 있는 조건부 평균 흐름 기반 증류 프레임워크인 STMD 의 제안.
2. 이론적 수렴: 2-워asserstein 거리에서 평균 흐름 알고리즘의 수렴 경계 유도. 저자들의 지식에 따르면, 이는 이 메트릭에서 평균 흐름 알고리즘에 대한 최초의 수렴 분석입니다. 그들은 이를 확장하여 학습된 조건부 전이 지도가 기본 역 SDE 의 전이 커널로 수렴함을 증명합니다.
3. 실증적 검증: MNIST, CIFAR-10, CelebA 데이터셋에 대한 이미지 생성 및 인페인팅 작업에서 방법론을 시연하여, 확률성을 유지하면서 베이스라인과 경쟁력 있는 성능을 보임.

실험 결과

저자들은 STMD 를 세 가지 데이터셋에서 평가했습니다:

• MNIST: U-Net 백본을 사용하여, 원래 평균 흐름 알고리즘 및 탈노이즈 확산 확률 모델 (DDPM) 과 비교했습니다. 저자들은 FID 가 MNIST 에서는 신뢰할 수 없으므로 분류기의 잠재 공간에서 프레이체 거리 (FD) 를 활용했습니다. 결과는 STMD 가 확률적 샘플링을 지원하면서도 베이스라인과 비교 가능한 성능을 보임을 보여주었습니다.
• CIFAR-10: U-Net 을 사용하여 유사한 실험이 수행되었습니다. 모델은 고정된 훈련 반복 횟수에 대해 DDPM 및 평균 흐름 베이스라인과 비교하여 경쟁력 있는 프레이체 인셉션 거리 (FID) 점수를 달성하면서도 확산 샘플러의 전체 확률성을 유지했습니다.
• CelebA: 실험은 잠재 DiT 아키텍처를 활용했습니다. 모델은 4 개의 추론 단계와 2 개의 평균 흐름 단계를 사용하여 8.28 의 FID 로 무조건부 샘플을 생성했으며, 이미지 인페인팅 작업을 성공적으로 수행했습니다.

모든 경우, 모델은 공정한 비교를 보장하기 위해 베이스라인과 동일한 네트워크 백본 및 반복 횟수로 처음부터 훈련되었습니다.

중요성 및 주장

이 논문은 STMD 가 빠른 결정론적 증류와 하위 작업에서 필요한 확률적 추론 사이의 격차를 해결한다고 주장합니다. 흐름 지도가 아닌 전이 지도를 증류함으로써, STMD 는 확산 과정의 확률적 구조를 보존합니다.

저자들은 그들의 접근 방식이 교사 없으며, 이중 수준 최적화나 궤적 캐싱의 필요성을 제거하여 훈련 과정을 더 효율적이고 확장 가능하게 만든다고 강조합니다. 2-워asserstein 수렴 경계를 확립하는 이론적 기여는 평균 흐름 접근법에 대한 엄격한 기초를 제공하여 이론적 보장과 실제 응용 사이의 격차를 메워줍니다.

이 연구는 결정론적 증류 방법이 부족한 역문제 및 사후 샘플링과 같은 확률적 추론이 필요한 작업을 위한 해결책으로 자신을 위치시킵니다. 저자들은 현재 성능에 대해 겸손하게 언급하며, 제한된 훈련 반복 횟수로 인해 최첨단 결과를 달성하지는 못했지만, 비교 연구를 통해 기존 베이스라인에 비해 방법론의 유효성이 검증되었음을 지적합니다.