e3nn: Euclidean Neural Networks

In plain words

Many 3D problems have a built-in symmetry. Rotate a molecule and its properties rotate with it, or stay put. A standard neural network does not know this, so it has to learn the symmetry from many rotated copies of the data, or it gets the answer wrong for orientations it has not seen.

e3nn is a set of building blocks in which every number carries a type that says how it changes when the scene is rotated or reflected through a point. Scalars stay fixed, vectors turn with the scene, and every operation is built so that these types are kept. The paper explains the mathematics behind these blocks and shows how several earlier 3D networks can be assembled from them. It is a framework paper with no benchmark numbers. Its value is a clear set of primitives that can be composed.

TL;DR

• e3nn is an open PyTorch library of E(3)-equivariant operations for 3D data §3.
• Its core is a general, weighted tensor product between typed representations (irreps), plus spherical-harmonic features and transforms to and from a sphere grid §6, §7.
• The paper gives the maths and the path-normalization rules, but no experiments of its own. Results are not reported §9.

Problem

Rotation, translation and inversion are symmetries of 3D geometry. A model that ignores them either needs data augmentation or must learn invariance itself. The paper cites a roughly 500-fold increase in brute-force training as the cost of learning 3D patterns in arbitrary orientations without a symmetry-aware model §1. Invariant models, which use only scalars, cannot predict quantities that rotate with the input, such as vectors §1, §4.3.

Implementing equivariant networks by hand is technical and error-prone. The paper's stated aim is to provide the algebra so that users can write modules for their own tasks §1. It assumes familiarity with group representations and their irreducible decompositions §4.

Method

Overview. Each tensor in the network is typed by how it transforms under the group E(3). A type is a direct sum of irreducible representations (irreps). An irrep has an order l (dimension 2l+1) and a parity p (even e or odd o). The string 1x0e + 1x2e denotes one even irrep of order 0 and one even irrep of order 2, and a symmetric 3×3 Cartesian tensor has this type §4.4. The paper's example of an odd vector is a velocity; a cross product of two vectors is an even pseudovector §4.5.

