ICDP：一种用于原子模拟的不可约化笛卡尔动态注意力势

《Computational Materials Science》：ICDP: an irreducible Cartesian dynamic-attention potential for atomistic simulations

　　

- ICDP采用紧凑的不可约笛卡尔张量编码原子环境。
- 双分支动态注意力在保持等变性的同时重新加权邻居。
- 克莱布施-戈丹耦合捕获高阶局域几何相互作用。
- 消融实验证实了不可约张量与注意力带来的互补增益。
- ICDP能够准确描述结构、输运和振动性质。

## 引言

分子动力学（MD）模拟的预测质量主要取决于所选的原子间模型能多么忠实地再现底层能量面。传统的经验势函数计算效率高，但受限于预定义的函数形式和人工参数化，这限制了其准确描述复杂多体相互作用和化学环境变化的能力。第一性原理方法精度更高，但通常不适用于大规模、长时间的原子模拟。机器学习原子间势（MLIPs）从量子力学参考数据构建计算高效的替代能量模型，它们以接近第一性原理的精度预测原子能量和力，从而实现更大长度和时间尺度的模拟 [1], [2]。

早期MLIPs主要依赖预定义的旋转不变描述符。Behler-Parrinello神经网络将总能量分解为局域原子能量，并使用原子中心对称函数来表示局域环境；而高斯近似势（GAP）对邻近原子排列的表示应用高斯过程推断，生成的能量模型随着更多量子力学训练数据的增加可以逐步优化 [3], [4], [5]。随后，SchNet通过连续滤波卷积直接从原子种类和原子间距离学习表示，减少了对手工设计特征的依赖。DimeNet进一步纳入了键角和方向性消息传递，以提高局域几何关系的分辨率 [6], [7]。这些不变模型相对简单，自然保持能量的旋转不变性，并保持计算高效。然而，方向信息通常需要间接地通过距离、键角或高阶描述符来编码。当局域环境包含复杂的多体几何相关性时，仅传播标量不变量可能会压缩方向信息，而显式构造三元组和高阶组合则会增加计算和内存成本。

等变图神经网络的发展大幅缓解了这些局限性。可极化原子相互作用神经网络（PaiNN）在消息传递过程中将标量通道与向量特征耦合，实现方向信息的高效传递；然而其对原子几何的描述主要依赖于有限张量阶的特征 [8]。神经等变原子间势（NequIP）使用E(3)不可约表示和等变张量积传播不同角动量阶的特征，实现了高数据效率和预测精度。消息传递原子团簇展开（MACE）将等变消息与升高的体阶纳入基于团簇的表示，仅需少数交互层即可捕获复杂的多体相关性。Allegro采用严格局域等变表示，以提高大规模并行原子模拟的可扩展性 [9], [10], [11]。这些研究表明，将物理对称性直接嵌入神经网络可以提高几何表达力和数据效率。然而，基于球谐函数、克莱布施-戈丹系数和高阶张量积的等变运算随着最大角动量和耦合路径数量的增加而变得越来越昂贵。等变球面通道网络（eSCN）等方法通过将SO(3)卷积转换为SO(2)操作来减少部分开销，但对于高阶等变特征如何平衡精度、速度和内存消耗仍然是模型设计中的重要挑战 [12]。

除了球张量表示之外，笛卡尔张量为描述原子环境提供了一种更直观的途径。TensorNet使用二阶及以下的笛卡尔张量及其不可约分量，在低张量阶实现高效的等变消息传递。笛卡尔原子团簇展开（CACE）在笛卡尔空间中构建具有显式体阶的原子团簇展开。笛卡尔原子矩势（CAMP）使用径向函数、边方向的笛卡尔张量积、原子矩和高阶矩，在图神经网络中系统地编码高阶多体相互作用 [13], [14], [15]。笛卡尔表示简洁、几何直观，并且易于与常规矩阵操作兼容。然而，传统的可约笛卡尔张量包含多个旋转阶和冗余的迹分量。在密集存储下，秩-?张量包含3^?个笛卡尔分量，导致高阶特征的存储和计算成本快速增长。将模型限制为低秩张量可以控制这种成本，但可能会限制其向更高角分辨率的系统扩展。

