Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Equivariance by Local Canonicalization: A Matter of Representation

Gerrit Gerhartz, Peter Lippmann, Fred A. Hamprecht

TL;DR

The paper presents a framework to translate existing equivariant tensor-field networks into the lightweight local canonicalization approach, achieving exact equivariance with substantially reduced runtime. It introduces LoCaFormer, a PyTorch Geometric-based architecture that uses an EDGE-like message-passing layer and flexible local-frame transformations to enable efficient, tensor-aware communication. Through experiments on QM9 and a 10k N-methylacetamide dataset, the authors demonstrate competitive accuracy and 4–5x speedups over the traditional Equiformer, while providing a modular software package (tensor_frames) for broad integration. The study also analyzes representations (Cartesian vs irreps) and radial/angular embeddings, highlighting data-efficiency advantages of built-in equivariance and detailing the trade-offs for tensorial targets and computational costs.

Abstract

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.

Equivariance by Local Canonicalization: A Matter of Representation

TL;DR

The paper presents a framework to translate existing equivariant tensor-field networks into the lightweight local canonicalization approach, achieving exact equivariance with substantially reduced runtime. It introduces LoCaFormer, a PyTorch Geometric-based architecture that uses an EDGE-like message-passing layer and flexible local-frame transformations to enable efficient, tensor-aware communication. Through experiments on QM9 and a 10k N-methylacetamide dataset, the authors demonstrate competitive accuracy and 4–5x speedups over the traditional Equiformer, while providing a modular software package (tensor_frames) for broad integration. The study also analyzes representations (Cartesian vs irreps) and radial/angular embeddings, highlighting data-efficiency advantages of built-in equivariance and detailing the trade-offs for tensorial targets and computational costs.

Abstract

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.

Paper Structure

This paper contains 17 sections, 18 equations, 3 figures, 6 tables.

Table of Contents

  1. Introduction
  2. Background: group representations and equivariance
  3. Representations of $\boldsymbol{\mathrm{O}(3)}$.
  4. Equivariance via local canonicalization.
  5. Learning representations
  6. Transforming tensor field based architectures to local canonicalization
  7. Discussion and results on molecular property prediction
  8. Software contribution
  9. Connection between irreducible and Cartesian tensor representation
  10. Computational complexity of Cartesian tensors and irreducible representations
  11. Radial and angular embedding of the molecule geometry
  12. Experimental details
  13. Hyperparameter choices.
  14. Architectural design.
  15. Tensorial property prediction.
  16. ...and 2 more sections

Figures (3)

  • Figure 1: LoCaFormer attention block.
  • Figure 2: Bessel embedding functions. Here, Bessel functions of the first kind are depicted with ten different frequencies. The Bessel functions are also multiplied by the envelope function with a cutoff radius of $1$.
  • Figure 3: Equivariance increases data efficiency compared to data augmentation on QM9. However, in the low data regime the data augmentation yields superior accuracy.