1. Types and parity. Use irreps of O(3), the rotation group together with inversion. Parity must be tracked because it changes sign for odd quantities §4.5. Learned weights must be scalars, which do not transform under the group, so the weights do not break equivariance §4.6.
2. Real basis. Irreps are unique only up to a change of basis. The paper uses real-valued matrices. It obtains the complex matrices from QuTiP and changes basis with a fixed map that carries a factor (-i)^l §5, §7.1. The code builds the same change of basis, and its comment says the factor makes the Clebsch-Gordan coefficients real [code: raw/code/2207.09453/e3nn/o3/_wigner.py#L31-L41].
3. Spherical harmonics. The paper defines Y^l as an equivariant map from the sphere to irrep l, unit-normalised on the sphere. The sign of each order is fixed by Y^{l+1} proportional to C^{l+1,l,1} Y^l Y^1 §6.1. Spherical harmonics extend to polynomials on R^3 §6.1, fig. 1. The normalisation options are named norm, integral and component in the code _spherical_harmonics.py.
4. Sphere transforms. A scalar function on the sphere is written as a sum of spherical harmonics with coefficient vectors v^l §6.2. Truncating at a maximum order L removes high frequencies §6.2. The paper uses a grid made of rings around the y axis §6.2, fig. 2. The code's grid uses beta_i = (i + 0.5) pi / N and alpha_j = 2 pi j / M [code: raw/code/2207.09453/e3nn/o3/_s2grid.py#L93-L96].
5. Tensor product. The tensor product is the bilinear, equivariant operation between two typed inputs that outputs irreps §7. Each allowed triple (l1, l2, l3) with parity p1 p2 = p3 and |l1 - l2| <= l3 <= l1 + l2 gives one path §7. A tensor product is a weighted sum of paths, and each output irrep sums the paths that write to it §7, §7.2. Learned weights may be attached to any path, and paths can be included, excluded or reweighted §7, fig. 3.
6. Connection modes. Two modes are described in the paper §7.2:
   • uvw: every pair of input channels is connected and each triple has its own weight. See the pseudocode below.
   • uvu: the first input's channel index passes through and is not summed. Other modes such as uuu and uvuv exist in the library but are not described §7.2.
7. Convenience classes. FullyConnectedTensorProduct wraps the general class with all paths §7 listing. ReducedTensorProduct decomposes a tensor with index-permutation symmetry into irreps, with a basis that satisfies the symmetries §7.4.
8. Generality. The paper argues that the general tensor product can represent any bilinear equivariant map between two sets of irreps. The argument rests on the uniqueness of the irrep decomposition §7.3.
9. Modules. Point convolution in the style of Tensor Field Networks, voxel convolution in the style of 3D Steerable CNNs, and an SE(3)-Transformer-style module are built from these operations §8. The paper also gives a polynomial module over a set of positions. It is equivariant to rotations, translations and inversion, and its listing uses spherical harmonics of the neighbour offsets, two tensor products and neighbourhood sums scaled by inverse square roots of the neighbour and node counts §8 listing.
10. Initialization. All weights are drawn from N(0, 1). A layer computes z = phi(W x / sqrt(d)), where d is the fan-in. The activation phi is rescaled so that its second moment under a standard Gaussian is 1. Components are normalised so that the squared norm of a vector equals its dimension ("component" normalisation). The paper says this matches Maximal Update Parametrization, shifted by theta = 1/2 §8.1.

Tensor-product path, mode uvw §7.2:

# Path with connection mode uvw  [§7.2]
# Input:   x (B, m1, 2*l1+1), y (B, m2, 2*l2+1)
# Weights: w (m1, m2, m3) learned, C (2*l1+1, 2*l2+1, 2*l3+1) fixed
# Output:  (B, m3, 2*l3+1)

def uvw_path(x, y, w, C, m1, m2):
    # 1. Contract both inputs with C and the weights  [§7.2]
    raw = einsum("ijk,bui,bvj,uvw->bwk", C, x, y, w)
    # 2. Scale by 1 / sqrt(m1 * m2), as stated for each path  [§7.2]
    return raw / sqrt(m1 * m2)

Tensor-product path, mode uvu §7.2:

# Path with connection mode uvu  [§7.2]
# Input:   x (B, m1, 2*l1+1), y (B, m2, 2*l2+1)
# Weights: w (m1, m2) learned, C fixed as above
# Output:  (B, m1, 2*l3+1)

def uvu_path(x, y, w, C, m2):
    # 1. Contract with C and the weights; the first channel index u is kept
    #    rather than summed  [§7.2]
    raw = einsum("ijk,bui,bvj,uv->buk", C, x, y, w)
    # 2. Scale by 1 / sqrt(m2)  [§7.2]
    return raw / sqrt(m2)

Scalar layer with variance scaling §8.1:

# Fully connected layer on a scalar vector  [§8.1]
# Input:   x (d,)  with components of variance about 1
# Weights: W (n, d) with entries drawn from N(0, 1)
# Output:  z (n,)

def layer(x, W, phi, d):
    # 1. Divide the matrix-vector product by sqrt(d), the fan-in  [§8.1]
    h = W @ x / sqrt(d)
    # 2. Apply phi, which is rescaled so E[phi(g)^2] = 1 for g ~ N(0, 1)
    #    [§8.1]
    return phi(h)

Key displayed formulas in the paper, cited by section: the symmetric-tensor transformation and its irrep decomposition §4.4, the Clebsch-Gordan equivariance §7.1, the path count and admissible-triple condition §7, the output scale for the sum of paths, out = sqrt(1/number of paths) times the sum over paths §7.2, and the reduced-tensor-product linear system §7.4. The paper's sign rule for spherical harmonics is given in §6.1.

Hyperparameters

The paper has almost no numeric hyperparameters. Values that appear below come from its listings and from the example code, and are marked as such.

Results

Author-reported. The paper contains no benchmark results, tables or experiments of its own §9, §10. Its quantitative content is limited to:

• the claim that a model with order-2 internal features shows a better learning curve than one with only orders 0 and 1, citing other works §1, §9;
• Figure 4, which illustrates a change of the power-law exponent of the learning curve for equivariant networks. The caption states the claim but not the data behind it fig. 4;
• the statement that the paper has no theoretical explanation for that change of exponent §9.

Setup: none. Baselines: none reported. Use the cited works for numbers.

Paper vs code

Checked against the official repository at commit 2aa7f58440a06b15352a2cbce01fa4c26f824969, which is the e3nn repository linked in the paper's bibliography. The code has moved on since 2022 (the e3nn API is newer than the paper's listings).

