面向时空动力学的等变图神经网络中时间尺度与环境的解耦

《Pattern Recognition》：Decoupling time scales and environments in equivariant graph neural networks for spatio-temporal dynamics

　　

•不变量标量控制将原子态、局部环境和时间振幅解耦，实现基于E(3)-等变坐标预测器的独立表征。
•环境感知的校准在保持用于注意力和门控的学习环境归属关系的同时，对齐无关标量通道的尺度变化。
•局部趋势-残差时间模块在短期历史窗口内分离持续性运动和残差运动，并保留它们的不变量振幅。
•混合键合与非键合图将校准后的标量状态与基于相对位移的等变消息传递相结合。
•在MD17和MD22数据集上的实验表明，在评估协议下一步坐标预测得到改善，且短期自回归误差累积减少。

**引言**

分子动力学（MD）模拟[1]是计算化学、材料科学和生物物理领域的核心工具。受MD启发的动力学已被应用于聚类和全局优化等模式识别任务[2]。为在降低计算成本的同时接近第一性原理精度，机器学习模型被广泛用于分子建模以及近似势能面和力场[3]、[4]。等变表示学习是模式识别领域的活跃研究方向[5]，包括用于三维配准的旋转稳定描述子[6]、旋转不变的点云学习[7]，以及对等变性何时最有益的深入分析[8]。图神经网络和图Transformer也在模式识别和计算机视觉领域得到广泛研究[9]。E(3)-等变图神经网络通过构造保持欧几里得对称性，代表性架构包括张量场网络[10]、SE(3)-Transformer[11]和EGNN[12]。对于动力学任务，近期变体纳入了显式时间结构，例如基于神经ODE的SEGNO[13]和基于时空注意力的ESTAG[14]。相关的分子图表示学习也出现在模式识别文献中[15]、[16]。

等变MD预测的常见范式将每一帧表示为分子图，利用节点标量特征结合短坐标历史来回归下一步坐标或速度[13]、[14]。大多数模型之间的差异在于如何实现等变性和时间聚合，包括球谐函数和张量场[17]、基于相对位移的坐标更新以及与连续时间积分的耦合。一旦施加了空间等变约束，性能往往受限于三个瓶颈：以坐标为中心的状态设计、时间尺度混合，以及环境依赖的标量缩放[18]。

首先，许多模型以坐标为中心。它们将大部分容量分配给等变坐标更新，而将节点标量视为辅助输入，通常仅包含原子序数和键类型等粗略标识符。结果是，动态状态被迫集中到坐标通道中，同一坐标必须同时表示静态结构和时变运动。

其次，分子轨迹具有明显的时间多尺度结构。局部键振动快速且振幅小，而扭转、折叠事件和协同重排则演变得更长的时间尺度[15]。在固定采样间隔和直接一步回归的情况下，训练信号被普遍存在的快速模式所主导，而慢速构象组分获得的有效监督较弱。这种不平衡可能降低优化稳定性，并可能导致迭代展开时误差累积，即使一步误差很小。缺乏显式的时间趋势-残差分解与对齐机制[19]时，高频和低频分量在整个时空处理流程中保持耦合。

第三，原子环境高度异质。相同元素在不同配位、刚度、极化率和电荷转移条件下可能表现出显著不同的局部势能和振动振幅[20]。在许多等变动力学模型（如[12]、[14]）中，标量输入仅限于原子序数和键类型等基本标识符，标量归一化通常仅在粗粒度图或层级别施加。因此，隐藏标量通道在不同原子和不同分子间可能发展出环境依赖的尺度和分布偏移，导致更新不平衡和跨系统泛化性能下降。

为解决这些瓶颈，我们提出了一种特征解耦的E(3)-等变模型，其中标量状态、时间信息和相对几何各承担独立角色。内在原子属性和轨迹统计表示为不变量标量。局部时间历史被分解为平滑趋势、归一化残差及其不变量振幅。空间图结合了半径邻域与共价键，下游更新使用标量门控和注意力权重来调制相对位移向量。这种设计将环境信息和时间尺度信息保留在不变量通道中，同时保持等变坐标更新。

