Constructing 3D Rotational Invariance and Equivariance with Symmetric Tensor Networks

Symmetry-aware architectures are central to geometric deep learning. We present a systematic approach for constructing continuous rotationally invariant and equivariant functions using symmetric tensor networks. The proposed framework supports inputs and outputs given as a tuple of Cartesian tensors of different rank as well as spherical tensors of different type. We introduce tensor network generators for invariant maps and obtain equivariant maps via differentiation. Specifically, we derive general continuous equivariant maps from vector inputs to Cartesian or spherical tensor output. Finally, we clarify how common equivariant primitives in geometric graph neural networks arise within our construction.

💡 Research Summary

The paper tackles a fundamental challenge in geometric deep learning: how to build neural networks that respect three‑dimensional rotational symmetry (SO(3)) both in an invariant and equivariant sense. The authors propose a unified, mathematically rigorous framework based on symmetric tensor networks. A symmetric tensor is defined as a tensor that remains unchanged under the action of the rotation group on each of its indices. By constructing networks solely from such tensors, the whole network automatically inherits the desired symmetry because tensor contraction preserves symmetry.

The core of the construction is the notion of tensor‑network generators for invariant polynomials. For any set of inputs, the algebra of continuous SO(3)‑invariant functions can be approximated arbitrarily well by polynomials of a finite set of generators (Stone‑Weierstrass theorem). The authors explicitly identify these generators for three families of inputs:

1. Vector inputs (x₁,…,xₙ∈ℝ³). The generators are the inner products xᵢ·xⱼ and the triple‑product (xᵢ×xⱼ)·xₖ. Every invariant polynomial can be expressed as a polynomial in these quantities.

2. Cartesian‑tensor inputs of arbitrary rank. Here the only building blocks needed are the Kronecker delta δᵢⱼ and the Levi‑Civita tensor εᵢⱼₖ. Even‑rank tensors are linear combinations of networks built from δ only; odd‑rank tensors require exactly one ε together with δ’s. This reproduces the classical theory of isotropic tensors.

3. Spherical (irreducible) tensor inputs. Each spherical tensor of type ℓ is first embedded into the reducible space (1)⊗ℓ via a specially constructed symmetric tensor P_ℓ. The P_ℓ tensors are isometric, permutation‑invariant, and annihilated by contraction with ε, guaranteeing that the embedded objects can be treated exactly like the Cartesian case. Consequently, invariant generators are again built from δ, possibly a single ε, and the embedded tensors P_ℓ(x).

Having obtained a set of invariant generators {g₁,…,g_m}, any continuous invariant function f can be written as f(x)=q(g₁(x),…,g_m(x)) where q is an ordinary (non‑symmetric) function. In practice, one can feed the scalar generators g_i into a standard feed‑forward neural network (which need not respect any symmetry) to learn q. The overall composition is guaranteed to be exactly SO(3)‑invariant because the generators themselves are invariant.

To obtain equivariant mappings, the authors observe that differentiation of an invariant function with respect to its vector inputs yields an equivariant map. Concretely, starting from the invariant polynomial expressed via the generators, one differentiates with respect to the input vectors; the resulting expression automatically transforms as a vector (or higher‑rank tensor) under rotations. This provides a systematic recipe for constructing general continuous equivariant maps from vectors to Cartesian or spherical tensors. The familiar equivariant primitives used in geometric graph neural networks—vector addition, cross product, tensor product, tensor contraction—appear as special cases of this differentiation procedure.

The paper also emphasizes the graphical representation of the tensor networks. Nodes correspond to tensors (δ, ε, P_ℓ, or input tensors), edges to contracted indices. This visual language simplifies both theoretical proofs (e.g., showing symmetry preservation) and practical implementation, as tensor contractions can be efficiently executed on GPUs.

In summary, the contributions are:

• Introduction of tensor‑network generators for SO(3)‑invariant polynomials covering vectors, arbitrary Cartesian tensors, and spherical tensors.
• Proof that any continuous invariant function can be approximated by a composition of a generic neural network with these generators (Stone‑Weierstrass based).
• A general construction of continuous equivariant maps via differentiation of invariant generators, yielding explicit formulas for vector‑to‑tensor mappings.
• Demonstration that common equivariant operations in existing geometric GNNs are special instances of the proposed framework.
• A clean graphical formalism that unifies theory and implementation.

This work provides a complete, mathematically grounded toolkit for building rotationally symmetric models, potentially improving data efficiency and generalization in domains such as molecular modeling, materials science, and any application involving 3D geometric data.