Stage-wise Distortion-Perception Traversal in Zero-shot Inverse Problems with Diffusion Models

기술 요약: 확산 모델을 활용한 제로샷 역문제에서의 단계별 왜곡–지각 탐색

문제 제기
원본 신호를 열화된 관측치로부터 재구성하는 것을 목표로 하는 베이지안 역문제는 근본적으로 왜곡–지각 (D–P) 트레이드오프에 의해 제약받습니다. 왜곡 성능은 샘플별 재구성 충실도 (예: MSE 또는 PSNR 을 통해) 를 측정하는 반면, 지각 성능은 시각적 품질과 실제 데이터에 대한 분포적 유사성을 반영합니다. 이론적 연구는 이러한 목표들이 본질적으로 긴장 관계에 있음을 규명했습니다. 즉, 한쪽을 개선하면 종종 다른 쪽이 저하됩니다. 확산 모델은 작업별 재학습 없이 제로샷 역문제 해결에서 최첨단 (SOTA) 결과를 달성해 왔으나, 기존 방법들은 추론 시 D–P 곡선을 탐색하기 위한 원칙적이고 계산 효율적인 전략이 부족합니다. 샘플링 단계 조정이나 샘플 평균화와 같은 현재 접근법들은 종종 휴리스틱이거나 계산 비용이 많이 들며, 확산 기반 제로샷 알고리즘에서 이 트레이드오프를 제어하는 메커니즘은 여전히 충분히 규명되지 않았습니다.

방법론
저자들은 단일 사전 훈련된 확산 모델을 사용하여 유연한 D–P 탐색을 가능하게 하는 MAP-RPS(최대 사후 확률 - 재노이즈 사후 샘플링) 라는 2 단계 프레임워크를 제안합니다. 이 방법은 이후 잠재 공간으로 확장되어 대규모 잠재 확산 백본을 활용하는 LMAP-RPS로 발전되었습니다.

1. 1 단계: 낮은 왜곡 초기화 (MAP 추정)
   첫 번째 단계는 D–P 곡선상의 최적 왜곡 지점에 해당하는 최소 평균 제곱 오차 (MMSE) 해를 근사합니다. 저자들은 강하게 로그 오목한 (strongly log-concave) 사후 분포를 분석하여, MAP 추정기가 유계 근사 오차를 가진 MMSE 해의 유효한 대리 변수가 됨을 입증했습니다. 그들은 특정 시간 단계 $t_1$에서 사전 훈련된 확산 모델의 스코어 함수를 사용하여 로그 사전 확률에 대한 실용적인 기울기를 유도하여, 확률적 경사 기반 최적화를 통한 효율적인 MAP 추정을 가능하게 합니다. 이를 통해 낮은 왜곡의 시작점 ($x_{MAP}$) 이 도출됩니다.

2. 2 단계: 지각 품질 개선 (재노이즈 사후 샘플링)
   두 번째 단계는 점진적으로 지각 품질을 향상시킴으로써 D–P 곡선을 탐색합니다. MAP 추정은 지정된 시간 단계 $t_0$까지 재노이즈되며, 이 노이즈가 추가된 초기화 상태에서 사후 샘플링이 수행됩니다. 이론적으로 저자들은 생성된 분포와 실제 데이터 분포 간의 지각 오차 지표인 Wasserstein-2 ($W_2$) 거리가 재노이즈 시간 단계 $t_0$에 의해 제어됨을 보여줍니다. $t_0$를 변화시킴으로써, 이 방법은 낮은 왜곡의 MAP 해 ($t_0=0$) 와 표준 사후 샘플링 ($t_0=T$) 사이의 연속적인 보간을 가능하게 합니다.

주요 기여

