기술 요약: 모드 탐색 손실을 통한 확산 모델 기반 역문제 해결

문제 정의

본 논문은 사전 학습된 무조건부 확산 모델(unconditional diffusion models)을 사용하여 **역문제(inverse problems)**를 해결하는 과제를 다룹니다. 역문제란 알려진 열화 연산자(degradation operator) $A$와 가우시안 노이즈 $\eta$에 의해 제어되는 방정식 $y = A(x_0) + \eta$를 통해, 부분적으로 손상된 측정값 $y$가 주어졌을 때 원래의 데이터 샘플 $x_0$를 추정하는 과정입니다.

사전 학습된 확산 모델은 제로샷(zero-shot, plug-and-play) 방식으로 이러한 문제를 해결할 수 있지만, 기존 접근 방식들은 다음과 같은 상당한 한계에 직면해 있습니다:

1. 사후 확률 샘플링 방법(Posterior Sampling Methods): 이들은 사후 확률 $p(x_0|y)$에 비례하는 샘플을 추출하는 것을 목표로 합니다. 그러나 이들은 종종 모델링 근사(예: 조건부 스코어에 대한 가우시안 근사)나 순차적 몬테카를로(sequential Monte-Carlo) 방법에 의존하며, 이는 계산 비용이 많이 들 수 있고 실제 정답 검증 없이 엄격하게 진정한 사후 분포를 따르지 못할 수 있습니다.
2. MAP 추정 방법(MAP Estimation Methods): 이들은 최대 사후 확률(MAP) 추정치(사후 분포의 모드)를 찾는 것을 목표로 합니다. 기존 솔버들은 휴리스틱한 근사에 의존하거나, 변수 분리(예: HQS, ADMM)를 사용하거나, 일관성 모델(consistency models)을 필요로 합니다. 이러한 접근 방식은 빈번하게 값비싼 그래디언트 계산을 포함하며, 모든 선형 연산자에 대해 유효하지 않을 수 있는 모델링 가정을 수반합니다.

저자들은 많은 이미지 복원 작업이 전체 사후 분포를 포괄하기보다는 측정값과 일치하는 그럴듯한(plausible) 이미지를 생성하는 것을 우선시한다고 주장합니다. 이는 MAP 추정과 더 밀접하게 맞닿아 있으며, 원칙적이고 근사 없는 MAP 전략의 필요성을 뒷받침합니다.

방법론: 변분 모드 탐색 손실 (Variational Mode-Seeking Loss, VML)

핵심 기여는 역 확산 과정 동안 생성된 샘플을 MAP 추정치로 유도하도록 설계된 새로운 목적 함수인 **변분 모드 탐색 손실(VML)**입니다.

이론적 기초

저자들은 각 역 확산 시간 단계 $t$에서 확산 사후 분포 $p(x_0|x_t)$와 측정 사후 분포 $p(x_0|y)$ 사이의 쿨백-라이블러(Kullback-Leibler, KL) 발산을 최소화함으로써 VML을 변분 관점에서 도출합니다:
$$\text{VML}(x_t, t) = D_{KL}(p(x_0|x_t) \parallel p(x_0|y))$$

직관적으로, $t \to 0$으로 갈수록 분포 $p(x_0|x_t)$는 $x_t$ 주변으로 점점 더 집중됩니다. 각 단계에서 이 KL 발산을 최소화하는 최적의 $x_t^*$를 찾음으로써, 이 과정은 샘플을 사후 분포 $p(x_0|y)$의 모드로 안내합니다.

선형 연산자를 위한 폐쇄형 도출 (Closed-Form Derivation)

주요 이론적 성과는 어떠한 모델링 근사 없이 선형 역문제($A(x) = Hx$)의 경우에 대한 VML의 폐쇄형 표현식을 도출한 것입니다. 전체 VML 표현식에는 다음이 포함됩니다:

1. 측정 일치성 항(measurement consistency term): $\|y - HD(x_t, t)\|^2 / (2\sigma_y^2)$.
2. 사전 항(prior terms): 디노이저 $D(x_t, t)$, 노이즈 스케줄 $\sigma_t$, 그리고 사후 공분산 $Cov[x_0|x_t]$를 포함합니다.

단순화 및 재가중된 목적 함수

전체 VML은 계산 및 미분이 매우 까다로운 공분산 행렬 $Cov[x_0|x_t]$를 포함하는 고차 항들을 포함하고 있습니다. 저자들은 이러한 고차 항들이 $t \to 0$일 때 균등하게 0(또는 상수)으로 수렴함을 증명합니다. 결과적으로, 이들은 측정 일치성과 사전 항만을 유지하고 이러한 항들을 제거한 **단순화된 VML ($VML_S$)**을 제안합니다.

최적화 안정성을 더욱 향상시키기 위해, 저자들은 **재가중된 단순화 VML ($VML_S^\rho$)**을 도입합니다:
$$VML_S^\rho(x_t, t) = \frac{\|y - HD(x_t, t)\|^2}{2\sigma_y^2} - \rho_t \log p(x_t) - \rho_t \frac{\|D(x_t, t) - x_t\|^2}{2\sigma_t^2} + C$$
여기서 $\rho_t$는 $t \to 0$일 때 $\rho_t \to 1$이 되는 사전 가중치 스케줄입니다. 이 공식은 표준 자동 미분을 사용하여 효율적인 그래디언트 계산을 가능하게 합니다.

알고리즘

VML 최소화를 기반으로 저자들은 다음을 제안합니다:

