Source-linked AI summary
Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds
Nathaniel Thomas, Tess Smidt, Steven Kearnes, Lusann Yang, Li Li, Kai Kohlhoff, Patrick Riley
TL;DR
The paper develops neural networks equivariant to 3D rotations, translations, and point permutations, addressing the need to recognize geometric features across orientations and locations. Using spherical-harmonic filters and equivariant point convolutions, it supports scalar, vector, and higher-order tensor fields, with demonstrations in geometry, physics, and chemistry.
Problem
Existing approaches may require rotational data augmentation to identify 3D features across orientations, motivating networks that directly respect 3D Euclidean symmetries.
Method
The paper constructs tensor field networks from spherical-harmonic filters and equivariant point-convolution components whose layers input and output scalar, vector, and higher-order tensor fields.
Results
The networks perform tasks in geometry, physics, and chemistry, including rotated shape classification, agreement with Newtonian gravity and moment-of-inertia radial functions, and point replacement in molecular structures.
Takeaways & Limitations
Rotation and translation equivariance lets the network identify local features across orientations and locations with the same filters, reducing reliance on rotational data augmentation.
Takeaways & Limitations
Maintaining equivariance constrains the architecture to use the same weights across all m values for a given rotation order, with biases permitted only for l = 0.
Abstract
from arXiv · showhide
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.
1 Motivation
The paper proposes tensor field networks for atomic systems and other 3D data, with equivariance to 3D rotations and translations. This design avoids rotational augmentation, supports geometric tensor features, and targets applications in geometry, physics, and chemistry.
- The paper presents networks with equivariance to the symmetries of 3D Euclidean space, including rotations and translations.
- Equivariance can reduce computation and training costs compared with data augmentation for rotation-invariant outputs.
- Tensor field networks operate on point clouds using continuous convolutions over 3D coordinates and point features.
- Their filters combine learnable radial functions with spherical harmonics, requiring compatibility with the algebra of geometric tensors.
- Every layer inputs and outputs tensor fields containing scalars, vectors, and higher-order tensors at geometric points.
2 Related work
Related work established equivariance and rotation-aware modeling in 2D, finite groups, spherical signals, and atomic systems. The paper distinguishes its approach by addressing continuous 3D rotations and geometric tensors for 3D point clouds and atomic systems.
- Harmonic Networks achieve 2D rotation equivariance, while SchNet uses continuous convolutions in a rotation-invariant network; tensor field networks can emulate both.
- Three-dimensional rotation equivariance is mathematically more complicated than the 2D case because 3D rotations do not generally commute.
- Prior approaches include finite symmetry groups, spherical signals with spherical harmonics and Wigner D-matrices, and permutation-group tensors.
- Atomic-system models based on pairwise distances or angular information generally lack the paper’s general equivariance.
- Voxelized approaches to 3D atomic systems incur significant expense, lack guaranteed smooth rotation transformations, and introduce voxelization edge effects.
3 Group representations and equivariance in 3D
The paper formalizes equivariance through group representations and layer conditions for point-cloud transformations. It focuses on rotations while ensuring permutation and translation equivariance through set-based processing and relative coordinates.
- Group representations: A representation D maps each group element to a square matrix while preserving group composition.
- Equivariance: A function is equivariant when transforming its input before applying it matches applying the corresponding output transformation afterward.
- Equivariance: The relevant symmetry operations include 3D isometries and point permutations, with invariance as the special case of an identity output transformation.
- Network construction: Composing equivariant layers produces an equivariant network, so proving equivariance layer by layer is sufficient.
- Point-cloud transformations: Permutation equivariance follows from treating point clouds as unordered sets, while translation equivariance follows from using coordinate differences.
- Rotation equivariance: Rotation equivariance is handled through SO(3) representations, whose irreducible components have dimensions 2l + 1 and include scalars, vectors, and symmetric traceless matrices for l = 0, 1, 2.
4 Tensor field network layers
Tensor field network layers represent point-associated features as SO(3) tensors and use spherical-harmonic filters with tensor-product composition to preserve equivariance. Point convolutions combine relative geometry and input features into multiple rotation orders, while self-interaction and nonlinear layers maintain equivariance through shared representation-wise transformations.
- Layer representations: Each layer associates every point with a vector in an SO(3) representation, decomposed into rotation-order channels indexed by l.Channels contain arrays indexed by point, channel, and representation indices m, with shape [|S|, n_l, 2l + 1].
- Point convolution: Point convolutions apply the same operation at each point using relative locations and the feature values of all other points.The relative vector is r⃗ = r⃗_a − r⃗_b, with its direction and length used by the convolution.
- Scope: The architecture is strictly more general than standard CNNs, which can be represented as point clouds on regularly spaced grids.
- Spherical harmonics and filters: Filters are restricted to a learned radial function multiplied by a spherical harmonic, inheriting spherical harmonics’ SO(3) transformation property.The radial function depends on distance and is scalar in the representation index m, while the spherical harmonic depends on the normalized direction.
- Representation products: Tensor products combine input and filter representations into output representations whose rotation orders range from |l_i − l_f| through l_i + l_f.Clebsch-Gordan coefficients implement this decomposition, including dot and cross products as the 1 ⊗ 1 → 0 and 1 ⊗ 1 → 1 cases.
- Equivariance: The resulting point convolutions are equivariant because both the filters and Clebsch-Gordan coefficients are equivariant.A point convolution can produce 2 min(l_i, l_f) + 1 rotation orders, although a network may omit some outputs.
- Equivariant channel operations: Self-interaction layers mix feature channels like 1x1 convolutions, while shared weights across m and scalar transforms preserve equivariance.Biases are permitted for l = 0 channels; nonlinearities act along the m dimension as scalar transforms.
5 Demonstrations
The demonstrations show tensor field networks handling equivariant geometry, physics, and chemistry tasks, including rotated shape classification, physical quantities, and missing-point prediction.
- The demonstrations span geometry, physics, and chemistry tasks designed to show the power and flexibility of 3D equivariance.The authors state these tasks are unnatural or impossible in existing models.
- 5.1 Geometry: shape classification: Rotation- and translation-equivariant shape classification matches training performance on randomly transformed test shapes without rotational data augmentation.The network classifies eight 3D Tetris shapes with perfect accuracy and performs as well on randomly rotated and translated shapes as on training data.
- 5.1 Geometry: shape classification: The shape classifier distinguishes mirrored 3D Tetris shapes that distance- or angle-only networks cannot distinguish.Its architecture uses scalar outputs because the shape classes are rotation-invariant, then sums point features for classification.
- 5.2 Physics: vectors and tensors in classical mechanics: Single-layer tensor field networks recover the Newtonian gravity inverse-square law and moment-of-inertia radial functions.Gravity uses an l = 1 convolution, while the moment-of-inertia task uses l = 0 and l = 2 convolutions.
- 5.3 Chemistry: toward geometrically generating molecular structures: The missing-point network predicts relative displacement vectors, point types, and confidence values for reconstructing removed molecular points.It trains on QM9 molecular structures and aggregates location votes using confidence-weighted outputs.
- 5.3 Chemistry: toward geometrically generating molecular structures: After 225 epochs on 1000 molecules, the missing-point model generalizes well to test sets containing larger molecules.The training molecules contain 5–18 atoms.
6 Future work
The authors identify future applications for tensor field networks in atomic systems, physics, and 3D-oriented scientific and engineering domains.
- For atomic systems, the authors propose predicting properties of large heterogeneous systems, learning molecular dynamics, and calculating electron densities.They also envision hypothesizing stable structures and designing materials, drugs, and chemicals.
- For physics, proposed applications include complex fluid flows, particle-physics detector events, and configurations of stars and galaxies.
- Additional proposed applications include 3D perception, robotics, computational geometry, and bioimaging.
A Proofs of general equivariance propositions
The proofs establish general equivariance propositions for transformations, point convolutions, and compositions of equivariant functions and networks.
- Equivariance to transformations g and h implies equivariance to their composition.The proposition is stated for all g, h ∈G and x ∈X.
- Composing equivariant functions L1 and L2 yields an equivariant function L2 ◦L1.
- Point clouds can be represented as continuous functions that are non-zero at finitely many points using Dirac δ functions.
- A point convolution is formulated as an integral transform with kernel F, producing the corresponding transformed output.
C Proof of equivariance of point convolution layer
The point-convolution layer achieves rotation equivariance by combining equivariant spherical-harmonic filters with equivariant Clebsch–Gordan coefficients.
- The layer’s rotation-equivariance condition is established through the transformation properties of spherical harmonics.
- The filter construction inherits spherical-harmonic transformation behavior because its radial function is scalar in m.
- Clebsch–Gordan coefficient equivariance combines with filter equivariance to make point convolutions equivariant.
D Details for gravitational accelerations and moment of inertia tasks
The physics demonstrations test whether tensor field filters can learn radial functions for gravitational acceleration and moment of inertia from point-mass configurations.
- D.1 Moment of inertia radial functions: The moment of inertia tensor is expressed using an l = 2 spherical-harmonic component associated with a symmetric traceless tensor.
- D.1 Moment of inertia radial functions: The learned l = 0 and l = 2 radial functions show excellent agreement with the analytical moment-of-inertia solution.
- D.2 Point generation details and radial hyperparameters: The demonstrations generate point counts uniformly from 2 through 10, with masses sampled uniformly from 0.5 to 2.0.
- D.2 Point generation details and radial hyperparameters: The models use 30 Gaussian basis functions with evenly spaced centers between 0 and 2, and Gaussian variance set to half the center spacing.
- D.2 Point generation details and radial hyperparameters: The gravity experiment excludes distances below 0.5 because the inverse-square law diverges as r → 0 and such examples are difficult to sample sufficiently.
E Proof of weighted point-averaging layer equivariance
The weighted point-averaging layer aggregates point votes while preserving rotation equivariance, but its use of absolute coordinates makes translation equivariance non-manifest.
- Each point produces a scalar and vector offset, with the scalar converted by softmax into a probability and the offset defining a vote location.
- The averaging operation uses absolute point coordinates, so translation equivariance is not manifest at this stage.
- The voting scheme remains rotation-equivariant because it sums 3D vectors weighted by probabilities.
- The loss is translation-invariant through vector differences and rotation-invariant through vector dot products.
F Missing point task accuracies and MAE by epoch
The missing-point evaluation reports accuracy and distance MAE by atom type and training epoch, with the main results recorded after 225 epochs.
- Table S1 breaks missing-point prediction accuracy and distance MAE down by atom type.
- The relevant atom-type comparison uses the number of examples in which that atom type was removed, rather than the total dataset size.
- Tables 1 and S1 report results after 225 training epochs.
- Figures S6 and S7 track accuracy and distance MAE as functions of training epochs.