Equivariant Polynomial Functional Networks
Thieu N. Vo, Viet-Hoang Tran, Tho Tran Huu, An Nguyen The, Thanh Tran, Minh-Khoi Nguyen-Nhat, Duy-Tung Pham, Tan Minh Nguyen
TL;DR
This work introduces MAGEP-NFN, a permutation and scaling equivariant polynomial neural functional network that uses a weight-sharing mechanism with nonlinear polynomial layers. Central to the approach is a newly defined class of stable polynomial terms that remain compatible with monomial group actions, enabling expressive invariant and equivariant layers with low memory and fast runtime. The authors provide a formal characterization of all invariant and equivariant layers within this stable polynomial framework and instantiate them into practical MAGEP-NFNs. Empirical results across implicit neural representations, CNN generalization from weights, and weight-space editing show competitive or superior performance to existing baselines while preserving efficiency. The work advances NFNs by combining strong symmetry guarantees with enhanced expressivity through higher-order weight relationships, with broad applicability to architectures using known weight-space symmetries.
Abstract
Neural Functional Networks (NFNs) have gained increasing interest due to their wide range of applications, including extracting information from implicit representations of data, editing network weights, and evaluating policies. A key design principle of NFNs is their adherence to the permutation and scaling symmetries inherent in the connectionist structure of the input neural networks. Recent NFNs have been proposed with permutation and scaling equivariance based on either graph-based message-passing mechanisms or parameter-sharing mechanisms. However, graph-based equivariant NFNs suffer from high memory consumption and long running times. On the other hand, parameter-sharing-based NFNs built upon equivariant linear layers exhibit lower memory consumption and faster running time, yet 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 (Monomial mAtrix Group Equivariant Polynomial 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.