[论文解读] ForceNet: A Graph Neural Network for Large-Scale Quantum Calculations

Weihua Hu, Muhammed Shuaibi+6 more|arXiv (Cornell University)|Mar 2, 2021

Machine Learning in Materials Science参考文献 53被引用 24

一句话总结

ForceNet 是一个可扩展的以力为中心的图神经网络，在 OC20 数据集的原子力预测方面超越基于物理的 GNN，同时在训练和推理方面提供显著的加速；ForceNet-large 与 DimeNet++ 相比在计算量显著更少的情况下达到或超过其性能。

ABSTRACT

With massive amounts of atomic simulation data available, there is a huge opportunity to develop fast and accurate machine learning models to approximate expensive physics-based calculations. The key quantity to estimate is atomic forces, where the state-of-the-art Graph Neural Networks (GNNs) explicitly enforce basic physical constraints such as rotation-covariance. However, to strictly satisfy the physical constraints, existing models have to make tradeoffs between computational efficiency and model expressiveness. Here we explore an alternative approach. By not imposing explicit physical constraints, we can flexibly design expressive models while maintaining their computational efficiency. Physical constraints are implicitly imposed by training the models using physics-based data augmentation. To evaluate the approach, we carefully design a scalable and expressive GNN model, ForceNet, and apply it to OC20 (Chanussot et al., 2020), an unprecedentedly-large dataset of quantum physics calculations. Our proposed ForceNet is able to predict atomic forces more accurately than state-of-the-art physics-based GNNs while being faster both in training and inference. Overall, our promising and counter-intuitive results open up an exciting avenue for future research.

研究动机与目标

• 使用大规模数据对昂贵的量子物理计算在原子力方面进行快速、准确的近似。
• 开发一个灵活、表达力强的以力为中心的 GNN，使其在体系结构中不依赖显式的物理约束。
• 利用基于物理的数据增强来对旋转相关的物理属性进行隐式约束。
• 研究在 OC20（200M+ 结构的量子数据集）上的可扩展性和效率，包括非平衡结构。

提出的方法

• 设计 ForceNet，采用编码器–解码器 GNN 架构，其中原子为节点，边通过条件滤波卷积捕获三维相互作用。
• 使用边特征 e_st，编码完整的三维相对位置及旋转协变方向信息。
• 引入基函数 B，将原始边特征转换为更丰富的表示（包括球谐函数变体），并在边与节点的 MLP 中使用表达力强的非线性激活（Swish）。
• 在训练中应用旋转数据增强，以促进旋转协变预测，而无需强制实现显式的架构旋转约束。
• 通过解码器 MLP 直接从最终节点嵌入构建力预测，避免显式的能量守恒约束以提升可扩展性。
• 在 OC20 S2F 任务上与 SchNet、DimeNet++、DimeNet++-large 和 GNS 进行对比，同时对训练/推理预算进行严格控制。
• 评估一系列基函数（单位、线性+激活、高斯、正弦、球谐函数）和激活函数（Swish 效果最好），以研究它们对性能的影响。

实验结果

研究问题

• RQ1 一个以力为中心的 GNN，在没有显式旋转约束的情况下，是否在预测大规模非平衡数据集的原子力方面达到或超过基于物理的 GNN 的最新水平？
• RQ2 结构选择（基函数、激活、边条件）和数据增强如何影响力预测精度与旋转协变性？
• RQ3 在 OC20 上扩展 ForceNet 与现有模型相比，训练/推理效率的权衡是什么？
• RQ4 旋转增强在非旋转不变模型中在多大程度上能够实现旋转协变行为？

主要发现

• ForceNet 在 OC20 上的原子力预测方面，在可比的计算量下超越基于物理的 GNN，且训练和推理更快。
• ForceNet-large 在力预测方面达到或优于 DimeNet++-large，同时需要的计算量显著更少（训练约少 8.2x–18.9x，推理约快 2.4x–3.0x，视对比而定）。
• ForceNet–large 在平均力 MAE 上优于 DimeNet++，并提供显著更快的训练和推理；早期结果显示与 SchNet 相比达到类似性能仅需约少 6.3x 的计算量，而与 DimeNet++ 相比推理快至约 6.5x。
• 旋转增强和大规模数据有助于降低旋转稳定性问题，使 ForceNet 的预测更接近旋转协变行为。
• 消融研究表明，表达力强的边级计算以及含 Swish 激活的球面基函数对 ForceNet 的性能至关重要；移除这些要素会降低精度。