기술 요약: 확산 역문제(Diffusion Inverse Problems)를 위한 스케일 일관적 사후 역학(Scale-Consistent Posterior Dynamics)

1. 문제 정의

본 논문은 노이즈가 포함된 선형 측정값 $y = Ax_0 + n$으로부터 깨끗한 이미지 $x_0$를 재구성하는 것을 목표로 하는 **확산 역문제(diffusion inverse problems)**를 다룹니다. 베이지안 프레임워크에서 이는 $p(x_0 | y) \propto p(y | x_0)p_0(x_0)$인 사후 분포로부터 샘플링하는 것을 요구하며, 여기서 $p_0$는 사전 학습된 확산 생성기(diffusion generator)에 의해 모델링된 사전 확률(prior)입니다.

핵심적인 어려움은 중간 단계의 우도 스코어(likelihood score), 즉 $\nabla_{x_t} \log p_t(y | x_t)$의 계산 불가능성(intractability)에 있습니다. 정확한 사후 샘플링을 위해서는 이 스코어를 평가해야 하며, 이는 $p(x_0 | x_t)$에 대한 조건부 분포를 통해 깨끗한 이미지의 우도를 적분하는 과정을 포함합니다. 기존 방법들은 디노이징된 점 추정치(예: DPS), 연산자 특화 분해(예: DDRM), 또는 변분 근사(variational approximations)를 사용하여 이 적분을 근사합니다. 이러한 접근 방식들은 사전 확률과 우도 사이의 스케일 불일치, 악조건(ill-conditioned)인 연산자가 존재하는 경우의 수치적 경직성(numerical stiffness), 또는 왜곡(distortion)과 지각적 품질(perceptual quality) 사이의 트레이드오프 문제를 겪는 경우가 많습니다.

2. 방법론

저자들은 문제를 세 가지 개념적 계층, 즉 이상적인 참조 모델(ideal reference model), 다룰 수 있는 연속적 대리 모델(tractable continuous surrogate), 그리고 이산적 수치 알고리즘(discrete numerical algorithm)으로 분리하는 프레임워크를 제안합니다.

2.1. 이상적 사후 SDE 패밀리

저자들은 먼저 하나의 매개변수 $\zeta \geq 0$를 갖는 정확한 사후 확률 SDE(Stochastic Differential Equations) 패밀리를 정의합니다. 이 패밀리는 동일한 시간 마진(true posterior)을 공유하지만, 확률성(stochasticity) 매개변수 $\zeta$에 따라 달라집니다.

• $\zeta = 0$: 확률 흐름 수송(probability-flow transport, 결정론적)에 해당합니다.
• $\zeta = 1$: 표준 역시간 SDE(reverse-time SDE)에 해당합니다.
• $\zeta > 1$: 타겟 분포를 변경하지 않으면서 확률적 탐색(stochastic exploration)을 증가시킵니다.
  이 공식화는 정확한 조건부 스코어를 사용할 수 있다면, 탐색(exploration)과 수송(transport) 사이의 균형을 타겟 분포를 바꾸지 않고도 조절할 수 있음을 입증합니다.

2.2. 스케일 일관적 대리 모델 및 연속화(Continuation)

정확한 스코어를 구할 수 없으므로, 저자들은 다음 두 단계를 통해 다룰 수 있는 대리 모델을 구축합니다.

1. 스케일 정렬(Scale Alignment): 우도를 재스케일링된 "깨끗한 이미지" 좌표 $z = x_t / \sqrt{\alpha_t}$로 표현합니다. 이는 학습된 사전 확률과 데이터 충실도(data-fidelity) 항 사이의 시간 의존적 스케일 불일치를 방지합니다. 이 좌표계에서 사후 근사치(posterior proxy)는 로그 SNR($\lambda$)에 의해 인덱싱되는 **연속화 경로(continuation path)**를 형성하며, 가우시안 평활화(Gaussian smoothing)를 점진적으로 제거하여 깨끗한 사후 분포에 접근합니다.
2. 노이즈 조건부 공분산(Noise-Conditioned Covariance): 엔드포인트(endpoint)가 아닌 곳에서는 재스케일링된 상태가 여전히 노이즈를 포함하고 있다는 점을 처리하기 위해, 측정 공분산을 투영된 확산 불확실성($\Gamma_\lambda = \sigma_d^2 I + s_\lambda AA^\top$)에 의해 보강합니다. 이는 높은 노이즈 수준에서 신뢰할 수 없는 측정 방향을 완화하고, 엔드포인트에서 정확한 우도를 회복합니다.

2.3. 교정된 수송 역학(Corrected Transport Dynamics)

저자들은 타겟의 엔드포인트 일관성이 근사된 스코어에 의해 구동되는 수송 과정이 해당 타겟을 추적하는 것을 보장하지 않는다는 점을 관찰했습니다. 이를 해결하기 위해, 수송 SDE와 **고정된 타겟 랑제방 교정기(frozen-target Langevin corrector)**를 교차(interleave)시킵니다.

• **예측기(Predictor)**는 변화하는 연속화 경로를 따라 샘플을 전진시킵니다.
• 교정기(Corrector)(보조 시간에서 실행됨)는 현재의 고정된 타겟을 향해 샘플을 평형 상태로 만듭니다.
  이를 통해 교정 강도 $\gamma_\lambda$가 움직이는 타겟의 추적을 제어하는 연속적인 대리 SDE가 생성됩니다.

2.4. 이산화: 분산 매칭 분할 단계 IMEX(Variance-Matched Split-Step IMEX)

연속 모델은 **Lie–Trotter 분할(splitting)**과 결합된 **분산 매칭 분할 단계 IMEX(Implicit-Explicit)**를 사용하여 이산화됩니다.

