[論文レビュー] Solving Linear Inverse Problems Provably via Posterior Sampling with Latent Diffusion Models

Litu Rout, Negin Raoof+4 more|arXiv (Cornell University)|Jul 2, 2023

Advanced Neuroimaging Techniques and Applications被引用数 5

ひとこと要約

本稿では、微調整を伴わずに事前学習済み潜在拡散モデルを用いて線形逆問題を解く最初のフレームワークである、潜在拡散モデルを用いた事後分布サンプリング（PSLD）を提案する。VAEボトルネックを通じて再構成忠実度を維持する追加の勾配更新をDPSアルゴリズムに拡張することで、PSLDは証明可能なサンプル回復を達成し、補完、ノイズ除去、ぼかし除去、超解像、ストライプ除去の各タスクで最先端の性能を示す。定量的・定性的な両方の指標において、先行手法を上回っている。

ABSTRACT

We present the first framework to solve linear inverse problems leveraging pre-trained latent diffusion models. Previously proposed algorithms (such as DPS and DDRM) only apply to pixel-space diffusion models. We theoretically analyze our algorithm showing provable sample recovery in a linear model setting. The algorithmic insight obtained from our analysis extends to more general settings often considered in practice. Experimentally, we outperform previously proposed posterior sampling algorithms in a wide variety of problems including random inpainting, block inpainting, denoising, deblurring, destriping, and super-resolution.

研究の動機と目的

• Stable Diffusionのような事前学習済み潜在拡散モデルを活用できない既存の逆問題ソルバのギャップを埋めること。
• タスク固有の微調整を回避して、潜在拡散モデルを用いて線形逆問題のための事後分布サンプリングを可能にする手法を開発すること。
• 2段階の拡散プロセスを想定した線形モデル設定下で、アルゴリズムの理論的分析を行い、証明可能なサンプル回復を確立すること。
• 補完、ノイズ除去、ぼかし除去、超解像、ストライプ除去を含む多様な逆問題において、実験的に最先端の性能を示すこと。

提案手法

• VAEボトルネックを通じて再構成忠実度を維持する追加の勾配更新ステップをDPSアルゴリズムに導入し、符号化器-復号器マップが概ね恒等写像となるような潜在変数のサンプリングを促進する。
• 潜在空間における再構成誤差を最小化するために、VAEの符号化器と復号器を活用するグリューリング目的関数を用い、サンプルされた潜在変数が正確に再構成可能であることを保証する。
• 事前学習済み潜在拡散モデルの無条件スコア関数と、線形劣化演算子から導出される測定値整合性項を組み合わせた、変更された逆SDEを適用する。
• 2段階の拡散プロセスを採用：まず、変更された逆SDEを用いて潜在空間でサンプリングを行い、次にピクセル空間に復号する。グリューリング目的関数が忠実度を保証する。
• 微調整や再学習の必要がないため、事前学習済みモデルの潜在空間内でのみ動作する。
• 理論的分析により、線形モデル設定下でサンプル回復が証明可能であり、スコアベース生成モデルの非漸近的解析から収束保証が導かれる。

実験結果

リサーチクエスチョン

• RQ1 事前学習済み潜在拡散モデルを用いた事後分布サンプリングは、生成モデルの微調整を伴わず線形逆問題を解くのに利用可能か？
• RQ2 逆問題における事後分布サンプリング中に、潜在拡散モデルの再構成忠実度をどのように維持できるか？
• RQ3 線形逆問題設定下で、潜在拡散モデルを用いた事後分布サンプリングに対して、どのような理論的保証を確立できるか？
• RQ4 提案手法は、DPS や DDRM といった既存の事後分布サンプリングアルゴリズムを、多様な逆問題において上回ることができるか？
• RQ5 潜在空間にグリューリング目的関数を組み込むことで、標準的な事後分布サンプリングと比較して再構成品質がどのように向上するか？

主な発見

• FFHQ 256検証セットにおけるランダム補完タスクで、PSLDはPSNR 30.31、SSIM 0.851を達成し、DPS（PSNR 25.23、SSIM 0.851）を顕著に上回った。
• 4倍超解像において、PSLDはPSNR 30.73、SSIM 0.867を達成し、DPS（PSNR 25.67、SSIM 0.852）とDDRM（PSNR 25.36、SSIM 0.835）を上回った。
• 分布外のImageNetサンプルにおけるガウスぼかし除去では、PSLDはPSNR 30.10、SSIM 0.843を達成し、DPSの24.25（PSNR）、0.811（SSIM）を上回った。
• 運動ぼかし除去では、LPIPSがDPSの0.613からPSLDの0.244に低下し、より優れた知覚的品質を示した。
• ストライプ除去では、水平方向でLPIPS 0.244、垂直方向で0.255を達成し、DPSの0.613（水平）と0.597（垂直）を大きく下回った。
• 理論的分析により、2段階の拡散プロセスを用いた線形モデル設定下で、スコア推定誤差が有界である条件下で、証明可能なサンプル回復が確認された。