ShuffleFlow: Scalable Posterior Inference for Bayesian Inverse Imaging

技術要約：ShuffleFlow

問題提起
ベイズ逆問題イメージングは、ノイズを含む不完全な観測値 $y = A(x) + n$ から、潜在的な画像 $x$ の事後分布 $p(x|y)$ を推定することで、潜在画像を復元することを目的としている。変分推論（Variational Inference; VI）は、最適化後に効率的なサンプリングを可能にする原理的な事後分布推論の手法であるが、大規模な画像再構成への適用においては、スケーラビリティの低さが課題となる。フローベースのネットワークに依存する従来のVI手法は、通常、画像分布全体を直接モデル化するため、計算コストが画像の次元に対して極めて悪く（例：$N \times N$ ピクセルの場合 $O(N^4)$）、最適化におけるメモリおよび時間の予算を圧迫する。パッチ分割やピクセル化といった素朴な戦略による既存の改善試行は、多くの場合、画像の品質や事後分布の表現力を犠牲にしている。対照的に、拡散モデルは強力ではあるものの、学習に大規模なデータセットを必要とし、反復的なサンプリングの性質上、高解像度かつ大量のサンプルによる事後分布の生成において高い計算コストを要する。

手法：ShuffleFlow
著者らは、事後分布モデリングのタスクを「座標サンプリング」、「特徴量エンコーディング」、「事後分布推定」の3つのモジュール成分に分解するように設計された、スケーラブルなVIフレームワークであるShuffleFlowを提案する。

1. 座標サンプリング（ピクセル・アンシャッフル）： 全体の $N \times N$ 画像や独立したパッチをモデル化する代わりに、ShuffleFlowはダウンサンプリング係数 $S$ を用いたピクセル・アンシャッフル操作により、画像を $S^2$ 個の相関のあるサブ画像のスタックへと分割する。各サブ画像は $(N/S) \times (N/S)$ の次元を持つ。これにより、問題の次元数は $N^2$ から $N^2/S^2$ へと削減される。
2. 特徴量エンコーディング（ニューラルフィールド）： フローネットワークに全画像の再構成タスクを強いることなく空間構造を捉えるため、本フレームワークはニューラルフィールド（NF）エンコーダ（具体的には座標ベースのMLPであるSIREN）を採用している。サンプリング位置の2D座標オフセットがNFに入力され、各サブ画像に対する特徴量埋め込み（$h_i$）が生成される。
3. 事後分布推定（条件付き正規化フロー）： 共有された条件付き正規化フロー（Conditional Normalizing Flow; CNF）が、サブ画像スタックの結合分布をモデル化する。CNFは、各サブ画像の場所に対応する特徴量埋め込み（$h_i$）によって条件付けられる。決定的な点として、本フレームワークはド・フィネッティの定理に着想を得たアプローチを採用し、すべてのサブ画像に対して単一の基底潜在変数 $z_0$ を共有する。これにより、結合分布を以下のように因数分解できる：
   $$q(x) = \left( \prod_{i=1}^{S^2} q(x_i | z_0; c_i) \right) q(z_0)$$
   ここで $z_0$ は標準正規分布から抽出される。この共有メカニズムは、サブ画像間の相関を維持してピクセルの独立性に起因するノイズアーティファクトを防ぎつつ、フローネットワークのサイズを $O(N^4)$ から $O(N^4/S^4)$ へと大幅に縮小させる。

最適化と推論
CNFとNFは、変分分布とベイズ事後分布の間のカルバック・ライブラー（KL）ダイバージェンスを最小化することにより、共同で訓練される。損失関数は、データ尤度、事前分布（古典的な全変動事前分布または重みを固定した学習済みスコアベース事前分布）、および多様性を促進するための負のエントロピー項の3つの項で構成される。推論時には、NFがサンプリング座標の埋め込みを生成し、CNFがこれらの埋め込みと共有された基底変数に基づいてサンプルを生成する。これらはその後、ピクセル・シャッフルを通じて完全な画像へと再構成される。

主な貢献

1. 問題の分解： 著者らは、大規模な画像事後分布モデリングを、座標サンプリング、部分的な画像表現、および軽量な不確実性定量化へと分解する、新しいモジュール型VIフレームワークを導入した。この構造は特定のネットワークアーキテクチャに依存せず、ニューラルフィールドや生成的事前分布の進展を統合することを可能にする。
2. スケーラビリティと効率性： ピクセル・アンシャッフルと共有潜在変数を利用することで、ShuffleFlowはモデルサイズと計算リソースの大幅な削減を実現する。理論的および経験的分析によれば、本手法はピクセル単位の推論と同等の効率でスケールし、$256 \times 256$ の画像（$S=8$ の場合）に対して、フル画像フローモデルと比較してフローネットワークのサイズを4096倍削減する。
3. 高解像度事後分布の生成： 本手法は、高サンプル数の事後分布を迅速に生成する能力を示す。サンプリング時間がサンプル数に対して線形に増加する拡散サンプラーとは異なり、ShuffleFlowは短い訓練フェーズの後、数千のサンプルを効率的に生成できる。

実験結果
著者らは、線形問題（モーションデブラー、圧縮センシングMRI）および非線形問題（フーリエ位相回復）の両方でShuffleFlowを評価している。

• 線形問題： ShuffleFlowは、画像品質（PSNR, SSIM, LPIPS）および不確実性の較正（NLL, ECE）の観点において、VI手法の中で最先端（SOTA）の性能を達成し、従来のフローベースVI手法（DPI）や素朴なパッチ分割戦略を上回った。画像品質においては主要な拡散サンプラー（例：DAPS）に匹敵する性能を示しながら、計算オーバーヘッドは大幅に低い。
• 非線形問題： 非常に困難なフーリエ位相回復タスク（ upright と rotated の画像という二峰性の事後分布を示す）において、ShuffleFlowは二峰性を捉えることに成功した。対照的に、CF-NeRFや素朴なパッチ分割のような手法は、これらの複雑なグローバルな相関をモデル化することに失敗した。
• 効率性： 10,000個の事後サンプルを必要とするシナリオにおいて、ShuffleFlowは16分以内に高解像度の二峰性事後分布を生成した。同じ時間枠内で、拡散サンプラー（DAPS, DPS）ははるかに少ないサンプルしか生成できず、第二のモードを明確に特徴付けることもできず、同等の解像度を得るには約20倍の時間を要した。

意義と主張
本論文は、ShuffleFlowが、事後分布の精度を損なうことなく、科学的イメージングにおける変分推論の決定的なスケーラビリティのボトルネックに対処することを主張している。本手法は、拡散モデルが学習できないようなデータ不足の領域や、複雑で高解像度な事後分布の迅速な探索（例：非線形逆問題における二峰性分布の特定）が必要とされるアプリケーションにおいて、VIが拡散モデルに代わる重要な選択肢であることを位置づけている。著者らは、本手法は非常に効率的である一方で、訓練の不安定さにより非線形設定において現在やや過信（overconfidence）を示す傾向があることを認めており、これを今後の課題として挙げている。本フレームワークは、事後分布の複雑さが広範なサンプル生成を正当化するかどうかを判断できる、探索的なツールとして提示されている。