E3x: $\mathrm{E}(3)$-Equivariant Deep Learning Made Easy (2401.07595v3)

Abstract: This work introduces E3x, a software package for building neural networks that are equivariant with respect to the Euclidean group $\mathrm{E}(3)$, consisting of translations, rotations, and reflections of three-dimensional space. Compared to ordinary neural networks, $\mathrm{E}(3)$-equivariant models promise benefits whenever input and/or output data are quantities associated with three-dimensional objects. This is because the numeric values of such quantities (e.g. positions) typically depend on the chosen coordinate system. Under transformations of the reference frame, the values change predictably, but the underlying rules can be difficult to learn for ordinary machine learning models. With built-in $\mathrm{E}(3)$-equivariance, neural networks are guaranteed to satisfy the relevant transformation rules exactly, resulting in superior data efficiency and accuracy. The code for E3x is available from https://github.com/google-research/e3x, detailed documentation and usage examples can be found on https://e3x.readthedocs.io.

• The paper introduces a novel deep learning framework that ensures E(3)-equivariance using irreducible representation (irrep) features for 3D data.
• It presents a unique tensor layer that couples irrep features through a learnable tensor product, improving connectivity in geometric processing.
• The framework integrates advanced modules like message-passing and attention, making it versatile for applications in quantum chemistry and computer vision.

Equivariant Neural Networks with E3x

Introduction to E3x

DeepMind's E3x framework presents itself as a pivotal development in the arena of deep learning for 3D objects and datasets. The E3x toolkit fundamentally addresses the challenge of creating neural network models that are equivariant to Euclidean group E(3) transformations, encompassing translations, rotations, and reflections in three-dimensional space. The adoption of E(3)-equivariant neural networks is paramount in scenarios where input or output data are inherently connected to three-dimensional geometries, as these networks inherently uphold transformation rules precisely, vastly enhancing data efficiency and model accuracy.

E3x Software Architecture

The success of E3x hinges on its alignment with familiar neural network structures while infusing E(3)-equivariance properties. At its core, E3x introduces the concept of “Irrep features,” built from irreducible representations of the O(3) group that describe tensors and pseudotensors up to a certain maximum degree. These features encapsulate geometric information that adjusts under coordinate transformations, retaining the correct equivariant relationships. The framework extrapolates traditional neural network components, like activation functions and dense layers, to be compatible with and maintain the equivariance of these irrep features. Furthermore, E3x furnishes a novel “Tensor layer” that binds different irrep features through a learnable tensor product coupling, akin to a mixing function that does not have a direct analogue in non-equivariant networks.

Advanced Components and Utilities

Beyond the foundational components, E3x is equipped with more sophisticated layers like message-passing and attention mechanisms that support equivariance. It provides a suite of utilities tailored for 3D datasets, such as functionalities for dealing with point clouds and performing operations indexed by neighbor lists. Additionally, the E3x library is equipped with tools for evaluating spherical harmonics and managing rotations, facilitating the administration of three-dimensional data.

Significance and Utilization

E3x serves as a versatile tool, readily integrating with existing model-building processes. By enabling the expression of complex behavior under rotational and reflective transformations, it holds the potential to significantly influence research and applications across disciplines such as quantum chemistry and computer vision. The framework’s scalability, ranging from configurations equivalent to ordinary neural networks to complex equivariant architectures, posits the E3x library as a significant contribution to the field of geometric deep learning. The availability of the E3x codebase on public repositories allows for broad accessibility and potential for future extensions by the community.