[論文レビュー] Enhanced sampling of robust molecular datasets with uncertainty-based collective variables

Aik Rui Tan, Johannes C. B. Dietschreit+1 more|arXiv (Cornell University)|Feb 6, 2024

Computational Drug Discovery Methods被引用数 4

ひとこと要約

本論文では、1つの機械学習モデルからの不確実性を集合変数（CV）として用いることで、分子動力学における強化サンプリングを誘導する手法を提案する。予測不確実性が高くなる領域にシミュレーションをバイアスすることで、希少かつ高エネルギー状態の効率的探索が可能となり、アラニンジペプチドを用いた実験で、構造的配置空間の広範なカバーが達成された。これにより、機械学習型原子間ポテンシャル（MLIP）のデータセットの多様性が向上し、外挿的誤差が低減された。

ABSTRACT

Generating a data set that is representative of the accessible configuration space of a molecular system is crucial for the robustness of machine learned interatomic potentials (MLIP). However, the complexity of molecular systems, characterized by intricate potential energy surfaces (PESs) with numerous local minima and energy barriers, presents a significant challenge. Traditional methods of data generation, such as random sampling or exhaustive exploration, are either intractable or may not capture rare, but highly informative configurations. In this study, we propose a method that leverages uncertainty as the collective variable (CV) to guide the acquisition of chemically-relevant data points, focusing on regions of the configuration space where ML model predictions are most uncertain. This approach employs a Gaussian Mixture Model-based uncertainty metric from a single model as the CV for biased molecular dynamics simulations. The effectiveness of our approach in overcoming energy barriers and exploring unseen energy minima, thereby enhancing the data set in an active learning framework, is demonstrated on the alanine dipeptide benchmark system.

研究の動機と目的

• 機械学習型原子間ポテンシャル（MLIP）の訓練データが疎で代表的でないという課題に対処し、モデルの一般化能力を向上させること。
• 従来のサンプリング手法が低エネルギー極小点に偏るという制限を克服し、モデルの耐性を高めるために不可欠な希少で高エネルギー状態を無視しないこと。
• MLIPの予測が最も不確実な配置空間領域を能動的に標的とすることで、データセットのカバー範囲とモデルの信頼性を向上させること。
• システム固有の反応座標を事前に定義する必要を排除し、モデルの不確実性を普遍的かつ適応的CVとして用いること。
• アンサンブルベースの不確実性推定よりも単一モデルからの不確手性推定を活用することで、計算コストを削減しながらも、有効な探索を維持すること。

提案手法

• 1つの訓練済みニューラルネットワーク原子間ポテンシャル（NNIP）の予測結果から不確実性を推定するためにガウス・ミックスチャネル・モデル（GMM）を用い、これを集合変数（CV）として扱う。
• 拡張系適応的バイアス力（eABF）とガウス加速分子動力学（GaMD）を組み合わせ、不確実性CVを用いたバイアス付き分子動力学シミュレーションを実行する。
• 活性学習の各イテレーションにおいて、バイアス強度（γ）と不確実性のしきい値（u_cutoff）を動的に調整し、探索と活用のバランスを取る。
• 高不確実性領域から新たな配置を繰り返し収集し、NNIPを再訓練することで、段階的にモデルの精度とカバー範囲を向上させる。
• 予測の信頼性のある誤差境界を保証するため、適合予測（conformal prediction）を用いて不確実性推定をキャリブレーションする。
• 不確実性をCVとして用いる手法と、不確実性をバイアスエネルギーとして用いる手法を比較し、配置空間カバー範囲と予測誤差低減の両面で評価する。

Figure 1: (a) , Structure of the alanine dipeptide molecule with carbon (C), nitrogen (N), oxygen (O), and hydrogen (H) atoms labeled in grey, blue, red, and white, respectively. Four backbone dihedral angles $\phi$ , $\psi$ , $\omega_{1}$ , and $\omega_{2}$ are annotated. (b) , Potential mean force

実験結果

リサーチクエスチョン

• RQ1 1つの訓練済みMLIPからの不確実性が、分子系における強化サンプリングを誘導する有効で普遍的な集合変数として機能するか？
• RQ2 不確実性をCVとして用いることで、従来のサンプリング法や不確実性ベースのバイアス法と比較して、希少または高エネルギー領域における配置空間カバー範囲が向上するか？
• RQ3 バイアス強度（γ）の選択が、新規配置の探索と物理的に不自然な歪みの回避のバランスにどのように影響するか？
• RQ4 不確実性をCVとして用いる手法が、ポテンシャルエネルギー表面の異なる領域におけるMLIPの一般化能力および外挿的精度をどの程度向上させるか？
• RQ5 この手法により、アンサンブルベースの不確実性推定を必要とせず、データセット強化のための有効なアクティブラーニングを維持できるか？

主な発見

• バイアス強度γ = 0.005が最も効果的な探索を実現し、γ = 0.01よりも予測不確実性が高く、カバー範囲も広がった。これは、低いγ値が段階的かつ障壁に配慮したサンプリングを可能にしていることを示唆している。
• 不確実性をCVとして用いる手法は、非バイアスシミュレーションや不確実性をバイアスエネルギーとして用いる手法よりも、φ < 0°のφ–ψ二面角空間領域に効果的に到達できた。
• 活性学習の第4世代以降、カバー率の上昇が止まっているにもかかわらず、エネルギーおよび力の予測における平均絶対誤差（MAE）は継続的に低下を示し、モデルの外挿的能力の有効な強化が確認された。
• 非平面的ω₁–ω₂二面角の領域についても、従来のeABF-GaMDシミュレーションで不十分にカバーされていた領域を効果的に探索し、複雑なトランスレーション運動のサンプリングが向上した。
• 配置空間の多様性とモデルの耐性の両面で、不確実性をCVとして用いる手法は、不確実性をバイアスエネルギーとして用いる手法を上回った。物理的に不自然な構造が原因でシミュレーションが早期に終了することもなかった。
• 単一モデルからの不確実性推定が、アンサンブルベースの手法に比べて計算コストを削減しながらも、効果的なサンプリングを可能にすることが実証された。

Figure 2: Cumulative exploration of configuration space projected onto the $\phi$ - $\psi$ (left column) and $\omega_{1}$ - $\omega_{2}$ (right column) plane of 10 NVT simulations at 300 K. Top) No biasing of any sort, middle) uncertainty-guided eABF, and bottom) uncertainty-guided eABF-GaMD.