[论文解读] Uncertainty Estimation for Molecules: Desiderata and Methods

Tom Wollschläger, Nicholas Gao+3 more|arXiv (Cornell University)|Jun 20, 2023

Machine Learning in Materials Science被引用 5

一句话总结

本文提出了一种基于高斯过程的不确定性估计方法——局部神经核（Localized Neural Kernel, LNK），用于图神经网络（GNN）驱动的分子力场，该方法满足六项关键期望——三项物理信息相关（对称性、能量守恒、局域性）和三项应用导向（速度、可解释性、分布外检测）。LNK 实现了可靠的单前向传播不确定性量化，同时保持了高预测精度，在分布外（OOD）检测中，其 AUC-ROC 误差相比蒙特卡洛丢弃法和证据回归法分别降低了 2.5 倍和 2.1 倍。

ABSTRACT

Graph Neural Networks (GNNs) are promising surrogates for quantum mechanical calculations as they establish unprecedented low errors on collections of molecular dynamics (MD) trajectories. Thanks to their fast inference times they promise to accelerate computational chemistry applications. Unfortunately, despite low in-distribution (ID) errors, such GNNs might be horribly wrong for out-of-distribution (OOD) samples. Uncertainty estimation (UE) may aid in such situations by communicating the model's certainty about its prediction. Here, we take a closer look at the problem and identify six key desiderata for UE in molecular force fields, three 'physics-informed' and three 'application-focused' ones. To overview the field, we survey existing methods from the field of UE and analyze how they fit to the set desiderata. By our analysis, we conclude that none of the previous works satisfies all criteria. To fill this gap, we propose Localized Neural Kernel (LNK) a Gaussian Process (GP)-based extension to existing GNNs satisfying the desiderata. In our extensive experimental evaluation, we test four different UE with three different backbones and two datasets. In out-of-equilibrium detection, we find LNK yielding up to 2.5 and 2.1 times lower errors in terms of AUC-ROC score than dropout or evidential regression-based methods while maintaining high predictive performance.

研究动机与目标

• 解决分子力场中可靠不确定性估计的关键需求，尤其是在标准 GNN 失效的分布外（OOD）预测场景下。
• 识别六项关键期望——三项物理信息相关（对称性、能量守恒、局域性）和三项应用导向（速度、可解释性、OOD 检测能力）——以实现分子系统中有效的不确定性估计。
• 调研分子机器学习中的现有不确定性估计方法，并表明目前无一方法能同时满足全部六项期望。
• 提出 LNK，一种新颖的单前向传播不确定性估计方法，将高斯过程回归与 GNN 主干网络结合，以满足所有期望。
• 在两个数据集上，使用多种主干网络和不确定性方法对 LNK 进行实证评估，重点关注 OOD 检测与预测性能。

提出的方法

• 定义分子力场中不确定性估计的六项期望：三项物理信息相关（对欧几里得变换与排列对称性的不变性、通过无旋力场实现的能量守恒），以及三项应用导向（速度、可解释性、OOD 检测能力）。
• 提出局部神经核（LNK），一种可微分的、基于高斯过程的 GNN 扩展，通过在分子图上学习局部核函数来建模不确定性。
• 将 LNK 集成到现有 GNN 主干网络（如 DimeNet++）中，实现无需重训练主干网络的不确定性估计，仅需一次前向传播。
• 采用可微分核机制，尊重物理对称性，并通过构造确保力场为保守场。
• 通过基于变分的方法实现对能量（u_E）和力（u_F）的不确定性量化，其中 u_F 由核函数的迹计算得出。
• 使用结合能量/力预测与不确定性校准的概率损失函数，端到端训练模型，确保不确定性估计具有良好校准性。

实验结果

研究问题

• RQ1 分子力场中不确定性估计的必要物理与应用特定要求是什么？
• RQ2 现有分子机器学习中的不确定性估计方法在多大程度上满足这些期望？
• RQ3 能否设计一种基于高斯过程的方法，以尊重分子对称性与守恒定律，同时实现快速、单前向传播的不确定性估计？
• RQ4 所提出的 LNK 方法在 OOD 检测与预测性能方面，相较于成熟技术（如蒙特卡洛丢弃法、证据回归）表现如何？
• RQ5 将 LNK 集成到 GNN 主干网络中，是否能在显著提升不确定性校准性的同时，保持高预测精度？

主要发现

• 现有分子机器学习中的不确定性估计方法均未满足全部六项所提期望，尤其在物理一致性或推理速度方面存在缺陷。
• 在 QM7X 数据集的非平衡态检测中，LNK 的 AUC-ROC 误差相比蒙特卡洛丢弃法和证据回归法分别降低了 2.5 倍和 2.1 倍。
• LNK 保持了高预测性能，能量 MAE 为 0.059 kcal/mol，力 MAE 为 0.007 kcal/mol（在 1% 丢弃率下），与最先进模型相当。
• LNK 比基于方差的方法（如 MCD 和集成）快逾 4.6 倍，QM7X 数据集上单样本推理时间仅为 59.3±2.68 ms，与证据回归速度相当。
• MCD-LNK（在推理阶段使用丢弃的变体）在增加丢弃率时表现出稳定性能，优于标准蒙特卡洛丢弃法，在 OOD 检测与预测精度方面均更优。
• 集成方法在 AUC-ROC 方面表现最佳（力的 AUC-ROC 达 0.996），但代价是训练时间增加五倍，且推理时间显著更高（274±8.25 ms），使其在主动学习或实时应用中不切实际。