Mixture-of-Gaussians-Guided Schedule Design for Brownian Bridge Diffusion Models

기술 요약: 가우시안 혼합 모델(Mixture-of-Gaussians) 기반의 브라운 브리지 확산 모델(BBDM) 스케줄 설계

문제 정의

브라운 브리지 확산 모델(BBDian, BBDM)은 표준 DDPM/DDIM처럼 순수 노이즈를 향해 확산하는 대신, 깨끗한 신호($x_0$)에서 저하된 관측값($y$)으로 직접 확률적 브리지를 구축함으로써 이미지 복원 및 역문제(inverse problems)를 해결하는 독특한 프레임워크를 제공한다. 유망한 방식임에도 불구하고, BBDM의 브리지 스케줄(보간 계수와 분산의 궤적) 선택은 일반적으로 휴리스틱에 의존해 왔다. 복잡한 신호 분포와 일반적인 선형 저하(linear degradation)에 대해 원칙적이고 분석적인 프레임워크를 설계하는 데에는 한계가 있었다. 기존의 분석적 연구들은 단일 가우시안 사전 확률(single-Gaussian prior)이나 저하 연산자와 신호 공분산 사이의 공유된 고유 기저(eigenbasis)와 같은 제한적인 가정에 의존하여, 실제 환경의 역문제에 적용하기에는 제약이 많았다.

방법론

본 연구는 가우시안 혼합(Mixture-of-Gaussians, MoG) 사전 확률 하에서 BBDM 스케줄 설계를 위한 방법론에 독립적인 분석적 프레임워크를 개발한다. MoG는 단일 가우시안보다 풍부하면서도 정확한 사후 확률 공식(posterior formulas)을 허용한다. 본 방법론은 다음과 같은 몇 가지 핵심적인 이론적 단계를 거친다:

1. 정확한 사후 확률 도출: 저자들은 MoG 사전 확률 하에서의 정확한 사후 분포 $p(x_0 | x_s, y)$와 그에 대응하는 베이즈 최적 MMSE 디노이저(denoiser)를 도출한다. 단일 가우시안 사례와 달리, 사후 평균은 현재 잠재 상태 $x_s$에 따라 결정되는 가중치(성분 확률)를 가진 아핀 추정치(affine estimators)의 가중 합으로 나타난다.
2. 전역적 아핀 구조의 상실: 정확한 MoG 디노이저를 BBDM 역방향 샘플러에 삽прав입하면, 체인이 전역적 아핀 구조를 상실함을 알 수 있다. 역방향 업데이트의 계수가 사후 성분 확률을 통해 현재 상태에 의존하게 되므로, 재구성 법칙(reconstruction law)을 폐쇄형(closed-form)으로 펼치는 것이 불가능해진다.
3. 선택된 라벨 대리 모델(Selected-Label Surrogate): 계산 가능성을 회복하기 위해, 저자들은 선택된 라벨 근사법을 도입한다. 측정 사후 확률 $P(C=r|y)$로부터 이산 라벨 $J$를 샘в링하고 이를 역방향 체인 전체에서 고정한다. 이는 비선형성을 동결시켜 역방향 역학을 아핀 가우시안 재귀(affine Gaussian recursion)로 변환한다.
   • 이 대리 모델은 **평균-정확(mean-exact)**임이 입증되었다(즉, 기대 재구성 값이 성분 조건부 사후 평균과 일치한다).
   • 유일한 차이점은 **공분산 결핍(covariance deficit)**으로, 대리 모델의 분산은 실제 사후 분산보다 엄격히 낮게 제한된다.
4. 스케줄 설계 목적 함수: 계산 가능한 대리 모델을 사용하여 스케줄 파라미터 $\Theta = \{(m_s, \delta_s)\}$를 최적화하기 위한 두 가지 상보적인 목적 함수를 정립한다:
   • 와서스타인($W_2$) 목적 함수: 대리 모델의 분포와 실제 사후 분포 사이의 거리를 최소화함으로써 지각적 품질(perceptual quality)을 목표로 한다. 이는 사후 확산(variance)을 맞추는 데 유리하다.
   • MSE 목적 함수: 기대 평균 제곱 오차(MSE)를 최소화함으로써 재구성 충실도(reconstruction fidelity)를 목표로 한다. 이는 사후 평균 주변으로의 집중을 선호한다.
