选择性差异性注意力机制增强了笛卡尔原子矩机器学习方法在原子间势能计算中的应用，并提升了跨系统间的通用性

《Computational Materials Science》：Selective differential attention enhanced cartesian atomic moment machine learning interatomic potentials with cross-system transferability

SA-CAMP模型通过选择性微分注意力机制与Cartesian原子矩展开结合，有效解决传统方法O(L^6)复杂度和远距离交互建模问题，在多个体系验证其高效性与高精度，代码开源。

　　

Xuan Li|Pandi Teng|Yunna Ou|Zhao Niu|Shu Zhan|Jiajia Xu

合肥工业大学计算机与信息工程学院，中国合肥

摘要

机器学习原子间势（MLIPs）的出现通过将第一性原理的精确性与传统力场方法的快速性相结合，彻底改变了分子动力学模拟。然而，当前的方法面临两个关键挑战：与球谐描述符相关的过高O(L^6)阶缩放问题，以及消息传递神经网络（MPNNs）在捕捉远距离原子相互作用方面的固有困难。本文介绍了选择性差分注意力笛卡尔原子矩势（SA-CAMP），它将笛卡尔原子矩势框架与选择性差分注意力机制相结合。该机制仅作用于标量特征，同时保留了完整的张量信息以确保几何精度。我们的方法使用笛卡尔坐标和切比雪夫多项式构建原子矩，然后仅对标量超矩（零秩张量）应用多头差分注意力。SA-CAMP有效地捕捉了全局相关性，并通过标量差分注意力消除了注意力噪声，同时通过向量和张量特征保持了局部几何关系，同时确保了旋转、平移和排列对称性的完全保留。该模型在各种系统上表现出色，例如 bulk water、双层石墨烯和碳同素异形体。我们的代码可在以下链接公开获取：https://github.com/lixuan000/SA-CAMP

引言

机器学习原子间势（MLIPs）代表了在计算精度（如从头算分子动力学AIMD）和计算效率（如经典分子动力学方法）之间的一个重要进步[1]、[2]。通过捕获密度泛函理论计算中的结构配置与能量属性之间的联系，MLIPs 的计算速度比传统从头算方法快几个数量级，同时保持了接近量子力学的精度[3]、[4]。这种平衡使研究人员能够探索更大的分子系统和更长时间尺度的动态过程，从而促进新材料发现和分子设计的新方法。

近年来，基于图神经网络的MLIPs取得了显著进展，等变性成为其核心设计原则[5]、[6]。这些模型不仅提高了预测精度，还通过显式保持基本物理对称性显著提高了数据效率[7]、[8]。在各种等变表示方法中，球谐函数因其作为SO(3)群不可约表示的自然基础而被广泛采用[9]、[10]。然而，基于球谐函数的高阶几何表示通常需要复杂的克莱布施-戈登系数计算和张量积运算，导致计算复杂性呈指数级增长，即O(L^6)，与最大角动量量子数L成正比[11]、[12]。

相比之下，笛卡尔坐标系中的原子环境表示在简单性和效率方面具有潜在的优势。笛卡尔原子矩势（CAMP）框架直接在笛卡尔空间中构建原子矩张量，消除了球谐展开的复杂性，同时保留了局部原子结构的完整特征[13]。同样，笛卡尔原子簇展开（CACE）方法通过消除多项式依赖特征，展示了用笛卡尔坐标描述原子环境的有效性[14]。这些工作表明，笛卡尔表示不仅在计算效率上具有某些优势，还可以系统地改进多体相互作用的建模。

尽管笛卡尔表示在处理局部相互作用时具有计算优势，但MLIPs 在处理非局部电子效应时仍面临重大挑战[2]、[15]。传统的消息传递神经网络（MPNNs）由于采用了局部邻域假设，通常难以有效描述超出截止半径的长程相互作用[16]、[17]。在描述涉及长程相关性的化学现象（如共轭系统、氢键网络和色散相互作用）时，这一限制尤为明显[18]、[19]。虽然简单地增加截止半径或堆叠更多消息传递层可以在一定程度上扩展感受野，但这种方法往往以牺牲计算效率为代价，并未能从根本上解决长程相互作用的问题。SO3KRATES模型通过引入球谐坐标（SPHCs）和球形自注意力机制，成功实现了任意长度尺度上的非局部量子效应建模，展示了注意力机制在编码远距离结构依赖性方面的有效性[20]。同时，基于Transformer架构的分子建模方法也显示出巨大潜力，Geoformer通过设计的原子间位置编码（IPE）将几何信息有效整合到Transformer架构中，在多个分子属性预测基准测试中取得了优异性能[21]。

