Equivariance by Local Canonicalization: A Matter of Representation
This work addresses the computational bottleneck for researchers and practitioners using equivariant networks in molecular and geometric data applications, though it is incremental as it builds on existing tensor field networks.
The authors tackled the computational inefficiency of equivariant neural networks by introducing a framework that transfers tensor field networks into the local canonicalization paradigm, preserving equivariance while significantly improving runtime, with empirical results showing speed-ups and competitive accuracy.
Equivariant neural networks offer strong inductive biases for learning from molecular and geometric data but often rely on specialized, computationally expensive tensor operations. We present a framework to transfers existing tensor field networks into the more efficient local canonicalization paradigm, preserving equivariance while significantly improving the runtime. Within this framework, we systematically compare different equivariant representations in terms of theoretical complexity, empirical runtime, and predictive accuracy. We publish the tensor_frames package, a PyTorchGeometric based implementation for local canonicalization, that enables straightforward integration of equivariance into any standard message passing neural network.