GRepsNet: A Simple Equivariant Network for Arbitrary Matrix Groups
Primary Area: general machine learning (i.e., none of the above)
Code Of Ethics: I acknowledge that I and all co-authors of this work have read and commit to adhering to the ICLR Code of Ethics.
Keywords: Group equivariant networks, geometric deep learning
Submission Guidelines: I certify that this submission complies with the submission instructions as described on https://iclr.cc/Conferences/2024/AuthorGuide.
TL;DR: We provide a simple equivariant architecture for arbitrary matrix groups
Abstract: Group equivariance is a strong inductive bias useful in a wide range of domains including images, point clouds, dynamical systems, and partial differential equations (PDEs). But constructing efficient equivariant networks for general groups and domains is difficult. Recent work by Finzi et al. (2021) directly solves the equivariance constraint for arbitrary matrix groups to obtain equivariant MLPs (EMLPs). However, this method does not scale well and scaling is crucial to get the best from deep learning. This necessitates the design of group equivariant networks for general domains and groups that are simple and scalable. To this end, we introduce Group Representation Networks (GRepsNets), a simple equivariant network for arbitrary matrix groups. The key intuition for our design is that using tensor representations in the hidden layers of a neural network along with appropriate mixing of various representations can lead to expressive equivariant networks, which we confirm empirically. We find GRepsNet to be competitive to EMLP on several tasks with group symmetries such as O(5), O(1, 3), and O(3) with scalars, vectors, and second-order tensors as data types. To illustrate the simplicity and generality of our network, we also use it for image classification with MLP-mixers, predicting N-body dynamics using message passing neural networks (MPNNs), and for solving PDEs using Fourier neural operators (FNOs). Surprisingly, we find that using simple first-order representations itself can yield benefits of group equivariance without additional changes in the architecture. Finally, we illustrate how higher-order tensor representations can be used for group equivariant finetuning that outperforms the existing equivariant finetuning method Basu et al. (2023b).
Anonymous Url: I certify that there is no URL (e.g., github page) that could be used to find authors' identity.
No Acknowledgement Section: I certify that there is no acknowledgement section in this submission for double blind review.
Submission Number: 4426
Loading