Permutation Equivariant Neural Functionals
This work addresses a foundational challenge in machine learning for researchers and practitioners by providing a unifying principle for processing neural network weights, which is incremental as it builds on symmetry principles to improve architecture design.
The paper tackles the problem of designing neural networks that can process the weights or gradients of other neural networks, known as neural functional networks (NFNs), by introducing a permutation equivariant framework that encodes symmetries as an inductive bias. The result shows that these networks are effective on tasks like predicting classifier generalization, producing sparsity masks, and classifying or editing implicit neural representations, with code provided for reproducibility.
This work studies the design of neural networks that can process the weights or gradients of other neural networks, which we refer to as neural functional networks (NFNs). Despite a wide range of potential applications, including learned optimization, processing implicit neural representations, network editing, and policy evaluation, there are few unifying principles for designing effective architectures that process the weights of other networks. We approach the design of neural functionals through the lens of symmetry, in particular by focusing on the permutation symmetries that arise in the weights of deep feedforward networks because hidden layer neurons have no inherent order. We introduce a framework for building permutation equivariant neural functionals, whose architectures encode these symmetries as an inductive bias. The key building blocks of this framework are NF-Layers (neural functional layers) that we constrain to be permutation equivariant through an appropriate parameter sharing scheme. In our experiments, we find that permutation equivariant neural functionals are effective on a diverse set of tasks that require processing the weights of MLPs and CNNs, such as predicting classifier generalization, producing "winning ticket" sparsity masks for initializations, and classifying or editing implicit neural representations (INRs). In addition, we provide code for our models and experiments at https://github.com/AllanYangZhou/nfn.