[论文解读] Equivariant Graph Neural Operator for Modeling 3D Dynamics

Minkai Xu, Jiaqi Han+7 more|arXiv (Cornell University)|Jan 19, 2024

Machine Learning in Materials ScienceMaterials Science被引用 3

一句话总结

该论文提出了一种新型方法——等变图神经算子（EGNO），通过基于傅里叶的时序卷积实现SE(3)等变神经算子，将3D物理动力学建模为连续轨迹。通过将动力学学习为时间上的函数，EGNO在粒子模拟、人体动作捕捉和分子动力学任务中均取得优越性能，相较于EGNN等基线模型，平均性能提升高达52%。

ABSTRACT

Modeling the complex three-dimensional (3D) dynamics of relational systems is an important problem in the natural sciences, with applications ranging from molecular simulations to particle mechanics. Machine learning methods have achieved good success by learning graph neural networks to model spatial interactions. However, these approaches do not faithfully capture temporal correlations since they only model next-step predictions. In this work, we propose Equivariant Graph Neural Operator (EGNO), a novel and principled method that directly models dynamics as trajectories instead of just next-step prediction. Different from existing methods, EGNO explicitly learns the temporal evolution of 3D dynamics where we formulate the dynamics as a function over time and learn neural operators to approximate it. To capture the temporal correlations while keeping the intrinsic SE(3)-equivariance, we develop equivariant temporal convolutions parameterized in the Fourier space and build EGNO by stacking the Fourier layers over equivariant networks. EGNO is the first operator learning framework that is capable of modeling solution dynamics functions over time while retaining 3D equivariance. Comprehensive experiments in multiple domains, including particle simulations, human motion capture, and molecular dynamics, demonstrate the significantly superior performance of EGNO against existing methods, thanks to the equivariant temporal modeling. Our code is available at https://github.com/MinkaiXu/egno.

研究动机与目标

• 解决现有图神经网络在捕捉3D物理动力学中长程时间相关性方面的局限性。
• 通过将完整解轨迹建模为时间上的连续函数，突破仅预测下一步的限制。
• 在学习3D空间动力学的时间依赖性时，保持SE(3)-等变性——这对物理泛化至关重要。
• 实现在任意时间离散化下的高效、并行推理而无需微调，提升部署灵活性。
• 开发一种可泛化的框架，可与任意EGNN架构集成，适用于多样化的物理动力学应用。

提出的方法

• 将3D动力学表述为从初始状态到完整轨迹的映射函数，通过神经算子实现连续时间建模。
• 引入在傅里叶域中参数化的等变时序卷积层，以在时序建模过程中保持SE(3)-等变性。
• 通过设计保持频率空间中旋转和平移不变性的核积分算子，利用傅里叶变换与逆傅里叶变换的等变性。
• 将基于傅里叶的时序层与标准EGNN组件堆叠，构建EGNO，实现几何与时间动态的端到端学习。
• 使用时间嵌入实现对任意时间分辨率的零样本推理，通过插值训练时间步之间实现。
• 将几何特征——节点特征（h）、位置（x）和速度（v）——整合到时序卷积中，以捕捉丰富的动态模式。

Figure 1: Illustration of EGNO . EGNO blocks (green) can be built with any EGNN layers (blue) and the proposed equivariant temporal convolution layers (yellow). Consider discretizing the time window $\Delta T$ into $P$ points $\{\Delta t_{1},\dots,\Delta t_{P}\}$ . Given a current state ${\mathcal{G

实验结果

研究问题

• RQ1 神经算子框架是否能在建模连续时间3D动力学的同时保持SE(3)-等变性？
• RQ2 与自回归的下一步预测相比，在傅里叶域中引入时序卷积在建模长程时间依赖性方面有何改进？
• RQ3 EGNO在无需微调的情况下，对未见过的时间离散化步长的泛化能力如何？
• RQ4 不同几何特征（h, x, v）对EGNO中时序建模性能的贡献有多大？
• RQ5 在轨迹精度和跨多样化物理系统的泛化能力方面，EGNO相较于现有基于GNN的模型表现如何？

主要发现

• 在Aspirin分子动力学基准测试中，EGNO相较于EGNN在F-MSE指标上实现了36%的相对性能提升。
• 在人体动作捕捉（Mocap-Run）数据集上，EGNO相较于EGNN在预测精度上实现了平均52%的提升。
• 消融实验表明，将速度（v）引入时序卷积可显著提升性能，全量EGNO变体（使用h, x, v）在N体模拟中相比消融版本最高可提升15%。
• EGNO在未微调的情况下实现对时间分辨率翻倍（2P）的零样本泛化，插值低分辨率时间步之间时可生成平滑且准确的轨迹。
• 傅里叶模态数量（I）对性能具有非单调影响，最优性能出现在I=2或I=3，超过此范围可能出现过拟合。
• EGNO支持在单次前向传播中并行解码任意时间窗口的未来状态，而自回归模型则需顺序推理。

Figure 2: Ablation studies on the number of modes $I$ on N-body simulation and Mocap- Run datasets.