[논문 리뷰] A Variational Perspective on Solving Inverse Problems with Diffusion Models

Morteza Mardani, Jiaming Song외 2명|arXiv (Cornell University)|2023. 05. 07.

Generative Adversarial Networks and Image Synthesis인용 수 36

한 줄 요약

본 논문은 모든 확산 타임스텝을 활용하는 KL 기반 목표와 노이즈 제거 기반 정규화를 사용하여 확산 사전분포를 갖는 역문제 해결을 위한 변분 샘플러 RED-diff를 제시한다.

ABSTRACT

Diffusion models have emerged as a key pillar of foundation models in visual domains. One of their critical applications is to universally solve different downstream inverse tasks via a single diffusion prior without re-training for each task. Most inverse tasks can be formulated as inferring a posterior distribution over data (e.g., a full image) given a measurement (e.g., a masked image). This is however challenging in diffusion models since the nonlinear and iterative nature of the diffusion process renders the posterior intractable. To cope with this challenge, we propose a variational approach that by design seeks to approximate the true posterior distribution. We show that our approach naturally leads to regularization by denoising diffusion process (RED-Diff) where denoisers at different timesteps concurrently impose different structural constraints over the image. To gauge the contribution of denoisers from different timesteps, we propose a weighting mechanism based on signal-to-noise-ratio (SNR). Our approach provides a new variational perspective for solving inverse problems with diffusion models, allowing us to formulate sampling as stochastic optimization, where one can simply apply off-the-shelf solvers with lightweight iterates. Our experiments for image restoration tasks such as inpainting and superresolution demonstrate the strengths of our method compared with state-of-the-art sampling-based diffusion models.

연구 동기 및 목표

• 확산 사전을 이용하여 다양한 역문제의 보편적이고 학습 없이 해결을 촉진한다.
• 가정된 p(x0|y)를 근사하기 위해 변분 추론을 도입하여 불가피하게 계산이 어려운 후방 분포의 문제를 극복한다.
• 확산 타임스텝 전반에 걸친 노이즈 제거기를 활용하는 노이즈 제거 기반 정규화(RED-diff) 관점을 개발한다.
• 경량 반복을 갖는 확률적 최적화를 통한 실용적이고 효율적인 샘플링 메커니즘을 제공한다.

제안 방법

• q(x0|y)와 p(x0|y) 간의 KL 최소화로 후방 추론을 형식화한다.
• 확산 사전을 통해 점수 매칭 정규화 항을 얻는 변분 경계식을 도출한다.
• 여러 확산 단계에서의 노이즈 제거를 사용하는 가중된 점수 매칭 목적함수를 도입한다 (RED-diff).
• 적절한 가중치 함수를 사용할 때 정규화 항의 기울기를 점수 네트워크를 역전파하지 않고도 계산할 수 있음을 보인다.
• 재구성 항과 정규화 항을 결합한 단순하고 합리적으로 계산 가능한 손실을 갖는 mu에 대한 확률적 최적화로 샘플링을 근사한다.
• 이 접근법을 REGULARIZATION-BY-DENOISING (RED) 프레임워크와 연결하고 최적화상의 이점을 논의한다.

Figure 1: The schematic diagram of our proposed variational sampler (RED-diff). The forward denoising diffusion process gradually adds noise to the estimate $\mu$ . The denoisers of the backward diffusion process apply score-matching regularization to the measurement matching loss. The refined estim

실험 결과

연구 질문

• RQ1 확산 모델 기반 역문제에서 y가 주어졌을 때 x0에 대한 원칙적 후방 추론을 변분 KL 프레임워크가 제시할 수 있는가?
• RQ2 적절한 가중치를 가진 모든 확산 타임스텝의 노이즈 제거 스코어를 통합하는 것이 기존 방법들보다 복원 품질을 향상시키는가?
• RQ3 샘플링을 점수 야코비안(score Jacobians)을 필요로 하지 않는 경량 반복의 효율적 확률적 최적화로 공식화할 수 있는가?
• RQ4 최신 샘플러와 비교했을 때 RED-diff가 선형 및 비선형 이미지 복구 작업에서 어떤 성능을 보이는가?
• RQ5 타임스텝 가중치와 샘플링 전략이 충실도와 지각 품질에 미치는 실질적 영향은 무엇인가?

주요 결과

• RED-diff는 PSNR, SSIM, KID, LPIPS 및 top-1 정확도 측면에서 이미지 인페이팅에서 최첨단 샘플러(DPS, Pi-GDM, DDRM)를 능가한다.
• 이 방법은 점수 야코비안 없이 경량 반복으로 메모리 효율과 GPU 친화성을 개선한다.
• 노이즈 제거 SNR에 기반한 가중 메커니즘이 초기 확산 단계를 상향, 이후 단계를 하향 가중해 더 나은 정규화를 제공한다.
• T에서 1로의 내림차 timestep 샘플링은 다른 샘플링 전략에 비해 충실도와 지각 품질을 향상시킨다.
• Ablation 연구는 노이즈 제거기 가중치와 타임샘플링 전략이 콘텐츠와 디테일의 균형 및 안정성을 조절하는 핵심 매개변수임을 보여준다.
• 비선형 역문제(위상 회수, HDR, 디블러링 등)에서 RED-diff는 DPS에 비해 우수한 지표를 보여주고, 적용 가능한 경우 Pi-GDM 및 DDRM보다 우월하게 성능을 발휘한다.

Figure 2: Comparison of the proposed variational sampler with alternatives for inpainting representative ImageNet examples. Each sampler is tuned for the best performance.