我们的贡献有三点。首先，我们从归一化的理化属性构建不变量标量状态，并使用跨节点和环境感知的统计量校准其通道尺度。学习到的环境归属关系被保留为注意力和门控的条件变量。其次，我们引入了一种基于可学习单边核的时间趋势-残差分解与对齐模块。它提供了慢速和快速等变向量分支以及不变量振幅标量。第三，我们将校准后的标量、时间振幅和混合键合与非键合图集成到一个等变时空预测器中，其坐标更新为相对位移的线性组合。在MD17和MD22上的实验，连同短期展开评估以及关于校准、标量状态设计和时间分解的消融实验，证明了更好的一步坐标预测和更少的展开误差累积，且在更大更异质的系统上增益更显著。

**相关工作**

我们回顾了与本方法相关的三条研究线索：用于分子系统的E(3)-等变GNN、用于多尺度分子动力学的时间建模，以及用于异质原子环境的表示和归一化方法。我们将这些线索与实验中使用基线方法进行关联，并突出了我们的设计所填补的空白。

**任务与对称性**

考虑一个包含N个原子的分子，其离散时间三维轨迹为$X(t)=\{x_i(t)\in\mathbb{R}^3\}_{i=1}^{N}$，$t\in\{0,\ldots,T-1\}$，采样间隔为$\Delta t>0$。给定历史窗口$X(0:T-1)$，任务是在欧几里得群$E(3)$的对称性下预测下一步坐标$\hat{X}(T)$。等变性要求对于任何刚体运动$g=(R,b)\in E(3)$：

$$\hat{X}_{T}(g\cdot X(0:T-1))=g\cdot\hat{X}_{T}(X(0:T-1))$$

其中$(g\cdot x_i)(t)=Rx_i(t)+b$。

这一约束编码了分子势能和力仅依赖于相对这一事实。

**概述**

整体框架如图1所示。该架构通过不变量标量控制与基于相对位移的等变更新的耦合分解，实例化了第3.3节中定义的允许预测器。标量模块构建校准后的不变量状态和软环境归属关系。时间模块将坐标历史转换为慢趋势向量分支、归一化残差向量分支和不变量振幅标量。空间模块……

**数据集与设置**

我们在MD17[3]和MD22[34]分子动力学基准数据集上进行评估，涵盖从小型有机分子到较大的生物分子和超分子系统。表1总结了这些分子及其规模。

MD17包含8个小分子（9–21个原子）：阿司匹林、苯、乙醇、丙二醛、萘、水杨酸、甲苯和尿嘧啶。MD22包含7个更大且更异质的系统（42–370个原子），覆盖肽、脂质、核酸片段、碳水化合物和超分子系统。

**结论**

本工作研究了分子动力学中基于E(3)-等变图神经网络的有监督坐标预测。我们聚焦于在欧几里得对称性施加后仍然存在的两个困难来源：在以坐标为中心的表示中局部残差运动与持续趋势运动的耦合，以及当节点特征受限于粗粒度原子标签时隐藏标量状态中环境依赖的尺度变化。

我们提出了一种不变量标量控制的特征解耦……

**CRediT作者贡献声明**

李建伟：可视化、方法论、数据整理。陈静：写作—审阅与编辑、方法论、概念化。徐文强：可视化。徐海涛：写作—审阅与编辑。邓智吉：可视化。蒋哲华：写作—审阅与编辑。

**利益冲突声明**

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

**致谢**

本工作部分由中国教育部基础学科与交叉学科突破计划（编号：JYB2025XDXM108）、国家自然科学基金（编号：62572160）以及浙江省科技计划项目（编号：2025C01065）资助。

李建伟 | 陈静 | 徐文强 | 徐海涛 | 邓智吉 | 蒋哲华