• 이론적 프레임워크: 이 논문은 두 단계 모두에 대한 이론적 분석을 제공합니다. 강하게 로그 오목한 가정 하에서 MAP 추정기로 MMSE 해를 근사할 때의 오차 상한을 확립하고, 재노이즈 사후 샘플링 단계에 대한 $W_2$ 거리의 상한을 유도하여 지각 오차가 $t_0$ 조정을 통해 제어될 수 있음을 증명합니다.
• 알고리즘 설계: MAP-RPS 와 그 잠재 공간 변형인 LMAP-RPS 의 도입은 D–P 탐색을 위한 계산 효율적인 메커니즘을 제공합니다. 반복 샘플링이나 복잡한 하이퍼파라미터 튜닝이 필요한 기존 방법들과 달리, 이 접근법은 하나의 MAP 추정 절차와 하나의 사후 샘플링 과정만 요구합니다.
• 잠재 공간 확장: 프레임워크를 잠재 확산 모델 (LMAP-RPS) 로 확장함으로써, 저자들은 실제 세계의 고해상도 작업에 대한 보다 광범위한 적용성을 위해 강력한 사전 훈련된 백본 (예: Stable Diffusion) 을 활용합니다.

실험 결과
FFHQ 데이터셋 (D–P 곡선 분석용) 과 MS-COCO 데이터셋 (실제 세계에 가까운 역문제용) 에서 광범위한 실험이 수행되었습니다.

• D–P 탐색: FFHQ 작업 (잡음 제거, 초해상도, 인페인팅, 디블러링) 에서 MAP-RPS 와 LMAP-RPS 모두 기존 방법들 (VSDPS 및 PSCGAN 등) 보다 이상적인 왼쪽 아래 모서리 (최적의 왜곡과 지각) 에 더 가까운 곡선을 일관되게 달성했습니다.
• 실제 세계 성능: MS-COCO 에서 LMAP-RPS 는 9 가지 최근 잠재 확산 기반 베이스라인과 비교하여 거의 모든 작업과 지표 (PSNR, SSIM, LPIPS, FID) 에서 SOTA 성능을 달성했습니다.
• 효율성: LMAP-RPS 는 상당한 계산적 이점을 보여주었습니다. $t_0=0$일 때 모든 베이스라인보다 훨씬 빨랐으며, $t_0=600$에서도 여전히 계산 효율적이었으며, 종종 우수한 재구성 품질을 유지하면서 다른 방법들보다 속도가 더 빨랐습니다.

의의 및 주장
이 논문은 MAP-RPS 와 LMAP-RPS 가 제로샷 역문제에서 D–P 탐색을 실현하기 위한 새로운 관점을 제공한다고 주장합니다. 저자들은 그들의 방법이 기존 휴리스틱 전략에 대한 원칙적이고 이론적으로 근거하며 계산 효율적인 대안을 제공한다고 강조합니다. 왜곡과 지각 목표를 두 개의 명확한 단계로 분리함으로써, 이 프레임워크는 모델을 재학습하지 않고도 특정 응용 요구 사항에 가장 잘 맞는 재구성을 유연하게 선택할 수 있게 합니다. 저자들은 또한 LMAP-RPS 의 강력한 실용적 잠재력을 강조하며, 현재 최첨단 알고리즘보다 낮은 계산 복잡도로 복잡한 실제 세계 역문제에서 SOTA 결과를 제공할 수 있음을 지적합니다.

한계 및 향후 과제
저자들은 그들의 이론적 보장이 강하게 로그 오목한 (거의 단일 모드인) 사후 분포에 대한 가정에 의존함을 인정합니다. 경험적 결과는 더 일반적인 설정에서도 이 방법이 효과적임을 시사하지만, 복잡한 다중 모드 사후 분포 하에서의 거동은 추가적인 조사가 필요합니다. 또한, 현재 구현은 2 단계에 특정 사후 샘플링 알고리즘 (DPS 등) 에 의존하며, 이는 추정 오차를 유발할 수 있습니다. 향후 연구는 D–P 탐색을 더 정교하게 만들기 위해 점근적으로 일관된 사후 샘플링 방법을 통합할 수 있습니다.