[论文解读] Fast, Expressive SE$(n)$ Equivariant Networks through Weight-Sharing in Position-Orientation Space

Erik J. Bekkers, Sharvaree Vadgama+3 more|arXiv (Cornell University)|Oct 4, 2023

Advanced Image and Video Retrieval Techniques被引用 5

一句话总结

本文提出了一种快速且表达能力强的 SE(n) 等变网络，通过在位置-方向空间（R³×S²）上共享权重，形式化了在 SE(n) 变换下点对等价类的权重共享。通过将消息函数条件化于唯一标识这些等价类的几何属性，该方法在 3D 点云基准测试中实现了最先进（SOTA）的准确率与速度，涵盖原子间能量预测、N体轨迹预测以及基于扩散模型的分子生成任务。

ABSTRACT

Based on the theory of homogeneous spaces we derive geometrically optimal edge attributes to be used within the flexible message-passing framework. We formalize the notion of weight sharing in convolutional networks as the sharing of message functions over point-pairs that should be treated equally. We define equivalence classes of point-pairs that are identical up to a transformation in the group and derive attributes that uniquely identify these classes. Weight sharing is then obtained by conditioning message functions on these attributes. As an application of the theory, we develop an efficient equivariant group convolutional network for processing 3D point clouds. The theory of homogeneous spaces tells us how to do group convolutions with feature maps over the homogeneous space of positions $\mathbb{R}^3$, position and orientations $\mathbb{R}^3 { imes} S^2$, and the group $SE(3)$ itself. Among these, $\mathbb{R}^3 { imes} S^2$ is an optimal choice due to the ability to represent directional information, which $\mathbb{R}^3$ methods cannot, and it significantly enhances computational efficiency compared to indexing features on the full $SE(3)$ group. We support this claim with state-of-the-art results -- in accuracy and speed -- on five different benchmarks in 2D and 3D, including interatomic potential energy prediction, trajectory forecasting in N-body systems, and generating molecules via equivariant diffusion models.

研究动机与目标

• 将消息传递网络中的权重共享形式化，推广至传统卷积设置之外的场景。
• 利用齐性空间理论，推导出适用于 SE(n) 等变消息函数的几何最优属性。
• 开发一种高效且表达能力强的 3D 点云处理架构，同时保持 SE(n) 等变性。
• 在多样化的 3D 学习基准上展示最先进性能，涵盖分子系统与物理系统。

提出的方法

• 将权重共享形式化为通过识别 SE(n) 变换下点对等价类的属性来条件化消息函数。
• 利用齐性空间理论推导几何属性，识别 R³×S² 为方向信息与效率的最优特征空间。
• 在 R³×S² 上构建可分离的群卷积，实现方向与通道上的并行化，提升计算效率。
• 使用球面卷积与批量通道混合技术，加速推理与训练过程。
• 将消息函数条件化于不变属性，如相对位置、相对方向以及交互特征（例如电荷乘积）。
• 实现一个完全卷积的架构，采用 ConvNeXt 块，并引入可学习的球面嵌入以处理输入与输出特征。

实验结果

研究问题

• RQ1 如何将消息传递网络中的权重共享从平移等变卷积推广至完整的 SE(n) 等变性？
• RQ2 哪些几何属性能唯一标识 SE(n) 变换下点对的等价类？
• RQ3 为何位置-方向空间（R³×S²）在 SE(3) 等变学习中比 R³ 或完整 SE(3) 更高效且更具表达力？
• RQ4 纯卷积的、SE(n)-等变架构能否在多样化的 3D 点云基准上实现最先进性能？
• RQ5 为何将消息函数条件化于几何不变量能提升等变消息传递中的表达力与效率？

主要发现

• 所提方法在原子间势能预测任务中实现了最先进性能，准确率与推理速度均优于先前方法。
• 在 N 体轨迹预测任务中，该模型实现了最先进结果，且推理时间显著短于基线等变模型。
• 在基于扩散模型的分子生成任务中，该架构生成了高质量、物理上合理的分子，且采样效率得到提升。
• 使用位置-方向空间（R³×S²）相比完整 SE(3) 索引，降低了计算成本，同时保持了方向敏感性并提升了性能。
• 理论框架证明，仅通过几何属性与共享消息函数，即可实现对 SE(n)-等变函数的通用逼近。
• 实证结果证实，将消息函数条件化于推导出的属性，可提升所有基准上的泛化能力与更快的收敛速度。