MAD: Manifold Attracted Diffusion

技術要約：Manifold Attracted Diffusion (MAD)

問題提起
スコアベースの拡散モデルは、画像合成において、サンプルの品質と多様性の両面でGANやVAEを凌駕することが多く、最先端の手法として確立されています。しかし、訓練データが測定誤差、圧縮アーティファクト、またはデータ収集プロセスによるノイズによって汚染されている場合、重大な制限が生じます。このようなノイズを含むデータセットで直接訓練された標準的な拡散モデルは、その汚染を再現するように学習してしまい、結果として訓練データのノイズ特性を保持したサンプルを生成してしまいます。ノイズを含む訓練に対する既存のアプローチの多くは、訓練プロセス自体の変更（例：損失関数の適合や、クリーンなデータとノイズ混じりのデータの混合の必要性）を必要としますが、確立された訓練フレームワークや学習済みネットワークを使用せずに、ノイズを含む分布からクリーンなサンプルを生成できる効率的な手法は不足しています。

手法
提案手法である**Manifold Attracted Diffusion (MAD)は、訓練プロセスではなく、スコアベースモデルの推論手順を修正することでこの問題に対処します。このアプローチは、意味のあるデータは高次元の周囲空間内にある低次元の多様体（マニフォールド）上に存在する、という多様体仮説（manifold hypothesis）**に基づいています。著者らは、ノイズをこの構造に対して直交する（多様体から外れた）方向への低振幅の変動と解釈し、一方で意味のあるデータの変動は多様体に沿って（多様体上で）起こると解釈しています。

技術的な核心となるイノベーションは、**拡張スコア（extended score）**の導入です。

• 定義: 確率測度 $p$ とパラメータ $\gamma > 0$ に対して、拡張スコア $H_\gamma p$ は $(1 + \gamma)S(p * g_\gamma) + \gamma \frac{d}{d\gamma}S(p * g_\gamma)$ と定義されます。ここで、$S$ は標準的なスコア演算子（$\nabla \log p$）、$g_\gamma$ はガウスカーネルです。
• 特性: 標準的なスコアは、特定の方向において分散がゼロである（低次元の多様体上のデータを表す）ディラックのデルタ分布に対しては定義されませんが、拡張スコアは定義可能です。決定的なことに、ディラックのデルタに対して、拡張スコアは単位分散を持つガウス分布のスコアと同様に振る舞います。これにより、推論プロセスにおいて、低分散（多様体から外れた）方向を非ゼロの分散を持つものとして扱い、実質的に「ソフト閾値処理」効果を適用することが可能になります。
• 推論アルゴリズム: 著者らは、拡張スコアの近似を利用した反復推論アルゴリズムを提案しています。学習済みのデノイザーネットワーク $D_\theta$（標準的なスコアを近似するもの）が与えられたとき、拡張スコアは、わずかに摂動を与えたノイズレベルでネットワークを評価し、$\gamma$ に関する微分を有限差分によって近似することで算出されます。
• 実装: アルゴリズムは、「多様体への吸引力（manifold attraction）」の強さを制御するハイパーパラメータ ($a, b, p$) を導入しています。これらを調整することで、推論プロセスは多様体から外れた方向の小さな変動（ノイズ）を抑制しつつ、多様体上の方向の大きな変動（信号）を保持します。本手法は、Karrasら (2022) による確率流ODE定式化などの確立されたフレームワークと互換性があり、学習済みネットワークを利用できます。

主な貢献

1. 形式的な定義と分析: 本論文は、拡張スコアの厳密な数学的定義を提供し、標準的な分布とディラック測度の両方におけるその挙能を分析することで、退化した分布に対してスコア演算子を一般化できることを示しています。
2. 効率的な推論アルゴリズム: 著者らは、特別な訓練手順を必要とせずに生成サンプルのノイズを低減する推論アルゴリズムを提示しています。これにより、既存の訓練パイプラインや学習済みモデルとの互換性が確保されます。
3. 実証的検証: 本アプローチの有効性は、トイ問題、合成データ、およびFFHQ、AFHQ、ImageNet、クライオ電子顕微鏡（Cryo-EM）データを含む実世界のデータセットを通じて実証されています。

結果

• トイデータおよび合成データ: 2次元の例示的な例において、拡張スコアは低分散のガウス分布の軌跡をその平均へと収束させる（ディラックのデルタを模倣する）一方で、高分散成分の広がりを保持することに成功しました。ガウスノイズとピクセル除去が加えられた合成データセットにおいて、MADは標準的なスコア推論を大幅に上回り、元のクリーンな形状を復元したサンプルを生成しました。
• 実世界の画像データ: FFHQ、AFHQv2、およびImageNetを用いた実験では、多様体への吸引パラメータ ($b$) を大きくすることで、背景がよりクリーンになり、主要な特徴がより鮮明になることが示されました。これは、高周波ノイズや二次的な変動を効果的に閾値処理していることを意味します。
• 定量的指標: ガウスノイズ、ショットノイズ、およびピクセル除去によって汚染されたCIFAR-10およびFFHQデータセットにおいて、MADは標準的な推論と比較して大幅に低いFréchet Inception Distance (FID) を達成しました。特筆すべきは、標準的な推論のFIDがノイズ混じりの訓練データのFIDと密接に一致している点であり、これは標準的なモデルが汚染を再現してしまう一方で、MADが効果的にデノイジングを行っていることを裏付けています。
• Cryo-EMへの適用: EMPIAR-11618データセット（ノイズの多い単一粒子Cryo-EM画像）に適用した結果、ネットワークがノイズの多いデータのみで訓練されていたにもかかわらず、MADは従来の計算手法によって得られる精緻な形状によく似たサンプルを生成しました。

意義と主張
本論文は、MADが「ブラインド生成デノイジング（blind generative denoising）」に対する実用的な解決策を提供すると主張しています。これは、具体的な劣化モデルが未知であり、ノイズを含むデータのみが利用可能な状況を指します。本手法の意義は、生成フェーズにおいて、複雑な訓練時の適応に頼るのではなく、データの多様体の幾何学的構造を利用してノイズをフィルタリングできる点にあります。

著者らは、その主張の範囲について慎重な姿勢をとっています：

• 特定の画像（与えられた単一の画像など）の逆問題を解く（デノイジングする）ことを目的としているのではなく、ノイズを含むデータセットから、基礎となる分布からなる新しいクリーンなサンプルを生成することを目的としています。
• 特定のドメイン（Cryo-EMなど）において最先端の性能に到達するには、ドメイン固有の事前知識を取り入れる必要があることを認めており、それは本研究の範囲外としています。
• 手法は効果的であるものの、計算コストは標準的な推論の最大2倍程度であること（追加のネットワーク評価による）を指摘しています。ただし、経験的な観察によれば、推論ステップの後半部分のみで拡張スコアを使用すれば十分であることも示唆されています。

本研究は、拡張スコアを、生成の質を向上させるための幾何学的に根ざしたメカニズムとして位置づけており、これはクリーンなデータで訓練されたモデルに対しても適用可能である可能性があり、適応的なパラメータ選択や条件付き拡散モデルへの統合といった将来の研究への道を開くものです。