e3nn: Euclidean Neural Networks

Mario Geiger; Tess Smidt

Source-linked AI summary

e3nn: Euclidean Neural Networks

Mario Geiger, Tess Smidt

arXiv:2207.09453v1Jul 18, 2022cs.LGcs.AIcs.NE

TL;DR

3D geometric data are difficult for machine learning because coordinate systems are sensitive to spatial symmetries, motivating equivariant models. e3nn provides composable E(3)-equivariant operations and general tensor-algebra building blocks, with high-order representations reported as more data-efficient than scalar-and-vector-only models.

Problem
3D geometry and geometric tensors are challenging because coordinates and coordinate systems are sensitive to rotations, translations, and inversion, while equivariant implementation is technical and error prone.
Method
e3nn is a general framework of composable E(3)-equivariant operations, including flexible TensorProduct-based tensor algebra and building blocks for trainable modules.
Results
Models built with e3nn using high-order representations (l > 1) are more data-efficient than models limited to scalars and vectors.
Takeaways & Limitations
e3nn lets users compose expressive equivariant functions and sophisticated trainable modules without imposing specific training operations.
Takeaways & Limitations
The authors report no theoretical explanation for the observed change in the learning-curve exponent.

Abstract

from arXiv · show

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.

1 Introduction

3D data are difficult for machine learning because coordinate choices interact with rotations, translations, and inversion. e3nn addresses this by providing composable E(3)-equivariant operations and a framework for building such models.

Motivation: 3D geometry and geometric tensors are challenging because their coordinates and coordinate systems respond to 3D rotations, translations, and inversion.Incorporating symmetry can eliminate the need for brute-force orientation augmentation, described as a 500-fold increase in training.
Motivation: Invariant models are simpler because they operate on scalars, but their scalar-only representations restrict what they can express.Representations indexed by l describe transformation behavior: scalars have l = 0, vectors l = 1, and higher-order quantities have larger l.
Motivation: Reframing tasks equivariantly can expose hidden symmetry-breaking biases, including arbitrary eigenvector choices in degenerate subspaces.
Motivation: Models with internal l = 2 representations have been observed to display better learning curves than models containing only l ≤1 internal representations.
Framework: e3nn is a Python framework that handles geometric tensor algebra while providing general mathematical primitives composable into diverse E(3)-equivariant architectures.Its operations include a unified parameterizable TensorProduct in later versions.

2 Related work

Prior work includes invariant, restricted equivariant, kernel, and general-group approaches. e3nn aims to express the most general equivariant operations across 3D data and unify these methods through its operations.

Invariant methods: Molecular and materials machine learning commonly uses invariant representations such as symmetry functions, ACE, and SOAP for per-atom energy predictions.
Invariant methods: Most symmetry-aware neural networks in molecules and materials operate on scalars, while DimeNet incorporates precomputed angular information as scalar inputs.
Equivariant methods: Restricted equivariant methods use scalar and vector operations but do not span the full space of equivariant operations and representations.Such restricted operations can nevertheless be adequate for some learning tasks.
Other approaches: Other approaches place equivariant operations inside kernels or derive equivariant linear layers for general groups from group generators.
Empirical motivation: Many studies report improved training-task accuracy when moving from invariant to equivariant models, even within the same framework.
e3nn's goal: e3nn targets a flexible framework for expressing the most general equivariant operations for all forms of 3D data, including methods discussed above.Applications span symmetry-breaking order parameters, molecular properties, phonons, interatomic potentials, imaging, flow, and other tasks.

3 Our Contribution

e3nn provides parity-equivariant, general tensor-decomposition, representation-conversion, and foundational operations for constructing equivariant modules of arbitrary complexity.

Core capabilities: e3nn supports parity equivariance, making models equivariant to O(3) rather than only SO(3).
Core capabilities: Its TensorProduct class enables flexible trainable bilinear operations between irreducible representations.
Core capabilities: ReduceTensorProduct decomposes arbitrary-rank tensors, including tensors with index-symmetry constraints, into direct sums of O(3) irreducible representations.
Core capabilities: The framework includes fast conversion between irreducible representations and spherical signals plus base operations for modules of arbitrary complexity.
Scope: These methods are intended for tensor computations beyond deep learning.

