テクニカル・サマリー：拡散モデルの推論時における分布制御のための平均場フレームワーク

問題提起

拡散モデルは、生成プロセスを推論時に特定の目的を満たすように誘導する、制御可能なサンプラーとしてますます活用されています。既存の手法は点別報酬（個々のサンプルに対して定義される $r: \mathbb{R}^d \to \mathbb{R}$）を効果的に扱いますが、多くの重要なアプリケーションでは分布報酬（生成された測度に対して定義される $R: \mathcal{P}(\mathbb{R}^d) \to \mathbb{R}$）を必要とします。これらの分布的目標には、実験的な観測結果への集団の較正、多様性の促進、あるいはモードのバランス調整などが含まれます。

分布的なステアリングに関する現在のアプローチは、多くの場合、報酬の勾配項を推論ダイナミクスに組み込みます（例：反発ポテンシャルやアンサンブル・ライクリフッド）。しかし、初期の点別ステアリング手法と同様に、これらのヒューリスティックな手法は、サンプリングされる分布に関する理論的な保証を欠いています。具体的には、単に勾配項を加えるだけでは、粒子が望ましい傾斜測度（tilted measure） $\mu^*$ をトレースすることを保証できません。ここで $\mu^*$ は、以下の KL 正則化された目的関数を最大化するものです：
$$\arg \max_{\mu} \left\{ R(\mu) - \text{KL}(\mu \| p_1) \right\}$$
点別のケースでは、傾斜されたターゲットは明示的な指数型傾斜 $\mu^* \propto p_1 e^r$ となりますが、分布的なケースでは、傾斜ポテンシャルが測度自体に依存するため、ターゲットは暗黙的になります。この暗黙的な分布をターゲットにするための、ダイナミクスと重みを補正する理論的に根拠のあるフレームワークは、現在存在しません。

手法

著者らは、推論時の分布的制御を定式化し解決するために、**平均場フレームワーク（Mean-Field Framework）**を提案しています。

1. 理論的定式化

この問題は、各時刻 $t \in [0,1]$ において、以下のように暗黙的に定義される傾斜測度 $\mu^*_t$ をターゲットにするものとして定式化されます：
$$\mu^*_t(dx) = \frac{1}{Z_t} e^{\Psi_t(x, \mu^*_t)} p_t(x) dx$$
ここで $\Psi_t(x, \mu) = \frac{\delta R_t}{\delta \mu}(x, \mu)$ は、時刻依存の報酬 $R_t$ の第一変分（汎関数微分）です。これにより、ターゲットとなる測度がポテンシャルに依存し、そのポテンシャルがさらに測度に依存するという、自己整合的な関係が生じます。

2. 重み付き McKean-Vlasov ダイナミクス

このパスを辿るために、著者らは重み付き McKean-Vlasov システムを導出しています。標準的な勾配ステアリングとは異り、ステアリング項 $\nabla_x \Psi_t(X_t; \mu_t)$ は、進化する経験的測度 $\mu_t$ に依存するため、粒子間に相互作用を誘発します。

粒子 $X_t$ とその対数重み $A_t$ の理想的なダイナミクスは以下の通りです：
$$dX_t = \left[ b_t(X_t) + \epsilon_t s_t(X_t) + \epsilon_t \nabla_x \Psi_t(X_t; \mu_t) \right] dt + \sqrt{2\epsilon_t} dB_t$$
$$dA_t = \left[ b_t(X_t) \cdot \nabla_x \Psi_t(X_t; \mu_t) + \dot{\Psi}_t(X_t) \right] dt$$
ここで、$\dot{\Psi}_t(x) = \frac{d}{dt} \Psi_t(x, \mu_t)$ は極めて重要な補正項です。点別のケースでは $\dot{\Psi}_t$ は単なる $\partial_t r$ ですが、ここでは測度 $\mu_t$ の時間発展が含まれます。著者らは、$\dot{\Psi}_t$ が報酬汎関数の第二変分を含む暗黙的な積分方程式を満たすことを示しています。

3. 有限粒子アルゴリズム

これを実装するために、著者らは重み付き相互作用粒子スキーム（Algorithm 1）を提案しています：

• 位置の更新: 粒子は、ベースとなるドリフト、スコア、および現在の経験的測量 $\hat{\mu}_t$ に関する報酬の勾配を用いて更新されます。
• 対数重みの更新: ステアリングされたダイナミクスとターゲットとなる傾斜パスとの間の不一致を補正するために、重みが更新されます。これには $\dot{\Psi}_t$ の計算が必要です。
• $\dot{\Psi}_t$ の算出: 論文では、この暗黙的な時間微分項を計算するための2つの方法を提供しています：
  1. 暗黙的線形システム: 積分方程式から導出された線形システムを解く方法（Proposition 3.3）。
  2. 不動点反復法: 重みの更新を近似する軽量なピカール反復スキーム。
