Tensor Field Networks and Geometric Deep Learning

The TFN Framework

In TFN, each node i carries a feature field — a collection of irreducible representations:

Where c_l is the number of channels at degree l. This is like having separate “colour channels” for each geometric degree:

• c_0 channels of scalars (l=0)
• c_1 channels of 3D vectors (l=1)
• c_2 channels of 5D quadrupolar features (l=2)
• etc.

The TFN Layer

Breaking this down:

• Y^{l_f}(r̂_{ij}): spherical harmonics of degree l_f evaluated at the edge direction
• f_j^{(l_in)} ⊗_{l_f} Y^{l_f}: CG tensor product combining node features (degree l_in) with geometric features (degree l_f) to produce output degree l_out

• W^{…}(

  r_{ij}

  ): radial weight function (depends only on distance, so invariant)

Aggregation and update:

The Geometric Deep Learning Blueprint

The TFN paper, together with Bronstein et al. (2021) “Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges,” established a unified framework:

Every neural network architecture = symmetry group + representation + invariant/equivariant layers

This unification shows that architectural choices are really choices about which symmetries to encode — and which geometric domain the data lives in.

From TFN to NequIP and MACE

NequIP (Batzner et al., 2022): extends TFN with:

• Message passing framework (not just pair interactions)
• Gate nonlinearity: f^{(l)} ← f^{(l)} · σ(f^{(0)}) (scalars gate higher-order features)
• Achieves state-of-the-art in molecular force fields with very few training points

MACE (Batatia et al., 2022): extends NequIP with:

• Higher-body-order interactions (not just pairwise — triplets, quadruplets)
• Many-body basis functions via tensor product pooling
• State-of-the-art on MD17 (molecular dynamics benchmark)

Equivariant Nonlinearities

Standard MLPs (ReLU, sigmoid) break equivariance when applied to l>0 features — the result is not equivariant. Two equivariant nonlinearity designs:

Gate activation: multiply l>0 features by a gating scalar (l=0):

Scalars are gated by nonlinear functions; higher-order features are gated by scalars (maintaining equivariance).

Norm nonlinearity: apply nonlinearity to the norm of each feature vector:

Norm is invariant; normalised direction is equivariant. Applying σ to the norm and scaling preserves equivariance.

Summary

TFN’s contribution is not just an architecture — it is the mathematical language in which geometric deep learning is now written. Understanding spherical harmonics, CG products, and irreducible representations is prerequisite knowledge for reading the current state-of-the-art in equivariant GNNs.

References

• Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., & Riley, P. (2018). Tensor Field Networks: Rotation- and Translation-Equivariant Neural Networks for 3D Point Clouds. arXiv 2018 (TFN: the original SE(3)-equivariant MPNN using spherical harmonics and Clebsch-Gordan tensor products for arbitrary-order geometric features).
• Batzner, S., Musaelian, A., Sun, L., Geiger, M., Mailoa, J. P., Kornbluth, M., Molinari, N., Smidt, T. E., & Kozinsky, B. (2022). E(3)-Equivariant Graph Neural Networks for Data-Efficient and Accurate Interatomic Potentials. Nature Communications 2022 (NequIP: TFN + MPNN architecture achieving state-of-the-art accuracy and data efficiency on molecular force fields).
• Batatia, I., Kovacs, D. P., Simm, G., Ortner, C., & Csányi, G. (2022). MACE: Higher Order Equivariant Message Passing Neural Networks for Fast and Accurate Force Fields. NeurIPS 2022 (MACE: many-body interactions via equivariant tensor products enabling higher-order correlations).