[论文解读] TorchMD-NET: Equivariant Transformers for Neural Network based Molecular Potentials

Philipp Thölke, Gianni De Fabritiis|arXiv (Cornell University)|Feb 5, 2022

Machine Learning in Materials Science被引用 58

一句话总结

TorchMD-NET 引入了一个等变 Transformer，用于预测分子的量子力学性质，在 QM9、MD17 和 ANI-1 基准测试上实现了最先进的精度与效率，并通过注意力分析提供洞见。

ABSTRACT

The prediction of quantum mechanical properties is historically plagued by a trade-off between accuracy and speed. Machine learning potentials have previously shown great success in this domain, reaching increasingly better accuracy while maintaining computational efficiency comparable with classical force fields. In this work we propose TorchMD-NET, a novel equivariant transformer (ET) architecture, outperforming state-of-the-art on MD17, ANI-1, and many QM9 targets in both accuracy and computational efficiency. Through an extensive attention weight analysis, we gain valuable insights into the black box predictor and show differences in the learned representation of conformers versus conformations sampled from molecular dynamics or normal modes. Furthermore, we highlight the importance of datasets including off-equilibrium conformations for the evaluation of molecular potentials.

研究动机与目标

• 开发用于预测分子能量和力的旋转等变 Transformer（ET）架构。
• 在标准量子化学基准 QM9、MD17、ANI-1 上实现高精度和计算效率。
• 通过注意力权重分析和跨数据情形的表示比较，为模型提供可解释性。
• 展示非平衡构象在评估神经网络势中的重要性。

提出的方法

• 通过每个原子两个学习嵌入来对原子类型和局部邻域进行嵌入。
• 使用通过 RBF 展开边数据将原子间距离并入的改进多头注意力机制。
• 在注意力值通路中引入连续滤波图卷积，将距离信息注入特征。
• 应用门控等变块来计算标量和向量更新，通过对坐标对能量求导实现力的预测。
• 使用 Adam 以均方误差训练，包含学习率预热和基于平台期的衰减；在预测能量和力时应用指数平滑。

实验结果

研究问题

• RQ1 等变 Transformer 是否能在 QM9、MD17 和 ANI-1 基准测试上超越现有的神经网络势？
• RQ2 将边缘（距离）信息并入注意力机制如何影响预测精度和效率？
• RQ3 哪些架构组件对性能影响最大（例如邻居嵌入、更新层、向量特征）？
• RQ4 注意力权重是否揭示静态与非平衡构象下有意义的化学与动力学结构？
• RQ5 训练数据多样性（非平衡构象）对能量和力预测有多重要？

主要发现

• 等变 Transformer（ET）在 QM9 目标上的 MAE 达到最先进水平，相对于不变和协变架构，并且与专门的等变模型相媲美或优于之。
• 在 MD17 上，ET 在大多数分子中实现能量和力的领先精度，通常优于先前方法，消融实验显示邻居嵌入和更新层对性能贡献显著。
• 在 ANI-1 上，ET 在大规模非平衡数据集中表现出色，在测试 MAE 上优于若干基线。
• 注意力分析表明，训练数据类型（静态 QM9 与动态 ANI-1/MD17）影响注意力焦点，动态数据集中氢原子得到更多关注，静态数据集中核心原子被强调。
• 非平衡构象对于准确的能量预测很重要，这在注意力模式和消融结果中有所体现。
• 较小的 ET 变体仍与最先进模型竞争，表明在参数数量关系下具有有利的数据效率和可扩展性。