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

TL;DR

E3x is a software package that simplifies the development of neural networks with Euclidean group equivariance, improving data efficiency and accuracy for 3D data by ensuring models respect geometric transformations.

Contribution

The paper introduces E3x, a user-friendly software library for building E(3)-equivariant neural networks, making advanced geometric deep learning more accessible.

Findings

E3x enables easy implementation of E(3)-equivariant models.

Models built with E3x show improved data efficiency.

E3x facilitates accurate learning of 3D geometric data transformations.

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.