기술 요약: 선형 역문제 해결을 위한 정확한 사후 스코어 추정 (Exact Posterior Score Estimation)

문제 정의

선형 역문제(Linear inverse problems)는 알려진 순방향 연산자(forward operator) $A$와 관측 노이즈 $\eta$가 포함된 노이즈 섞인 선형 측정값 $y = Ax_0 + \eta$로부터 미지의 신호 $x_0$를 복구하는 과정을 포함합니다. 이러한 문제들은 인페인팅(inpainpainting), 초해상도(super-resolution), 디블러링(deblurring) 등 영상 처리 분야에서 널리 퍼져 있으며, 종종 부정치 문제(ill-posed problem)를 가집니다. 즉, 동일한 관측값에 대해 여러 개의 신호가 존재할 수 있습니다. 목표는 단일 점 추정치를 찾는 것을 넘어, 사후 분포 $p(x_0|y)$로부터 샘플링하여 불확실성을 포착하고 데이터 충실도(data fidelity)와 사전 확률의 타당성(prior plausibility) 사이의 균형을 맞추는 것입니다.

확산 모델(diffusion) 및 플로우 기반 모델(flow-based models)은 가우시안 오염 과정을 역전시키는 법을 학습함으로써 강력한 데이터 사전 확률(data priors)을 제공하지만, 본질적으로 무조건적 스코어(unconditional score) $\nabla_{x_t} \log p(x_t)$만을 제공합니다. 역문제를 해결하기 위해서는 사후 스코어(posterior score) $\nabla_{x_t} \log p(x_t|y)$가 필요합니다. 기존의 접근 방식은 다음과 같은 중대한 한계를 가진 두 가지 범주로 나뉩니다:

1. 학습이 필요 없는 방법(Training-free methods) (예: DPS, DDNM, $\Pi$GDM): 사전 학습된 디노이저를 고정시킨 채 각 역과정 단계에서 근사적인 측정치 매칭 보정(예: 그래디언트 단계 또는 투영)을 적용합니다. 이는 편향(bias)을 유발하여 흔히 과도한 평활화(oversmoothing), 환각 구조(hallucinated structures) 생성, 또는 부적절하게 보정된 불확실성을 초래합니다.
2. 학습 기반 방법(Training-based methods) (예: Palette, InDI): 새로운 조건부 모델을 처음부터 학습시키거나 증류(distillation)를 통해 학습시킵니다. 표현력은 뛰어나지만, 이들은 사전 확률의 디노이징 구조를 포기하는 경향이 있습니다. 즉, 네트워크가 연산자 의존성을 엔드투엔드로 학습해야 하므로, 사전 학습된 체크포인트로부터 효율적인 웜 스타트(warm-starting)가 불가능하며 학습 비용이 증가합니다.

방법론: 정확한 사후 스코어 (Exact Posterior Score, EPS)

본 논문은 선형 가우시안 역문제 및 일반적인 가우시안 보간법(Gaussian interpolants) 하에서 정확한 사후 스코어에 대한 폐쇄형 식(closed-form expression)을 도출합니다. 핵심 통찰은 사후 샘플링이 여전히 디노이징 문제이지만, **측정치를 인식하는 입력(measurement-aware input)**과 **연산자에 의존하는 비등방성 노이즈 공분산(operator-dependent anisotropic noise covariance)**을 가진다는 점입니다.

이론적 도출

저자들은 시간 $t$에서의 사후 스코어가 다음과 같이 표현될 수 있음을 보여줍니다:
$$\nabla_{x_t} \log p(x_t|y) = \frac{1}{\beta_t^2} \left( \alpha_t D_{\Sigma^\star(t)}(\mu^\star(x_t, y, t)) - x_t \right)$$
여기서:

