[论文解读] Spatial Attention Kinetic Networks with E(n)-Equivariance

Yuanqing Wang, John D. Chodera|arXiv (Cornell University)|Jan 21, 2023

Machine Learning in Materials Science被引用 11

一句话总结

SAKE 引入一种简单的空间注意力形式，该形式在 E(n) 不变/协变，并利用它构建一个快速、准确的 SAKE 模型，用于多体系统中的不变和协变任务，超越若干基线，同时速度显著更快。

ABSTRACT

Neural networks that are equivariant to rotations, translations, reflections, and permutations on n-dimensional geometric space have shown promise in physical modeling for tasks such as accurately but inexpensively modeling complex potential energy surfaces to guiding the sampling of complex dynamical systems or forecasting their time evolution. Current state-of-the-art methods employ spherical harmonics to encode higher-order interactions among particles, which are computationally expensive. In this paper, we propose a simple alternative functional form that uses neurally parametrized linear combinations of edge vectors to achieve equivariance while still universally approximating node environments. Incorporating this insight, we design spatial attention kinetic networks with E(n)-equivariance, or SAKE, which are competitive in many-body system modeling tasks while being significantly faster.

研究动机与目标

• 动机：在神经模型中引入几何感知的归纳偏置以用于物理系统的必要性。
• 提出一个简单的、神经参数化的空间注意力机制，能够实现 E(n)-不变性/协变性。
• 开发利用空间注意力来进行 n-body 动力学中速度和位置更新的 SAKE 架构。
• 在标准基准测试中，在不变性（势能）和协变性（动力学）任务上展示具有竞争力的性能和更高的速度。

提出的方法

• 将空间注意力定义为一个函数，该函数聚合对边向量的神经参数化线性组合，并取其范数来形成不变描述符。
• 采用类似多头的设置：在边嵌入上使用函数 lambda_i，在边向量上使用一个协变的 f，随后是处理连接范数的 mu。
• 用 SchNet 风格的边嵌入编码成对相互作用，将节点嵌入与距离信息通过径向基函数结合。
• 引入一个虚拟的速度更新和类似欧拉的位置更新，以创建一个几乎哈密顿的积分器，确保动力学的 E(n)-Equivariance。
• 包含语义注意力和基于距离截断的注意力，以促进相互作用中的各向异性。
• 提供基于消融的支撑性论证，表明空间注意力是性能的主要贡献者。

实验结果

研究问题

• RQ1 一个简单的神经参数化的空间注意力机制是否能够实现局部节点环境的 E(n)-不变性/协变性？
• RQ2 与最先进的模型在准确性和速度方面相比，SAKE 在不变性任务（如势能预测）和协变性任务（如动力学预测）上的表现如何？
• RQ3 空间注意力、语义注意力以及速度/位置更新对模型性能的影响是什么？

主要发现

• SAKE 在 MD17、ISO17、QM9 和合成的 N-body 动力学基准上，在不变性和协变性任务中实现具有竞争力或更优的表现。
• SAKE 通常比若干最先进的协变模型需要更少的训练和推理时间。
• 消融研究表明空间注意力有助于改善能量和力的预测，其他组件进一步提升。
• SAKE 展示了对未见化学空间（ISO17 未知分子）的强外推能力，以及有利的轨迹预测结果。
• 与使用球面调和函数的方法相比，SAKE 提供了一个更快的替代方案，在多种任务上具有竞争力的准确性。