4 Group Representations

The framework types data by finite-dimensional representations of the Euclidean group, using irreducible O(3) representations to preserve transformation behavior. Equivariance is maintained compositionally through network operations.

Groups and representations: The Euclidean group combines 3D rotations, translations, and inversion.
Groups and representations: e3nn types network data by how it transforms under the Euclidean group and restricts representations to finite dimensions.
Groups and representations: A representation maps each group element to a matrix while preserving the group's identity and composition structure.
Irreducible representations: Irreducible representations are minimal representations, and finite-dimensional O(3) representations can be decomposed into them.O(3) irreducible representations factor into SO(3) irreducible representations and parity representations.
Irreducible representations: SO(3) irreducible representations are indexed by l and have dimension 2l + 1; l = 0 corresponds to scalars and l = 1 to vectors.Higher l representations include the l = 2 decomposition of a symmetric rank-2 tensor.
Symmetric matrix example: A symmetric 3 × 3 matrix decomposes into one scalar trace and five symmetric-traceless components, represented in e3nn as 1x0e + 1x2e.
Equivariance: An equivariant function commutes with group actions on inputs and outputs, and composing equivariant functions preserves equivariance.
Equivariance: Learned parameters are required to be scalars, so they remain unchanged under E(3) transformations and coordinate-system changes.

5 Irreducible representations used in e3nn

e3nn uses real-valued bases for SO(3) irreducible representations, resolving basis degeneracy through generator diagonalization and basis changes.

Irreducible representations are unique up to an invertible change of basis.
Diagonalizing a rotation generator provides a preferred basis for each SO(3) irrep.
SO(3) representations are equivalent to their conjugates, allowing a change of basis into real-valued representations.
e3nn converts representations obtained from QuTip into the usual real basis for memory and computation reasons.

6 Spherical harmonics

Spherical harmonics provide equivariant polynomial bases on the sphere and support truncated signal representations that e3nn can convert efficiently to and from grids.

6.1 Definition: Spherical harmonics Y^l map the unit sphere to the SO(3) irrep D^l and contain 2l + 1 component functions.
6.1 Definition: They transform equivariantly under rotations and are normalized to unit norm on the sphere in e3nn’s norm convention.
6.1 Definition: Once the irrep basis is fixed, spherical harmonics on the unit sphere are unique up to a sign.
6.1 Definition: Extended to R3, spherical harmonics can be chosen as polynomials and form a basis for equivariant polynomials on the sphere.
6.2 Project to and from the sphere: A scalar spherical signal is expanded using coefficients v_l and harmonics Y^l, whose coefficients transform as irreducible representations.
6.2 Project to and from the sphere: Truncating at maximum L removes high frequencies, while e3nn provides efficient conversion between irrep coefficients and discretized grid signals.

7 Tensor Product

e3nn’s TensorProduct represents bilinear equivariant operations by combining allowed irreducible-representation paths, with ReducedTensorProduct extending this framework to permutation-constrained tensors.

7.1 Tensor product: The tensor product is an equivariant multiplication operation for two representations.
7.1 Tensor product: Tensor products decompose into independently bilinear and equivariant paths connecting input irreps to allowed output irreps.
7.1 Tensor product: e3nn TensorProduct specifies input and output irreps plus paths that identify connected inputs and their weights.
7.2 Different Paths in the Tensor Product: Each output is a weighted sum of paths formed with Clebsch–Gordan coefficients, which change basis from tensor products to output irreps.
7.3 Generality of Tensor Product: The TensorProduct class can represent any bilinear equivariant operation combining two sets of irreps into irreps.
7.4 Reducing Tensor Products into irreps: ReducedTensorProduct reduces arbitrary-rank tensors by combining rotation irreps with index-permutation constraints, whose commuting symmetries can be solved independently.