• $D_{\Sigma^\star(t)}$는 $E[x_0 | x_0 + \xi = \mu^\star]$를 추정하는 **비등방성 디노이저(anisotropic denoiser)**입니다 (단, $\xi \sim \mathcal{N}(0, \Sigma^\star(t))$).
• $\mu^\star(x_t, y, t)$는 **사후 피벗(posterior pivot)**으로, 현재의 노이즈 섞인 상태 $x_t$와 측정값 $y$를 정밀도 가중치로 융합한 것입니다.
• $\Sigma^\star(t)$는 역문제 파라미터($A, \sigma_y$)와 확산 스케줄($\alpha_t, \beta_t$)으로부터 유도된 **사후 공분산(posterior covariance)**입니다.

결정적으로, $\mu^\star$는 $(x_t, y)$에 대한 충분 통계량(sufficient statistic) 역할을 하며, $\Sigma^\star$는 비등방성 불확실성을 포착합니다. 즉, $A$에 의해 관측된 방향은 더 확실해지는 반면, 관측되지 않은 방향(A의 영공간(nullspace)에 있는 방향)은 여전히 불확실한 상태로 남습니다.

EPS 학습 목적 함수

저자들은 이 항등식을 활용하는 정확한 사후 스코어(EPS) 학습 목적 함수를 정의합니다:

1. 입력 변환: 표준적인 노이즈 섞인 입력 $x_t$ 대신, 피벗 $\mu^\star(x_t, y, t)$를 사용하여 네트워크를 학습시킵니다.
2. 비등방성 노이즈: 학습 타겟은 특정 비등방성 공분산 $\Sigma^\star(t)$ 하에서의 $\mu^\star$ 디노이징을 목표로 합니다.
3. 효율적인 시뮬레이션: 핵심적인 이론적 기여(Proposition 3)는 필요한 비등방성 오염이 $\mu^\star$의 폐쇄형 구성에 의해 자연스럽게 유도된다는 것을 보여줍니다. 따라서 표준적인 등방성 노이즈 샘플링($x_t = \alpha_t x_0 + \beta_t \epsilon$)을 사용하여 학습하고 $\mu^\star$를 분석적으로 계산할 수 있으므로, 비등방성 노이즈를 직접 샘플링할 필요가 없습니다.
4. 아키텍처: 네트워크 아키텍처는 표준적인 사전 학습된 디노이저와 동일하며, 입력 채널은 피벗 $\mu^\star$와 원시 관측값 $y$를 결합하도록 확장됩니다. 이를 통해 EPS는 처음부터 학습하거나, 사전 학습된 무조건적 디노이저로부터 더 효율적으로 **미세 조정(fine-tuning)**할 수 있습니다.

추론

추론 시, EPS는 기본 백본과 동일한 샘플러(예: Euler ODE)를 사용합니다. 유일한 수정 사항은 표준 디노이저 호출 $D_\theta(x_t, t)$를 $D_\theta(\mu^\star(x_t, y, t), y, t)$로 교체하는 것입니다. 이를 위해 라이클리후드(likelihood) 그래디언트, 투영, 또는 샘플링 중 내부 최적화 루프가 전혀 필요하지 않습니다.

주요 기여

1. 정확한 사후 스코어 항등식: 선형 가우시안 역문제에 대한 정확한 사후 스코어의 폐쇄형 식을 도출하여, 사후 샘플링이 이동된 피벗(shifted pivot)에서의 디노이징으로 귀결됨을 보여주었습니다.
2. EPS 알고리즘: 사전 학습된 구조의 입출력 형식을 보존하는 학습 목적 함수 및 샘플링 알고리즘을 제공합니다. 이는 효율적인 미세 조정을 가능하게 하며, 추론 시 백본의 샘플러를 변경 없이 그대로 사용할 수 있게 합니다.
3. 효율성 및 성능: EPS는 그래디언트 기반 샘플러보다 훨씬 적은 수의 디노이저 평가(NFE)로 최첨단(SOTA) 성능을 달성합니다. EPS는 약 20단계 내에 수렴하는 반면, DPS나 $\Pi$GDM과 같은 베이스라인은 유사한 성능에 도달하기 위해 100단계 이상의 단계가 필요합니다.
4. 고노이즈 극한(High-Noise Limit): 고노이즈 극한에서 사후 평균 $E[x_0|y]$에 대한 일단계 추정치(one-step estimator)를 제공하여, 반복적인 샘플링 없이도 빠른 MMSE-최적 점 추정치를 제공합니다.

