作者｜袁超豪 阿里巴巴达摩院实习生
导语
在学习和模拟物理动力学（如分子运动、蛋白质折叠）时，图神经网络（GNN）已展现出巨大潜力。为了提升模型的泛化能力，一系列“等变GNN”（Equivariant GNNs）被开发出来，它们将平移、旋转等物理对称性作为归纳偏置（inductive bias）融入模型。然而，现有的方法普遍忽略了物理动力学的一个关键特性：非平稳性（non-stationarity），即系统的联合分布会随时间动态变化。
更棘手的是，传统上用于处理非平稳时间序列的方法（如数据归一化）会破坏物理系统固有的对称性假设。为了在模拟非平稳物理动力学的同时保持模型的对称性，本文提出了一种非平稳等变图神经网络（NS-EGNN）。
NS-EGNN 的核心在于：
采用分块傅里叶变换（Patch Fourier Transform）来提取物理轨迹中随时间变化的频率信息，该模块提取的信息是不变 (Invariant) 的。
使用一阶和二阶差分来缓解数据中的非平稳趋势，此操作同样保持等变性。
通过捕捉变化的频率特性并缓解原始物理动力学中的线性和二次趋势，NS-EGNN 能够更好地建模物理动力学中的时间依赖性。
实验证明，NS-EGNN 在分子动力学、人体运动捕捉和蛋白质动力学等多种物理模拟任务上，均超越了现有的SOTA算法。
论文地址：https://openreview.net/pdf?id=0dp90Bq8r8
代码仓库：https://github.com/MaojiWEN/NS-EGNN
现存问题及挑战
精确模拟真实世界的物理动力学在分子动力学、药物发现和蛋白质折叠等众多领域至关重要。近年来，等变图神经网络（EGNNs）已成为主流方法，因为它们能将物理世界的基本对称性（如旋转、平移不变性）编码到模型中，使得模型能更准确、更泛化地学习粒子间的复杂相互作用。
尽管取得了成功，但现有模型在模拟复杂的真实物理系统时仍面临两大挑战：
非马尔可夫性（Non-Markovian）： 大多数模型采用单步预测（frame-to-frame），即只使用前一帧的状态来预测下一帧。但这在物理上是不充分的，因为未来的状态往往依赖于过去一段时间的轨迹，而非仅仅当前状态。
非平稳性（Non-stationary）： 这是本文关注的核心挑战。非平稳性意味着系统的统计特性（如均值、方差）是随时间变化的。例如，在分子动力学中，分子的势能不断变化，导致其振动幅度和粒子间的连接关系随之改变；在人体运动数据中，人的速度和生理状态也是动态的。
（图1：非平稳性与等变性图示。左侧不同时间段（G0:t，Gt:2t）的分子动态，其傅里叶频率（S1, S2）截然不同，表现出非平稳性。右侧展示了对系统进行旋转/平移后，模型应捕捉到相同的频率信息，即不变性。)
现有的深度学习模型很难处理这种时变的分布，导致泛化能力差。而经典的时间序列分析方法（如归一化（Normalization））虽然可以“稳定”数据，但它们在应用到三维坐标时会直接破坏 E(3) 等变性，这会造成深度模型的过拟合从而导致泛化性差。
NS-EGNN：兼顾非平稳性与等变性的新框架:
为了解决上述挑战，我们设计了 NS-EGNN 框架。其核心思想是：将非平稳性分解为动态变化的方差和动态变化的均值（趋势），并分别使用保持对称性的模块来处理它们。
(图2：NS-EGNN 整体框架。模型主要由两个核心模块构成：不变分块傅里叶变换（PFT）和等变非平稳时序池化（NS-Pooling）。)
不变分块傅里叶变换 (PFT): 捕捉动态方差
物理系统随时间变化的方差（如振动幅度）会体现在频域上。
传统的傅里叶变换（FT）会计算整个轨迹的全局频谱，这隐含地假设了信号是平稳的，无法捕捉动态变化。
PFT 的做法是：在轨迹上使用重叠的滑动窗口，并对每个窗口内的片段分别进行傅里叶变换。
通过这种方式，NS-EGNN 获得了随时间演变的局部频率信息。最关键的是，我们提取的频率特征（如功率谱）对于E(3)变换（旋转、平移）是不变的（invariant）。这意味着PFT模块可以在不破坏物理对称性的前提下，将时变的统计特性信息安全地馈入后续的EGNN骨干网络。
等变非平稳时序池化 (NS-Pooling): 缓解动态均值（趋势）
非平稳性的另一个来源是均值的时变趋势（例如，分子团簇的整体漂移）。
受经典时间序列分析（ARIMA）的启发，我们使用一阶和二阶差分（Differencing）来消除轨迹中的线性和二次趋势。
对于三维坐标 X，差分操作 ΔXi = Xi (t) – Xi (t-1)本身就是等变的（equivariant）。
NS-Pooling 模块通过对差分后的、更平稳的轨迹进行池化来预测未来。
综上所述，NS-EGNN 通过 PFT 捕捉动态方差（不变性），通过 NS-Pooling 处理动态均值（等变性），最后结合强大的 EGNN 骨干网络进行空间信息建模。这使其成为第一个能够端到端处理非平稳轨迹，同时严格保持E(3)等变性的框架。
实验效果与分析:
我们在三类经典的物理动力学数据集上验证了 NS-EGNN 的性能：MD17（分子动力学）、CMU Motion Capture（人体运动捕捉）和 AdK（蛋白质动力学）。
首先，我们使用 ADF 统计检验（Augmented Dickey-Fuller test）证实了这些数据集均表现出强烈的非平稳特性（如表1所示）。绝大多数的动力学轨迹都表现出了非常强的非平稳特性。
(表2 & 表3：MD17 数据集上的 ADE / FDE 误差对比。NS-EGNN 在所有分子上均取得了最低的平均误差和最终误差，显著优于 ESTAG, EqMotion 等前SOTA模型。)
分子动力学 (MD17) 如上表所示，在MD17数据集上，NS-EGNN 相比当前的SOTA模型（如 ESTAG），在平均位移误差（ADE）上实现了 29.77% 的相对提升，在最终位移误差（FDE）上实现了 33.33% 的相对提升。这证明了建模非平稳性的巨大优势。
人体运动捕捉 (CMU Motion) 在预测人体“行走”和“奔跑”的任务中，NS-EGNN 同样在大多数设置下取得了最佳性能。如下图所示，相比 ESTAG，NS-EGNN 预测的运动轨迹（蓝色）更接近真实轨迹（红色），尤其是在复杂的关节运动上，轨迹更稳定、更符合物理规律。
(图3：人体运动捕捉可视化。在“行走”和“奔跑”任务中，NS-EGNN（右侧）的预测轨迹比 ESTAG（左侧）更平滑、更接近真实值。)
消融实验 我们通过消融实验证明了每个组件的必要性：
去除 PFT（仅用全局FT）：性能显著下降。证明了捕捉动态时变的频率信息至关重要。
去除差分（NS-Pooling）：性能大幅下降。证明了缓解非平稳趋势（动态均值）是必要的。
去除等变性（使用普通GNN）：模型完全失效。证明了等变性是物理模拟的基石。
总结与未来工作
本文揭示了真实世界物理动力学中广泛存在的非平稳特性，并指出现有等变GNN模型在处理这一问题上的空白。
我们提出的 NS-EGNN 框架，通过创新的不变分块傅里叶变换（PFT）和等变时序池化（NS-Pooling），首次实现了在严格保持E(3)等变性的同时，对非平稳物理动力学进行建模。在多个基准数据集上的SOTA性能充分验证了该框架的有效性和泛化能力。
未来，我们将探索将 NS-EGNN 扩展到更多应用场景，并尝试将其他非平稳性分析工具（如小波变换）适配到等变框架中。
｜往期内容回看
