技術的サマリー：Stein Diffusion Guidance (SDG)

問題定義

拡散モデルは、複雑で高次元の分布をモデル化する上で顕著な成功を収めています。しかし、ラベルなしデータで訓練された場合、それらは主にデータ多様体の高密度領域からサンプリングを行い、高価値かつ稀なサンプル（例：新規の創薬候補）が存在する可能性がある低密度領域を見落としがちです。

この制限に対処する既存の手法は、一般的に以下の 2 つのカテゴリに分類され、いずれも重大な欠点を有しています：

1. 分類器ベースのガイダンス：これらの手法は、複数のノイズレベルでノイズの混入したデータに訓練された補助分類器を活用してサンプリングを誘導します。しかし、これらは訓練の複雑さを増大させ、分類器の勾配がノイズを含む場合、サンプルを誤って誘導し、特に低密度領域において生成多様体から逸脱させる可能性があります。
2. 確率的最適制御（SOC）：SOC は、最適コントローラーによって拡散プロセスを誘導することにより、事後サンプリングのための原理的な枠組みを提供します。理論的に堅牢であり、クリーンなデータで訓練された市販の分類器を活用できる一方で、最適制御の計算には、報酬信号を Neural-SDE 全体の軌道を通じて逆伝播させる必要があります。このため、SOC は効率的な推論に対して計算量が膨大すぎます。

最近の「訓練不要」アプローチは、SOC の計算コストを回避するためにTweedie の公式を用いて拡散事後分布を近似しようとしています。しかし、本論文では、この近似は本質的に偏りがあり、最適ではないと主張しており、特にスコアベースモデルの精度が最も低い低密度領域において、信頼性の低いガイダンスをもたらすと指摘しています。

手法：Stein Diffusion Guidance (SDG)

著者は、訓練不要手法の効率性と SOC の原理的な性質を両立させるStein Diffusion Guidance (SDG)を提案します。SDG は、代理 SOC 目的関数に基づいており、近似事後分布を精緻化するためのStein 補正メカニズムを導入します。

1. 代理 SOC 定式化

著者は、低密度拡散サンプリングを、新しいコスト汎関数 $\tilde{J}(u, x, t)$ を持つ確率的最適制御問題として定式化します。この汎関数には以下の要素が含まれます：

• 状態コスト：高密度領域を罰則し（$\alpha(t)$ というアニーリングスケジュールを通じて低密度領域の探索を促進）、
• 終端コスト：市販の分類器または報酬モデル $r(x_T)$ に基づく（$\beta(t)$ という強度を通じてサンプルを所望の特性へと誘導）。

彼らは、SOC 値関数 $V(x,t)$ に関する変分的上界を導出します。この上界は、Tweedie の公式に依存する標準的な訓練不要ガイダンスの最適化が、近似事後分布 $q$ と真の事後分布 $p$ 間のKullback-Leibler (KL) ダイバージェンスを無視したまま、上界の最初の 2 つの項のみを最小化することに相当することを明らかにしています。この無視が、最適ではないガイダンスを引き起こします。

2. 往復 Stein 補正

完全な SOC の計算コストを伴わずに KL ダイバージェンス項に対処するため、SDG は往復 Stein 補正メカニズムを導入します：

• 後方マッピング：拡散多様体 $M_t$ 上のノイズ粒子 $x_t$ を、Tweedie の公式を用いてデータ多様体 $M_T$ へ後方にマッピングし、近似事後分布サンプルの初期セット $x_T$ を取得します。
• Stein 補正：これらの粒子を、**Stein 変分勾配降下法（SVGD）**を用いて $M_T$ 上で進化させます。粒子は、近似事後分布と真の事後分布間の KL ダイバージェンスを最小化するように、最急降下方向に沿って更新されます。これは、モデルスコアを通じて近似された条件付き事後分布のスコア $\nabla_{x_T} \log p(x_T|x_t) \approx s_\theta(x_T) - \eta(t)s_\theta(x_t)$ を利用します。
• 前方マッピング：補正された粒子を前方にマッピングして現在の多様体 $M_t$ に戻し、次の拡散ステップで使用します。

このメカニズムは事後分布サンプルを効果的に正則化し、市販の分類器を活用する場合でも、それらが真の拡散ダイナミクスと生成多様体に近接したままであることを保証します。

3. アルゴリズム

