[论文解读] Fast and Uncertainty-Aware Directional Message Passing for Non-Equilibrium Molecules

Johannes Gasteiger, Shankari Giri+2 more|arXiv (Cornell University)|Nov 28, 2020

Machine Learning in Materials Science被引用 160

一句话总结

DimeNet++ 提速定向信息传递用于非平衡分子，在实现更高精度的同时，利用 COLL 进行非平衡研究，并进行不确定性量化分析。

ABSTRACT

Many important tasks in chemistry revolve around molecules during reactions. This requires predictions far from the equilibrium, while most recent work in machine learning for molecules has been focused on equilibrium or near-equilibrium states. In this paper we aim to extend this scope in three ways. First, we propose the DimeNet++ model, which is 8x faster and 10% more accurate than the original DimeNet on the QM9 benchmark of equilibrium molecules. Second, we validate DimeNet++ on highly reactive molecules by developing the challenging COLL dataset, which contains distorted configurations of small molecules during collisions. Finally, we investigate ensembling and mean-variance estimation for uncertainty quantification with the goal of accelerating the exploration of the vast space of non-equilibrium structures. Our DimeNet++ implementation as well as the COLL dataset are available online.

研究动机与目标

• 将机器学习分子建模扩展到非平衡反应和扭曲构型。
• 开发用于能量和原子力的快速、准确且具有广泛泛化性的 GNN。
• 创建 COLL，一个包含反应性分子构型的非平衡数据集。
• 研究对非平衡态能量和力预测的不确定性量化策略。

提出的方法

• 用 Hadamard 乘积替换 DimeNet 的昂贵双线性定向相互作用，并以 2-layer MLPs 为基表示进行补偿。
• 实现带有下投影和上投影层的嵌入层次结构，以减小嵌入尺寸并提升速度与表达能力。
• 在可能的情况下将网络深度从 6 层降至 4 层，并调整嵌入/张量表示以保持精度。
• 在更新的架构中保留定向信息传递，以快速、可扩展的方式建模键长和键角。
• 在 QM9 和 COLL 数据集上进行评估，以评估非平衡构型下的速度、准确性和鲁棒性。
• 探索不确定性量化方法（集成与均方差估计）及其在能量与力预测中的局限性。

实验结果

研究问题

• RQ1 DimeNet++ 是否能够在 QM9 等平衡基准上实现更快的运行时，同时提高或维持精度？
• RQ2 DimeNet++ 如何在 COLL 这类高度非平衡、具有反应性的构型上实现良好的泛化？
• RQ3 标准的不确定性量化方法（集成、均值-方差估计）在非平衡态下对能量和力的不确定性预测是否可靠？
• RQ4 影响 ML 模型训练与不确定性估计的非平衡分子数据的特征与挑战是什么？

主要发现

• DimeNet++ 比原始 DimeNet 快 8x，平均精度提高约 ~10%，对 QM9 的能量精度约提升 ~20%。
• 在 COLL 数据集上，DimeNet++ 显著优于 SchNet（对 E 的 MAE = 0.047 eV，F 的 MAE = 0.040 eV，DimeNet++，相对于 SchNet 的更高误差）。
• COLL 数据集（140k 配置）涵盖更广的非平衡空间，包含扭曲的键长/键角，带来比 QM9 更强的挑战。
• 通过集成和均值-方差估计的不确定性量化存在挑战：能量不确定性并不能可靠地作为力不确定性的代理，且集成带来显著的计算开销。
• 集成在力的不确定性估计方面提供更好的结果，但成本更高；而 MVE 对力不确定性的有用性有限。