기술 요약: DDIM을 통한 선형 역문제를 위한 증명 가능한 확산 기반 사후 샘플링

문제 정의

본 논문은 $Y = AX_0 + \varepsilon$ 형태의 **노이즈가 포함된 선형 역문제(noisy linear inverse problems)**를 해결하는 문제를 다룹니다. 여기서 $X_0 \in \mathbb{R}^d$는 미지의 신호이고, $A \in \mathbb{R}^{n \times d}$는 알려진 선형 측정 연산자이며, $\varepsilon \sim \mathcal{N}(0, \sigma^2 I_n)$은 가우시안 노이즈입니다. 목표는 관측값 $y$에 조건화된 베이지안 사후 분포(Bayesian posterior distribution) $p_{X_0|Y}$로부터 샘플을 생성하는 것입니다.

확산 모델(Diffusion models)은 복잡한 신호 분포를 포착하기 위한 강력한 데이터 기반 사전 확률(prior)로 부상했으나, 기존의 사후 샘플링 방식들은 대개 다음 두 가지 한계점 중 하나를 겪습니다:

1. 이론적 보증의 부재: 투영 기반 방식(projection-based methods)이나 근사적인 스코어 보정(approximate score corrections)과 같은 많은 휴리스틱 접근 방식들은 생성된 샘플이 실제 사후 분포로 수렴한다는 엄격한 증명이 부족합니다.
2. 계산적 오버헤드: 사후 스코어를 근사하거나 순차 몬테카를로(Sequential Monte Carlo, SMC)를 사용하려는 방법들은 스코어 함수의 그래디언트 계산이나 대량의 파티클(particle)을 요구하여 과도한 계산 비용을 초래합니다.

방법론: Posterior-DDIM

저자들은 표준 DDIM(Denoising Diffusion Implicit Models) 샘플러를 변형하여 사후 샘플링을 수행하는 간단하고 효율적인 알고리즘인 Posterior-DDIM을 제안합니다. 핵심 혁신은 신호 공간을 특이값 분해(SVD) 기반의 분할을 통해 나누고, 관측 신호 대 잡음비(SNR)와 확산 SNR 사이의 상대적 강도에 맞추어 업데이트를 조정하는 데 있습니다.

1. SNR 가이드 분할 (SNR-Guided Partitioning)

측정 행렬 $A = U\Sigma V^\top$의 SVD를 사용하여, 문제를 각 특이 방향 $v_s$를 따르는 스칼라 추론 문제로 분리합니다. 저자들은 확산 SNR($\alpha_t/\sigma_t$)이 관측 SNR($\Sigma_{ss}/\sigma$)과 같아지는 각 방향에 대한 임계 시간 $\tau_s$를 정의합니다. 이를 바탕으로, 임의의 시간 $t$에서 각 특이 방향은 다음 세 그룹으로 분할됩니다:

• 측정 지배 방향 ($S^{meas}_t$): 관측 SNR이 확산 SNR보다 높은 방향 ($t \ge \tau_s$).
• 사전 확률 지배 관측 방향 ($S^{prior}_t$): 관측 SNR이 확산 SNR보다 낮지만 여전히 관측되는 방향 ($t < \tau_s$).
• 미관측 방향 ($S^c$): 측정 정보가 존재하지 않는 $A$의 영공간(null space)에 있는 방향.

2. 좌표별 업데이트 규칙 (Coordinate-Wise Update Rules)

알고리즘은 이 그룹들에 대해 서로 다른 업데이트 규칙을 적용하며, 이는 표준 DDIM 업데이트에 가벼운 수정만을 요구합니다:

• $S^{meas}_t$ (측정 지배)를 위해: 알고리즘은 학습된 사전 확률을 건너뛰고 직접 교정된 관측치의 노이즈 버전을 주입합니다. 구체적으로, 관측 $y$에 의한 보조 순방향 과정(auxiliary forward process)을 사용하여 좌표 $v_s^\top \hat{X}_{t_{i+1}}$를 생성하며, 이는 해당 방향으로 투영된 순방향 확산 과정의 주변 분포(marginal distribution)와 일치하도록 합니다.
• $S^{prior}_t$ (사전 확률 지배 관측)를 위해: 알고리즘은 학습된 데이터 예측기 $\mu_t(\cdot)$에 의해 구동되는 표준 DDIM 역학을 따릅니다. 특이값이 동일하지 않은 일반적인 경우, 변화하는 특이값을 고려하여 예측기에 **보정 항(correction term)**을 추가함으로써 업데이트가 조건부 사후 분포와 일치하도록 보장합니다.
• $S^c$ (미관측)를 위해: 이 좌표들은 관측된 좌표들에 조건화된 학습된 확산 사전 확률로부터 완전히 샘플링됩니다. 미관측 부분 공간에 대한 충분한 탐색을 장려하기 위해 더 큰 확률성 파라미터 $\eta_1$을 사용하여, 업데이트가 더 랑제빈(Langevin)스러운 성격을 갖도록 합니다.

