Abstract: Neural networks efficiently encode learned information within their parameters. Consequently, many tasks can be unified by treating neural networks themselves as input data. When doing so, recent studies demonstrated the importance of accounting for the symmetries and geometry of parameter spaces. However, those works developed architectures tailored to specific networks such as MLPs and CNNs without normalization layers, and generalizing such architectures to other types of networks can be challenging. In this work, we overcome these challenges by building new metanetworks — neural networks that take weights from other neural networks as input. Put simply, we carefully build graphs representing the input neural networks and process the graphs using graph neural networks. Our approach, Graph Metanetworks (GMNs), generalizes to neural architectures where competing methods struggle, such as multi-head attention layers, normalization layers, convolutional layers, ResNet blocks, and group-equivariant linear layers. We prove that GMNs are expressive and equivariant to parameter permutation symmetries that leave the input neural network functions unchanged. We validate the effectiveness of our method on several metanetwork tasks over diverse neural network architectures.

What if neural nets could process, analyze, interpret, and edit other neural nets' weights? These models, termed metanetworks, are extremely versatile and useful. For instance, metanets have been used for learned optimization (map [params, gradients] → new params), 3D data generation (generating INR weights), analyzing parameter spaces (predict hyperparameters from weights), and more.

We develop Graph Metanetworks (GMNs). GMNs are metanetworks that operate on neural architectures while respecting parameter symmetries. To achieve this we,

1. Convert input neural networks to graphs, where neurons are nodes and parameters are edges.
2. Process this with a graph neural network or Graph Transformer

Crucially, our design allows a single metanet to operate on a diverse set of neural architectures.

We develop procedures to turn many types of layers into graphs, such as: multihead attention, convolutions, linear, residual, group equivariant linear, spatial grid, and normalization layers. We experimentally process small ViTs, ResNets, DeepSets, and other input nets.

Theoretically, we prove that our GMNs are equivariant/invariant to computation graph automorphisms of the inputs, which correspond to permutation parameter symmetries (e.g. permuting hidden neurons in MLPs, hidden channels in CNNs).

Propositions 1&2 (informal): Graph Metanets are equivariant to parameter permutation symmetries.

We prove that our GMNs are able to express existing metanets.

Proposition 3: On MLP inputs (where parameter graphs and computation graphs coincide), Graph Metanets can express StatNN and NP-NFN.

They can even simulate the forward pass of any network defined as a computation graph. This means they are at least as expressive as the input networks they operate on.

Proposition 4: On computation graph inputs, graph metanets can express the forward pass of any input feedforward neural network.

We empirically show that our GMNs improve over previous metanets on several tasks:

• Predicting CIFAR-10 test accuracy of diverse input image classifiers (ViTs, ResNets, DeepSets, etc) from their weights alone
• Editing 2D neural implicits
• Self-supervised learning of representations of neural implicits

Citation

Derek Lim, Haggai Maron, Marc T. Law, Jonathan Lorraine, James Lucas
Graph Metanetworks for Processing Diverse Neural Architectures. ICLR, 2024.