기술 요약: 역문제 해결을 위한 듀얼 어센트 확산 (Dual Ascent Diffusion)

문제 정의

역문제 (Inverse problems) 는 전방 연산자 $A(\cdot)$에 의해 지배되는 잡음이 섞이고 손상된 관측치 $y$로부터 알려지지 않은 신호 또는 이미지 $x$를 추정하는 것을 포함합니다. 이러한 문제들은 천체물리학 및 의료 영상에서 단백질 구조 결정에 이르기까지 다양한 분야에서 광범위하게 존재합니다. 역문제는 일반적으로 잘 정의되지 않아 (ill-posed) 해 공간을 제약하기 위해 사전 정보 (priors) 를 사용해야 합니다.

현대적 접근법은 사전 훈련된 확산 모델을 강력한 사전 정보로 활용하지만, 최대 사후 확률 (MAP) 문제를 해결하거나 사후 분포에서 샘플링하는 기존 방법들은 중대한 한계를 안고 있습니다. 많은 방법들이 관측 모델의 가능도 (likelihood) 와 확산 사전 정보를 결합하기 위해 계산적 근사법에 의존합니다. 이러한 근사법들은 종종 부정확한 샘플, 비최적의 재구성 품질, 또는 관측치를 충실히 반영하지 못함 (즉, 높은 잔차) 을 초래하며, 이는 생성된 해에서 "환각 (hallucinations)"으로 나타날 수 있습니다. 또한, 많은 최첨단 방법들은 높은 수준의 측정 잡음에 취약하며, 오차의 누적이나 보조 마르코프 연쇄 몬테 카를로 (MCMC) 단계의 필요성으로 인해 계산 비효율성을 겪습니다.

방법론: 듀얼 어센트 확산 (DDiff)

저자들은 사전 훈련된 확산 모델을 사용하여 MAP 문제를 해결하기 위한 새로운 프레임워크인 DDiff를 제안합니다. 사후 분포에서 증명 가능하게 샘플링하려는 방법들과 달리, DDiff 는 MAP 문제를 정확하고 효율적으로 해결하는 최적화 전략을 도출하는 데 중점을 둡니다. 이 방법은 승수 교대 방향법 (ADMM) 에서 영감을 받았지만, 확산 모델에 특화된 새로운 듀얼 어센트 (dual ascent) 구조를 도입했습니다.

핵심 알고리즘 구성 요소

DDiff 프레임워크는 각 단계에서 3 가지 주요 업데이트를 반복하며, 주변 변수 ($x, z$) 와 듀얼 변수 ($u$) 사이를 교대로 업데이트합니다:

1. x-업데이트 (데이터 충실도): 제약된 최소제곱 문제를 정확하게 푸는 대신, DDiff 는 선형화된 ADMM 접근법을 사용합니다. 데이터 충실도 항 $-\log p(y|x)$에 대한 단일 경사 하강 단계를 수행하여, 선형이든 비선형이든 미분 가능한 모든 전방 모델 $A(\cdot)$에 대해 직접 적용 가능하도록 합니다.
   $$x \leftarrow v - \gamma \nabla_v \|y - A(v)\|_2^2, \quad \text{where } v = z - u$$

2. z-업데이트 (디노이징): 이 단계는 확산 모델을 단순한 디노이저로 적용하는 방식을 대체합니다. 저자들은 표준 스코어 모델 $s_\theta(x, t)$가 시간 $t$에서 확산 매니폴드 (diffusion manifold) 위에서만 정확하다는 점을 관찰했습니다. 변수 $x+u$가 일반적으로 이 매니폴드에 존재하지 않는다는 사실을 해결하기 위해, DDiff 는 디노이징 과정을 고정하는 재귀적 변수 $x_t$를 구성합니다. 이 업데이트는 "예측된 $x_0$"가 현재 주변 추정치 $x$인 DDIM 스타일의 역방향 확산 단계를 포함하며, 듀얼 변수 $u$는 신호 수준과 일관성을 유지하기 위해 재조정되어 추가됩니다.
   $$z \leftarrow \frac{1}{\sqrt{\bar{\alpha}_t}} (x_t + (1-\bar{\alpha}_t)s_\theta(x_t, t))$$
   여기서 $x_t$는 듀얼 변수 $u$를 통합한 DDIM 단계를 통해 업데이트됩니다.

3. 듀얼 업데이트: 듀얼 변수 $u$는 주변 변수 간의 잔차 오차를 누적하여 제약 조건 $x=z$를 강제하도록 업데이트됩니다:
   $$u \leftarrow u + x - z$$

잠재 공간 확장

이 프레임워크는 잠재 확산 모델 (LDMs) 과 함께 사용하기 위해 LatentDDiff로 확장되었습니다. 이 변형은 메모리 및 계산 요구 사항을 줄이기 위해 컴팩트한 잠재 표현에서 작동합니다. 데이터 일관성을 강제하기 위해 두 가지 모드를 지원합니다:

