[论文解读] Spherical Channels for Modeling Atomic Interactions

C. Lawrence Zitnick, Abhishek Das+6 more|arXiv (Cornell University)|Jun 29, 2022

Machine Learning in Materials Science被引用 36

一句话总结

介绍了球面通道网络（SCN），一种使用球面谐波通道来建模原子能量和力的图神经网络，具备可旋转引导的嵌入和变体等变性，在 energy 和 force 预测上在 OC20 上达到最新研究成果。

ABSTRACT

Modeling the energy and forces of atomic systems is a fundamental problem in computational chemistry with the potential to help address many of the world's most pressing problems, including those related to energy scarcity and climate change. These calculations are traditionally performed using Density Functional Theory, which is computationally very expensive. Machine learning has the potential to dramatically improve the efficiency of these calculations from days or hours to seconds. We propose the Spherical Channel Network (SCN) to model atomic energies and forces. The SCN is a graph neural network where nodes represent atoms and edges their neighboring atoms. The atom embeddings are a set of spherical functions, called spherical channels, represented using spherical harmonics. We demonstrate, that by rotating the embeddings based on the 3D edge orientation, more information may be utilized while maintaining the rotational equivariance of the messages. While equivariance is a desirable property, we find that by relaxing this constraint in both message passing and aggregation, improved accuracy may be achieved. We demonstrate state-of-the-art results on the large-scale Open Catalyst dataset in both energy and force prediction for numerous tasks and metrics.

研究动机与目标

• 动机并解决对原子能量与力的 DFT 高计算成本的问题。
• 开发通过球面通道编码角信息的图神经网络，以提高预测精度。
• 探索旋转等变性及其放宽在提升模型表达能力与性能方面的作用。
• 在 Open Catalyst 2020 (OC20) 数据集的能量与力任务上展示最新性能。
• 评估相对于现有最先进模型的样本效率与可扩展性。

提出的方法

• 在基于图的模型中，将每个原子表示为一个节点，邻居表示为边。
• 将节点嵌入编码为球面上的一组球面通道，利用球谐函数扩展至阶数 L。
• 将边缘定向嵌入旋转以与边的方向对齐，在信息传递过程中传递角信息。
• 使用神经网络计算边消息，并将其旋转回全局坐标系，必要时使用多次 z 轴旋转（taps）以促进等变性。
• 对球面网格应用逐点、无约束的非线性函数来聚合信息，然后转换回球谐函数以进行更新。
• 通过对取向进行球面积分的神经网络来估计每个原子能量，并通过能量梯度或在球面方向上的直接力估计来计算力。
• 使用 AdamW 进行训练，利用数据并行和混合精度，在适用时使用球谐函数的多分辨率（带数）以在内存和精度之间取得平衡。

实验结果

研究问题

• RQ1 SCN 的球面通道嵌入是否能在 OC20 上优于现有 GNN 提高能量和力的预测？
• RQ2 放宽严格的旋转等变性（超出 m=0）如何影响模型精度和训练效率？
• RQ3 在力预测中使用旋转嵌入和多次 taps 对角信息的影响是多少？
• RQ4 与此前的最先进模型相比，SCN 在 OC20 任务中的样本效率如何？

主要发现

• SCN 在 S2EF 任务的力 MAE 和力余弦上相比若干基线取得了最新的改进。
• 在 OC20 上，m 在 [-1,1] 或 [-2,2] 的 SCN 变体相对于严格的 m=0 有提升，显示放宽严格等变性的好处。
• 将 SCN 扩展到更大的 L 和更深的 K 可获得强劲的性能，在 IS2RE 能量预测和总体力精度方面有显著提升。
• 在 F_e 计算中对多次旋转（taps）进行采样可显著改善力预测，虽然会在吞吐量上带来一定成本；而使用两个带（L, L-1）则在吞吐量影响适中的前提下提升能量和力预测。
• 直接力预测（与基于梯度的能量力相比）在多种情况下提供更优的速度/精度权衡，同时仍保持竞争力的精度。
• SCN 显示出比若干强基线更高的样本效率，这意味着它在较少的训练样本下也能达到相当的精度。