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

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.

Figures (3)

• Figure 1: Visualization of the spherical harmonics $Y_{\ell}^{m}$ up to degree $\ell=3$ (positive values shown in blue, negative values in red, the $x$--, $y$--, and $z$--axes are depicted as red, green, and blue arrows). All $Y_{\ell}^{m}$ are evaluated on the unit sphere, but its radius at point $\boldsymbol{r}$ is scaled by $|Y_{\ell}^{m}(\boldsymbol{r})|$, leading to distinct shapes for the different $Y_{\ell}^{m}$.
• Figure 2: (A) Color-coded visualization of the memory layout of (randomly drawn) features $\boldsymbol{x}\in{\mathbb R}^{2\times(L+1)^2\times F}$ with $L=2$ and $F=8$. (B) Visualization of features as three-dimensional shapes (negative components are drawn in red and positive components in blue).
• Figure 3: For better memory efficiency, E3x supports representing features where all pseudotensor components are zero (A) without the need for explicitly storing the zero components (B).