[论文解读] 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

一句话总结

本论文提出利用单一机器学习模型的不确定性作为集体变量（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训练至关重要的罕见高能构型的局限性。
• 开发一种主动针对MLIP预测不确定性最高的构型空间区域的方法，从而提升数据集覆盖范围与模型可靠性。
• 通过使用模型不确定性作为通用自适应CV，消除对预定义、体系特异性反应坐标的依赖。
• 通过利用单模型不确定性而非基于集合的不确定性估计，降低计算成本，同时保持有效的探索能力。

提出的方法

• 使用高斯混合模型（GMM）从单个训练好的神经网络势能（NNIP）的预测中估计不确定性，并将此不确定性视为集体变量（CV）。
• 采用扩展系自适应偏置力（eABF）结合高斯加速分子动力学（GaMD），利用不确定性CV执行偏置分子动力学模拟。
• 在主动学习迭代过程中，动态调整偏置强度（γ）和不确定性阈值（u_cutoff），以平衡探索与利用。
• 从高不确定性区域迭代收集新构型并重新训练NNIP，逐步提升模型准确度与覆盖范围。
• 通过置信预测校准不确定性估计，确保预测结果具有可靠的误差边界。
• 将基于不确定性的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 单一训练好的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.