[论文解读] Equivariant Spatio-Temporal Attentive Graph Networks to Simulate Physical Dynamics

Liming Wu, Zhichao Hou+3 more|arXiv (Cornell University)|May 21, 2024

Advanced Graph Neural Networks被引用 4

一句话总结

该论文提出ESTAG，一种等变时空图网络，通过利用历史轨迹数据来建模非马尔可夫物理动力学。它集成了等变离散傅里叶变换（EDFT）以提取周期性模式，结合等变空间模块（ESM）与等变时间模块（ETM），并采用前向注意力机制和等变池化，实现了在分子、蛋白质及宏观尺度动力学基准上的最先进性能。

ABSTRACT

Learning to represent and simulate the dynamics of physical systems is a crucial yet challenging task. Existing equivariant Graph Neural Network (GNN) based methods have encapsulated the symmetry of physics, \emph{e.g.}, translations, rotations, etc, leading to better generalization ability. Nevertheless, their frame-to-frame formulation of the task overlooks the non-Markov property mainly incurred by unobserved dynamics in the environment. In this paper, we reformulate dynamics simulation as a spatio-temporal prediction task, by employing the trajectory in the past period to recover the Non-Markovian interactions. We propose Equivariant Spatio-Temporal Attentive Graph Networks (ESTAG), an equivariant version of spatio-temporal GNNs, to fulfill our purpose. At its core, we design a novel Equivariant Discrete Fourier Transform (EDFT) to extract periodic patterns from the history frames, and then construct an Equivariant Spatial Module (ESM) to accomplish spatial message passing, and an Equivariant Temporal Module (ETM) with the forward attention and equivariant pooling mechanisms to aggregate temporal message. We evaluate our model on three real datasets corresponding to the molecular-, protein- and macro-level. Experimental results verify the effectiveness of ESTAG compared to typical spatio-temporal GNNs and equivariant GNNs.

研究动机与目标

• 解决现有物理动力学模拟方法中马尔可夫假设的局限性，这些方法无法捕捉未观测到的环境相互作用。
• 将动力学模拟重新表述为基于历史轨迹数据的时空预测任务，以建模非马尔可夫依赖关系。
• 设计一种保持欧几里得对称性（旋转、平移）的同时建模时空依赖关系的等变架构。
• 通过一种新颖的等变离散傅里叶变换（EDFT）实现周期性模式建模，以捕捉物理系统中的重复动力学。
• 证明结合时间记忆与等变性相较于标准时空GNN与非等变模型的优势。

提出的方法

• 提出一种新颖的等变离散傅里叶变换（EDFT），从历史节点轨迹中提取频域特征，同时保持E(3)等变性。
• 设计一种等变空间模块（ESM），通过等变更新节点坐标与特征，在空间图上执行消息传递。
• 引入一种等变时间模块（ETM），结合前向注意力机制与等变池化，以在保持几何对称性的同时聚合跨帧的时间消息。
• 从历史系统状态构建时空图，以建模超越帧间假设的长程时间依赖关系。
• 通过精心设计消息传递与聚合层，确保所有学习到的表示与更新在平移、旋转和反射下保持等变性。
• 在ETM中使用可学习注意力机制，以动态加权过去帧的重要性，实现自适应时间建模。

实验结果

研究问题

• RQ1 将物理动力学建模为时空预测任务（而非帧间预测），是否能更好地捕捉由未观测环境相互作用引起的非马尔可夫行为？
• RQ2 如何设计一种保持E(3)对称性的等变时空GNN，同时有效建模长程时间依赖关系？
• RQ3 通过EDFT实现周期性模式建模对物理动力学模拟性能的贡献有多大？
• RQ4 与标准STGNN或等变GNN相比，结合时间记忆与几何等变性在泛化能力与准确性方面提升了多少？
• RQ5 EDFT、注意力机制、等变性、时间池化等各组件对整体模型性能的贡献分别是什么？

主要发现

• ESTAG在MD17、蛋白质与宏观尺度数据集上均达到最先进性能，在Aspirin数据集上的平均MSE为0.063，在Benzene数据集上为0.003，显著优于非等变与非时间基线模型。
• 消融实验表明，移除EDFT后MSE平均增加25%，证明其在建模空间周期性中的关键作用。
• 在ETM中移除注意力机制后，MSE平均增加38%，表明其在捕捉动态时间相互作用中的重要性。
• ESTAG的非等变版本在MSE上平均增加10倍，证明几何等变性对鲁棒物理模拟至关重要。
• 将等变时间池化替换为可学习系数后，MSE增加35%，证实结构化时间聚合的优势。
• 在Ethanol数据集上，当L=2层时，ESTAG在准确率与效率之间达到最优平衡，MSE为0.990×10⁻⁴。