8 Equivariant modules

e3nn composes TensorProduct and spherical harmonics into equivariant modules, including convolutions, transformers, and position-dependent equivariant polynomials, with an initialization scheme designed to stabilize learning.

8 Equivariant modules: TensorProduct and spherical harmonics implement point convolutions, voxel convolutions, and equivariant transformers.
8 Equivariant modules: An e3nn polynomial module remains E(3)-equivariant for every choice of its learned weights.
8.1 Initialization: Component normalization targets ||x||² = d so average component magnitudes are approximately one.
8.1 Initialization: The initialization scheme gives preactivations mean 0 and variance 1, post-activation second moment 1, and learning at infinite width.
8.1 Initialization: For a fully connected layer, scaling Wx by 1/√n restores unit variance when weights have variance 1.

9 Data Efficiency

Equivariance can improve data efficiency beyond reducing the amount of training data by changing the learning-curve exponent. The theoretical explanation for this change remains unavailable.

Observed scaling: The data-efficiency of equivariant neural networks affects the power-law exponent rather than only contributing a multiplicative constant.This is the central interpretation associated with Figure 4.
Data-efficiency interpretation: Removing orientation-specific data augmentation would predict a horizontal learning-curve shift, requiring c times less training data for equal performance.This would reflect a multiplicative data-efficiency gain rather than a change in learning-curve shape.
Observed scaling: Equivariant neural networks have a different test-error learning-curve exponent, visible as a changed slope in a log-log plot.The learning curve relates test error to the number of training data.
Open question: No theoretical explanation is provided for the observed change in exponent.The paper explicitly identifies this as unresolved.

10 Expressing previous works with e3nn

e3nn primitives can express several prior equivariant architectures, while preserving their differing choices for filters, tensor products, nonlinearities, and transforms. These constructions include Tensor Field Networks, 3D Steerable CNNs, Spherical CNNs, Clebsch-Gordan Networks, and SE(3) Transformers.

Tensor Field and Steerable CNNs: 3D Steerable CNNs and Tensor Field Networks differ in whether the filter or the tensor-product operation is represented through spherical harmonics.The distinction reflects voxel compatibility in 3D Steerable CNNs and point-cloud computational considerations in Tensor Field Networks.
Tensor Field and Steerable CNNs: Tensor Field Networks reduce point-cloud overhead by omitting tensor-product weights and applying a linear operation afterward, implemented as ‘uvu’ versus ‘uvw’.The passage compares this design to matrix decomposition and SchNet.
Spherical CNNs: Spherical CNNs can operate in real space or irrep space, where convolution becomes linear layers and activations require forward and reverse Fourier transforms.The representation choice changes how convolution and activation functions are implemented.
Other architectures: Clebsch-Gordan Networks use TensorSquare as a nonlinearity followed by a Linear layer.This architecture can be constructed from e3nn primitives.
Other architectures: An SE(3)-equivariant transformer, along with the other described models, can be constructed using e3nn primitives.The cited prior architectures are represented through the framework’s reusable operations.

11 Conclusion

e3nn provides composable E(3)-equivariant operations for learning on 3D data and beyond. Its expressive building blocks support complex equivariant functions and enable higher-order representations to improve data efficiency over scalar-and-vector-limited models.

Framework scope: e3nn is a general framework for composing E(3)-equivariant operations for learning on 3D data and beyond.The framework is designed around composable operations rather than a single fixed model.
Framework design: Expressive general operations let users build complex equivariant functions and sophisticated trainable modules without imposed training operations.The framework provides building blocks analogous to a flexible autodifferentiation API.
Framework design: Distilling equivariant operations into a few core building blocks concentrates optimization efforts on those classes.This design centralizes optimization within the framework’s core operations.
Data efficiency: Models built with e3nn using l > 1 representations are more data-efficient than models limited to scalars and vectors.The conclusion attributes this comparison to prior evidence.