• Path normalization default. The paper states that each output is sqrt(1/number of paths) times the sum of paths §7.2. The code's default is element, which scales each path by the sum of input variances over all paths into the same output [code: raw/code/2207.09453/e3nn/o3/_tensor_product/_tensor_product.py#L226] [code: raw/code/2207.09453/e3nn/o3/_tensor_product/_tensor_product.py#L305-L320]. The path option multiplies by the number of paths into the output, which is closer to the paper's stated constant [code: raw/code/2207.09453/e3nn/o3/_tensor_product/_tensor_product.py#L307-L309]. Pass the option explicitly if you want the paper's constant.
• Spherical harmonic normalization. The definition uses unit norm §6.1, but the code default is integral [code: raw/code/2207.09453/e3nn/o3/_spherical_harmonics.py#L116].
• Activation normalization. The paper states the constant as an exact Gaussian integral §8.1. The code computes it by Monte Carlo with one million samples and a fixed seed of 0 [code: raw/code/2207.09453/e3nn/math/_normalize_activation.py#L9-L17] [code: raw/code/2207.09453/e3nn/math/_normalize_activation.py#L43]. The result is close to the exact value but not identical.
• Polynomial module. The paper's listing feeds the second tensor product with the middle features for both inputs, and takes the output irreps as an argument §8 listing. The repository's example feeds the spherical harmonics of the edges as the second input, and its output is 0o + 6x0e tetris_polynomial.py. The pooling and normalization steps follow the same pattern.
• Linear layers. The paper's variance-scaled layer §8.1 matches the code in the case where the variance arguments are 1 [code: raw/code/2207.09453/e3nn/nn/_fc.py#L41-L44]. The code multiplies the output by sqrt(var_out), which the paper does not describe.
• Reduced tensor product. The paper prints the decomposition of the rank-4 tensor with formula ijkl=jikl=ijlk=klij as 2x0e+2x2e+1x4e §7.4. The repository's test checks a total dimension of 21 for the formula ijkl=jikl=klij, with input 1e or 1o [code: raw/code/2207.09453/tests/o3/reduce_tensor_test.py#L69-L70]. I did not run the printed example, so the exact string is not checked.
• Learning rate. The paper says the learning rate scales with the number of features §8.1 but gives no constant. The example uses a fixed Adam rate of 1e-2 tetris_polynomial.py.

I did not run any code. Statements about behaviour come from reading the sources.

Pitfalls

• The source is a draft. The LaTeX begins with a line stating that the document is a draft §1. Many passages are commented out and do not appear in the compiled paper, so do not implement a claim from a commented block. This card follows the LaTeX source fetched for the arXiv version. Compare against the PDF before you rely on a section number.
• Equation numbers. The source uses several display environments without numbers in the text, and some align blocks have lines without labels. Cite by section, as this card does, unless you have checked the PDF numbering.
• Path normalization is a choice, not a detail. The default in code is not the constant the paper states. Set path_normalization explicitly and state which one you used.
• Weights must be scalars. The equivariance condition in §4.6 requires the learned weights to be invariant under the group. A weight that transforms with the input breaks equivariance.
• Parity must be tracked. The types 1o and 1e differ under inversion. An operation that is equivariant to SO(3) but not to O(3) will fail on pseudovector or mirror-related inputs §4.5.
• Basis dependence. Spherical harmonics depend on the chosen basis of each irrep §6.1, remark. Use one consistent basis, the real basis of §5, across the network. Mixing it with a complex or differently signed basis changes the outputs without changing equivariance.
• Sign convention. The sign of each harmonic order is fixed by the recursion in §6.1. I did not check the sign convention in the code, so verify it if your model depends on the sign of individual outputs.
• Monte Carlo constant. The activation rescaling is estimated from samples in the code [code: raw/code/2207.09453/e3nn/math/_normalize_activation.py#L17]. Results will differ slightly from the exact integral in §8.1, and the difference is not quantified.
• Reduced tensor product. The solver step uses an eigenvalue cut-off of 1e-9 to select the null space §7.4, listing. The tolerance is sensitive to dtype and conditioning, and the paper gives no guidance on it. The paper also argues that a single component of each irrep suffices to find the solution space, without a formal proof §7.4.
• Costs the paper does not quantify. The paper gives no runtime, memory or scaling numbers. The tensor product carries all paths, so its cost grows with the number of allowed paths and with multiplicities. Plan for this yourself.
• Data efficiency is observed, not explained. The paper reports an observed change in the learning-curve exponent and says it has no theoretical explanation §9. Do not treat it as a guaranteed effect in new settings.
• Learning rate is not stated. The paper says the rate is proportional to the number of features but gives no constant §8.1. The example's fixed rate is a repository choice, not the paper's rule [code: raw/code/2207.09453/examples/tetris_polynomial.py#L114].

Lineage

• Builds on: Tensor Field Networks (arXiv:1802.08219) §1, §2, 3D Steerable CNNs §2, Clebsch-Gordan Nets §2, and the SE(3)-Transformer §2, §8. The paper says these models can be written with e3nn primitives §10. None of these papers has a card on this wiki yet, so they are named, not linked.
• Initialization: the paper uses the Maximal Update Parametrization of Yang and Hu, arXiv:2011.14522 §8.1.
• Compared against: invariant descriptors (Behler-Parrinello symmetry functions, ACE, SOAP), invariant networks (SchNet, DimeNet), restricted equivariant message-passing methods that use only scalar and vector operations, equivariant kernel methods, and Spherical CNNs §2.
• Superseded by: not stated in this paper; no later card on this wiki records it.