Equivariance with Learned Canonicalization Functions

Sékou-Oumar Kaba, Arnab Kumar Mondal, Yan Zhang, Yoshua Bengio, Siamak Ravanbakhsh
Proceedings of the 40th International Conference on Machine Learning, PMLR 202:15546-15566, 2023.

Abstract

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.

Cite this Paper

BibTeX
@InProceedings{pmlr-v202-kaba23a, title = {Equivariance with Learned Canonicalization Functions}, author = {Kaba, S\'{e}kou-Oumar and Mondal, Arnab Kumar and Zhang, Yan and Bengio, Yoshua and Ravanbakhsh, Siamak}, booktitle = {Proceedings of the 40th International Conference on Machine Learning}, pages = {15546--15566}, year = {2023}, editor = {Krause, Andreas and Brunskill, Emma and Cho, Kyunghyun and Engelhardt, Barbara and Sabato, Sivan and Scarlett, Jonathan}, volume = {202}, series = {Proceedings of Machine Learning Research}, month = {23--29 Jul}, publisher = {PMLR}, pdf = {https://proceedings.mlr.press/v202/kaba23a/kaba23a.pdf}, url = {https://proceedings.mlr.press/v202/kaba23a.html}, abstract = {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.} }
Endnote
%0 Conference Paper %T Equivariance with Learned Canonicalization Functions %A Sékou-Oumar Kaba %A Arnab Kumar Mondal %A Yan Zhang %A Yoshua Bengio %A Siamak Ravanbakhsh %B Proceedings of the 40th International Conference on Machine Learning %C Proceedings of Machine Learning Research %D 2023 %E Andreas Krause %E Emma Brunskill %E Kyunghyun Cho %E Barbara Engelhardt %E Sivan Sabato %E Jonathan Scarlett %F pmlr-v202-kaba23a %I PMLR %P 15546--15566 %U https://proceedings.mlr.press/v202/kaba23a.html %V 202 %X 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.
APA
Kaba, S., Mondal, A.K., Zhang, Y., Bengio, Y. & Ravanbakhsh, S.. (2023). Equivariance with Learned Canonicalization Functions. Proceedings of the 40th International Conference on Machine Learning, in Proceedings of Machine Learning Research 202:15546-15566 Available from https://proceedings.mlr.press/v202/kaba23a.html.

Related Material