[论文解读] e3nn: Euclidean Neural Networks

Mario Geiger, Tess Smidt|arXiv (Cornell University)|Jul 18, 2022

Medical Imaging and Analysis被引用 106

一句话总结

tldr: 一个用于构建 E(3) 等变神经网络的 Python 框架，使用 irreps、张量积和球谐函数来处理 3D 几何。

ABSTRACT

We present e3nn, a generalized framework for creating E(3) equivariant trainable functions, also known as Euclidean neural networks. e3nn naturally operates on geometry and geometric tensors that describe systems in 3D and transform predictably under a change of coordinate system. The core of e3nn are equivariant operations such as the TensorProduct class or the spherical harmonics functions that can be composed to create more complex modules such as convolutions and attention mechanisms. These core operations of e3nn can be used to efficiently articulate Tensor Field Networks, 3D Steerable CNNs, Clebsch-Gordan Networks, SE(3) Transformers and other E(3) equivariant networks.

研究动机与目标

• 动机：为 3D 数据使用对称性意识的模型，以避免过度的数据增强并提升泛化能力。
• 提供一个灵活、统一的框架，用于构建遵守 3D 欧几里得对称性 (E(3)) 的等变神经网络。
• 解释数学基元（irreps、张量积、球谐函数）以及如何将它们组合成卷积和 transformer 等模块。
• 展示这些构建块如何使 Tensor Field Networks、3D Steerable CNNs、Clebsch-Gordan Networks 和 SE(3) Transformers 等体系结构成为可能。

提出的方法

• 引入 O(3) 的 irreps 和奇偶性来建模数据变换。
• 开发一个通用的 TensorProduct 运算，通过 Clebsch–Gordan 系数将双线性输入分解为 irreps。
• 在张量积中定义多条“路径”以捕获所有双线性等变操作，并允许加权组合。
• 提供 ReducedTensorProduct 实用工具，以在考虑指标对称性的情况下将相互作用分解为 irreps。
• 解释初始化方案和组件归一化，以确保稳定训练并维持等变性。

实验结果

研究问题

• RQ1 我们如何系统地构建所有在输入 irreps 之间、映射到输出 irreps 的双线性等变操作？
• RQ2 如何使用统一的 TensorProduct 框架高效实现并组合等变模块（如卷积、注意力）？
• RQ3 在 3D 几何数据学习中，强制 E(3) 等变性在数据效率和泛化方面有哪些好处？

主要发现

• 该框架围绕单一 TensorProduct 基元统一了各种 E(3) 等变结构。
• 球面谐函数和 Clebsch–Gordan 系数被用来在实值基底中将张量积分解为 irreps。
• 基于多路径的张量积构造允许完整的、带可学习权重的双线性等变操作。
• ReducedTensorProduct 提供系统化地将高秩张量在保持指标对称性的同时降维为 irreps。
• 本文讨论等变模型在初始化和数据效率方面的优势，包括对学习曲线斜率的影响。