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

技術的サマリー：拡散モデルを用いたゼロショット逆問題における段階的歪み - 知覚トラバース

問題定義
元の信号から劣化した観測値を復元することを目的とするベイズ逆問題は、本質的に歪み - 知覚（D-P）のトレードオフによって制約されている。歪み性能はサンプルごとの復元忠実度（例えば MSE や PSNR による）を測定し、知覚性能は実データとの視覚的品質および分布の類似性を反映する。理論的研究により、これらの目的は本質的に緊張関係にあることが確立されている。すなわち、一方を改善することはしばしば他方を劣化させることを必要とする。拡散モデルはタスク固有の再学習なしにゼロショット逆問題解決において最先端（SOTA）の結果を達成してきたが、既存の手法には推論時に D-P 曲線を原理的かつ計算効率的にトラバースするための戦略が欠けている。サンプリングステップの調整やサンプル平均化などの現在のアプローチは、しばしばヒューリスティックであるか計算コストが高く、拡散ベースのゼロショットアルゴリズムにおけるこのトレードオフを制御するメカニズムは十分に特徴付けられていない。

手法
著者は、単一の事前学習済み拡散モデルを用いて柔軟な D-P トラバースを可能にする、MAP-RPS（Maximum A Posteriori - Re-noised Posterior Sampling）と呼ばれる 2 段階のフレームワークを提案する。この手法は後に潜在空間へ拡張され、大規模な潜在拡散バックボーンを活用する LMAP-RPS として発展させられる。

1. 段階 1：低歪み初期化（MAP 推定）
   最初の段階は、D-P 曲線上の最適な歪み点に対応する最小平均二乗誤差（MMSE）解を近似する。著者は強く対数凹性の事後分布を分析し、MAP 推定量が有界な近似誤差を伴う MMSE 解の有効な代理となることを確立する。彼らは、特定のタイムステップ $t_1$ における事前学習済み拡散モデルのスコア関数を用いて対数事前分布の扱いやすい勾配を導出し、確率的勾配ベースの最適化による効率的な MAP 推定を可能にする。これにより、低歪みの出発点（$x_{MAP}$）が得られる。

2. 段階 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 トラバースのための計算効率的なメカニズムを提供する。反復サンプリングや複雑なハイパーパラメータ調整を必要とする手法とは異なり、このアプローチは 1 回の MAP 推定手順と 1 回の事後サンプリングプロセスのみを必要とする。
• 潜在空間への拡張: フレームワークを潜在拡散モデル（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 トラバースを実現するための新たな視点を提供すると主張する。著者は、この手法が既存のヒューリスティックな戦略に対する、原理的かつ理論的根拠があり、計算効率的な代替手段を提供することを強調する。歪みと知覚の目的を 2 つの明確な段階に分離することで、このフレームワークはモデルの再学習なしに、特定のアプリケーション要件に最も適合する復元を柔軟に選択することを可能にする。著者はまた、LMAP-RPS の実用的な可能性を強調し、現在の最先端アルゴリズムよりも低い計算複雑性で複雑な実世界の逆問題において SOTA 結果を提供する能力に言及している。

限界と今後の課題
著者は、理論的保証が強く対数凹性（ほぼ単峰性）の事後分布という仮定に依存していることを認めている。経験的結果は、より一般的な設定でもこの手法が有効であることを示唆しているが、複雑な多峰性事後分布における挙動についてはさらなる調査が必要である。さらに、現在の実装は段階 2 において DPS などの特定の事後サンプリングアルゴリズムに依存しており、これは推定誤差を導入する可能性がある。今後の研究では、D-P トラバースをさらに洗練させるために、漸近的に整合性のある事後サンプリング手法を統合することが考えられる。