e3nn: Euclidean Neural Networks

TL;DR

e3nn introduces a flexible framework for constructing Euclidean neural networks that are equivariant under 3D rotations and translations, enabling advanced geometric learning in 3D systems.

Contribution

The paper presents a unified framework with core equivariant operations for building various E(3) equivariant neural network architectures.

Findings

Supports tensor field networks and 3D steerable CNNs

Enables efficient construction of SE(3) transformers

Provides a modular toolkit for geometric deep learning

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.