3. 확률성 파라미터 (Stochasticity Parameters)

방법론은 $S^{prior}_t$와 $S^c$의 DDIM 업데이트에서 노이즈 주입을 제어하는 두 개의 확률성 파라미터 $\eta_0$와 $\eta_1$을 활용합니다.

• $\eta_0$는 사전 확률이 주요 가이드 역할을 하는 관측 방향에서 결정론적 업데이트를 유지하기 위해 작게(종종 0으로) 유지됩니다.
• $\eta_1$은 분포가 오직 사전 확률에 의해서만 결정되는 미관측 부분 공간이 효과적으로 탐색될 수 있도록 충분히 크게 설정됩니다.

주요 기여

1. 알고리즘의 단순성: 제안된 방법은 사후 샘플링을 단순한 좌표별 DDIM 업데이트 시퀀스로 축소합니다. 이는 복잡한 투영 단계, 가능도(likelihood) 그래디언트 계산, 또는 확산 모델의 재학습 없이도 측정 모델을 명시적으로 통합합니다.
2. 이론적 보증: 본 논문은 **점근적 사후 정확성(asymptotic posterior correctness)**을 확립합니다. 정확한 확산 사전 확률과 충분히 미세한 시간 이산화(time discretization)를 가정할 때, 저자들은 Posterior-DDIM에 의해 생성된 샘플의 분포가 총 변동 거리(total variation distance) 측면에서 실제 베이지안 사후 분포 $p_{X_0|Y}$로 수렴함을 증명합니다. 이는 동일한 특이값을 갖는 특수 사례와 임의의 특이값을 갖는 일반적인 경우 모두에 적용됩니다.
3. 실험적 성능: 이미지 복원 작업(inpaint, super-resolution, Gaussian deblurring, compressed sensing)에 대한 광범위한 실험을 통해, Posterior-DDIM이 대부분의 평가 지표(PSNR, SSIM, LPIPS, FID)에서 기존의 확산 기반 샘플러들(DPS, DDRM, DDNM+, ProjDiff 포함)을 능가함을 입증했습니다.

실험 결과

저자들은 CelebA 및 ImageNet 데이터셋에 대해 알고리즘을 평가했습니다. 주요 결과는 다음과 같습니다:

• 우수성: Posterior-DDIM은 대부분의 작업에서 최소 3개의 4개 지표에서, 그리고 고려된 모든 작업에서 최소 2개의 지표에서 최고의 성능을 달성했습니다.
• 강건성: 이 방법은 다양한 노이즈 수준과 열화 유형에 걸쳐 일관되게 베이스라인들을 능가했습니다.
• 효율성: 알고리즘은 안정적인 성능에 도달하기 위해 100회의 함수 호출(NFE)만을 필요로 하며, 이는 종종 훨씬 더 많은 단계(예: DPS의 1000 NFE)를 요구하는 다른 방법들과 비교했을 때 비슷하거나 더 빠릅니다.
• 하이퍼파라미터 민감도: 실험 결과 $\eta_0 = 0$(사전 확률 지배 방향에 대한 결정론적 업데이트)과 충분히 큰 $\eta_1$(예: 16 또는 32)을 설정하는 것이 최적의 결과를 낳는 것으로 나타났습니다.

의의 및 주장

본 논문은 확산 사전 확률과 선형 측정 모델을 통합하기 위한 원칙적이고 증명 가능한 프레임워크를 제공한다고 주장합니다. 복잡한 사후 샘플링 문제를 단순한 방향별 DDIM 업데이트로 변환함으로써, 저자들은 이론적 엄밀함과 실용적 효율성 사이의 간극을 메웁니다.

이 연구의 의의는 다음과 같은 능력에 의해 강조됩니다:

• 많은 휴리스틱 방법에서 결여된 수렴 보증 문제를 해결함으로써, 역문제를 위한 DDIM 유형 샘플러에 대한 엄격한 이론적 지원을 제공합니다.
• 측정 모델을 명시적으로 처리하면서도 표준 확산 샘플러의 계산 효율성을 유지하는 구현하기 쉬운 솔루션을 제공합니다.
• 순차 몬테카를로의 계산적 제약이나 투영 기반 방법의 휴리스틱 근사 없이도 **사후 일관성(posterior consistency)**을 달성할 수 있음을 보여줍니다.

저자들은 현재 연구가 선형 역문제에 집중하고 있지만, 이 프레임워크를 비선형 설정으로 확장하고 수렴 속도를 규명하는 것이 향후 연구의 자연스러운 방향임을 언급했습니다.