不可约笛卡尔张量提供了一种将笛卡尔表示的直观性与不可约表示的紧凑性相结合的方法。对秩-?笛卡尔张量施加置换对称性和零迹条件后，留下2?+1个自由度，它们在SO(3)下按照与角动量?的球面张量相同的不可约表示进行变换。Cartesian-nj进一步建立了不可约笛卡尔张量分解、张量积和缩并的统一框架，以及笛卡尔张量空间与不可约系数空间之间的基变换工具 [16]。该框架明确分离了不同的旋转阶并去除了迹和重复分量，同时保留了笛卡尔构造。然而，如何将这种紧凑表示与CAMP中的原子矩和高阶矩层级以及受限耦合路径相结合，并在不违反旋转等变性的前提下自适应调节邻居消息，仍有待探索。

传统消息聚合可以通过依赖元素的径向滤波器区分不同的原子对，但通常缺乏根据当前局域隐藏状态动态重新分配邻居消息的显式机制。因此，注意力机制越来越多地被引入图神经网络中。MLANet使用几何感知的双路径动态注意力来调节消息传递，而SA-CAMP对CAMP的标量高阶矩应用选择性微分注意力，以增强标量特征之间的相互作用 [17], [18]。这些研究表明动态注意力可以改善环境依赖的信息选择。然而，从高阶张量分量直接构建非不变注意力分数可能会增加计算成本并违反等变消息的变换律。

基于以上考虑，我们提出了不可约笛卡尔动态注意力势（ICDP）。ICDP基于CAMP的原子矩和高阶矩层级 [15] 以及现有的不可约表示工具 [16]，重点在基于CAMP的不可约矩框架内实现不变边重新加权。受几何感知注意力 [17] 的启发，标量环境分数与不可约查询/键特征的不变缩并共同决定原子矩聚合之前邻居消息的权重。由此得到的加权原子矩随后进入高阶矩构造，使得学习到的边权重可以进一步调节每个交互层内形成的多邻居组合。我们通过五种子消融实验评估骨干修改和注意力机制的效果。特别地，四种模型配置共享相同的不可约骨干，并分别检验无注意力、标量注意力、几何注意力及其组合。基准结果表明，ICDP在体相水、LiPS和GAP17数据集上具有具有竞争力的能量和力预测精度。分子动力学（MD）模拟和声子计算进一步评估了模型描述所研究系统结构、动力学和振动性质的能力。

## 节选段落

**整体架构**

ICDP使用CAMP的原子矩和高阶矩层级 [15] 构建局域多体表示。边方向张量使用Cartesian-nj [16] 的不可约笛卡尔张量分解工具获得的不可约系数来表示。在该框架内，ICDP在原子矩聚合之前为邻居消息分配旋转不变权重。这些权重结合了标量环境分数与不可约查询/键特征的不变缩并。

**实现细节**

ICDP使用PyTorch和PyTorch Lightning实现，所有模型在单块NVIDIA RTX 4090 GPU上训练。对于一个包含??个原子构型的迷你批次，第??个构型包含????个原子。总能量和原子力用作联合监督信号，训练目标是能量和力均方误差的加权和：

? = ??? + ???

能量损失按每个构型中的原子数进行归一化，而力损失则取平均。

## 结论

本研究提出了ICDP，它结合了笛卡尔表示的几何直观性、不可约表示的紧凑性以及局域消息聚合的环境适应性。ICDP将边方向的高阶笛卡尔张量积投影到维度为2?+1的不可约系数空间，明确分离不同的旋转阶并去除冗余的迹分量。满足角动量选择规则的克莱布施-戈丹耦合……

## CRediT作者贡献声明

庄怀娟：资源、方法论、形式分析。李璇：论文撰写——审阅与编辑、论文撰写——原稿、可视化。李阳：验证。金晶：验证、监督、资源、资金获取。赵冲：验证、资源。杨帆：监督。牛钊：资源。

## 利益冲突声明

作者声明不存在任何已知可能影响本文工作的竞争性经济利益或个人关系。

## 致谢

本研究得到了教育部产学合作协同育人项目（项目编号：250603873093026）和中央高校基本科研业务费专项资金（项目编号：PA2025GDSK0036）的资助。计算在合肥工业大学高性能计算平台上完成。

庄怀娟 | 李璇 | 李阳 | 金晶 | 赵冲 | 杨帆 | 牛钊