[Paper Review] Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds

Nathaniel Thomas, Tess Smidt+5 more|arXiv (Cornell University)|Feb 22, 2018

3D Shape Modeling and AnalysisEngineering19 references564 citations

TL;DR

Introduces tensor field networks that are rotation- and translation-equivariant for 3D point clouds, enabling scalar, vector, and higher-order tensor outputs via spherical-harmonic-based filters and Clebsch-Gordan tensor products.

ABSTRACT

We introduce tensor field neural networks, which are locally equivariant to 3D rotations, translations, and permutations of points at every layer. 3D rotation equivariance removes the need for data augmentation to identify features in arbitrary orientations. Our network uses filters built from spherical harmonics; due to the mathematical consequences of this filter choice, each layer accepts as input (and guarantees as output) scalars, vectors, and higher-order tensors, in the geometric sense of these terms. We demonstrate the capabilities of tensor field networks with tasks in geometry, physics, and chemistry.

Motivation & Objective

• Motivate learning on 3D Euclidean space with full rotation and translation equivariance to reduce data augmentation and improve efficiency.
• Develop a universal architecture operating on point clouds with continuous local convolutions and geometric tensor outputs.
• Provide rotation-equivariant building blocks (point convolution, self-interaction, nonlinearity) compatible with SO(3) representations.
• Demonstrate applications in geometry, physics, and chemistry including shape classification, dynamics, and molecular structure generation.

Proposed method

• Define layers that take a set of 3D points with tensor-valued features and output tensor-valued features at each point.
• Constrain convolution filters to F_{cm}^{(l_f,l_i)}(r) = R^{(l_f,l_i)}_{c}(r) Y_m^{(l_f)}( hat) using spherical harmonics Y_m^{(l)}, ensuring SO(3) equivariance.
• Combine input and filter representations via tensor products using Clebsch-Gordan coefficients to produce outputs transforming under SO(3) irreducible representations.
• Prove equivariance of point convolutions, self-interactions, and nonlinearities under permutations, translations, and rotations.
• Implement l=0,1,2 representations corresponding to scalars, vectors, and symmetric traceless tensors, with layer-wise channels per rotation order.
• Demonstrate on tasks: geometry (shape classification without data augmentation), physics (accelerations and inertia tensors), and chemistry (missing-point generation in QM9).

Experimental results

Research questions

• RQ1 How can a neural network be locally equivariant to 3D rotations, translations, and point permutations at every layer?
• RQ2 Can filters built from spherical harmonics and radial functions realize SO(3) equivariance and support scalar, vector, and higher-order tensor outputs?
• RQ3 What tasks in geometry, physics, and chemistry demonstrate the practical benefits of tensor field networks over non-equivariant or partially equivariant models?
• RQ4 How does the network perform on rotation- and translation-invariant outputs across diverse 3D data modalities?

Key findings

• Tensor field networks achieve 3D rotation- and translation-equivariance at every layer.
• With rotation-equivariant filters, the model outputs tensors of consistent transformation order (scalars, vectors, higher-order tensors) per point.
• In geometry tasks, the network classifies 3D shapes in a single orientation and remains accurate under random rotations/translations (no rotational data augmentation).
• In physics tasks, the model learns Newtonian accelerations and moment of inertia tensors with correct radial behavior using minimal layers.
• In chemistry, the network can predict missing atoms in QM9 structures with high accuracy and good generalization across molecule sizes.
• The approach scales to hundreds to thousands of points and complements existing molecular/directed-graph architectures.