5. 유니버설 스케줄: 스케줄을 제한된 4-파라미터 패밀리로 제한함으로써, 저자들은 두 목적 함수에 대한 **유니버설 스케줄(universal schedules)**의 존재를 증명한다. 이 스케줄들은 특정 저하 연산자, 노이즈 수준, 또는 사전 확률 파라미터에 독립적이다. 이들은 내재적인 트레이드오프를 드러낸다: MSE 성능을 높이면 $W_2$ 성능이 감소하고, 그 반대도 마찬가지이다.

주요 기여

• 정확한 MoG 분석: 단일 가우시안 분석을 확장하여, MoG 사전 확률 하에서의 BBDM에 대한 정확한 사후 확률 및 MMSE 디노이저를 도출하였다.
• 선택된 라벨 대리 모델: 혼합 라벨을 고정하여 폐쇄형 표현식과 최종 재구성 및 스케줄 간의 의존성을 제공하는 효과적인 근사법을 도입하였다. 이 대리 모델은 평균-정확하며 공분산 결핍이 제한적임이 증명되었다.
• 이중 스케줄 기준: 왜곡(distortion)과 지각(perception) 사이의 트레이드오프를 명확히 하는 두 가지 상보적인 설계 기준($W_2$ 및 MSE)을 정립하였다.
• 보편적 문제 불가지론적 스케줄: 특정 저하나 사전 확률을 알 필요 없이 유니버설 스케줄을 도출할 수 있음을 증명하였으며, 이를 다양한 역문제에 대해 검증하였다.

실험 결과

본 논문은 세 단계의 실험을 통해 이론과 실무적 가치를 검증한다:

1. 제어된 MoG 설정: 공유 공분산을 가진 합성 환경에서, 선택된 라벨 근사법이 "오라클(oracle)" MoG BBDM 체인의 동작을 매우 밀접하게 모사함을 보여준다. 주변 차원이 증가함에 따라 선택된 라벨이 안정화되며, 이는 근사의 이론적 자기 일관성(self-consistency)을 확인시켜 준다.
2. MNIST 검증: 푸리에 저역 통과 측정(Fourier low-pass measurements)이 적용된 MNIST 데이터셋에서, 고정된 대리 모델은 훈련된 BBDM 모델과 일치하는 스케줄링 경향을 성공적으로 포착한다. 대리 모델의 MSE 목적 함수를 최적화하면 기본 스케줄에 비해 훈련된 모델의 PSNR 및 SSIM이 향상된다.
3. FFHQ 실제 데이터셋: 인페인팅(inpainting), 디블러링(deblurring), 초해상도(super-resolution) 작업에 대한 FFHQ(256x256) 실험을 통해 제안된 휴리스틱의 실용적 유용성을 입증한다:
   • MSE 지향 스케줄은 일관되게 최고의 PSNR 및 SSIM(왜곡 지표)을 달성한다.
   • $W_2$ 지향 스케줄은 일관되게 최고의 FID 및 LPIPS(지각 지표)를 달성하며, 지각적 품질 측면에서 기본 BBDM 및 조건부 DDIM 베이스라인을 능가하는 경우가 많다.
   • 결과는 내재된 트레이드오프를 확인시켜 준다: 왜곡을 최소화하는 스케줄은 지각적 품질을 희생하고, 그 반대도 마찬가지이다.

의의 및 주장

본 논문은 일반적인 MoG 사전 확률 하에서 BBDM 스케줄 설계를 위한 최초의 원칙적인 분석 프레임워크를 제공한다고 주장한다. 휴리스틱을 넘어, 본 연구는 확산 기반 역문제에서의 왜곡-지각 트레이드오프에 대한 이론적 근거를 제시한다. 저자들은 도출된 유니버설 스케줄이 값비싼 재학습이나 각 역문제마다 복잡한 하이퍼파라미터 튜닝 없이도 BBDM 성능을 최적화할 수 있는 신뢰할 수 있고 직접적인 방법을 제공한다고 단언한다. 선택된 라벨 대리 모델은 정확한 베이즈 추론과 계산 가능한 최적화 사이의 간극을 메우며, 스케줄 설계를 가이드하는 계산 효율적인 도구 역할을 한다.

결론적으로, 정확한 MoG 역방향 과정은 직접적인 펼치기(unrolling)가 어렵지만, 선택된 라벨 근사법은 필수적인 스케줄링 역학을 포착하여 다양한 이미지 복원 작업 전반에서 왜곡 충실도와 지각 품질 사이의 균형을 효과적으로 맞추는 유니버설 스케줄을 도출할 수 있게 한다.