Equivariant Polynomial Functional Networks

Thieu Vo, Hoang V. Tran, Tho Tran Huu, An Nguyen The, Thanh Tran, Minh-Khoi Nguyen-Nhat, Duy-Tung Pham, Tan Minh Nguyen

TL;DR: We design a family of monomial matrix-equivariant neural functional networks based on a parameter-sharing mechanism that achieves lower memory consumption and running time while preserving expressivity.

Abstract: A neural functional network (NFN) is a specialized type of neural network designed to process and learn from entire neural networks as input data. Recent NFNs have been proposed with permutation and scaling equivariance based on either graph-based message-passing mechanisms or parameter-sharing mechanisms. Compared to graph-based models, parameter-sharing-based NFNs built upon equivariant linear layers exhibit lower memory consumption and faster running time. However, their expressivity is limited due to the large size of the symmetric group of the input neural networks. The challenge of designing a permutation and scaling equivariant NFN that maintains low memory consumption and running time while preserving expressivity remains unresolved. In this paper, we propose a novel solution with the development of MAGEP-NFN (**M**onomial m**A**trix **G**roup **E**quivariant **P**olynomial **NFN**). Our approach follows the parameter-sharing mechanism but differs from previous works by constructing a nonlinear equivariant layer represented as a polynomial in the input weights. This polynomial formulation enables us to incorporate additional relationships between weights from different input hidden layers, enhancing the model's expressivity while keeping memory consumption and running time low, thereby addressing the aforementioned challenge. We provide empirical evidence demonstrating that MAGEP-NFN achieves competitive performance and efficiency compared to existing baselines.

Lay Summary: A Neural Functional Network (NFN) learns from entire neural networks as input. While efficient NFNs use shared parameters to handle permutations and scaling, they often lack expressivity due to symmetry constraints. We propose MAGEP-NFN, which introduces a nonlinear, polynomial-based layer to capture richer relationships between weights. This approach maintains low memory and fast runtime while improving expressivity, achieving strong performance in practice.

Application-Driven Machine Learning: This submission is on Application-Driven Machine Learning.

Primary Area: Deep Learning

Keywords: neural functional network, equivariant model, permutation equivariance, scaling equivariance

Submission Number: 3839