실험 결과

저자들은 FFHQ 및 ImageNet 데이터셋에 대해 다섯 가지 선형 역문제(랜덤 인페인팅, 박스 인페인팅, 4$\times$ 초해상도, 가우시안 디블러링, 모션 디블러링)에 대해 EPS를 평가했습니다.

• 지표: EPS는 점별 충실도(PSNR, SSIM), 지각적 품질(LPIPS, FID), 그리고 분포적 보정(CRPS, MMD) 측면에서 학습이 필요 없는 방법(DPS, DAPS, DDNM, $\Pi$GDM) 및 학습 기반 방법(Palette) 모두를 능가합니다.
• 샘플링 효율성: EPS는 약 20 NFE 내에 성능 정점에 도달합니다. 반면, 샘플링 기반 베이스라인들은 100 NFE까지 천천히 개선되며 종종 EPS의 점근선(asymptote)에 도달하지 못합니다.
• 학습 효율성: 사전 학습된 체크포인트로부터 미세 조정할 때, EPS는 다른 학습 기반 방법(예: Palette)에 필요한 데 비해 매우 적은 반복 횟수만으로 수렴하는데, 이는 사전 학습된 디노이징 기하학을 보존했기 때문입니다.
• 런타임: EPS는 비교된 방법들 중 가장 빠른 역문제 해결사입니다. 그래디언트 기반 방법에서 요구되는 역전파(backward pass)를 피하기 때문에, 사실상 무조건적 사전 학습된 디노이저와 거의 동일한 실제 실행 시간(약 1.006$\times$)을 가집니다.
• 어블레이션(Ablations): 실험을 통해 이동된 피벗 $\mu^\star$가 성능의 결정적 요인임을 확인했습니다. 또한, 제로샷 피보팅(사전 학습된 디노이저를 미세 조정 없이 사용하는 것)은 성능이 저조하므로, EPS가 학습한 비등방성 기하학이 필수적임을 입증했습니다.

의의 및 주장

본 논문은 EPS가 선형 역문제에 대해 원칙적이고 보정된(principled and calibrated) 접근 방식을 제공한다고 주장합니다. 정확한 사후 기하학을 도출함으로써, EPS는 학습이 필요 없는 방법의 효율성과 학습 기반 방법의 표현력 사이의 간극을 메웁니다.

• 구조적 정렬: 새로운 조건부 경로를 처음부터 학습하도록 강요하는 조건부 학습 방법과 달리, EPS는 측정에 의해 유도된 특정 비등방성 기하학에 맞춰 사전 학습된 Prior를 적응시킵니다. 이를 통해 빠른 미세 조정과 강력한 무조건적 Prior의 재사용이 가능합니다.
• 투명성: 폐쇄형 피벗과 공분산은 샘플러를 분석할 수 있는 투명한 핸들을 제공하며, 측정된 방향과 관측되지 않은 방향을 명시적으로 분리합니다.
• 실용성: 추론 시 복잡한 샘플링 루프나 그래디언트 계산을 요구하지 않으므로 저지연(low-latency) 애플리케이션에 적합합니다.

저자들은 한계점도 인정합니다. 정확한 도출은 선형 순방향 연산자와 가우시안 관측 노이즈를 가정합니다. 비선형 연산자나 비가우시안 노이즈의 경우 국소 선형화(local linearization)나 실제 라이클리후드에 대한 학습이 필요하며, 이 경우 폐쇄형 식이 직접적으로 적용되지 않습니다. 또한, 잠재 확산 백본(latent diffusion backbones)의 경우, 디코더가 픽셀 공간의 선형 연산자를 잠재 공간에서 비선형으로 만들 수 있으므로 주의가 필요합니다.