기술 요약: 확산 모델을 위한 국소 MAP 샘플링

문제 정의

확산 사후 샘플링 (DPS) 은 사후 분포 $p(x_0 | y)$에서 샘플링하여 역문제 해결을 위한 강력한 프레임워크로 부상했으며, 원칙적인 불확실성 정량화와 다중 모드성을 제공합니다. 그러나 이미지 처리 및 과학 분야의 많은 실제 응용에서는 전체 분포 샘플링보다 정확한 점 추정을 우선시합니다. $\arg \max p(x_0 | y)$로 정의된 최대 사후 확률 (MAP) 추정량은 여전히 표준 재구성 목표입니다.

기존 최적화 기반 방법들 (예: Resample, DiffPIR, DCDP) 은 소음 제거 단계와 최적화 단계를 번갈아 수행함으로써 재구성 성능을 향상시키려 시도합니다. 그러나 이러한 반복 절차, 전역 MAP 해, 그리고 사후 샘플링 사이의 이론적 연결은 여전히 불분명합니다. 이러한 방법들이 전역 MAP 에 수렴하는지, 또는 DPS 와 일관성이 있는지 여부는 입증되지 않았습니다. 또한, 이러한 방법들에서 최적화 단계의 구체적인 성격, 즉 전역 또는 국소 목표를 해결하는지 여부에 대한 명확화가 필요합니다.

방법론: 국소 MAP 샘플링 (LMAPS)

저자들은 확산 궤적을 따라 국소 최대 사후 확률 (MAP) 부분 문제를 반복적으로 해결하는 새로운 추론 프레임워크인 **국소 MAP 샘플링 (LMAPS)**을 소개합니다.

1. 이론적 프레임워크

LMAPS 는 단일 전역 문제가 아닌 일련의 국소 MAP 문제를 해결함으로써 작동합니다. 각 확산 단계 $t$에서, 소음 상태 $x_t$와 측정값 $y$가 주어졌을 때, 이 방법은 다음을 찾습니다:
$$x^*_0(t, x_t, y) := \arg \max_{x_0} p(x_0 | x_t, y)$$
이후 표준 DDIM 전이 단계가 따릅니다:
$$x_{t-\Delta t} = g(x^*_0, x_t, \epsilon)$$

기존 방법과의 차이점:

• 전역 MAP 대비: 전역 MAP 는 전체 사후 분포 $p(x_0 | y)$의 모드를 찾습니다. 반면 LMAPS 는 각 단계에서 조건부 사후 분포 $p(x_0 | x_t, y)$의 모드를 찾습니다. 저자들은 일련의 국소 최대화자가 반드시 $p(x_0 | y)$의 전역 최대화자로 수렴하지는 않는다고 주장합니다.
• DPS 대비: DPS 는 업데이트 시 조건부 평균 $E[x_0 | x_t, y]$를 사용합니다. LMAPS 는 이 평균을 조건부 모드로 대체합니다. 조건부 사후 분포가 가우시안인 경우에만 이 둘은 동일하며, 그렇지 않으면 LMAPS 는 모드 탐색 편향을 도입하여 모드 간의 저밀도 영역을 회피합니다.
• DMAP 대비: DMAP 또한 국소 MAP 관점을 활용하지만, 잠재 확산 변수 공간 ($x^*_{t-\Delta t} = \arg \max p(x_{t-\Delta t} | x_t, y)$) 에서 최적화를 수행합니다. LMAPS 는 문제 를 직접 깨끗한 신호 공간 ($x^*_0$) 에서 제기하여 측정 가능도 $p(y|x_0)$와 명시적으로 연결되는 보완적인 관점을 제공합니다.

2. 실용적 알고리즘 및 근사

일반적인 비선형 역문제에 대해 국소 MAP 문제를 실용적으로 만들기 위해 저자들은 특정 근사 및 재구성을 제안합니다:

• 공분산 근사: 저자들은 가우시안 사전 가정에서 영감을 받은 등방성 형태를 사용하여 사전 분포 $p(x_0 | x_t)$의 조건부 공분산 $\Sigma_{0|t}$를 근사합니다. 소음 수준 $\sigma_t \to 0$으로 갈 때, 사후 공분산은 점근적으로 등방성이 됩니다. 그들은 $\Sigma_{0|t} \approx \frac{1}{SNR} I$로 근사하며, 여기서 $SNR = \alpha_t^2 / \sigma_t^2$입니다.