• Latent-DC: 디코더를 통한 역전파가 필요하며 잠재 공간에서 직접 일관성을 강제합니다.
• Pixel-DC: 잠재 변수를 픽셀 공간으로 디코딩한 후 데이터 일관성 보정을 적용하고 다시 인코딩합니다. 저자들은 Pixel-DC 가 일관되게 더 높은 품질의 재구성을 산출한다고 지적합니다.

주요 기여

1. 확산 사전 정보를 위한 듀얼 어센트 프레임워크: 이 논문은 ADMM 에서 영감을 받은 듀얼 어센트를 확산 모델과 통합하는 새로운 최적화 전략을 제시합니다. 이 접근법은 기존 사후 샘플링 방법에서 흔히 발생하는 근사법을 피합니다.
2. 고정점 수렴 증명: 저자들은 확산 기반 사후 최적화 방법에 대한 최초의 고정점 수렴 증명을 확립했습니다. 그들은 유계 디노이저와 경사 가정 하에, 볼록한 사전 정보를 가정하지 않더라도 주변 및 듀얼 변수 시퀀스가 고정점 $(x^\star, z^\star, u^\star)$으로 수렴함을 입증했습니다.
3. 향상된 강건성과 충실도: 이 방법은 최첨단 방법들에 비해 잔차가 0 에 더 가까워 관측치를 더 충실히 모델링하는 해를 생성하도록 설계되었습니다. 이는 생성적 사전 정보로 인해 발생하는 환각 가능성을 줄여줍니다.
4. 효율성: DDiff 는 (DAPS 와 같은 방법에서 요구되는) 보조 MCMC 평형화 단계를 제거하고 고품질 결과를 위해 필요한 신경망 함수 평가 횟수 (NFEs) 를 줄임으로써 더 빠른 추론 시간을 달성합니다.

실험 결과

저자들은 초해상도, 인페인팅, 가우시안 및 모션 디블러링, 위상 복원, 비선형 디블러링, 고동적 범위 (HDR) 영상 등 다양한 선형 및 비선형 역문제에 대해 DDiff 를 평가했습니다. 실험은 FFHQ 및 ImageNet 데이터셋에서 수행되었습니다.

• 정량적 성능: DDiff 는 대부분의 작업에서 최첨단 기준 (DAPS, DMPlug, DCDP, DPS, DiffPIR, DDRM 포함) 을 능가했습니다. 최대 신호 대 잡음비 (PSNR), 구조적 유사성 지수 (SSIM), 학습된 지각적 이미지 패치 유사성 (LPIPS) 에서 우수한 점수를 기록했습니다. 특히, 더 나은 데이터 일관성을 나타내는 가장 낮은 잔차 오차를 달성했습니다.
• 잡음에 대한 강건성: 이 방법은 높은 수준의 측정 잡음에 대해 훨씬 더 큰 강건성을 보여주었습니다. 일부 기준 방법들은 잡음이 증가함에 따라 PSNR 이 급격히 저하되는 반면, DDiff 는 안정적인 성능을 유지했습니다. 고잡음 시나리오 (예: $\sigma=0.3$인 위상 복원) 에서 DDiff 는 다른 방법들이 손상된 아티팩트를 생성한 곳에서 일관된 전역 구조와 의미를 복원했습니다.
• 효율성: 시간 효율성 측면에서 DDiff 는 DAPS 에 비해 더 적은 NFE 로 더 나은 지각적 품질을 달성했습니다. 선형 작업의 경우 약 절반의 샘플링 시간이, 비선형 작업의 경우 약 6 분의 1 의 시간이 동등하거나 더 우수한 결과를 얻는 데 필요했습니다.
• 애블레이션 연구: 실험은 듀얼 변수 업데이트와 특정 노이즈 단계 (논문 내 식 12) 가 모두 결정적임을 확인했습니다. 듀얼 변수를 제거하거나 (DiffPIR 을 모방) 노이즈 단계를 제거하면 고주파 아티팩트 및 데이터 일관성 측면에서 성능이 크게 저하되었습니다.

중요성 및 주장

이 논문은 DDiff 가 사후 샘플링 근사에서 견고한 듀얼 어센트 최적화 프레임워크로의 전환을 통해 확산 사전 정보를 활용한 역문제 해결에서 중요한 진전을 이루었다고 주장합니다. 저자들은 그들의 방법이 다음과 같음을 강조합니다:

• 관측치를 충실히 모델링: 잔차를 최소화함으로써 생성적 역문제 해결기에서 흔히 보이는 "환각" 효과를 줄입니다.
• 계산적으로 효율적: MCMC 단계의 병목 현상을 제거하고 경사 계산의 계산 부담을 줄입니다.
• 이론적 기반: 제공된 수렴 증명은 많은 기존 확산 기반 최적화 접근법에서 결여되어 있는 방법의 안정성에 대한 이론적 보장을 제공합니다.

저자들은 DDiff 가 저선량 CT 재구성처럼 측정 잡음이 두드러지는 시나리오를 포함한 대규모 복잡한 역문제에 대한 실용적인 해결책이라고 결론지었습니다. 향후 작업은 3D 볼륨 및 비디오와 같은 고차원 데이터로 프레임워크를 확장하는 것을 제안합니다.