e3nn: Euclidean Neural Networks

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.

• The paper presents a framework (e3nn) that constructs E(3) equivariant neural networks using TensorProduct and spherical harmonics for precise 3D data handling.
• The framework maintains equivariance under rotations, translations, and inversion, reducing the need for extensive data augmentation.
• e3nn's modular design enhances data efficiency and supports diverse applications, including molecular science, medical imaging, and fluid dynamics.

e3nn: Euclidean Neural Networks

Euclidean Neural Networks (e3nn) represent a comprehensive framework for constructing E(3) equivariant neural networks, which inherently operate on 3D geometric data and tensors. These networks are specifically designed to transform predictably under coordinate system changes, eliminating the need for extensive data augmentation typically necessary to learn patterns from 3D data.

By leveraging core operations such as TensorProduct and spherical harmonics, e3nn provides a robust structure for articulating various advanced architectures, including Tensor Field Networks, 3D Steerable CNNs, and SE(3) Transformers.

Core Concepts and Framework

Equivariance in Neural Networks

Equivariance is a fundamental property that ensures neural network operations remain consistent under transformations such as rotations and translations. This property is crucial for handling 3D geometric data, as it helps eliminate symmetry-breaking biases inherent in conventional approaches. In the context of e3nn, equivariance is maintained by employing operations like TensorProduct which can handle geometric tensor algebra with precision and accuracy.

TensorProduct and Spherical Harmonics

One of the pivotal elements in e3nn is the TensorProduct operation, which enables the bilinear and equivariant combination of irreducible representations (irreps) of the rotation group. Through efficient manipulation of Clebsch-Gordan coefficients, TensorProduct can decompose complex tensor interactions into simpler, more manageable forms.

Spherical harmonics are another core component, facilitating the projection of functions onto the unit sphere, thereby capturing the angular aspects of 3D data. They serve as a bridge between spatial representations and rotational equivariance, enabling the construction of sophisticated neural network modules.

Figure 1: Left Grid on the sphere used in e3nn. Right Random signal on the sphere (cutoff at L=5).

Implementation Strategy

Building E(3) Equivariant Models

e3nn allows for the construction of a wide range of equivariant models using its comprehensive set of primitives. The primary aim is to provide a flexible framework that supports varied applications, from molecular dynamics to medical imaging. The use of operations like TensorProduct and module composition ensures that models remain expressive and computationally efficient.

For practical deployments, the library supports various implementation-specific requirements, such as:

• Parity Equivariance: Ensures models are invariant not just to rotations but to inversion operations as well.
• Reducing Tensor Products: Facilitates the decomposition of complex tensors into irreps, catering to advanced tensor manipulations.
• Universality of Operations: Guarantees that all necessary equivariant transformations are encapsulated within the provided primitives, ensuring model robustness across different domains.

Data Efficiency and Learning Dynamics

A significant advantage of e3nn is its impact on data efficiency. By using models with higher-order representations (e.g., $l > 1$), the framework improves learning dynamics by altering the power law exponent of the learning curve. This translates into better generalization from fewer training samples, overcoming what traditional data augmentation achieves merely through quantity.

Figure 2: The data-efficiency of equivariant neural network is more than a multiplicative constant, it affects the power law exponent of the learning curve.

Practical Applications and Prospects

The applications of e3nn are vast and varied, encompassing:

• Molecular and Materials Science: Prediction and analysis of molecular properties and interatomic potentials.
• Medical Imaging: Advanced processing and interpretation of diffusion MRI data.
• Fluid Dynamics: Modeling and simulation of complex fluid interactions around 3D surfaces.

The framework enables scientists and engineers to construct highly specialized models quickly, focusing more on scientific inquiry than on intricate implementation details.

Conclusion

e3nn stands as a powerful framework facilitating the construction and deployment of 3D equivariant neural networks. By abstracting complex geometric transformations into generalized primitives, it empowers users to efficiently express sophisticated operations while maintaining high data efficiency and expressivity. This positions e3nn as a versatile tool, essential for advancing AI-driven insights in scientific and engineering domains.