[論文レビュー] Learning Diffusion Priors from Observations by Expectation Maximization

François Rozet, Gérôme Andry+2 more|arXiv (Cornell University)|May 22, 2024

Gaussian Processes and Bayesian Inference被引用数 5

ひとこと要約

この論文は、不完全/ノイズのある観測からベイズ逆問題の事前分布として拡散モデルを学習し、EMフレームワークを用いて無条件拡散モデルのモーメント一致後方サンプリング方式を導入する。

ABSTRACT

Diffusion models recently proved to be remarkable priors for Bayesian inverse problems. However, training these models typically requires access to large amounts of clean data, which could prove difficult in some settings. In this work, we present DiEM, a novel method based on the expectation-maximization algorithm for training diffusion models from incomplete and noisy observations only. Unlike previous works, DiEM leads to proper diffusion models, which is crucial for downstream tasks. As part of our methods, we propose and motivate an improved posterior sampling scheme for unconditional diffusion models. We present empirical evidence supporting the effectiveness of our approach.

研究の動機と目的

• ノイズありの観測しか得られない場合に、ベイズ逆問題のための事前分布を学習する動機づけ。
• 不完全データから拡散モデルを学習するためのEMベースのパイプラインを開発。
• 得られた拡散前処理が下流のベイズタスクに適切で使用可能であることを保証。
• 無条件拡散モデルの新規後方サンプリング方式を提供し、サンプリング品質を向上。

提案手法

• 連続時間のSDEフレームワークで拡散モデルを定式化し、デノイジングスコアマッチング目的を用いる。
• Monte Carlo EMを用いて事前分布 q_theta(x) を学習し、 posterior q_theta(x|y) から反復的にサンプリングして再学習する。
• 観測特有のフォワードを扱う前方モデル p(y|x,A)=N(y|Ax,Sigma_y) を定義。
• モーメント一致後方サンプリング（MMPS）を導入し、後方スコアを grad log p(x_t) + grad log p(y|x_t) に分解して用いる。
• 尤度スコアは共分散 V[x|x_t] を考慮したガウス近似で推定し、Tweedeieの共分散公式を用いる。
• V[x|x_t] を Tweedieの公式で効率的に計算し、線形系を共役勾配法で解く際に大しい行列を作成せずに済ませる。
• 収束を速めるために EM で識別されたガウス事前分布で q_0(x) を初期化。
• デノイザー網 d_theta(x_t,t) を用いたデノイジングスコアマッチングとして Eq. (4) の形で拡散モデル q_theta(x) を訓練。

実験結果

リサーチクエスチョン

• RQ1 EMフレームワークを用いて不完全/ノイズのある観測から拡散モデルを事前分布として学習できるか。
• RQ2 MMPS は経験ベイズ設定において無条件拡散モデルの後方サンプルを安定かつ正確に提供できるか。
• RQ3 実世界の破損データタスク（例：破損 CIFAR-10、加速MRI）に対して、既存手法と比較して EM ベースの拡散前処理はどう機能するか。
• RQ4 後方共分散 V[x|x_t] の正確なモデリングが後方品質と下流推論にどのような影響を与えるか。

主な発見

手法	Deleted	FID ↓	IS ↑
0.4	0.4	18.85	7.45
AmbientDiffusion [77]	0.6	28.88	6.88
0.8	0.8	46.27	6.14
Ours w/ Tweedie	0.75	19.56	7.60
Ours w/ (I+Σ_t^{-1})^{-1}	0.75	49.45	6.99
Ours w/ Σ_t	0.75	125.58	3.98

• EMベースのパイプラインは、制御実験下でground-truth latent distributions に近い分布へ収束する拡散前処理を得られる。
• Tweedie共分散を用いたモーメント一致後方サンプリング（MMPS）は、ヒューリスティック共分散よりも後方 q(x_t|y) をより正確に推定し、安定性とカバレッジを向上させる。
• 破損 CIFAR-10 では、Tweedie共分散を用いた提案手法が同等の破損レベルで AmbientDiffusion を上回るか同等で、ヒューリスティック共分散は性能が低い。
• 加速MRI では、拡散前処理によって高い加速因子（最大 R=32）で妥当な再構成と詳細な後方サンプルが得られる。
• EMで識別されたガウス事前分布による初期化は収束までの反復回数を削減し、効率性を向上。
• A V[x|x_t] A^T を扱うのに共役勾配ソルバを用いることで高次元でも計算を実現可能にし、わずか 1-3 回の CG 反復で利益を得られる。