• 내재적 단계(Implicit Step): 선형 우도 항은 경직된 결정론적 과도 응답(deterministic transients)을 댐핑하기 위해 내재적(Backward Euler)으로 해결됩니다.
• 외재적 단계(Explicit Step): 학습된 사전 확률은 외재적으로 처리됩니다.
• 분산 매칭(Variance Matching): 결정적으로, 확률적 혁신(stochastic innovation)은 내재적 우도 해법(implicit likelihood solve) 이후에 추가됩니다. 이는 새로운 확률적 방향이 우도 해법(likelihood resolvent)에 의해 감쇄되지 않도록 하여, 올바른 유한 단계 공분산을 유지하도록 보장합니다.
• 교정기(Corrector): 불변 법칙 편향(invariant-law bias)을 최소화하기 위해 교정 단계에는 Crank–Nicolson 업데이트가 사용됩니다.

3. 주요 기여

1. 이상적 참조 모델: 정확한 하나의 매개변수 사후 SDE 패밀리를 공식화하여, 확률성 매개변수가 사후 마진을 변경하지 않고 탐색/수송의 균형을 어떻게 제어하는지 입증했습니다.
2. 스케일 일관적 대리 모델: 노이즈 조건부 공분산 경로로 구성된 깨끗한 이미지 좌표계에서의 대리 모델을 도출하여, 타겟이 총 변동(total variation) 관점에서 정확한 사후 분포로 수렴함을 보장했습니다.
3. 교정된 대리 SDE: 모델링 에러(대리 스코어)와 수치적 에러(이산화)를 분리하여, 타겟의 일관성만으로는 샘플러의 수렴을 보장할 수 없다는 문제를 해결하는 연속 교차 SDE를 구축했습니다.
4. 수치 스킴: 경직된 측정 방향에서 확률적 탐색을 유지하기 위해 내재적 해법 이후에 혁신을 주입하는 분산 매칭 분할 단계 IMEX 예측기를 개발했습니다.
5. 이론적 보장: log-Sobolev 및 수송 결함(transport-defect) 조건 하에서 대리 모델의 사후 수렴을 증명하였으며, 이산 알고리즘에 대한 1차 약 오차(first-order weak error) 상한을 제시했습니다.

4. 실험 결과

본 방법은 FFHQ 및 ImageNet 데이터셋에 대해 초해상도(super-resolution), 가우시안 디블러링(Gaussian deblurring), 모션 디블러링(motion deblurring) 작업에 대해 고정된 100 NFE(Number of Function Evaluations) 예산으로 평가되었습니다.

• 성능: 제안된 방법(VM-IMEX)은 대부분의 작업에서 경쟁력 있거나 최첨단(SOTA)의 재구성 충실도(PSNR/SSIM)를 달enc성했으며, 특히 ImageNet의 모션 디블러링에서 뛰어난 성능을 보였습니다. 또한 왜곡-지각 트레이드오프에서도 강력한 성능을 입증했습니다.
• 절제 연구(Ablation Studies):
  • 스케일 일관성(SC): 가장 지배적인 요소로 확인되었으며, 스케일링되지 않은 우도와 비교했을 때 FFHQ에서 약 10 dB, ImageNet에서 약 13.8 dB의 PSNR 향상을 가져왔습니다.
  • 노이즈 조건부 연속화(NC): 작업에 따라 효과가 달랐는데, 모션 디블러링에는 크게 도움이 되었으나 가우시안 디블러링에는 영향이 미미했습니다.
  • 교정기 할당(PC): NFE를 교정 단계에 재할당하면 왜곡 지표의 약간의 손실과 함께 지각적 품질(LPIPS)이 개선되었습니다.
  • 혁신 배치(VM): 노이즈가 포함된 디블러링 작업에서는 내재적 해법 이후에 혁신을 배치하는 효과가 작았습니다. 그러나 노이즈 없는 박스 인페인팅(noiseless box inpainting) 스트레스 테스트(특이 제약 조건)에서는, 분산 매칭 접근 방식이 직접적인 Backward Euler보다 훨씬 우수한 성능을 보여 경직된 우도 해법에 의해 확률적 혁신이 필터링되는 것을 방지했습니다.
• 확률성($\zeta$): 실험 결과 최적의 $\zeta$는 작업에 따라 달라지며 비단조적(non-monotonic)임을 보여주었습니다. 즉, 탐색을 늘리는 것이 항상 충실도나 지각 품질을 개선하는 것은 아닙니다.

5. 의의 및 주장

본 논문은 확산 역문제에서 흔히 혼동되는 세 가지 문제를 엄격하게 분리하는 프레임워크를 제공한다고 주장합니다.

1. 타겟 경로가 실제 사후 분포에 접근하는지 여부.
2. 연속 샘플러가 해당 타겟을 추적하는지 여부.
3. 유한 단계 이산화가 연속 역학을 근사하는지 여부.

저자들은 자신들의 접근 방식이 모델링 에러(대리 스코어로부터 기인)와 수치적 에러(이산화로부터 기인)를 분리한다는 점을 강조합니다. 본 연구의 의의는 스케일 일관적 구축과 분산 매칭 이산화에 있으며, 이는 샘플러가 경직된 선형 제약 조건이 존재하는 상황에서도 필요한 확률적 탐색을 유지할 수 있도록 보장합니다. 이 연구는 좌표를 정렬하고, 우량 공분산을 조건화하며, 분할 단계 적분을 통해 확률 법칙을 보존함으로써 제한된 NFE 내에서 고충실도 재구성이 가능하다는 것을 보여줍니다.