• 목적 함수 재구성: 국소 MAP 목적 함수는 사전 거리와 데이터 충실도의 가중 합을 최소화하는 형태로 구성됩니다:
  $$x^*_0 = \arg \min_{x_0} \left\{ \frac{SNR}{k} \|x_0 - m_{0|t}\|^2 + \frac{1}{\sigma_y^2} \|y - H(x_0)\|^2 \right\}$$
  수치적 안정성과 해석 가능성을 개선하기 위해 이를 척도 불변의 볼록 결합으로 재구성합니다:
  $$x^*_0 = \arg \min_{x_0} \left\{ (1 - \mu_t) \frac{1}{2} \|x_0 - m_{0|t}\|^2 + \mu_t k_2 \|y - H(x_0)\|^2 \right\}$$
  여기서 $\mu_t = \sigma_t^2 / (\sigma_t^2 + k_1^2)$입니다. 이 공식은 자동 어닐링을 가능하게 합니다: $t$가 감소함에 따라 ($\sigma_t \to 0$), 목적 함수는 측정 주도에서 사전 주도 방향으로 전환됩니다.

• 선형 대 비선형: 선형 역문제에 대해서는 국소 MAP 해를 폐쇄형으로 유도할 수 있으며, 이는 특수한 경우로 Tweedie Moment Projected Diffusion (TMPD) 을 복원합니다. 비선형 문제의 경우, 해는 경사 하강법을 통해 얻어집니다.

주요 기여

1. 이론적 통합: 본 논문은 LMAPS 를 공식화하여 최적화 기반 역문제 방법과 Tweedie Moment Projected Diffusion (TMPD) 을 단일 확률 프레임워크 하에 통합합니다. 최적화 기반 방법들이 본질적으로 전역 MAP 나 전체 사후 샘플링이 아닌 국소 MAP 문제를 해결함을 명확히 합니다.
2. 방법론적 발전: 저자들은 가우시안 기반 공분산 근사와 자동 어닐링을 통해 수치적 안정성과 해석 가능성을 향상시킨 재구성된 목적 함수를 소개합니다.
3. 실증적 검증: 이 프레임워크는 10 가지 이미지 복원 작업 (선형, 비선형, 미분 불가능 연산자 포함) 과 3 가지 과학적 역문제에서 검증되었습니다.

실험 결과

저자들은 이미지 복원을 위해 FFHQ 및 ImageNet 데이터셋에서, 그리고 과학적 역문제 (선형 역산란, 압축 센싱 MRI, 블랙홀 이미징) 벤치마크에서 LMAPS 를 평가했습니다.

• 이미지 복원: LMAPS 는 FFHQ 및 ImageNet 에서 PSNR, SSIM, LPIPS 지표를 기준으로 60 개 사례 중 43 개에서 최첨단 성능을 달성했습니다. 이는 밀접하게 관련된 DMAP 베이스라인을 일관되게 능가하여, 모션 디블러링, JPEG 복원, 양자화 등의 작업에서 2dB 이상의 PSNR 개선을 보여주었습니다.
• 과학적 역문제: LMAPS 는 모든 과학적 작업에서 일관되게 가장 높은 PSNR 을 달성했습니다. 특히, 이전 방법들에 비해 세 가지 선형 역산란 (LIS) 작업에서 1.5dB 이상의 PSNR 향상을 달성했습니다.
• 효율성: LMAPS 는 DAPS 에 비해 우수한 효율성을 보여주었으며, 더 적은 계산 자원 (예: 확산 단계당 더 적은 최적화 단계) 으로 더 나은 결과를 달성했습니다.

중요성 및 주장

본 논문은 LMAPS 가 최적화 기반 방법에 대한 원칙적인 확률적 해석을 제공하며, 사후 샘플링 (DPS) 과 점 추정 (MAP) 사이의 간극을 연결한다고 주장합니다. 명시적으로 국소 MAP 부분 문제를 해결함으로써, 이 방법은 기존 최적화 기반 솔버의 동작을 설명하는 통합된 관점을 제공합니다.

저자들은 LMAPS 가 전체 사후 분포에서 샘플링하지는 않으므로 (따라서 DPS 와 동일한 방식으로 다중 모드성을 포착하지는 않지만), 정확한 점 추정이 주요 목표인 시나리오를 위해 특별히 설계되었다고 강조합니다. 이 프레임워크는 재구성 충실도와 확산 모델의 구조적 사전 지식을 균형 있게 조화시키는 견고한 대안을 제공하며, 표준 이미지 복원과 까다로운 과학적 재구성 작업 모두에서 최첨단 결과를 달성합니다. 이 연구는 미래 연구가 여전히 열린 과제로 남아 있는 확산 사전 지식을 활용한 전역 MAP 추정의 유용성을 더 깊이 탐구할 수 있음을 시사합니다.