• リサンプリング: 重みの退化を防ぐために、標準的な逐次モンテカルロ（SMC）リサンプリングが採用されています。

このフレームワークは、$R(\mu) = \int r(x) d\mu(x)$ である場合において、標準的な点別報酬ステアリングを特殊なケースとして回収します。

主な貢献

1. 分布的ステアリングの定式化: 著者らは、分布レベルのステアリングを、傾斜が測度自体に依存する傾斜測度をターゲットにすることとして定式化し、点別報酬の傾斜ターゲットに対する分布的な類似性を確立しました。
2. 重み付き McKean-Vlasov ダイナミクスの導出: 重み付きの相互作用粒子が望ましい傾斜パス $\mu^*_t$ を正確に辿ることを保証する、補正的な Feynman-Kac 再重み付けを伴う原理的なダイナミクス・システムを導出しました。
3. 理論的回収と基盤: 本フレームワークは、標準的な点別ステアリングを特殊なケースとして回収し、既存のヒューリスティックなバッチレベルのステアリング手法（例：多様性サンプリング、実験的ガイダンス）に対して理論的な基礎を提供します。
4. 収束保証: 粒子数 $N \to \infty$ のとき、有限粒子システムが理想的なダイナミクスに収束することを証明しています。

実験結果

著者らは、低次元および高次元の設定の両方で本フレームワークを検証しています。

• 低次元（1D Mixture）:

  • タスク: 二峰性ガウス混合分布を、二乗最大平均ディスタンス（MMD）報酬を用いて、異なるモードの重みを持つターゲット分布へと傾斜させる。
  • 知見: 標準的な勾配ステアリングは粒子をターゲットへと押し出すものの、傾斜された分布 $\mu^*$ の正確な形状を一致させることには失敗します。対照的に、平均場メソッドは $\mu^*$ を正確にターゲットとし、ステアリング強度 $\lambda$ が目的関数のパラメータと一致する場合、最適性ギャップ $J(\mu) - J(\mu^*)$ を最小化します。これは様々なノイズスケジュールにおいて成立します。

• 高次元（タンパク質コンフォメーション）:

  • タスク 1 (HIV-1 Protease): 実験的に観測された残基距離分布（DEERデータ）に一致するように、Boltz-2 で生成されたアンサンブルをステアリングする。ベースモデルは「閉じた（closed）」状態に偏っており、目標は「半開（semi-open）」状態の割合を回復することである。
    • 結果: 平均場ステアリングは実験的な残基距離分布を正確に一致させましたが、勾配のみのステアリングは、わずかに異なる形状とより強いピークを持つ分布を生じさせました。
  • タスク 2 (Adenylate Kinase): 分子動力学シミュレーションから得られた結合角の2次元分布に向けてステアリングする。
    • 結果: 平均場ステアリングは、勾配のみの手法と比較して、ターゲット分布へのより強力なシフトを生み出しました。
  • タスク 3 (X-ray Crystallography): ヘテロなコンフォメーションを示す電子密度マップ（4OLE構造）に適合するように Protenix の生成をステアリングする。
    • 結果: 同じステアリング強度において、平均場ステアリングはターゲット密度との高いコサイン類似性を達成し、勾配のみのステアリングよりも多くのサンプルを代替コンフォメーション（Conformation B）に配置しました。これは、勾配のみの手法が真の逆問題ターゲットに対してアンダーステアリング（不足した誘導）している可能性を示唆しています。

意義と主張

本論文は、分布的制御のための原理的な理論的基礎を提供することにより、点別制御と分布的制御の間の溝を埋めるものであると主張しています。

• 正確性: 主要な主張は、提案された手順が無限粒子数の極限において、規定された報酬傾斜分布 $\mu^*$ を正確にターゲットとすることであり、ヒューリスティックな勾配ステアリングに内在する不一致を補正することです。
• 解釈可能性: 本フレームワークは、既存のバッチレベルのステアリング手法が、実質的に勾配ステップを実行しているものの、真の事後分布（または傾斜測度）をサンプリングするために必要な補正的な再重み付けを欠いていることを明らかにしています。
• 実用性: 本手法は追加の計算コスト（$\dot{\Psi}_t$ の解決）を導入しますが、著者らは、高次元の設定（例：タンパク質）におけるネットワーク評価と比較して、これはしばしば無視できる程度であると述べています。また、不動点ソルバーは効率的です。
• 限界: 著者らは、リサンプリングへの依存が画像のようなモダリティにおいて制限となる可能性があること（重複サンプルの発生など）、および微分可能な報酬関数が必要であることを認めています。これらを緩和するために、分散減少技術（例：ポアソン方程式ソルバー）との組み合わせを提案しています。

本研究は、既存の手法を置き換えるものではなく、なぜ特定のヒューリスティックが機能するのかを説明する厳密な基礎を提供し、それらが正しい分布をターゲットにすることを保証するメカニズムを提供するものであると位置づけられています。