Uncertainty Estimation for Molecular Diffusion Models

技術要約：分子拡散モデルにおける不確実性推定

問題提起
拡散モデルは、3D分子生成における支配的なパラダイムとなっています。しかし、これらのモデルには、生成された分子が化学的に無効であるか、あるいは不安定であるかを識別するための原理的なシグナルが欠けているという、極めて重要な実用上のボトルネックが存在します。予測モデルは通常、信頼性の低い出力をフラグ立てするために不確実性推定を使用しますが、既存の手法は主に画像領域に焦点を当ててきました。ダウンストリームの評価（ドッキングやウェットラボでの検証など）が高コストである分子領域において、粗いルールベースのフィルタリングや、モデル自身の尤度（多くの場合、較正が不十分である）に依存し続けることが依然として標準となっています。分子拡散モデルに対して、サンプル品質のシグナルを提供する類似の研究は不足しています。

手法
著者らは、事前学習済みの分子拡散モデルを再学習することなく、事後的に（post-hoc）各サンプルの不確実性を推定する手法を提案しています。このアプローチは、「モデルが内部的に不確実性を感じているサンプルは、デノイジングの軌跡がより変動しやすい」という直感に基づいています。

1. ラプラス近似（Laplace Approximation）: あらゆる事前学習済みモデルとの互換性を維持するため、本手法はデノイジングネットワークのパラメータに対してラプラス近似を適合させます。具体的には、事前学習済みのパラメータ $\hat{\theta}$ の周囲に事後分布 $q(\theta) \approx \mathcal{N}(\theta | \hat{\theta}, \Sigma)$ を構築します。共分散 $\Sigma$ は、スケーラビリティを確保するために（最終層のみに適用した）経験的フィッシャー情報行列と事前精度項を用いて近似されます。
2. ノイズ予測の変動性: 生成の軌跡の中で、選択されたタイムステップのサブセット $\mathcal{T}$ に対し、事後分布 $q(\theta)$ から $M$ 個の重みサンプル $\theta_m$ を抽出します。各サンプルについて、デノイザーはノイズ $\epsilon^m_t = f_{\theta_m}(\hat{x}_t, t)$ を予測します。
3. 不確実性スコアの集計: 各タイムステップにおける中間的な不確実性 $u_t$ は、これら $M$ 個のサンプル間におけるノイズ予測の要素ごとのサンプル分散として定義されます。このステップごとの不確実性は、選択されたタイムステップ、すべての原子数（$N_a$）、およびすべての特徴次元（$D$）にわたって集計され、生成された分子の単一のスカラー不確実性スコア $u(\hat{x})$ を算出します。
   $$u(\hat{x}) := \frac{1}{|\mathcal{T}| \cdot N_a \cdot D} \sum_{t \in \mathcal{T}} \sum_{n=1}^{N_a} \sum_{d=1}^{D} [u_t]_{n,d}$$
   経験的に、著者らは生成プロセスの最終10%（$t$ が小さい段階）においてのみ不確実性を測定することが、パフォーマンスに対して十分であることを発見しており、これにより計算オーバーヘッドを制限しています。

主な貢献

• 分子拡散への初の適用: 本研究は、拡散の不確地性アプローチを画像領域から分子生成へと拡張し、ダウンストリーム評価が高コストな科学領域におけるサンプル品質のシグナルの必要性に対処しています。
• 事後的なフレームワーク（Post-hoc Framework）: 本手法は、基礎となる拡散モデルの再学習を必要とせず、EDMやGeoLDMのような既存の最先端アーキテクチャに適用可能です。
• テスト時スケーリング: 本論文は、不確実性スコアをテスト時スケーリング（オーバーサンプリングとフィルタリング）の設定におけるフィルタリング基準としてどのように使用できるかを実証し、生成される分子の品質を向上させることを示しています。

結果
著者らは、EDMおよびGeoLDMモデルを用いて、QM9およびGEOM-Drugsデータセット上で彼らの手法を評価しました。

• 品質との相関: QM9データセットにおいて、提案された不確実性スコアは、確立された品質指標（分子の安定性、原子の安定性、および妥当性）と統計的に有意な負の相関を示しました。これは、不確実性スコアが高いほど、品質の低いサンプルに対応することを意味します。
• 尤度との比較: 不確実性スコアは、サンプルの品質を予測する指標として、拡散モデルの負の対数尤度（NLL）を一貫して上回りました。例えば、EDM/QM9において、不確実性スコアは分子の安定性に対して$-0.284$のスペアマン相関を達成しましたが、NLLは$-0.150$でした。
• テスト時スケーリングの性能: 2万個の分子を生成し、不確実性が最も低い1万個を保持するようにフィルタリングすることで、著者らはEDMモデルにおいて分子の安定性（約10%）、原子の安定性（約1%）、および妥当性（約5%）の大幅な向上を観察しました。これらの向上は、基礎となるモデルアーキテクチャをEDMからGeoLDMへ切り替えた際に得られる改善と同等のものでした。
• アブレーション研究:
  • フィッシャー近似: フィッシャー近似を除去（共分散を等方性ガウス分布に置き換え）しても、ほぼ同様の性能が得られました。これは、スコアが厳密なベイズ的エピステミック不確実性というよりも、小さな摂動に対する感度の尺度として機能していることを示唆しています。
  • タイムステップの選択: シグナルは生成軌跡の「クリーン」な端（$t$ が小さい段階）で最も強く、これは不確実性の測定を最終ステップに集中させるという選択を正当化しています。
• 限界: 本手法は、より大規模で複雑なGEOM-Drugsデータセットでは意味のある改善をもたらさず、このアプローチが現在のところ、より高いコンフォメーションの複雑さを持つデータセットには転移しないことを示しています。

意義と主張
本論文は、拡散の不確実性アプローチを分子ドメインへと拡張するための「第一歩」を踏み出したと主張しています。主な意義は、生成された分子をフィルタリングするための実用的な事後ツールを提供し、それによってダウンストリームの評価パイプラインの効率を向上させることにあります。著者らは、不確実性スコアを、適応的なテスト時スケーリング手法におけるNLLベースの検証器の自然な代替物として位置づけています。

著者らは、自身の手法が現在のところ大規模で複雑なデータセット（GEOM-Drugs）では失敗すること、および基礎となるメカニズムが真のベイズ的エピステミック不確実性ではなく感度の尺度であるように見えることを明記しており、自らの主張に対して謙虚な姿勢を保っています。彼らは、今後の研究として、高次元の画像生成における最近の成功と同様に、事前学習済みの分子エンコーダを用いた「セマンティック」な特徴空間における不確実性の計算を探索すべきであると示唆しています。