G-RepsNet: A Lightweight Construction of Equivariant Networks for Arbitrary Matrix Groups
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. Removed analysis and claims from figure and table captions as reviewer requested.
2. Merged and moved some tables and figures that were referenced in the main paper, but were in the appendix.
3. Added points to the limitations section that clarified points raised by the reviewers regarding complexity of higher order tensors, tensor decomposition layers, and difficulty in defining and computing group action-invariant functions for arbitrary matrix groups.
4. Added runtime and number of parameters for the models in the paper in the appropriate sections.
5. Added some sentences to make additional connections and explain similarities between (a) merging layers in G-RepsNets and bilinearity layers in EMLPs and (b) G-RepsCNNs being a special case of Steerable CNNs.
6. Fixed typos related to method name, some technical details, and citations.
7. Improved the readability of model design under Section 5.1
8. Added an additional paragraph in the Background section 3 to define the orthogonal and Lorentz groups.
9. Clarified method of computing the invariant from $T_2$ features in G-RepsCNNs.
Assigned Action Editor: ~Grigorios_Chrysos1
Submission Number: 4026
Loading