확산 모델을 활용한 역문제 해결에 대한 변분적 관점

A Variational Perspective on Solving Inverse Problems with Diffusion Models

May 7, 2023

저자: Morteza Mardani, Jiaming Song, Jan Kautz, Arash Vahdat

cs.AI

초록

디퓨전 모델은 시각 영역에서의 기초 모델(foundation model)의 핵심 요소로 부상했습니다. 이 모델의 중요한 응용 분야 중 하나는 각 작업별로 재학습 없이 단일 디퓨전 사전(prior)을 통해 다양한 하위 역작업(inverse task)을 보편적으로 해결하는 것입니다. 대부분의 역작업은 측정값(예: 마스킹된 이미지)이 주어졌을 때 데이터(예: 전체 이미지)에 대한 사후 분포(posterior distribution)를 추론하는 문제로 공식화될 수 있습니다. 그러나 디퓨전 과정의 비선형성과 반복적 특성으로 인해 디퓨전 모델에서 사후 분포를 다루는 것은 매우 어려운 과제입니다. 이러한 문제를 해결하기 위해, 우리는 설계상 진정한 사후 분포를 근사화하는 변분적 접근법을 제안합니다. 우리의 접근법은 자연스럽게 디노이징 디퓨전 과정을 통한 정규화(RED-Diff)로 이어지며, 여기서 서로 다른 시간 단계의 디노이저(denoiser)들이 이미지에 대해 다양한 구조적 제약을 동시에 부과합니다. 서로 다른 시간 단계의 디노이저들의 기여도를 평가하기 위해, 우리는 신호 대 잡음비(SNR)를 기반으로 한 가중치 메커니즘을 제안합니다. 우리의 접근법은 디퓨전 모델을 사용한 역문제 해결을 위한 새로운 변분적 관점을 제공하며, 샘플링을 확률적 최적화 문제로 공식화할 수 있게 합니다. 이를 통해 경량화된 반복 과정으로 기존의 표준 솔버를 간단히 적용할 수 있습니다. 인페인팅(inpainting) 및 초해상도(superresolution)와 같은 이미지 복원 작업에 대한 실험 결과는 우리의 방법이 최신 샘플링 기반 디퓨전 모델들과 비교하여 우수한 성능을 보임을 입증합니다.

English

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.