技術的概要：拡散モデルのための局所 MAP サンプリング

問題定義

拡散事後サンプリング（DPS）は、事後分布 $p(x_0 | y)$ からサンプリングすることにより逆問題を解決する強力な枠組みとして登場し、原理的な不確実性の定量化と多様性を提供しています。しかし、画像処理や科学分野の多くの実用的な応用では、完全な分布サンプリングよりも正確な点推定が優先されます。$\arg \max p(x_0 | y)$ と定義される最大事後確率（MAP）推定量は、依然として標準的な再構成目的関数です。

既存の最適化ベースの手法（Resample、DiffPIR、DCDP など）は、ノイズ除去ステップと最適化ステップを交互に行うことで再構成性能の向上を試みています。しかし、これらの反復手順、グローバル MAP 解、および事後サンプリングとの理論的関連性は依然として不明確です。これらの手法がグローバル MAP に収束するかどうか、あるいは DPS と整合性があるかどうかは確立されていません。さらに、これらの手法における最適化ステップの具体的な性質、すなわちそれらがグローバル目的関数か局所目的関数を解くのかについては、明確化が必要です。

手法：局所 MAP サンプリング（LMAPS）

著者は、拡散軌道上で局所的な最大事後確率部分問題を反復的に解く新しい推論フレームワークである**局所 MAP サンプリング（LMAPS）**を導入します。

1. 理論的枠組み

LMAPS は、単一のグローバルな問題ではなく、一連の局所 MAP 問題を解くことで動作します。各拡散ステップ $t$ において、ノイズ状態 $x_t$ と測定値 $y$ が与えられたとき、この手法は以下を求めます：
$$x^*_0(t, x_t, y) := \arg \max_{x_0} p(x_0 | x_t, y)$$
これに続いて、標準的な DDIM 遷移ステップが行われます：
$$x_{t-\Delta t} = g(x^*_0, x_t, \epsilon)$$

既存の手法との違い：

• グローバル MAP 対比： グローバル MAP は完全な事後分布 $p(x_0 | y)$ のモードを求めます。一方、LMAPS は各ステップにおいて条件付き事後分布 $p(x_0 | x_t, y)$ のモードを求めます。著者らは、局所最大化者の系列が必ずしも $p(x_0 | y)$ のグローバル最大化者に収束するわけではないと主張します。
• DPS 対比： DPS は更新において条件付き平均 $E[x_0 | x_t, y]$ を使用します。LMAPS はこの平均を条件付きモードに置き換えます。これらは条件付き事後分布がガウス分布である場合のみ同等であり、そうでない場合、LMAPS はモード探索バイアスを導入し、モード間の低密度領域を回避します。
• DMAP 対比： DMAP も局所 MAP の視点を利用しますが、潜在拡散変数空間（$x^*_{t-\Delta t} = \arg \max p(x_{t-\Delta t} | x_t, y)$）で最適化を行います。LMAPS はクリーン信号空間（$x^*_0$）で直接問題を定式化し、測定尤度 $p(y|x_0)$ と明示的に関連付ける補完的な視点を提供します。

2. 実用的アルゴリズムと近似

一般的な非線形逆問題に対して局所 MAP 問題を扱いやすくするため、著者は特定の近似と再定式化を提案します。

• 共分散近似： 著者は、ガウス事前分布仮説に基づく等方性形式を用いて、事前分布 $p(x_0 | x_t)$ の条件付き共分散 $\Sigma_{0|t}$ を近似します。ノイズレベル $\sigma_t \to 0$ になると、事後共分散は漸近的に等方性になります。彼らは $\Sigma_{0|t} \approx \frac{1}{SNR} I$ と近似し、ここで $SNR = \alpha_t^2 / \sigma_t^2$ です。

• 目的関数の再定式化： 局所 MAP 目的関数は、事前距離とデータ忠実度の加重和を最小化する形式で定式化されます：
  $$x^*_0 = \arg \min_{x_0} \left\{ \frac{SNR}{k} \|x_0 - m_{0|t}\|^2 + \frac{1}{\sigma_y^2} \|y - H(x_0)\|^2 \right\}$$
  数値的安定性と解釈可能性を向上させるため、これをスケール不変な凸結合に再定式化します：
  $$x^*_0 = \arg \min_{x_0} \left\{ (1 - \mu_t) \frac{1}{2} \|x_0 - m_{0|t}\|^2 + \mu_t k_2 \|y - H(x_0)\|^2 \right\}$$
  ここで $\mu_t = \sigma_t^2 / (\sigma_t^2 + k_1^2)$ です。この定式化により自動アニーリングが可能となり、$t$ が減少する（$\sigma_t \to 0$）につれて、目的関数は測定駆動型から事前駆動型へとシフトします。

• 線形対非線形： 線形逆問題の場合、局所 MAP 解は閉形式で導出でき、Tweedie モーメント射影拡散（TMPD）を特殊ケースとして回復します。非線形問題の場合、解は勾配降下法によって得られます。

主要な貢献

1. 理論的統合： 本論文は、最適化ベースの逆問題手法と Tweedie モーメント射影拡散（TMPD）を単一の確率的枠組みの下で統合するように LMAPS を定式化しました。最適化ベースの手法は本質的にグローバル MAP や完全な事後サンプリングではなく、局所 MAP 問題を解いていることを明確にしました。
2. 方法論的進展： 著者は、ガウス分布に基づく共分散近似と、自動アニーリングを通じて数値的安定性と解釈可能性を高める再定式化された目的関数を導入しました。
3. 実証的検証： このフレームワークは、10 の画像復元タスク（線形、非線形、微分不可能な演算子を含む）および 3 つの科学逆問題で検証されました。

実験結果

著者は、画像復元に対して FFHQ および ImageNet データセットを、科学逆問題に対してベンチマーク（線形逆散乱、圧縮センシング MRI、ブラックホール画像）を用いて LMAPS を評価しました。

• 画像復元： LMAPS は、FFHQ および ImageNet における PSNR、SSIM、LPIPS 指標の 60 のケースのうち 43 で最先端の性能を達成しました。これは、運動ブラー除去、JPEG 復元、量子化などのタスクにおいて、2 dB 以上の PSNR 改善を示し、密接に関連する DMAP ベースラインを一貫して上回りました。
• 科学逆問題： LMAPS は、すべての科学タスクで最も高い PSNR を一貫して達成しました。特に、以前の手法と比較して 3 つの線形逆散乱（LIS）タスクで 1.5 dB 以上の PSNR 向上を達成しました。
• 効率性： LMAPS は DAPS に比べて優れた効率性を示し、より少ない計算リソース（例えば、拡散ステップあたりの最適化ステップ数の削減）でより良い結果を達成しました。

意義と主張

本論文は、LMAPS が最適化ベースの手法に対する原理的な確率的解釈を提供し、事後サンプリング（DPS）と点推定（MAP）の間のギャップを埋めると主張しています。局所 MAP 部分問題を明示的に解くことで、この手法は既存の最適化ベースソルバーの動作を説明する統合された視点を提供します。

著者は、LMAPS が完全な事後分布からサンプリングするものではなく（したがって DPS と同様に多様性を捉えるものではない）、正確な点推定が主要な目標であるシナリオのために特別に設計されていることを強調しています。このフレームワークは、再構成忠実度と拡散モデルの構造的事前分布のバランスを取る堅牢な代替案を提供し、標準的な画像復元および困難な科学再構成タスクの両方で最先端の結果を達成しています。この研究は、拡散事前分布を用いたグローバル MAP 推定の有用性をさらに探求する今後の研究の可能性を示唆しており、それは依然として未解決の課題です。