[论文解读] Equivariant Graph Hierarchy-Based Neural Networks

Jiaqi Han, Wenbing Huang+2 more|arXiv (Cornell University)|Feb 22, 2022

Protein Structure and Dynamics被引用 6

一句话总结

本文提出等变图层次神经网络（EGHN），一种基于层次消息传递的框架，通过引入广义等变矩阵消息传递（EMMP）、E-Pool实现节点特征的粗粒度化以保持欧几里得对称性，以及E-UnPool利用高层动态信息精炼低层表示，从而增强等变图神经网络。EGHN在多体物理系统（包括蛋白质动力学、动作捕捉和合成N体模拟）上实现了最先进性能，通过保持对称性的层次结构实现了更高的精度和子结构发现能力。

ABSTRACT

Equivariant Graph neural Networks (EGNs) are powerful in characterizing the dynamics of multi-body physical systems. Existing EGNs conduct flat message passing, which, yet, is unable to capture the spatial/dynamical hierarchy for complex systems particularly, limiting substructure discovery and global information fusion. In this paper, we propose Equivariant Hierarchy-based Graph Networks (EGHNs) which consist of the three key components: generalized Equivariant Matrix Message Passing (EMMP) , E-Pool and E-UpPool. In particular, EMMP is able to improve the expressivity of conventional equivariant message passing, E-Pool assigns the quantities of the low-level nodes into high-level clusters, while E-UpPool leverages the high-level information to update the dynamics of the low-level nodes. As their names imply, both E-Pool and E-UpPool are guaranteed to be equivariant to meet physic symmetry. Considerable experimental evaluations verify the effectiveness of our EGHN on several applications including multi-object dynamics simulation, motion capture, and protein dynamics modeling.

研究动机与目标

• 为解决现有等变图神经网络（EGNs）中平面消息传递的局限性，其无法捕捉复杂多体系统中的空间和动力学层次结构。
• 实现物理系统中子结构（如蛋白质二级结构或刚性机械部件）的端到端学习，同时保持欧几里得等变性。
• 设计在平移、旋转和反射下保持等变性的层次消息传递机制，确保跨尺度的物理一致性。
• 通过将网络架构划分为编码器-解码器层级并结合对称的池化与反池化操作，提升全局信息融合与局部动力学建模能力。
• 验证层次化设计、等变性以及专用消息传递在真实世界与合成物理系统建模任务中的必要性与有效性。

提出的方法

• 提出广义等变矩阵消息传递（EMMP），将传统基于向量的消息传递扩展至矩阵值消息，提升表达能力，超越标准EGNN。
• 引入E-Pool，一种可微分的E(n)-等变池化层，基于学习的注意力机制将低层节点聚类分组，保持几何对称性。
• 设计E-UnPool，一种对称的反池化操作，通过注意力机制与逆聚类从池化后的高层表示中重建细粒度节点特征。
• 在EMMP中为内部（局部）和外部（全局）消息传递分别使用独立的邻接矩阵，实现对簇内与簇间相互作用的独立建模。
• 在训练过程中引入连接性损失，以促使E-Pool保留簇内连接并抑制簇间边，提升结构一致性。
• 将EGHN架构设计为编码器-解码器框架：编码器通过应用E-Pool与EMMP实现从细粒度到粗粒度的尺度转换，解码器则通过E-UnPool与EMMP重建细尺度动力学。

实验结果

研究问题

• RQ1 与EGNs中的平面消息传递相比，层次化消息传递在建模复杂物理系统方面有何改进？
• RQ2 在池化与反池化操作中保持E(n)-等变性在多体动力学中对泛化能力与物理一致性的影响有多大？
• RQ3 EMMP的基于矩阵的消息传递相较于基于向量的替代方案，在捕捉拓扑与动力学复杂性方面有何贡献？
• RQ4 EMMP中分别使用全局与局部邻接矩阵对性能与结构学习的影响如何？
• RQ5 连接性损失与聚类数（K）对子结构发现质量与预测精度的影响是什么？

主要发现

• 在M-complex系统基准测试中，EGHN的预测误差显著低于基线EGN，证明其在层次化物理系统中具备更优的泛化能力与动力学建模性能。
• 在人体动作捕捉数据集上，EGHN在第35号受试者行走任务中达到8.5 MSE（×10⁻²），在跑步任务中达到25.9 MSE，优于所有基线模型，包括EGNN与MPNN。
• 在蛋白质动力学任务中，EGHN能准确追踪α-螺旋的折叠过程，可视化结果表明其与真实值的偏差显著小于MPNN与EGNN。
• 消融实验表明，若移除等变性（通过非等变MPNN实现），误差上升至19.7 MSE，证明等变性对性能至关重要。
• 若移除层次化设计（无层次结构），误差上升至21.9 MSE，证实层次结构对建模复杂运动学的关键作用。
• 若在所有EMMP层中使用单一邻接矩阵而非区分全局与局部矩阵，性能出现显著下降，验证了范围感知消息传递的重要性。