然而，现有的基于注意力机制的MLIPs在实际应用中仍面临几个重要挑战。首先，不加选择地将注意力机制应用于所有特征组件会显著增加计算复杂性，尤其是在处理高阶张量特征时[22]、[23]。其次，传统的softmax注意力机制存在固有的注意力噪声，这往往会将过多的注意力权重分配给无关的上下文信息，从而降低模型关注关键化学相互作用的能力[24]。第三，注意力操作可能会破坏物理对称性，而这些对称性对于确保模型的物理合理性至关重要[25]。最后，如何在保持注意力机制的表现力的同时，保留原子矩展开方法的可解释性和物理严谨性，仍然是一个亟待解决的问题[26]、[27]。

基于此背景，我们提出了选择性差分注意力笛卡尔原子矩势（SA-CAMP），这是一种创新的MLIPs模型，它策略性地将选择性差分注意力机制与笛卡尔原子矩表示相结合。我们的核心创新在于巧妙地将差分注意力机制专门应用于标量特征，从而在保持基本物理对称性的同时实现原子系统之间的高效全局信息交换。差分注意力通过计算两个独立注意力图之间的差异来有效消除注意力噪声，使模型能够更精确地关注关键化学相互作用。这种选择性设计不仅显著降低了计算复杂性，还确保了关键物理对称性（如旋转不变性、平移不变性和排列等变性）的严格保持。

SA-CAMP 的设计基于一个物理考虑：由于目标势能是一个旋转和平移不变的标量，等变消息传递逐渐将方向依赖的几何信息（由向量和高阶张量特征携带）总结为用于能量读出的不变标量通道，使标量特征成为与势能表面最相关的全球相关性的自然且对称安全的载体[28]、[29]。通过选择性地将差分注意力机制应用于标量特征，我们可以有效建模长程相互作用，并在避免其他特征组件的不必要计算开销的同时消除注意力噪声。所提出的方法在各种基准数据集上的表现与当前领先技术相当或更优。

架构概述

我们提出了SA-CAMP，这是一种新颖的机器学习原子间势框架。该模型将物理驱动的CAMP框架与选择性多头差分注意力机制相结合，有效保持了原子系统的固有物理对称性，同时保持了计算效率。我们使用无向图 $G\left(V,?\right)$ 表示原子结构，其中顶点集 对应于单个原子，边集 编码原子之间的成对相互作用。

实现细节

在本节中，我们比较了SA-CAMP模型与其他模型在水数据集、双层石墨烯数据集和GAP17数据集上的能量和力预测误差。对于水和GAP17数据集，我们使用均方根误差（RMSE）进行准确性评估，而对于双层石墨烯数据集，则使用平均绝对误差（MAE）作为评估指标。我们还进行了一系列实验来验证SA-CAMP在分子动力学模拟中的优异性能。

结论

本研究提出了SA-CAMP，它通过将笛卡尔原子矩表示与选择性差分注意力机制相结合，有效解决了球谐表示的O(L^6)计算复杂性和MPNNs在建模长程相互作用方面的局限性。SA-CAMP 的核心创新在于仅将多头差分注意力应用于标量特征，实现了高效的全局信息交换，同时保留了完整的张量信息。

CRediT作者贡献声明

Xuan Li：撰写 – 审稿与编辑、撰写 – 原稿、可视化、软件、方法论、研究、形式分析、数据整理、概念化。 Pandi Teng：资源、方法论、形式分析。 Yunna Ou：资源。 Zhao Niu：资源。 Shu Zhan：验证、监督、资源获取。 Jiajia Xu：验证、监督、资源。

利益冲突声明

作者声明他们没有已知的竞争性财务利益或个人关系可能影响本文报告的工作。

致谢

本研究得到了安徽省重大科技项目（项目编号：JZ2024AKKZ0025）的支持。计算工作在合肥工业大学的HPC平台上完成。