1. VML-MAP: 각 역 확산 단계에서 $VML_S^\rho$를 최소화하기 위해 경사 하강법을 사용하는 알고리즘입니다.
2. VML-MAP$_{pre}$: 조건수가 나쁜(ill-conditioned) 선형 연산자를 위한 프리컨디셔닝된 변형입니다. 이는 열화 행렬 $H$의 특이값 분해(SVD)를 활용하여 프리컨디셔너 $P$를 구축함으로써, 블러링(deblurring)과 같은 어려운 작업에 대한 수렴을 가속화합니다.
3. VML-MAP$^\tau$: 노이즈가 있는 이미지 복원을 위한 전략입니다. 낮은 노이즈 수준에서의 최적화 문제로 인해 발생할 수 있는 날카로운 아티팩트를 방지하기 위해, 이 알고리즘은 확산 시간 임계값 $\tau$까지만 VML을 최소화한 후, 최종 이미지를 생성하기 위해 표준 DDIM 샘플링으로 전환합니다.
4. LatentVML-MAP: 잠재 공간(latent space)에서 VML을 근사하는 잠재 확산 모델(LDM)에 대한 확장입니다.

주요 기여

1. 새로운 손실 함수: 확산 및 측정 사후 분포 사이의 KL 발산을 최소화함으로써 샘플을 MAP 추정치로 유도하는 VML을 도입했습니다.
2. 근사 없는 도출: 선형 역문제에 대해, VML은 휴리스틱한 모델링 근사 없이 폐쇄형으로 도출되어, 사전 학습된 무조건부 확산 모델만을 활용한 MAP 추정의 공식적인 이론적 토대를 제공합니다.
3. 수렴 분석: 완만한 가정 하에 단순화 및 재가중된 VML 목적 함수가 $t \to 0$일 때 음의 로그 사후 확률(negative log-posterior, 즉 MAP 목적 함수)로 수렴한다는 것을 공식적으로 증명했습니다.
4. 실용적인 알고리즘: 다양한 이미지 복원 작업에 대해 경험적으로 효과적인 솔루션을 제공하는 VML-MAP 및 그 프리컨디셔닝 변형을 개발했습니다.
5. 광범위한 검증: 여러 데이터셋(ImageNet, FFHQ, CelebA)과 작업(inpainting, super-resolution, deblurring)에 걸쳐 기존 베이스라인(예: DDRM, $\Pi$GDM, MAPGA, DAPS)보다 우수한 성능을 입증했습니다.

실험 결과

저자들은 ImageNet256, FFHQ256, CelebA에서 사전 학습된 모델을 사용하여 이미지 복원 작업에 대해 방법을 평가했습니다.

• 지표: 본 논문은 픽셀 단위 지표(PSNR, SSIM)보다 지각적 지표(perceptual metrics) (LPIPS 및 FID)를 강조합니다. 이는 후자가 불량한 조건의 역문제에서 그럴듯한 재구성의 지각적 품질을 포착하는 데 실패하는 경우가 많기 때문입니다.
• 성능:
  • VML-MAP는 인페인팅, 초해상도, 블러링 작업 전반에서 LPIPS 및 FID 측면에서 사후 확률 샘플링 방법(DDRM, $\Pi$GDM) 및 기타 MAP 추정기(MAPGA, DAPS)를 일관되게 능가했습니다.
  • VML-MAP$_{pre}$는 조건수가 나쁜 작업(예: 블러링)에서 비-프리컨디셔닝 버전보다 유의미한 개선을 보였으며, 더 낮은 LPIP 및 FID 점수를 달 Ach성했습니다.
  • 계산 효율성: VML-MAP는 많은 역 확산 단계를 요구하는 방법들에 비해 더 적은 실행 시간(신경 함수 호출 횟수)으로 더 나은 지각적 품질을 달성했습니다.
• 노이즈 복원: VML-MAP$^\tau$ 전략은 노이즈가 있는 측정값을 성공적으로 처리하여, 표준 VML-MAP가 아티팩트로 인해 어려움을 겪었던 부분에서 고품질의 재구성을 생성했습니다.
• 잠재 모델: LatentVML-MAP는 CelebA에서 효과를 입증했으나, 저자들은 디코더의 비선형성으로 인해 잠재 공간에서의 최적화가 여전히 더 어렵다는 점을 언급했습니다.

의의 및 주장

본 논문은 VML이 오직 사전 학습된 무조건부 확산 모델만을 활용하여 선형 역문제에 대한 최초의 근사 없는 MAP 추정 모델링임을 주장하며, 이를 공식적인 이론적 분석으로 뒷받침합니다.

저자들은 자신들의 접근 방식이 분포 포괄(distribution coverage)을 목표로 하는 사후 확률 샘플링 및 휴리스틱이나 일관성 모델에 의존하는 기존 MAP 솔버와 근본적으로 다르다는 점을 강조합니다. VML을 최소화함으로써, 이 방법은 사후 분포의 "모드(modes)"를 효과적으로 타겟팅하며, 저자들은 이것이 그럴듯하고 고품질의 이미지 복원을 생성하기 위해 원하는 동작이라고 주장합니다.

본 연구는 다음과 같은 한계를 인정하며 겸허한 태도를 보입니다:

• 효과성은 VML 목적 함수를 사용할 수 있는 실용적으로 효율적인 최적화 도구의 가용성에 달려 있습니다.
• 조건수가 나쁘거나 비선형인 연산자에 대해서는 최적화 문제가 지속되지만, 프리컨디셔닝이 도움을 줍니다.
• 과도한 계산 비용 없이 최적화 문제를 더욱 완화하기 위해 고차 최적화 도구를 개발하는 향후 연구가 필요합니다.

결론적으로, VML-MAP는 확산 모델을 이용한 역문제를 해결하기 위한 원칙적이고 이론적 근거가 확실하며 경험적으로 효과적인 프레임워크를 제공하여, 이론적 MAP 추정과 실질적인 이미지 복원 사이의 간극을 메웁니다.