[论文解读] Frame Averaging for Invariant and Equivariant Network Design

Omri Puny, Matan Atzmon+5 more|arXiv (Cornell University)|Oct 7, 2021

Machine Learning in Materials Science参考文献 62被引用 10

一句话总结

本文提出框架平均（Frame Averaging, FA），一种系统性框架，用一个小而依赖数据的群元素子集（称为框架）替代计算上不可行的群平均，从而在神经网络中实现精确不变性或等变性。FA 保留了主干架构的表达能力，在点云法向估计、图分离和 n 体动力学预测任务中达到最先进性能。

ABSTRACT

Many machine learning tasks involve learning functions that are known to be invariant or equivariant to certain symmetries of the input data. However, it is often challenging to design neural network architectures that respect these symmetries while being expressive and computationally efficient. For example, Euclidean motion invariant/equivariant graph or point cloud neural networks. We introduce Frame Averaging (FA), a general purpose and systematic framework for adapting known (backbone) architectures to become invariant or equivariant to new symmetry types. Our framework builds on the well known group averaging operator that guarantees invariance or equivariance but is intractable. In contrast, we observe that for many important classes of symmetries, this operator can be replaced with an averaging operator over a small subset of the group elements, called a frame. We show that averaging over a frame guarantees exact invariance or equivariance while often being much simpler to compute than averaging over the entire group. Furthermore, we prove that FA-based models have maximal expressive power in a broad setting and in general preserve the expressive power of their backbone architectures. Using frame averaging, we propose a new class of universal Graph Neural Networks (GNNs), universal Euclidean motion invariant point cloud networks, and Euclidean motion invariant Message Passing (MP) GNNs. We demonstrate the practical effectiveness of FA on several applications including point cloud normal estimation, beyond $2$-WL graph separation, and $n$-body dynamics prediction, achieving state-of-the-art results in all of these benchmarks.

研究动机与目标

• 为解决设计对复杂对称性（如排列与欧几里得运动）实现精确不变性或等变性的神经网络的挑战，同时保持表达能力与计算效率。
• 通过用一个小而可高效计算的群元素子集替代大群或连续群中完整的群平均，克服全群平均的计算不可行性。
• 提供一种通用、系统性的框架，适用于多种架构（如 GNN、点云网络与 MLP），实现对称性感知学习。
• 证明基于 FA 的模型可保留其主干架构的最大表达能力，进而在关键场景中实现通用近似。
• 通过实证验证 FA 在多个基准上的表现，证明其在不变性稳定性与模型性能方面优于近似或全群平均基线方法。

提出的方法

• 框架平均将对称群 G 中所有元素的不可行群平均，替换为对一个精心选择的有限子集（称为框架 F(X)）的平均，该子集依赖于输入 X。
• 框架 F(X) 的构造满足集合等变性性质，确保对 F(X) 的平均可保持在全群作用下的精确不变性或等变性。
• 该框架被应用于多种架构：MLP 用于排列不变性，GNN 用于图级别不变性，消息传递 GNN 用于欧几里得运动不变性/等变性。
• 对于欧几里得运动对称性（E(d)），框架源自稳定子子群，并利用几何不变量以实现高效计算。
• 该方法采用改进的消息传递机制，使特征在每一层上通过对框架进行平均，从而在保持表达能力的同时确保等变性。
• 理论分析证明，当主干网络具有通用近似能力时，基于 FA 的模型可实现通用近似，且框架的选择不会降低表达能力。

实验结果

研究问题

• RQ1 能否开发一种系统性框架，在不承担全群平均计算成本的前提下，强制神经网络实现精确不变性或等变性？
• RQ2 在何种条件下，一个小而依赖数据的群元素子集（即框架）可替代全群平均，同时保持精确对称性？
• RQ3 框架平均是否在对称性感知学习任务中保持或增强底层主干架构的表达能力？
• RQ4 与近似平均方法（如蒙特卡洛）或全群平均相比，FA 在不变性准确率与模型性能方面表现如何？
• RQ5 FA 能否推广至多种对称群（如排列、欧几里得运动）并应用于包括 GNN 与点云网络在内的广泛架构？

主要发现

• 框架平均通过在小而依赖输入的框架上进行平均，实现精确不变性与等变性，避免了全群平均的计算不可行性。
• 基于 FA 的模型保留了其主干架构的最大表达能力，使图学习与点云学习任务中可实现通用近似。
• 在 n 体动力学预测任务中，FA-GNN 的测试均方误差为 0.0057，优于最先进方法 EGNN（0.0071），性能提升超过 20%，且参数量相近。
• 在 beyond-2-WL 图分离基准上，FA-MLP 与 FA-GIN+ID 实现了完美分离，证明其具备通用表达能力。
• 即使仅使用 k=1 个框架样本的近似 FA，其不变性误差也显著低于使用 k=1 的全群平均（GA），表明其具有更优的稳定性与泛化能力。
• 实证结果表明，FA 在不变性与有效性方面均优于近似与全群平均方法，尤其在低样本量与高对称性场景中表现更优。