[論文レビュー] A Variational Perspective on Solving Inverse Problems with Diffusion Models

Morteza Mardani, Jiaming Song+2 more|arXiv (Cornell University)|May 7, 2023

Generative Adversarial Networks and Image Synthesis被引用数 36

ひとこと要約

本論文は、拡散事前分布を用いた逆問題を解くための変分サンプル RED-diff を提案し、KL ベースの目的とすべての拡散タイムステップを活用するノイズ除去に基づく正則化を用いる。

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 拡散モデルベースの逆問題において、変分 KL フレームワークは x0 を与 Y に対する principled な後方推論を生み出せるか。
• RQ2 適切な重み付けを用いた全拡散タイムステップのデノイジング・スコアを組み込むことで、既存手法より復元品質が改善されるか。
• RQ3 スコア Jacobians を必要とせず、軽量なイテレートで効率的な確率的最適化としてサンプリングを定式化できるか。
• RQ4 RED-diff は線形および非線形の画像復元タスクで、最先端のサンプラーと比較してどう性能を示すか。
• RQ5 タイムステップの重み付けとサンプリング戦略が、忠実度と知覚品質に及ぼす実務的な影響は何か。

主な発見

• RED-diff は、PSNR、SSIM、KID、LPIPS、およびトップ1精度の観点で、画像インペインティングにおいて最先端のサンプラ(DPS、Pi-GDM、DDR,M)よりも優れている。
• 勾配のスコア Jacobian を用いずに、軽量イテレートを実現し、メモリ効率と GPU への適性を改善した。
• デノイジング SNR に基づく重み付け機構が、拡散初期ステップを上げ、後半のステップを下げて正則化を改善。
• 降順のタイムステップサンプリング（T から 1 へ）は、他のサンプリング戦略と比較して忠実度と知覚品質を向上させる。
• アブレーションでは、デノイザーの重み付けと時刻サンプリング戦略が、内容と細部のバランスおよび安定性の鍵となる。
• 非線形の逆問題（位相回復、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.