SDG アルゴリズムは、この補正を標準的な逆拡散プロセスに統合します。各ステップにおいて、以下の処理を行います：

1. Tweedie の公式を通じて初期事後分布推定値を生成します。
2. Stein 補正ステップ（カーネル相互作用とスコアの差に基づく粒子更新）を適用します。
3. 補正されたサンプルを用いて、逆 SDE 更新のためのガイダンス信号（報酬勾配）を計算します。

主要な貢献

1. 低密度フレームワーク：新しいコスト・トゥ・ゴー関数を通じて低密度領域を明示的にターゲットとする、拡散ガイダンスのための新しい SOC ベースの定式化。
2. 理論的洞察：Tweedie の公式による拡散事後分布の近似が最適ではなく、近似事後分布と真の事後分布間の KL ダイバージェンスを考慮するための補正ステップが必要であることを示す理論的証明。
3. Stein 補正メカニズム：近似事後分布と真の事後分布間の KL ダイバージェンスを反復的に最小化し、低密度領域における堅牢なガイダンスを可能にする、計算効率的かつ訓練不要な Stein 補正の導入。
4. 汎用性：モデルの微調整や蒸留を必要とせず、さまざまな拡散事後分布サンプリング問題に統合可能なプラグアンドプレイフレームワーク。

実験結果

著者は、画像ガイダンスおよびタンパク質ドッキングのための分子サンプリングなど、多様なタスクにおいて SDG を評価しました。

• 画像ガイダンスタスク：画像ラベルガイダンス、ガウスぼかし除去、超解像、テキストから画像へのスタイル転送などのタスクにおいて、SDG（特に低密度サンプリングを行わない変種 SDG$\heartsuit$）は、標準的な訓練不要ベースライン（例：DPS、LGD）および反復戦略を備えた高度な手法（UGD、MPGD）を、ほとんどの指標で一貫して上回りました。これは、低密度領域以外の設定であっても近似事後分布を補正することの重要性を浮き彫りにしています。
• 分子サンプリング（低密度）：リガンド - タンパク質ドッキングタスク（Fa7、5ht1b、Jak2、Parp1）において、SDG は、厳格な基準（新規性、合成の容易さ、ドラッグライクネス、ドッキングスコア）を満たす新規の「ヒット」分子を生成する点で、ベースライン（GDSS、MOOD、FREED など）を大幅に上回りました。
  • アブレーション研究：Stein 補正を除去した場合（SDG$\clubsuit$）、性能が低下し、サンプルがヒット基準を満たさなかったり、生成多様体から逸脱したりすることが頻繁に起こりました。低密度サンプリングを除去した場合（SDG$\heartsuit$）も、結果は最適ではなく、新規発見のために低密度領域をターゲットとする必要性が確認されました。
  • 報酬の過大評価：本論文は、Stein 補正がない場合、報酬モデルが低密度領域で報酬を過大評価し、非現実的な分子をもたらすことを示しています。SDG は正則化メカニズムとして機能し、この過大評価を軽減してサンプルの妥当性を維持します。
  • 効率性：SDG は実行時間のオーバーヘッド（ソルバーとタスクに応じて約 36〜85%）を伴いますが、計算上実行可能であり、完全な SOC に比べて最小限のメモリオーバーヘッドで、最大 3000 までの大規模な粒子サイズにスケーリング可能です。

意義と主張

本論文は、SDG が、事後サンプリングのための prohibitively expensive（計算的に許容しがたいほど高価な）確率的最適制御手法に対する計算効率的かつ訓練不要な代替手段を提供すると主張しています。事後分布の補正の必要性を理論的に確立し、それを Stein 変分推論を通じて実装することで、SDG は高価値な発見（新規の創薬候補など）が存在する可能性が高い低密度領域における効果的なサンプリングを可能にします。

著者は、自らの手法が以下の点に優れていることを強調しています：

• Tweedie ベースの近似の偏りを克服する。
• サンプルが生成多様体から逸脱するのを防ぐ。
• 低密度領域における報酬の過大評価を軽減する。
• 高密度領域を超えたさまざまな事後分布サンプリング問題に適用可能な一般的な枠組みを提供する。

この研究は、近似事後分布の補正が、訓練不要の拡散ガイダンスにおいて重要でありながらしばしば見落とされている構成要素であることを示唆しており、SDG は科学分野における稀で高価値なサンプルの発見のために拡散モデルの可能性を解き放つものとしています。