G-RepsNet: A Lightweight Construction of Equivariant Networks for Arbitrary Matrix Groups

Sourya Basu, Suhas Lohit, Matthew Brand

Abstract:

Group equivariance is a strong inductive bias useful in a wide range of deep learning tasks. However, constructing efficient equivariant networks for general groups and domains is difficult. Recent work by Finzi et al. directly solves the equivariance constraint for arbitrary matrix groups to obtain equivariant MLPs (EMLPs). But this method does not scale well and scaling is crucial in deep learning.
Here, we introduce Group Representation Networks (G-RepsNets), a lightweight equivariant network for arbitrary matrix groups with features represented using tensor polynomials. The key insight in our design is that using tensor representations in the hidden layers of a neural network along with simple inexpensive tensor operations leads to scalable equivariant networks. Further, these networks are universal approximators of functions equivariant to orthogonal groups. We find G-RepsNet 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.
On image classification tasks, we find that G-RepsNet using second-order representations is competitive and often even outperforms sophisticated state-of-the-art equivariant models such as GCNNs and $E(2)$-CNNs. To further illustrate the generality of our approach, we show that G-RepsNet is competitive to G-FNO and EGNN on N-body predictions and solving PDEs respectively, while being efficient.

Submission Length: Long submission (more than 12 pages of main content)

Changes Since Last Submission:

1. Added author names and affiliations.
2. Added link to the codebase we have released for this work.

Code: https://github.com/merlresearch/G-RepsNets

Assigned Action Editor: ~Grigorios_Chrysos1

Submission Number: 4026