# Equivariance with Learned Canonicalization Functions

arXiv: https://arxiv.org/abs/2211.06489 · HTML: https://mlsanity.com/paper/2211.06489/

## Key facts

- Cited by 109 papers in the map (data through 14 Oct 2026).
- Builds on 55 of its 79 references that are in the map.
- First posted on arXiv 11 Nov 2022.

## Builds on (most cited first, 30 of 55)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift](https://mlsanity.com/paper/1502.03167/) (2015)
- [PointNet: Deep Learning on Point Sets for 3D Classification and Segmentation](https://mlsanity.com/paper/1612.00593/) (2016)
- [PointNet++: Deep Hierarchical Feature Learning on Point Sets in a Metric Space](https://mlsanity.com/paper/1706.02413/) (2017)
- [ShapeNet: An Information-Rich 3D Model Repository](https://mlsanity.com/paper/1512.03012/) (2015)
- [Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distances](https://mlsanity.com/paper/1306.0895/) (2013)
- [Dynamic Graph CNN for Learning on Point Clouds](https://mlsanity.com/paper/1801.07829/) (2018)
- [Estimating or Propagating Gradients Through Stochastic Neurons for Conditional Computation](https://mlsanity.com/paper/1308.3432/) (2013)
- [Reformer: The Efficient Transformer](https://mlsanity.com/paper/2001.04451/) (2020)
- [Spatial Transformer Networks](https://mlsanity.com/paper/1506.02025/) (2015)
- [Deep Sets](https://mlsanity.com/paper/1703.06114/) (2017)
- [Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges](https://mlsanity.com/paper/2104.13478/) (2021)
- [Group Equivariant Convolutional Networks](https://mlsanity.com/paper/1602.07576/) (2016)
- [E(n) Equivariant Graph Neural Networks](https://mlsanity.com/paper/2102.09844/) (2021)
- [Tensor field networks: Rotation- and translation-equivariant neural networks for 3D point clouds](https://mlsanity.com/paper/1802.08219/) (2018)
- [SE(3)-Transformers: 3D Roto-Translation Equivariant Attention Networks](https://mlsanity.com/paper/2006.10503/) (2020)
- [Input Convex Neural Networks](https://mlsanity.com/paper/1609.07152/) (2016)
- [Neural Relational Inference for Interacting Systems](https://mlsanity.com/paper/1802.04687/) (2018)
- [Vector Neurons: A General Framework for SO(3)-Equivariant Networks](https://mlsanity.com/paper/2104.12229/) (2021)
- [Learning SO(3) Equivariant Representations with Spherical CNNs](https://mlsanity.com/paper/1711.06721/) (2017)
- [Kornia: an Open Source Differentiable Computer Vision Library for PyTorch](https://mlsanity.com/paper/1910.02190/) (2019)
- [Efficient and Modular Implicit Differentiation](https://mlsanity.com/paper/2105.15183/) (2021)
- [Fast Differentiable Sorting and Ranking](https://mlsanity.com/paper/2002.08871/) (2020)
- [A Practical Method for Constructing Equivariant Multilayer Perceptrons for Arbitrary Matrix Groups](https://mlsanity.com/paper/2104.09459/) (2021)
- [Learning Latent Permutations with Gumbel-Sinkhorn Networks](https://mlsanity.com/paper/1802.08665/) (2018)
- [Frame Averaging for Invariant and Equivariant Network Design](https://mlsanity.com/paper/2110.03336/) (2021)
- [Point Convolutional Neural Networks by Extension Operators](https://mlsanity.com/paper/1803.10091/) (2018)
- [Investigating Bi-Level Optimization for Learning and Vision from a Unified Perspective: A Survey and Beyond](https://mlsanity.com/paper/2101.11517/) (2021)
- [MDP Homomorphic Networks: Group Symmetries in Reinforcement Learning](https://mlsanity.com/paper/2006.16908/) (2020)
- [Equivariance Through Parameter-Sharing](https://mlsanity.com/paper/1702.08389/) (2017)

## Built on by (most cited first, 30 of 109)

- [MiDi: Mixed Graph and 3D Denoising Diffusion for Molecule Generation](https://mlsanity.com/paper/2302.09048/) (2023)
- [A Hitchhiker's Guide to Geometric GNNs for 3D Atomic Systems](https://mlsanity.com/paper/2312.07511/) (2023)
- [Orb: A Fast, Scalable Neural Network Potential](https://mlsanity.com/paper/2410.22570/) (2024)
- [SymmCD: Symmetry-Preserving Crystal Generation with Diffusion Models](https://mlsanity.com/paper/2502.03638/) (2025)
- [Equivariant Frames and the Impossibility of Continuous Canonicalization](https://mlsanity.com/paper/2402.16077/) (2024)
- [Equivariant Adaptation of Large Pretrained Models](https://mlsanity.com/paper/2310.01647/) (2023)
- [A Canonicalization Perspective on Invariant and Equivariant Learning](https://mlsanity.com/paper/2405.18378/) (2024)
- [Learning Probabilistic Symmetrization for Architecture Agnostic Equivariance](https://mlsanity.com/paper/2306.02866/) (2023)
- [Symmetry Breaking and Equivariant Neural Networks](https://mlsanity.com/paper/2312.09016/) (2023)
- [Scale Equivariant Graph Metanetworks](https://mlsanity.com/paper/2406.10685/) (2024)
- [Lie Algebra Canonicalization: Equivariant Neural Operators under arbitrary Lie Groups](https://mlsanity.com/paper/2410.02698/) (2024)
- [Learning on LoRAs: GL-Equivariant Processing of Low-Rank Weight Spaces for Large Finetuned Models](https://mlsanity.com/paper/2410.04207/) (2024)
- [Laplacian Canonization: A Minimalist Approach to Sign and Basis Invariant Spectral Embedding](https://mlsanity.com/paper/2310.18716/) (2023)
- [Emergent Equivariance in Deep Ensembles](https://mlsanity.com/paper/2403.03103/) (2024)
- [Improving Equivariant Networks with Probabilistic Symmetry Breaking](https://mlsanity.com/paper/2503.21985/) (2025)
- [Expressive Sign Equivariant Networks for Spectral Geometric Learning](https://mlsanity.com/paper/2312.02339/) (2023)
- [Beyond Canonicalization: How Tensorial Messages Improve Equivariant Message Passing](https://mlsanity.com/paper/2405.15389/) (2024)
- [Symmetry-Informed Governing Equation Discovery](https://mlsanity.com/paper/2405.16756/) (2024)
- [Relaxed Equivariance via Multitask Learning](https://mlsanity.com/paper/2410.17878/) (2024)
- [Lorentz Local Canonicalization: How to Make Any Network Lorentz-Equivariant](https://mlsanity.com/paper/2505.20280/) (2025)
- [Efficient Equivariant Transfer Learning from Pretrained Models](https://mlsanity.com/paper/2305.09900/) (2023)
- [Towards fully covariant machine learning](https://mlsanity.com/paper/2301.13724/) (2023)
- [Using Multiple Vector Channels Improves E(n)-Equivariant Graph Neural Networks](https://mlsanity.com/paper/2309.03139/) (2023)
- [Weisfeiler Leman for Euclidean Equivariant Machine Learning](https://mlsanity.com/paper/2402.02484/) (2024)
- [A Generative Model of Symmetry Transformations](https://mlsanity.com/paper/2403.01946/) (2024)
- [Equivariant Symmetry Breaking Sets](https://mlsanity.com/paper/2402.02681/) (2024)
- [Symmetry-Aware Bayesian Flow Networks for Crystal Generation](https://mlsanity.com/paper/2502.03146/) (2025)
- [Equivariant Denoisers Cannot Copy Graphs: Align Your Graph Diffusion Models](https://mlsanity.com/paper/2405.17656/) (2024)
- [Latent Field Discovery In Interacting Dynamical Systems With Neural Fields](https://mlsanity.com/paper/2310.20679/) (2023)
- [Tilt your Head: Activating the Hidden Spatial-Invariance of Classifiers](https://mlsanity.com/paper/2405.03730/) (2024)

## Abstract (arXiv)

Symmetry-based neural networks often constrain the architecture in order to achieve invariance or equivariance to a group of transformations. In this paper, we propose an alternative that avoids this architectural constraint by learning to produce canonical representations of the data. These canonicalization functions can readily be plugged into non-equivariant backbone architectures. We offer explicit ways to implement them for some groups of interest. We show that this approach enjoys universality while providing interpretable insights. Our main hypothesis, supported by our empirical results, is that learning a small neural network to perform canonicalization is better than using predefined heuristics. Our experiments show that learning the canonicalization function is competitive with existing techniques for learning equivariant functions across many tasks, including image classification, N-body dynamics prediction, point cloud classification and part segmentation, while being faster across the board.
