An Unconditional Representation of the Conditional Score in Infinite-Dimensional Linear Inverse Problems

技術的概要：無限次元線形逆問題における条件付きスコアの無条件表現

問題提起
スコアベース拡散モデル（SDM）は、ベイズ逆問題における事後分布からのサンプリングのための強力なツールとなり、従来の手作業で設計された正則化項を上回る柔軟な事前分布を提供します。しかし、事後サンプリングのための既存の手法は、計算効率とスケーラビリティの間に重大なトレードオフに直面しています。

1. 補正ベースの手法（例：スコアへの尤度勾配の追加やデータ整合性の強制）は、サンプリング過程において前方演算子 $A$ の繰り返し評価を必要とします。これは、大規模・高次元の問題において、許容しがたい計算ボトルネックを生み出します。
2. 直接的条件付き学習手法は、測定値に条件付けられたスコア関数（$s(x, t; \mu_y)$）を学習します。これはサンプリング中の前方評価を回避しますが、ニューラルネットワークの入力次元を（$x$ から $(x, y)$ へ）劇的に増加させ、特に無限次元の設定において条件付きスコアを正確に近似するために、広範な学習データと複雑なアーキテクチャを必要とします。

本論文は、前方モデル評価の高いオンラインコストや、直接的条件付き学習の高いサンプル複雑性を伴うことなく、無限次元線形逆問題において効率的かつスケーラブルな事後サンプリングを行うという課題に取り組みます。

手法：条件付きスコアの無条件表現（UCoS）
著者は、UCoS を提案します。これは、条件付けの計算負荷をオンラインサンプリング段階からオフライン学習段階へシフトさせるフレームワークです。中核となる洞察は、ガウスノイズを伴う線形逆問題において、条件付きスコア関数は、修正されたタスク依存の拡散過程の無条件スコア関数から、アフィン変換を介して正確に導出できるという点です。

• 理論的基盤:

  • 本手法は、離散化不変性を保証する可分ヒルベルト空間 $\mathcal{H}$ において動作します。
  • $s(x, t; \mu_y)$ を近似するネットワークを学習する代わりに、著者は、事前分布と前方演算子 $A$ に依存するガウス項の混合分布となる新しい拡散過程 $\tilde{X}_t$ を定義します。
  • 彼らは（定理 3.7 において）、真の条件付きスコア $s(x, t; \mu_y)$ が、この新しい過程の無条件スコア $\tilde{s}$ と以下の恒等式によって関連付けられることを証明しています：
    $$s(x, t; \mu_y) = \lambda(t) \left( r(\xi_t(x, y)) - e^{-t/2}x \right)$$
    ここで、$r$ は無条件スコアの変換版、$\lambda(t)$ は時間依存のスカラー、$\xi_t(x, y)$ は測定値 $y$ と随伴演算子 $A^*$ を含むアフィン変換です。
  • 決定的な点は、項 $A^*\Gamma^{-1}y$（ここで $\Gamma$ はノイズ共分散）をサンプリング開始前に一度計算できることです。その後のサンプリング過程では、訓練済みネットワーク $r_\theta$ の評価と単純なベクトル加算のみが必要となり、リバース拡散ステップ中に $A$ や $A^*$ をさらに適用する必要がなくなります。

• 実装:

  • オフライン段階: 関数 $r$ を近似するためにニューラルネットワーク（具体的にはフーリエニューラルオペレーター、FNO）が訓練されます。学習には、修正された過程 $\tilde{X}_t$ に対するデノイジング・スコアマッチング目的関数が利用されます。この段階では $A$ と $A^*$ へのアクセスが必要ですが、特定の測定データ $y$ は不要です。
  • オンライン段階: 事後サンプルは、変換されたスコアを用いて時間逆転 SDE を解くことで生成されます。この過程は「行列フリー」であり、事前に計算されたシフト $A^*\Gamma^{-1}y$ とネットワーク評価のみを必要とします。

主要な貢献

1. 無限次元における正確な恒等式: 本論文は、無限次元ヒルベルト空間において、アフィン変換を介して条件付きスコアをタスク依存の無条件スコアに結びつける厳密な理論的恒等式（定理 3.7）を確立します。これは、特定の有限次元離散化を超えてその原理を一般化します。
2. 収束解析: 著者は、生成過程に対する定量的な誤差 bound（定理 4.1）を提供し、学習損失、数値的離散化、および初期化からの誤差寄与を明示的に特徴づけます。これにより、手法が真の事後測度に収束することが正当化されます。
3. スケーラビリティと効率性: 条件付けをサンプリングループから分離することにより、UCoS は、サンプリング中の前方評価を行わない無条件手法の計算効率を達成しつつ、条件付き手法の精度を維持します。これは、結合空間 $(x, y)$ 上でネットワークを訓練することに伴う次元の呪いを回避します。

実験結果
著者は、UCoS を画像修復、スパースビュー計算トモグラフィ（CT）、画像デブラリングの 3 つの逆問題で検証しました。

• 画像修復（ガウス事前分布）: 事後分布がガウス分布であり、スコアが閉形式で計算可能な設定において、UCoS は正確な条件付き手法および他のベースラインと同等の性能を発揮し、理論的恒等式を確認しました。
• CT 画像化およびデブラリング（学習済み事前分布）: それぞれ LIDC-IDRI データセットと CelebA データセットを使用し、UCoS を SDE ALD、Diffusion Posterior Sampling (DPS)、射影ベース手法（Proj）、および直接的条件付き学習と比較しました。
  • 精度: UCoS は、低複雑性のアーキテクチャ（例：層あたり 32 ノード）を使用する際、無条件手法と同等かそれ以上のバイアスおよび標準偏差を持つサンプルを生成し、条件付き手法よりも著しく優れた性能を示しました。
  • 効率性: UCoS は条件付き手法（どちらも前方評価を回避）と同等のサンプリング速度を達成しますが、繰り返し前方パスを必要とする補正ベースの手法（DPS、Proj、SDE ALD）を大幅に上回ります。
  • ロバスト性: 結合分布の近似の複雑さが増すため、ネットワークパラメータが減少すると条件付き手法の性能は顕著に低下しますが、UCoS はより単純なアーキテクチャでも高品質を維持します。

意義と主張
本論文は、UCoS が「両方の世界の最良」を提供すると主張しています。すなわち、サンプリング中の前方モデル評価という計算ボトルネック（補正ベース手法の限界）を排除しつつ、直接的条件付き学習の高いサンプル複雑性とアーキテクチャ上の要求を回避するという点です。

著者は、そのアプローチが離散化不変であることを強調しており、高次元の有限次元近似では性能保証がしばしば失敗する無限次元逆問題に適していると述べています。彼らは、この手法がタスク依存であり、前方演算子 $A$ が変更される場合は再学習が必要であること、また正確な対応関係は前方モデルの線形性と尤度のガウス性（「平方完成」ステップを可能にする）に依存していることを認めています。この研究は、理論的定式化において近似誤差を生じさせることなく、条件付けの複雑さを学習段階へシフトさせる新しいクラスの拡散ベースサンプリングの概念実証として提示されています。