Graph Metanetworks for Processing Diverse Neural Architectures
Authors: Derek Lim
Haggai Maron
Marc T. Law
Jonathan Lorraine
James Lucas
Venue: ICLR 2024 spotlight (2024)
Keywords: Metanetwork
graph
equivariance
expressivity
Primary area: learning on graphs and other geometries & topologies
TL;DR: We develop metanetworks that allow expressive, permutation equivariant processing of diverse neural network architectures.

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.

Published as a conference paper at ICLR 2024
GRAPH METANETWORKS FOR PROCESSING DIVERSE NEURAL ARCHITECTURES
Derek Lim
Haggai Maron
Mare T. Law NVIDIA marel @nvidia.com
MIT CSAIL
Technion / NVIDIA
dereklim@mis t. edu
hmaron nvi dia. com
Jonathan Lorraine NVIDIA jlorraine@nvicia.com
James Lucas NVIDIA jalucasenvidia.com
ABSTRACT
Neural networks efficiently encode learned information within their parameters. Consequently, many tasks can be unified by treating neural networks themselves as imput data. When doing so, recent studies demonstrated the importance of ac- counting for the symmetries and geometry of parameter spaces. However, those works developed architectures tailored specific networks such as MLPs and CNNs without normalization layers, and generalizing such architectures to other types of networks can be challenging. this work, we overcome these chal- lenges by building 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 multi-head attention layers, normal- ization layers, convolutional layers, ResNet blocks, and group-equivarian linear layers. We prove that GMNs are expressive and equivariant parameter permu- tation 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.
metanetworks neural networks that take weights from
INTRODUCTION
Neural networks are well-established for predicting, generating, and transforming data. newer paradigm is to treat the parameters of neural networks themselves as data. This insight inspired re- searchers to suggest neural architectures that can predict properties of trained neural networks (Eil- ertsen et al. 2020). generate new networks (Erkoç et al. 2023). optimize networks (Metz et al.] 2022). or otherwise transform them (Navon et al. 2023; Zhou et al.) 2023a). We refer to these neural networks that process other neural networks as metanetworks, metanets for short.
Metanets enable new applications, but designing them is nontrivial. A common approach is to flatten the network parameters into vector representation, neglecting the input network structure. More generally, a prominent challenge in metanet design is that the space of neural network parameters exhibits symmetries. For example, permuting the neurons in the hidden layers of a Multilayer Per- ceptron (MLP) leaves the network output unchanged (Hecht-Nielsen) 1990). Ignoring these symme- tries greatly degrades the metanet performance (Peebles al.) 2022 Navon al. 2023). Instead, cquivariant metanets respect these symmetries, so that if the input network is permuted then the metanet output permuted in the same way.
Recently, several works have proposed equivariant metanets that have shown significantly improved performance (Navon et al. 2023 Zhou et al. 2023a b). However, these networks typically require highly specialized, land-designed layers that can be difficult to devise. A careful analysis of its symmetries is necessary for any input architecture, followed by the design of corresponding equiv- ariant metanet layers. Generalizing this procedure to more complicated network architectures is
Work done while interning NVIDIA.

---

Published as a conference paper at ICLR 2024
Eraph Definition
Graph Notanatwork
Output
Linear
ReLU
Linear
Graph Neural Networks
Baiihh tormalization
L Lincar
Input
Figure I: Overview of Graph Metanetworks (GMNs) Our method converts neural network ar- chitectures into parameter graph where edges correspond to network parameters. The bias (b) and batch-normalization parameters are incorporated via additional nodes with edges to the relevant layer's neurons. The graph is processed by graph neural network operating edge attributes. Fixed-length (invariant) predictions can be extracted by pooling the output graph features.
time-consuming and nontrivial, so existing methods can only process simple input networks with linear and convolutional layers, and cannot process standard modules such as normalization layers or residual connections - let alone more complicated modules such as attention blocks. Moreover, these architectures cannot directly process input neural networks with varying architectures, such as those with different numbers of layers hidden units.
This work offers a simple and clegant solution to metanet design that respects neural network param- eter symmetries. As in the concurrent work of Zhang et al.) (2023), our technique's 's crux is represent- ing an input neural network as graph (see Figure[]). We show how to efficiently transform a neural network into graph such that standard techniques for learning on graphs e.g., Message Passing Neural Networks Gilmer et al. (2017); Battaglia et al. (2018) or Graph Transformers (Rampášek) al., [022 will be equivariant the parameter symmetries. One of our key contributions is in developing a compact parameter graph representation, which in contrast to established compu- tation graphs allows us handle parameter-sharing layers like convolutions and attention layers without scaling with the activation count. While past work is typically restricted to processing MLPs and simple Convolutional Neural Networks (CNNs) (LeCun et al. 1989), we validate ex- perimentally that our graph metanets (GMNs) generalize more complicated networks such as Transformers (Vaswani et al. 2017), residual networks (He et al. 2016), normalization layers (Toffe & Szegedy 2015; Ba et al. 2016 Wu & He 2018), general group- eouuue-quivaiiant architectures (Ravan- bakhsh et al. (017) like Deep Sets (Zaheer et al. 2017), and more.
We prove theoretically that metanets equivariant to permutation symmetries in the input network, which we formulate via neural graph automorphisms (Section (.2p. This generalizes the hidden neuron permutations in MLPs and channel permutations of CNNs covered in prior work to arbitrary feedforward neural architectures. We further prove that our metanets operating on com- putation graphs are least expressive prior methods meaning they can approximate them to arbitrary accuracy and can express the forward pass of any input feedforward neural network, generalizing a result of Navon al. (2023) (Section
Empirical evaluations show that our approach solves variety of metanet tasks with diverse neural architectures. As part of this effort, we trained new datasets of diverse image classifiers, including 2D CNNs, ID CNNs, DeepSets, ResNets, and Vision Transformers. Our method is easier to imple- ment than past equivariant metanets while being at least as expressive, and is applicable more general input architectures. Crucially, our GMNs achieve strong quantitative performance across all tasks we explored.
GRAPH AUTOMORPHISSN-AASEH METANETS
We first explain how neural networks can be encoded as Directed Acyclic Graphs (DAGs). There many choices in representing neural networks DAGS, perhaps the most common being
2

---

Published as a conference paper at ICLR 2024
computation graph (see Appendix C This work introduces a more compact representation, referred to as parameter graphs.
We then introduce one of the paper's main concepts - Neural DAG Automorphisms. This concept generalizes previously studied symmetry groups for MLPs and CNNs to arbitrary feedforward ar- DAGS. To conclude this section, deseribe GNN-based metanet that operates over these graphs and is equivariant to Neural DAG Automorphisms. A glossary of our notation is provided in Appendix Table
chitectures represented
Motivation. Certain permutations of parameters in neural networks do not change the function they parameterize. For example, consider simple MIP defined such that fo(=) == W2 W1 with one hidden layer, where == (W2, W1) are the parameters of the network, and o is nonlinear element-wise activation function. For any permutation matrix P, if we define O := (W₂P*, PW1). then for all inputs æ, we have fa(x) = fo(a)). This P corresponds to per- mutation of the order of the hidden neurons, which is well-known not to affect the network function (Hecht-Nielsen] 1990p. Likewise, permuting the hidden channels of CNN does not affect the network function (Navon et al. 2023: Entezari et al. 2022 Ainsworth et al. 2023).
While these permutation symmetries for MLPs and simple CNNs are easy to determine by manual inspection, it is more difficult to determine the symmetries of general architectures. For exam- ple, simple residual connections introduce additional neuron dependencies across layers. Instead of manual inspection, we show that graph automorphisms (i.e. graph isomorphisms from graph to itself) on DAGs representing feedforward networks correspond to permutation parameter symme- tries. From this observation, it can be shown that GNNs acting these DAGS equivariant to their permutation symmetries.
Overview of our approach. See Figure for an illustration. Given a general input feedforward neural network, we first encode it as a graph in which each parameter is associated with an edge and then process this graph with a GNN. The GNN outputs single fixed-length vector or predictions for each node or edge depending on the learning task. For instance, one graph-level task is to predict the scalar accuracy of an input neural network on some task. An edge-level task is to predict new weights for an input neural network to change its functionality somehow.
We now discuss the graph construction, the symmetries of these graphs, and the GNN we use.
2.1
GRAPH CONSTRUCTION FOR GENERAL FEEDFORWARD ARCHITECTURES
2.1.1
COMPUTATION GRAPHS
Every feedforward neural network defines computation graph DAG (Zhang al. 2023). where nodes are neurons and edges hold neural network parameter weight values (see Fig. and Fig. Zb. Thus, this gives a method to construct weighted graph. However, the com- putation graph approach has some downsides. For one, it may be expensive due to weight- sharing: for instance, l-input-and-outpee channel 2D-convolution layer with a kernel size of 2 has parameters, but the 4 parameters are used in the computation of many activations (e.g. 1024 activations for 32 x 32 input grid). Further, we may want to add input node and edge features such layer number to help performance and expressive powe Figure 2 illustrates example of (small) computation
Convolution Computation Graph
An example computation graph for a network with a single convolutional layer. The layer has a 2 x 2 filter kernel, a single input and output channel, and applies the filter with stride of 2. Even in this small case of 4 x 4 input image, the graph has 16 edges for only parameters.
Figure 2:
graph for convolutions (for visual clarity, we exclude bias terms). More details, including the exact formal correspondence between feedforward neural network functions and computation graphs, are given in Appendix C.I
In Scction our proofs rely these features show graph metanets express existing metanets.
3

---

Published as a conference paper at ICLR 2024
20 Convolution Layer
Attention Layer
Residual Layer
2x2 Nernel
Quary
Output
Key
Figure 3: Parameter subgraph constructions for assorted layers that implemented in
pirical evaluation. Their descriptions are given in Section 2.1.2 Further details are discussed in Appendix B
2.1.2 PARAMETER GRAPHS
To deal with the challenges of computation graphs, we propose alternate neural network graph con- structions
examples of which are shown in Figure that (a) efficient and (b) allow expressive metanets. We call these graphs parameter graphs because we design the graphs so that each parameter is associated with single edge of the graph (whereas a parameter may be associated to many edges in a computation graph).
We design modular subgraphs that be created for each layer and then stitched together. Our goal is to design parameter graphs with at most one edge for each parameter in the neural network. Additionally, they should capture the correct set of parameter permutation symmetries. Full details are in Appendix B but we discuss the major points behind the design of selection of parameter graphs here.
Linear layers. Figure I depicts three linear layers in sequence. Each linear layer's parameter subgraph coincides with computation graph, but even so, there are important design choices to be made. Bias parameters could be encoded node features done by /Zhang et al. (2023) as self-loop edges on each neuron. Instead, we include bias node for each layer and encode the bias parameters edges from that node to the corresponding neurons in the layer.
The bias node approach is preferable because the self-loop or node feature approaches to encoding biases can hinder the expressive power of the metanet. In particular, the results Navon et al. (2023) and Zhou et al. (2023a) show that the permutation equivariant linear metanet maps for MLP inputs include crsssage-pasiing like operations where the representation of bias parameter is updated via the representations of other bias parameters in the same layer. Using a message passing GNN on graph with bias nodes allows naturally express these operations in single graph metanet layer, as explained in Appendix D.1.3
Convolution layers. Convolutions and other group equivariant linear layers leverage parameter sharing, where the same parameter is used to compute many activations (Ravanbakhsh et al. 2017). Therefore, representing convolutions network graph to a choice of input size. To avoid this. we develop succinct and expressive param- eter graph representation of convolutions. This depicted the left of Figure (3 for convolution layer with a 2 x 2 filter, one input channel, and two output channels. Note that we have two output channels here, unlike the computation graph Figure 2 where there is only one.
computation graph introduces scaling issues and binds the
Our subgraph construction allocates a node for each input and output channel. We then have parallel edges between each input and output node for each spatial location the filter kernel - making this multigraph. Bias nodes are added as in the linear layers. This subgraph contains exactly one edge for each parameter in the convolution layer while capturing the parameter permutation symmetries graph automorphisms. The spatial position of each weight within the kernel is included by positional encoding in the corresponding edge feature.
In Section 4.1
our graph metanets process ID and 2D convolutional networks. well as DeepSets networks that consist of permutation equivariant linear layers (Zaheer et al. (017).
Multi-head attention layers. Attention layers (Vaswani et al. 2017) also exhibit parameter sharing across sequence length. As with convolutions, we design an efficient subgraph representation where each parameter appears single edge. There is node for each feature dimension of the input,
4

---

Published as a conference paper at ICLR 2024
Linear Layers
20 Convolution Layer
Residual Layer
Figure 4: Examples of neural DAG automorphisms for linear layers, convolutional layers, and resid- ual layers. Possible node permutations are illustrated using red and purple dashed arrows, the same color represents an identical transformation.
vectors used in attention computation, and output. There are two sets of edges: a set that corresponds to the query, key, and value maps, and second set corresponding to the final output mapping.
In the middle Figure 3 we show a single-headed self attention layer, with bias nodes excluded for visual clarity. Generalizing this to multi-head attention simply requires adding additional node features to the middle layer that indicate which head each node belongs to.
Residual layers. A residual connection does not introduce any additional parameters, but it does affect the permutation symmetries in the parameter space of the network. Therefore, is crucial to represent residual connections additional parameter-free edges within the parameter graph. The top-right of Figure 3 shows a residual connection bypassing a linear layer. The edges are drawn as dashed lines emphasize that there graph, we fix the weight of the residual edge to be 1.
associated parameters. As is natural in the computation
2.2 NEURAL DAG AUTOMORPHISMS
The prior section describes how to represent (feedforward) neural networks as DAGs, A natural question from an equivariant learning perspective is: what are the symmetries of this DAG repre- sentation? Specifically, we consider graph automorphisms, which are structure-preservigg transfor- mations of graph unto itself. A neural DAG automorphism of DAG (V, E)) associated with neural network is a permutation of nodes Φ V + V that preserves adjacency, preserves types of nodes (e.g. d cannot map hidden neurons input neurons), and preserves weight-sharing straints (i.e. tied weights must still be tied after permuting the endpoints with the automorphism); see Appendix C.2 for more details. Every automorphism also induces permutation of edges 6º E + E. where edge (i, j.) is mapped to d= ((i.j)) = (o(1).((j))
Intuitively, a neural DAG automorphism represents a permutation of the neural network parameters via the induced edge permutation, Φ" We write ((0) to represent this permutation on the parameters themselves, meaning (+0) (()(()())(()))
0 (s.j)- Hidden node permutations in MLPs and hidden channel permutations in CNNs are special cases of neural DAG automorphisms, which we explain Appendix C.3 To formalize these notions, we show that every neural DAG automorphism & of computation graph is permutation parameter symmetry, in the sense that the induced parameter permutation & does not change the neural network function. Figure Bigure/fillusee illustrates several neural DAG automorphisms.
Proposition 1. For any neural DAG automorphism ó of a computation graph, the neural network function is left unchanged: Vx E fo( (x)))(x).
proof is given in Appendix C.4 Recall that our goal is to design metanets equivariant to parameter permutation symmetries. Proposition I shows to achieve this it is necessary to design metanets that are equivariant to neural DAG automorphisms. Graph metanets achieve this exactly since GNNs are equivariant to permutation symmetries of graphs (Maron et al. 2019).
Proposition 2. Graph metanets are equivariant to parameter permutations induced by neural DAG automorphisms.
5

---

Published as a conference paper at ICLR 2024
These results formally justify using graph metanets computation graphs for equivariance to pa- rameter permutation symmetries. Now, noting the parameters are stored as edge features in our computation and parameter graphs, design graph neural networks that operate these DAGs.
2.3 FORMULATING METANETS GNNS
After constructing the input graphs. we use GNN as our metanet to learn representations and perform downstream tasks. We desire GNNs that learn edge representations since the input neural network parameters are placed on the edges of the constructed graph. While countless GNN variants have been developed in the last several years (Hamilton 2020), most Iearn node representations.
For simplicity, mostly loosely follow the general framework of Battaglia al. (2018), which defines general message-passing GNNs that update node, edge, and global features. For a graph, let v; Rd be the feature of node i, (G.)) E Rd feature of the directed edge (i.j), global feature associated to the entire graph, and let E be the set of edges in the graph. The directed edge (i.j) represents edge starting from j and ending at i. We allow multigraphs, where there can be several edges (and hence several edge features) between a pair of nodes (i,j); thus, we let
Rd be
F(..) denote the set of edge features associated with (i,j). Then, general GNN layer updating these features can be written as:
0. < MLP
((.)),
(1)
((i)) < MLP"(uy Ug ((i))) U)
(2)
u t MLP" Σ
e,
(3)
cE
Intuitively, node features are updated by message passing along neighbors (Equation D the features of adjacent nodes update the features of the edges connecting them (Equation 2). and the global feature is updated with aggregations of all features (Equation 3). While our graphs DAGs, are free to use undirected edges by ensuring that (i.j) E E implies (j. E with (i.j) = -(.)) We often choose to allow message passing between layers in both directions.
For parameter-level metanet tasks with per-parameter predictions, we let the prediction be the fi- nal feature ((i)) of the parameter's corresponding edge. We pool edge features for network-level metanet tasks where fixed-length vector the final prediction for each graph. We can pool node features but found pooling edge features sufficient. We do not use global features for empirical results, but we find them crucial for the expressive power results of Proposition 3
3 EXPRESSIVE POWER OF GRAPH METANETS (GMNS)
Ideally, one does not sacrifice expressive power when restricting neural architecture to satisfy group equivariance constraints. We want our metanet architecture to be powerful enough learn useful functions of network parameters. To this end, we first show that GMNs can express two existing equivariant metanets on MI.P inputs. Consequently, GMNs are at least as expressive as these approaches.
Proposition 3. On MLP inputs (where parameter graphs and computation graphs coincide), graph metanets can express StatNN [Unterthiner et al. 2020) and NP-NFN (Zhou al. 2023a
The proof is given in Appendix D.I StatNN is based on permutation-invarian
statistics of each weight or bias, which graph metanetworks can easily compute. The linear layers of NP-NEN consist of local weessage-passinglite operations and global operations that the global feature can capture.
Second, we show that GMNs can simulate the forward pass of an input neural network repre- sented by the DAG of its computation graph (Appendix C.T). This substantially generalizes result of Navon et al. (2023), who show that their DWSNets can simulate the forward pass of MLPs.
Proposition 4. On computation graph inputs, graph metanets can express the forward pass of any input feedforward neural network as defmeed in Section C.1

---

Published as a conference paper at ICLR 2024
Table I: Results for predicting the test accuracy of input neural networks trained on CIFAR-IO. The top results use a training set of 15 000 uniformly selected input networks, the middle results use 10% of this training set, and the bottom results only train on imput networks of low hidden dimension (while testing on networks with strietly higher hidden dimension). Our method performs best in all setups, with increasing benefits in the low-data and OOD regimes.
Varying CNNs
Diverse Architectures
Metanet
R
7
R²
T
DeepSets Zaheer 2017
.778=0022 .022
.6971.002
562:1.020
.5591.011
50%
DMC (Bilertsen al.
2020
948--00
.876+.003
9577+0009
.883-.007
GMN (Ours)
978 002
.15+.0066
.95-+.002
.908+.004
DeepSets (aheer et al
2017
692 .006
.6881.002
.1261.015
.2901.010
5%
DMC (Bilertsen et al. 2020
816= .038
7622.0014
8104.046
.758+.013
GMN (Ours)
876=.010
.797=.005
.18+.002
8284.005
DeepSets (aheer al. 2017
-015
.883
.128:1.71
.3801.014
OOD
DMC (Eilertsen
2020
387 .29
760=.024
-.134 .147
5665.0055
GMN (Ours)
.89 037
.800
.768-t.066
.7801t.03
Varying CNNs
Diverse Architectures
Diverse Architectures, Test
Test
Test
CMN
Density
Pisptibs
3
2
2
5
7
2
3
6
77
9
CIFAR- 10 Accuracy
Figure 5: Histograms of CIFAR-10 accuracies for Varying CNNs and Diverse Architectures datasets. Left and middle show train and test accuracy for the two datasets. Right shows test accu- racy of Diverse Architectures split by model type.
The above result applies only to GMNs operating on computation graphs. This makes it directly applicable the parameter graphs for MLPs they coincide with the computation graph. However, we expect that a similar result possible for more general parameter graphs. We leave formal proof of this future work.
EXPERIMENTS
4.1 PREDICTING ACCURACY FOR VARYING ARCHITECTURES
Task. As in prior works (Unterthiner al. 2020: Zhou al. 2023a train metanets predict the test accuracy of input neural networks, We consider image classification neural networks trained on the CIFAR-10 dataset (Krizhevsky 2009), and the metanet task is to take the parameters of an input network and predict the network's image classification accuracy on the CIFAR-IO test set.
Datasets. To demonstrate the flexibility our graph metanets, train image classifiers that signif- icantly vary in size and architecture. For our "Varying CNNs" dataset, we train about 30 000 basic 2D CNNs varying common architectural design choices like hidden dimension. number of convo- lution layers, number of fully connected layers, and type of normalization layers (BatchNorm (Toffe
Szegedy
2015) or GroupNorm (Wu & He 2018)).
For our "Diverse Architectures" dataset, we also train four other diverse image classifiers to test our model's generalization with multiple architectures: (i) basic ID CNNs treating images as raster or- dered sequences of pixels, (ii) DeepSets treating images pixel sets, which maintains pixel position information with positional encodings (Zaheer et al. 2017). (iii) ResNets (He et al. 2016), and (iv) Vision Transformers (Dosovitskiy al. (021)) This dataset also totals about 300 networks. The Vision Transformers' patch embedding module can be viewed as a convolution. which is how we encode this module graph (the Transformer layer graph encodings are described in Figure 7.
7

---

Published as a conference paper at ICLR 2024
Figure shows these trained networks span wide range of accuracies, between 10% and 77.5% for Varying CNNs and between 7.5% and 87.8% for Diverse Architectures. Also, the number of parameters in these networks ranges from 970 to 21 165 in Varying CNNs and from 970 87 658 for Diverse Architectures. These networks are significantly more diverse and achieve higher accuracy than the dataset of Unterthiner et al. (2020), who train small CNNs of fixed architecture that obtain at most 56% test accuracy on CIFAR-I0. Also, our Diverse Architectures dataset contains many more network types and modules than the datasets of Bilertsen et al. (2020) and Schürrollt et al. 2022), which are limited to CNNs.
Metanetworks. For our graph metanet, we consider a simple message passing GNN as in Sec- tion 2.3 that does not use a global graph feature. To obtain an invariant prediction, we mean-pool over edge representations.
We cannot apply competing permutation equivariant methods like DWSNet (Navon et al. 2023 or NFN (Zhou al.) (2023a), because they cannot process normalization layers, imput neural networks of different sizes, modules like self-attention. Instead, as a baseline, we consider the Deep Meta Classifier (DMC) from Bilertsen et al. (2020), which vectorizes an input network's 's parameters and processes it with ID CNN, allowing the use of differently sized networks. We also consider baseline that treats the parameters as a set and applies a DeepSets network (Zaheer et al., (017) to output scalar prediction. Note that the DeepSets bascline is invariant permutation parameter symmetries, but it is also invariant to permutations that do not correspond to parameter symmetries (which significantly outnumber permutation parameter symmetries), it has low expressive power.
We evaluate our method and the two baselines across six different data settings. Using both the Varying CNN dataset and the Diverse Architectures dataset. We explore training on about half of the input networks, training on only 10% of this previous split, and an out-of-distribune (OOD) generalization setting where we train on a reduced set of architectures that have smaller hidden dimension than the held-out architectures.
Results. See Table T for quantitative results and Figure g in the Appendix for scatterplots in the OOD setting on Diverse Architectures. We report the R-Squared value and the Kendall 7 coefficient of the predicted generalization against the true generalization. Our GMNs outperform both baselines in predicting accuracy for input networks across all six data settings. When we restrict to 10% of the full training set size, we see that GMNs generalize substantially better than the baselines. This performance gap is maintained in the more challenging OOD generalization setting, where the non- cquivariant DMC performs very poorly in R?. The improved GMN performance could be from high expressive power (which the DeepSets bascline lacks), with better generalization due to equivariance to parameter permutation symmetries (which DMC lacks).
4.2 EDITING 2D INRs
Table 2: Test MSE (lower is better) for edit-
Table 3: Results for self-supervised learning of
ing 2D INRs, following the methodology of (hhou et al.] (2023a). Results of baselines are from (Zhou et al. [2023a
neural net representations, in test MSE of linear regressor on the learned representations. Num- bers besides GMN from Navon et al. (2023).
Metanetwork
Contrast
Dilate
Metanctwork
Test MSE
MLP
.031
.306
MLP
7.39 1.19
MLP-Aug
.029
.307
MLP + Perm. aug
5.65 1.01
NEN-PT
.029
.197
MLP- + Alignment
4.47 1.15
NEN-HNP
.0204 00000
.0706
0005
VeR2 (Arch.)
3.86 4.32
NEN-NP
.0203 0000
.0693
= 0009
Transformer
5.12
NFT
.0200 + 0.00
.0510
:.0004
DWSNets
1.39 +.06
GMN (ours)
.0197 1.000
.0603 +.0010
GMN (ours)
1.13 4.08
Next, we empirically test the ability of GMNs to process simple MLP inputs to compare against fles-flexible permutation equivariant metanets such as NFN (Zhou et al., 2023a) and NET (Zhou et al. 2023b). For this, we train metanetworks the 2D INR editing tasks of (Zhou et al. 2023a), where the imputs are weights of an INR representing an image, and the outputs are weights of an INR representing the image with some transformation applied to it.
8

---

Published as a conference paper at ICLR 2024
Table 2 shows results. We see that our GMNs outperform most metanetworks these simple input networks. In particular, GMN outperforms all methods on the Contrast task, and only beat by NFT on the Dilate task.
4.3 SELP-SUPERVISED LEARNING WITH INRS
We also compare GMNs against another less-flexible permutation equivariant metanet, DWS- Nets (Navon et al., 2023 in the self-supervised learning experiments Navon et al. (2023). Here, the imput data MLPs fit to sinusoidal functions of the form x + assinbbz), where b R are varying parameters. The goal is to learn a metanet encoder that gives strong representations of input networks, using contrastive-learinn framework similar SimCLR (Chen al. 2020). Our GMNs learn edge representations, which are then mean-pooled to get a vector representation each input network. The downstream task for evaluating these representations fitting linear model on the metanet representations to predict the coefficients and b of the sinusoid that the input MLP represents.
Table 3 shows the results. GMNs outperform all baselines in this task. Thus, even in the most advantageous setting for competitors, which they are restricted to, GMNs are empirically successful at metanet tasks.
5 CONCLUSION
In this work, proposed Graph Metanetworks, an approach processing neural networks with theoretical and empirical benefits. Theoretically, our approach satisfies permutation parameter sym- metry equivariance while having provably high expressive power. Empirically, we can process di- verse neural architectures, including layers that appear in state-of-the-an models, and we outperform existing metanetwork baselines across all tasks that we evaluated.
Limitations We make substantial progress towards improving the scalability of GMNs by intro- ducing parameter graphs. However, large neural networks can have billions of parameters and pro- cessing them may be more difficult. We believe that standard scalable GNN methods can be used to scale to the billion parameter regime (as existing GNNs are capable of processing graphs with billions of edges (Hu et al. 2021) using modest computing resources), but we have not yet tried to use Graph Metanets at such a scale. Additionally, we argue that parameter graphs are easier to design than specialized architectures of prior work, but do not give formal constraints design. Further work investigating parameter graphs is promising; for instance, our theory of neural DAG automorphisms depends the more expensive computation graphs, but it could possibly be extended to parameter graphs. Moreover, our approach only accounts for permutation-bast rameter symmetries and does account for e.g. symmetries induced by scaling weights in ReLU networks (Dinh et al. 2017 Godfrey et al. 2022).
their
expeesive
pa-
Future work Our graph-based approach metanets promising for future development. This pa- mostly uses a basic pressage-passingg GNN architecture. which can be further improved using the many GNN improvements in the literature. Furthermore, our theoretical developments largely apply to computation graphs and we expect future work can extend these results to the more practical pa- rameter graphs. Since graph metanets can process modern neural network layers like self-attention and spatial INR feature grids, they can be used to process and analyze state-of-the-an neural net- works. Further, future work could try applying graph metanets to difficult yet impactful metanets tasks such as pruning, learned optimization, and finetuning pretrained models.
ACKNOWLEDGMENTS
We thank Matan Atzmon, Jiahui Huang, Karsten Kreis, Francis Williams, and Xiaohui Zeng for helpful comments. We would also like to thank Jun Gao, Or Perel, and Frank Shen for helpful input on some of our early INR experiments. DL is funded by NSF Graduate Fellowship. HM is the Robert J. Shillman Fellow, and supported by the Israel Science Foundation through a personal grant (ISF 264/23) and equipment grant (ISF 532/23).

---

Published as a conference paper at ICLR 2024
REFERENCES
Samuel Ainsworth, Jonathan Hayase, and Siddhartha Srinivasa. Git re-baiin: Merging models mod- ulo permutation symmetries. In The Eleventh International Conference on Learning Representa- tions, 2023. htLhhtts:://// /
net (
Bruno Andreis, Soro Bedionita, and Sung Ju Hwang. Set-based neural network encoding. arXiv preprint arXiv:2305. 16625, 2023.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer normalization. arXiv preprint
arXivv1607.064450, 2016.
Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sunchez-Gonzaler Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al. Relational inductive biases, deep learning, and graph networks. arXiv preprint arXiv: arXiv:/806.011 1806.0/26,, 2018.
Matthias Bauer, Emilien Dupont, Andy Brock, Dan Rosenbaum, Jonathan Schwarz, and Hyunjik Kim. Spatial functa: Scaling functa to imagenet classification and generation. arXiv preprint arXiv:2302.03130, 2023.
Léon Bottou and Patrick Gallinari. A framework for the cooperation of learning algorithms. Ad- vances neural information processing systems, 3, 1990.
Adriano Cardace, Pierluigi Zama Ramirez, Francesco Ballerini, Allan Zhou, Samuele Salti, and Luigi Di Stefano. Neural processing of tri-plane hybrid neural fields. arXiv preprint arXiv:2310.01 140, 2023.
Eric R Chan, Connor Z. Lin, Matthew A Chan, Koki Nagano, Boxiao Pan, Shalini De Mello, Orazio Gallo, Leonidas J Guibas, Jonathan Tremblay, Sameh Khamis, et al. Efficient geometry- aware 3d generative adversarial networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp. 16123-16133, 2022.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597-1607. PMLR, 2020.
Terrance Veeeies and Graham W Taylor. Improved regularization of convolutional neural networks with cutout. arXiv preprint arXiv:1708.04533 2017.
Cameron Diao and Ricky Loynd. Relational attention: Generalizing transformers for graph- structured tasks. In The Eleventh International Conference on Learning Representations, 2023.
Laurent Dinh, Razvan Pascanu, Samy Bengio, and Yoshua Bengio. Sharp minima can generalize for deep nets. In International Conference on Machine Learning, pp. 1019-1028. PMLR, 2017.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16 words: Transformers for image recogni- tion at scale. In International Conference on Learning Representations, 2021. URL hetps: //openzeview.co
cNTTY
Emilien Dupont, Hyunjik Kim, SM Eslami. Danilo Rezende, and Dan Rosenbaum. From data to functa: Your data point is a function and you can treat it like one. International Conference on Machine Learning (ICMI.,, 2022.
Gabriel Eilertsen, Daniel Jönsson, Timo Ropinski, Jonas Unger, and Anders Ynnerman. Classifying the classifier: dissecting the weight space of neural networks. Proceedings of the European Conference on Artificial Intelligence, 2020.
Rahim Entezari, Hanie Sedghi, Olga Saukh, and Behnam Neyshabur. The role of permutation in- variance in linear mode connectivity of neural networks. In International Conference on Learning Representations, 2022. URL https: htttt:s::///peeeevviiie
net wtt.ee/wwww....eeoooooooooaaiiaa
10

---

Published as a conference paper at ICLR 2024
Ziya Erkoç, Fangchang Ma, Qi Shan, Matthias Nießner, and Angela Dai. Hyperdiffusion: Generat- ing implicit neural fields with weight-space diffusion. In Proceedings of the IEEE/CVF Interna- tional Conference on Computer Vision (ICCV), pp. 14300-14310, October 2023.
Jirí Fiala and Jan Kratochvíl. Locally constrained graph homomopphimm--stnenn and applications. Computer Science Review, 2(2):97-111, 2008.
complexity,
Mare Finzi, Max Welling, and Andrew Gordon Wilson. practical method for constructing equiv- ariant multilayer perceptrons for arbitrary matrix groups. In International conference on machine learning, pp. 1-3888.. PMLR, 2021.
Cedric Gegout, Bernard Girau, and Fabrice Rossi. mathematical model for feed-forwart neural networks theoretical description and parallel applications. Research Report LIP RR-1995-23, Laboratoire I l'informatique du parallélisme, September 1995. URL BtLLhttps://// //hal lara.
archi Archives-oo wwweress...eess....oooe fr/hal 02101991 945
Justin Gilmer, Samuel Schoenholz, Patrick Riley, Oriol Vinyals, and George E Dahl. Neural message passing for quantum chemistry. In International conference on machine learning, pp- 1263-1272. PMLR, 2017.
Charles Godfrey, Davis Brown, Tegan Emerson, and Henry Kvinge. On the symmetries of deep learning models and their internal representations. In Alice Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), Advances in Neural Information Processing Systems, 2022.
William Hamilton. Graph representation learning. Morgan & Claypool Publishers, 2020.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog- nition. In Proceedings of the IEEE conference computer vision and pattern recognition, pp- 770-778, 2016.
Robert Hecht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Ad- vanced Neural Computers, pp. 129-135. Elsevier, 1990.
Kurt Hornik. Approximation capabilities of multilayer feedforward networks. Neural networks, (2):251-251) 1991.
Weihua Hu, Matthias Fey, Hongyu Ren, Maho Nakata, Yuxiao Dong, and Jure Leskovec. Ogb-Ise: A large-scale challenge for machine learning on graphs. arXiv preprint arXiv:2103.09430, 2021.
Sergey Toffe and Christian Szegedy. Batch normalization: Accelerating deep network training by reducing internal covariate shift. In International conference machine learning, pp. 448-456. pmlr, 2015.
Jinwoo Kim, Sacyoon Oh, and Seunghoon Hong. Transformers generalize deepsets and can be extended graphs & hypergraphs. Advances Neural Information Processing Systems, 34: 28016- 2021.
Jinwoo Kim, Dat Nguyen, Seonwoo Min, Sungjun Cho, Moontae Lee, Honglak Lee, and Seunghoon Hong. Pure transformers are powerful graph learners. Advances in Neural Information Processing Systems, 35: 15:44582-44595 2022a.
Jinwoo Kim, Sacyoon Oh, Sungjun Cho, and Seunghoon Hong. Equivariant hypergraph neural networks. In European Conference Computer Vision, pp. 86-103. Springer, 2022b.
Diederik Kingma and Jimmy Ba. Adam: method for stochastic optimization. arXiv preprint arXiv: 1412.694/266980 2014.
Boris Knyazev, Michal Drozdzal, Graham W Taylor, and Adriana Romero Soriano. Parameter prediction for unseen deep architectures. Advances in Neural Information Processing Systems, 34:29433-294488 2021.
Alex Krizhevsky. Learning multiple layers of features from tiny images. 2009. URL https: / 7/www. cs toroomooeedb///ri///eeaaa
-R. TR.

---

Published as a conference paper at ICLR 2024
Guillaume Leclerc, Andrew Ilyas, Logan Engstrom, Sung Min Park, Hadi Salman, and Alek- sander Madry. Ffev: Accelerating training by removing data bottlenecks. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pp. 12011-12020, 2023.
Yann LeCun, Bernhard Boser, John Denker, Donnie Henderson, R. Howard, Wayne Hubbard, and Lawrence Jackel. Handwritten digit recognition with a propapaaonn network. In
D. Touretzky (ed.), Advances in Neural Information Processing Systems, volume 2. Morgan- fttpc/pppppp__eepppppppeerr http:///...cceeee4444f66666666e0
Kaufmann, 1989. URL htullittps//prrceddngg...eeeee
Derek Lim, Joshua David Robinson, Lingxiao Zhao, Tess Smidt, Suvrit Sra, Haggai Maron, and Stefanie Jegelka. Sign and basis invariant networks for spectral graph representation learning. In The Eleventh International Conference on Learning Representations, 2023.
Or Litany, Haggai Maron, David Acuna, Jan Kautz, Gal Chechik, and Sanja Fidler. Feder- ated learning with heterogeneous architectures using graph hypernetworks. arXiv preprint arXiv:2201.08845. 2022.
Hanxiao Liu, Karen Simonyan, and Yiming Yang. DARTS: Differentiable architecture search. In International Conference on Learning Representations, 2019. URL httttps:://peeeeeiieerr.
net / wtw.:eroondeeeeeeeeecc.ceeee
Luca De Luigi, Adriano Cardace, Riccardo Spezialetti, Pierluigi Zama Ramirez, Samuele Salti, and Luigi di Stefano. Deep learning on implicit neural representations of shapes. In The Eleventh International Conference on Learning Representations, 2023. httstt::::/ //tppr/vi///e..iee..
net/Eorum?ide 000IW
3uadi
Liheng Ma, Chen Lin, Derek Lim, Adriana Romero-Soriano, Puneet K Dokania, Mark Coates, Philip Torr, and Ser-Nam Lim. Graph inductive biases in transformers without message passing. In ICML, 2023.
Haggai Maron, Heli Ben-Hamu, Nadav Shamir, and Yaron Lipman. Invariant and equivariant graph networks. In International Conference on Learning Representations, 2019. URL https:
openreview. net/ www.eeee/oouddssyyeeeee 22
Luke Metz, James Harrison, Daniel Freeman, Amil Merchant, Lucas Beyer, James Bradbury, Naman Agrawal, Ben Poole, Igor Mordatch, Adam Roberts, et al. Velo: Training versatile learned optimizers by scaling up. arXiv preprint urXiv:2211.097700 2022.
Thomas Müller, Alex Evans, Christoph Schied, and Alexander Keller. Instant neural graphics prim- itives with a multiresolution hash encoding. ACM Transactions on Graphics (ToG), 41(4::1-15, 2022.
Aviv Navon, Aviv Shamsian, Idan Achituve, Ethan Fetaya, Gal Chechik, and Haggai Maron. Equiv- ariant architectures for learning deep weight spaces. International Conference on Learning Representations, 2023.
Adam Paszke, Sam Gross, Soumith Chintala, Gregory Chanan, Edward Yang, Zachary DeVito, Zeming Lin, Alban Desmaison, Luca Antiga, and Adam Lerer. Automatic differentiation in pytorch. 2017.
William Peebles, Ilija Radosavovic, Tim Brooks, Alexei A Efros, and Jitendra Malik. Learning to learn with generative models of neural network checkpoints. arXiv preprint arXiv:2209. 12892, 2022.
Ladislav Rampášek, Michael Galkin, Vijay Prakash Dwivedi, Anh Tuan Luu, Guy Wolf, and Do- minique Beaini. Recipe for general, powerful, scalable graph transformer. Advances Neural Information Processing Systems, 35:14501-14515, 2022.
Siamak Ravanbakhsh. Jeff Schneider, and Barnabas Poczos. Equivariance through parameter- sharing. In International conference machine learning, pp. 2892-2901. PMLR, 2017.
12

---

Published as a conference paper at ICLR 2024
Konstantin Schürholt, Diyar Taskiran, Boris Knyazev, Xavier Giró-i Nieto, and Damian Borth. Model zoos: A dataset of diverse populations of neural network models. Advances in Neural Information Processing Systems, 35:38134-3814 2022.
Towaki Takikawa, Joey Litalien, Kangxue Yin, Karsten Kreis, Charles Loop, Derek Nowrouzezahrai, Alec Jacobson, Morgan McGuire, and Sanja Fidler. Neural geometric level of detail: Real-time rendering with implicit 3D shapes. 2021.
Veronika Thost and Jie Chen. Directed acyclic graph neural networks. International Conference on Learning Representutions, 2021.
Thomas Unterthiner, Daniel Keysers, Sylvain Gelly, Olivier Bousquet, and Ilya Tolstikhin. Predict- ing neural network accuracy from weights. arXiv preprint 20Xi::2002.11445, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural informa- tion processing systems, 30, 2017.
Clement Vignac, Igor Krawezuk, Antoine Siraudin, Bohan Wang, Volkan Cevher, and Pascal Frossard. Digress: Diserete denoising diffusion for graph generation. The Eleventh Interna- tional Conference Learning Representations, 2023. URL httpps://cpeneeeier net//
Eorum?id-UaAD Nu8 6WX
Yuxin Wu and Kaiming He. Group normalization. In Proceedings of the European conference on computer vision (ECCV), pp- 3-19, 2018
Saining Xie, Alexander Kirillov, Ross Girshick, and Kaiming He. Exploring randomly wired neural networks for image recognition. In Proceedings of the IEEE/CVF International Conference on Computer Vision, pp. 1284-1293, 2019.
Jiaxuan You, Jure Leskovec, Kaiming He, and Saining Xie. Graph structure of neural networks. In International Conference Machine Learning, pp. 10881-10891. PMLR, 2020.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Russ R Salakhutdinov, and Alexander Smola. Deep sets. Advances in neural information processing systems, 30, 2017.
David W. Zhang, Miltiadis Kofinas, Yan Zhang, Yunlu Chen, Gertjan J. Burghouts, and Cees G. M. Snoek. Neural networks are graphs! graph neural networks for equivariant processing of neural networks. In 2nd Anmual Topology, Algebra, and Geometry in Machine Learning Workshop, 2023.
Muhan Zhang, Shali Jiang, Zhicheng Cui, Roman Garnett, and Yixin Chen. D-vae: A variational autoencoder for directed acyclic graphs. Advances in Neural Information Processing Systems, 32, 2019.
Allan Zhou, Kaien Yang, Kaylee Burs,, Yiding Jiang, Samuel Sokota, J Zico Kolter, and Chelsea Finn. Permutation equivariant neural functionals. Advances in Neural Information Processing Systems, 2023a.
Allan Zhou, Kaien Yang, Yiding Jiang, Kaylee Burns, Winnie Xu, Samuel Sokota, J Zico Kolter, and Chelsea Finn. Neural functional transformers. arXiv preprint 2rxiv:2305. 2023b.
13

---

Published as a conference paper at ICLR 2024
Table 4: Glossary and notation
MLP
Multilayer Perceptron
CNN
Convolutional Neural Network
GNN
Graph Neural Network
GMN
Graph Metanetwork
INR
Implicit Neural Representation
2, 1/, 2,
R
Scalars
8. y, 2, ... E R
Vectors
X, Y. Z,
Raxm
Matrices
X X,Y, 2
The domain of æ; 2,
dmn du day duy E N
The size of the network input or node/edge/global features
1 E N
An index for edges A directed edge from index j to
(i.j) P
A permutation matrix
W
A weight matrix in neural network An element-wiss nonlinearity function
X Rd
The input for neural network
E O + X
The parameters of an NN A function parameterized by neural network
fo (
E
A symmetry of 0, which induces the same f
V
The set of vertices
E
The set of edges
S
Set of parameter-sharing constraints
DAG = (V,
A Directed Acyclic Graph
:: V +V e E Φ E - E
An automorphism An edge
The edge permutation for an automorphism
+ (()) Rd. 0. Rd " E Idd..
The network parameter permutation induced by de
The features for the edge from to j
The features for node A global feature
h (x( ER
Activation at node v for fo on input æ
A RELATED WORK
Metanetworks. Recently, several works developed metanets equivariant to permutation parameter symmetries of simple networks. DWSNets (Navon et al., 2023) and NFN (Zhou et al., (023a) derive the form of linear layers that are group-equivariant permutation parameter symmetries of simple MLPs (and NFN also handles CNNs). NFTs (Zhou et al. 2023b) use special parameter-permuuttion equivariant attention layers and layer-number features to build Transformer-like metanets.
Other types of metanets have been developed, including those based on statistics such as weight mean and standard deviation (Unterthiner al. 2020).. ID convolutional models on the flattened weight vector (Eilertsen et al.) 2020). hierarchical LSTMs and per- parameter MLPs (Metz et al., 2022]. and set functions applied chunked parameter sets (Andreis et al. 2023).
Special metanets have also been developed to process implicit neural representations (INRs) repre- senting shapes, images, or other types of continuous data (Dupont al. [022] Luigi et al. 2023 Bauer et al. (023). However, these approaches require special training procedures for the input neu- ral networks and/or access to forward passes of the input network. In contrast, graph metanets can process datasets of independently trained neural networks of diverse architectures without modify- ing the inputs or requiring forward passes through the inputs, as we demonstrate in our experiments.
The recent work of Zhang et al. (223 also proposes graph neural networks as metanets. They use the computation graph construction MLPs not our general parameter graphs - which allows us to scale efficiently for models with parameter sharing (like CNNs). Moreover, they adopt different graph construction for bias nodes, where our approach is inspired by expressive power considerations from the form of the
linear maps (Navon al.
14

---

Published as a conference paper at ICLR 2024
2023
Zhou al. 2023a). Finally, Zhang et al. (2023) evaluate their method empirically simple networks, but they have limited theoretical guarantees and empirical results on other networks.
NNs Graphs. Many works consider neural networks as graphs, including early works that aimed at developing unified frameworks for neural networks (Bottou & Gallinari 1990] [Gegout et al. 1995]. The graph perspective has been used in several application areas, including: neural architec- ture search (Liu et al. 2019, Xie et al., (019)), analyzing relationships between graph structure and performance (You et al. 2020), and federated learning with differing neural architectures (Litany et al. 2022). Several works process neural networks in some graph representation (Zhang et al,,
2019 Thost Chen 202T Knyazev al.) 2021, Litany et al., 2022). but these works generally operate on a much coarser network representation (e.g. where node can correspond to a whole convolutional layer), whereas our approach we desire each parameter to correspond to an edge.
GNNs and Edge Representations. There has been much work on graph neural networks (GNNs) in recent years, leading to many variants that can be used in our framework (Hamilton) 2020). GNNs learning edge representations instead of just node or whole-graph representations are particularly rel- evant. These include geesaag-passing-liee
architectures, architectures based on graph graph-permutation-
equivariant linear layers, and Transformers for graph data (Battaglia et al. 2018 Maron al. 2019
Kim et al.
2021 2022a
Diao & Loynd 2023 Vignac et al. 2023 Ma ct al. 2023).
B MORE GRAPH CONSTRUCTION DETAILS
We implemented a procedure to automatically build the parameter graph given a PyTorch (Paszke et al. (017) definition of a model that uses supported layers. Our approach converts all layers in the network into their parameter subgraph representation and iteratively stitches the subgraphs together to form the overall parameter graph.
This section details how we build the per-layer subgraphs we partially described in Section 2.1 We also include additional examples in Figure which we explain below.
Node and edge features For all subgraph constructions, we use node and edge features to improve the expressive power of our GMNs. Some examples include layer index, neuron type (input neuron index, output neuron index, or layer type), edge direction (when message-passinng in an undirected manner), and positional encodings (for parameters with associated spatial features, such convo- lutions or the spatial feature grids detailed below). In all cases except for the residual layers, edge features always include the parameter value. All nodes and edges always include integer features for the layer index and neuron type. Importantly, these added features are invariant under neural DAG automorphisms, adding them does not break the equivariance properties of our Graph Metanets.
Linear layers For each linear layer, we use node and edge features corresponding to the layer index and the node/edge type. The latter is an integer value identifying that each node/edge is either a linear weight or bias type.
For convolutions, we include additional edge features to identify each param- eter's spatial position in the filter kernel. We use integer coordinates appended to the other edge features that indicate layer index and edge type. The node features include only the layer index and node type.
Convolution layers
Equivariant linear layers We can also handle general group-equivariant linear layers, of which convolutions are a special case. The high-level idea the same: each layer has node for each "channel", and each node in the input layer is connected with do edges each node in the output layer, where d, is the dimension of the space of equivariant linear maps.
More specifically, let the input space of the layer be R" the output space be RT= and the symmetry group be G. Suppose the group acts via representations P1 (g) E GL(m1) on the input and p2(g) E GL(ng) the output. Then the set of all G-equivariant linear maps, i.c. the set of linear maps T R > R such that T a pi (g) = (2(g) T for all g E G, forms a vector space of dimension
that we denote as dg (Maron al. 2019 Finzi al. 2021
Let B Bi...., Ba, be basis for the space
15

---

Published as a conference paper at ICLR 2024
of equivariant linear maps. Then group equivariant linear layer takes the form
W. Bix,
(4)
for some learnable weights wi E IR. These weights wi are the parameters of the layer.
As is often done in practice, this can be extended to mul- tiple imput or output feature dimensions (also called chan- nels) as follows. We now let the input space be m and the output space be Ro2xd, where di is the num- ber of input channels and dy the number of output chan- nels. The group now acts independently on each chan- nel: via representations #1 (g) = Pi(g) @ la, xd, on the input, and P2(9) = (g) & Id xd, on the output; e.g. for an input matrix X R x the action given by g. X = (g)X. The basis of equivariant linear maps for this space is now has ded, d2 elements, and they take the form B, @ Est, where [dg]. j [d2], and I [d1]. An equivariant linear layer then takes the form
Group- Group-Equivaciane Equivariant Linear Layer (OcepSats)
Figure 6: Example of the graph con- struction for a group-equivariant linear map, for the specific case of the group S. of permutations acting on sets of n elements in Raxd (Zaheer al. (017)- There are two types of equivariant linear maps for this case. The depicted linear layer maps R= + ggx3
de X Σ BXXW11,
(5)
where W; E Rd xd2 are learnable weights. To construct a graph for this layer, let the input layer have d, nodes, and the output layer have d2 nodes. For j E [d2] and I [d1], we connect input node to imput node j with dg edges. The ith edge has as weight wil))
j.l Moreover, the ith edge has a feature denoting that it is part of B, - convolutions this is an encoding of the spatial position of the parameter in the kernel.
As example, consider the permutation equivariant DeepSets linear layer (Zaheer al. 2017) Here, the group G' of permutations on elements acts on Ik" by permuting coordinates. There are two G-equivariant linear maps from R R² I and 11 (so dn 2). Thus, the equivariant linear layer from naxi; + Raxx2 takes the form
X + XW + 11 x2)
(6)
See Figure 6 for an illustration of how we encode this layer, as is used for our experiments in Section
Multi-head attention layers For multi-head attention (Vaswani et al. 2017). we can include the head index as additional node/edge feature for the relevant nodes and edges (though we do not currently do this our experiments of Section 4.] where we only use 2-head attention layers). Otherwise, we treat the edges standard linear parameters.
No parameters are associated with the edges for residual layers. Consequently. the edge features do not include the parameter values (we replace them with constant value of 1). Otherwise, we include the layer index (for the starting layer) and an integer value identifying the edge residual connection.
Residual layers
Normalization layers Similar to bias nodes, normalization layers such as BatchNorm (Ioffe & Szegedy 2015), LayerNorm (Ba et al. 2016). and GroupNorm (Wu & He 2018) can be represented by two additional nodes corresponding to the learned mean and variance. Edges are drawn between these nodes and the neurons to be normalized. The middle Figure deuurtss a normalization layer with 3 neurons. We include the layer index and an integer value for both the node and edge features, identifying them as either the mean/variance and which normalization layer they correspond to. The different types of normalization layers distinguished via node and edge features: for instance, LayerNorm mean parameter nodes are given different node label from NatchNorm mean parameter nodes.
16

---

Published as a conference paper at ICLR 2024
Layers
Normalization Layer
Spatial Grid
(9
C2213
Figure 7: Additional examples of parameter subgraph constructions. We include linear layers in iso- lation for completeness (left), an example of a normalization layer (middle), and a spatial parameter grid (right).
Spatial parameter grids An increasingly common design pattern in Implicit Neural Representa- tions (INRs) introduces grid parameters that are accessed according to input coordinates, such triplanar grids (Chan et al. 2022), octrees (Takikawa et al. 2021) and hash grids (Müller al., 2022). A typical implementation of these parameter grids involves finding a set of features according to the input coordinates and then using linear interpolation to combine the features. The resulting feature vector is then fed into an MLP to decode the features target object.
We show subgraph construction for simple 2D parameter grid in the right of Figure 7 The 2D grid contains only four entries at the corners the domain, with two feature channels. We assume the two feature vectors are bilinearly interpolated based the input coordinates (X, Y). The resulting vector is concatenated with the input coordinates to produce a four-dimensional vector. In practice, can add positional encoding of the grid position the edge features for each element of the spatial parameter grid.
Note that the graph construction for the triplanar grid can largely be reused for dense voxel feature grids, or hash grids (Müller et al.) 2022). The same principles also apply to sparse feature grids, for example, those using octrees (Takikawa al.) (021). A recent work by [ardace al. (2023) develops Transformer-based metanet for processing INRs with triplanar grid features.
B.1 CODE FOR CONSTRUCTING PARAMETER GRAPHS
We plan to release the code for constructing parameter graphs at later date. Here, we outline some basics about how the code works. The code takes as input a PyTorch sequential module, which consists of several modules applied in sequence to some data. We iteratively build graph, starting from the first layer in the input. For each layer in the input (e.g. linear layer), we convert the layer to subgraph, using the constructions outlined in Section 2.1 and Appendix B We connect each of each these subgraphs to graph that we are iteratively building, until we have done this of the layers. We also have way to invert this process, i.e. to go from parameter graph to neural network module. This is useful for metanet applications where we wish to modify the weights of an input neural network. To do this, take architecture specification and the edge weights of a parameter graph, and we then iteratively build a PyTorch module the correct architecture and parameter values.
C THEORY: EQUIVARIANCE AND NEURAL GRAPH AUTOMORPHISMS
C.I DAGS AND COMPUTATION GRAPHS
Consider a DAG (V, E), with weights 0 E for each edge. We assume throughout that this DAG is connected. We may also have parameter-sharing constraints (e.g. e. the case of convolution layers), which is set S that partitions the edge set E, so that if edges (i1,51) and (iz, j2) are in the same equivalence class, then (() = (2) Further choose a nonlinearity a - R. Let there be d,n input nodes in set Vin dout output nodes in Vout, some number of bias nodes in Valas and every other node a set Vnidden- This defines a neural network function fo Rd - Rd with
computation graph given by the DAG. To define the function, for input æ d we first define
17

---

Published as a conference paper at ICLR 2024
the activation of node i V to be real value hx(x)) E R, which is defined recursively as follows:
æi
i Vin
i E Viius
h (z) =
E(iSER (1)))(()) Σ E(DEE O ()))))(())
8 E Vhidden
i E Vout
Note that the activations are well-defined because the graph is a DAG. Here, we slightly abuse notation and assume that nodes ....., dn are the input nodes. Finally, the neural network function given by the values of the activations at the output nodes:
fo( [ (x) .... jamm E
E Rd
(8)
Where Ji. Jdou: E Vou are the output nodes. In the main paper, we generally consider our more efficient parameter graphs instead of computation graphs, since computation graphs may be much larger in general due to parameter sharing, in which one parameter associated with many edges. Nonetheless, study the theory of computation graphs here, they more closely related to the neural network function fo. and they can be defined unambiguously for feedforward neural networks (whereas make some choices in definitions of parameter graphs). We discuss extensions of our theoretical results to parameter networks in Appendix
C.2 NEURAL DAG AUTOMORPHISMS
We define an automorphism (of Neural DAGs) to be bijection & V with the following three properties:
1. Edge preservation: (o(1), +()) E E if and only if (i,j) E E.
2. Node label preservation: imput nodes, output nodes, and bias nodes are fixed points, mean- ing they are mapped to themselves.
3. Weight sharing preservation: if ((a.d)) and (ia.da)) are constrained to be equal, then (((())))( and 0 (0(12))002) are also constrained be equal.
In particular, node label preservation means that if i and j are distinct input, output, bias nodes, then ((i) - and <(j) = j. and we cannot have Φ(i) j.
C.3 NEURAL DAG AUTOMORPHISM GENERALIZE KNOWN PARAMETER S SYMMETRIES
This result shows that every hidden neuron permutation in an MI.P is an automorphism, and also that every automorphism of an MLP DAG takes this form.
Proposition 5. A permutation & V + V of an MLP DAG is Neural DAG Automorphism if and only if o only permutes hidden neurons, and ((i) the same layer for each hidden neuron
Proof. ( (=) Suppose ø is Neural DAG Automorphism. Then because of the node label preser- vation property, we know that & maps each input neuron, output neuron, and bias neuron to itself, meaning that 6 only permutes hidden neurons. Morcover, since is graph isomorphism, it does not change path lengths: the maximum path length between and j is the same as that between ((i) and (()). Hence, preserves layer number of each hidden neuron.
(=) ) Suppose & only permutes hidden neurons and preserves the layer of each hidden neuron. We need only show edge preservation to show that ø is Neural DAG Automorphism. If (i,j) E, layer(ø(ii) 1. As MLPs are fully
then we know that layer(j) = layer(i) - 1, so layer(ø(j)) connected between adjacent layers, this means that (o(i), (())) E. Analogously, if (ø(i), o()) e E. then layer(j) = layer(i) 1, so (i.j) E E. We have thus showed edge preservation and are done.
For this next proposition, will analyze the symmetries of simple CNN parameter spaces. In particular, we will assume the CNN only consists of convolutional layers (and no pooling, or fully connected layers). Every neuron has spatial position and channel index. For instance,
18

---

Published as a conference paper at ICLR 2024
input neurons in 32 32 input with 3 channels indexed by spatial position (z.y)
(1....,32) 32} x {1.....33 32}, and channel с E {1, 2,233,, 3). We say that a permutation ø V V permutes hidden channels if it only permutes hidden neurons, o(1) has the same spatial position as i, and channel(i) = () if and only channel (d(i)) = (hf)n
Here, show that any permutation of hidden channels that preserves neuron layers is Neural DAG Automorphism, though we do not show the converse. Nonetheless, this direction shows that Neural DAG Automorphisms capture the known permutation parameter symmetries of CNNs, and thus our Graph Metanetwork framework can handle these symmetries.
Proposition 6. Let Φ V V be a permutation of a CNN computation graph that permutes hidden channels and preserves neuron layers. Then & is Neural DAG Automorphism.
Proof. Suppose d only permutes hidden channels and preserves the layer of each hidden neuron. Then @ satisfies node label preservation, and also weight sharing preservation. Now, consider an layer(i) 1, and also the spatial positions of and j fit within the
edge (i,j) E E, so layer(j) convolution filter of this layer. Since @ preserves layers, we have layer(j) = layer(i) 1. Further, since & only permutes hidden channels, the spatial position of <ot) is the same i, and likewise the spatial position of Φ()) is the same as j. Thus, (o(1), (()) E E. A similar argument show that ((()) being in E implies that (i, j) E, are done.
C.4 NEURAL DAG AUTOMORPHISMS PRESERVE NETWORK FUNCTIONS
Proposition 7. For any neural DAG automorphism of computation graph, the neural network function is left unchanged: Va E X, fe = (z)
Proof. We will show a stronger statement: namely that the activations of the network with permuted parameters permuted version of the activations of the original network. In other words, will show that
h 0000) (((v)
(x)
for all nodes V. This is sufficient to show the proposition for the following reason. Let the output nodes be , Since a neural DAG automorphism & maps each output node to itself, if equation hoolds, then
foco (x)
-[h (b(0) (x),
(0)
(x)]
(10)
Josns
1h [h)())) (x), (x)..... afjant
(()
(x))
(11)
[h (z),... ...,. he
(x)
(12)
Joons
= fxr)
(13)
The second line follows as ó leaves the output nodes unchanged. The third line follows from equa- tion [ Thus, it suffices show equation 9 We proceed by induction the layer number of u, i.e. the maximum path distance from an input node to v.
For the base case, if is an input node, then Φ(v) v.
h (x) Ey (00) (x)
h (:())
h(m) (x))
(14)
Bias nodes always have activation 1, so this combined with the fact that automorphisms map bias nodes to themselves show equation 9 for bias nodes.
19

---

Published as a conference paper at ICLR 2024
Now, consider a hidden node at layer > O, and suppose that equation 9 holds for all nodes of lower layers. I is hidden node, then we have
h (()) (x) (())
100) ((0) (d()),01)))) aG)
(x)
(15)
(
Σ
0 (1,) 200 (a))
definition of (00)
(16)
(A())
Σ
0 (x)
induction (17)
(
Σ
(x))))))))
edge preservation property of Φ
(18)
(1.))E
h (m:)
(19)
If is an output node, then the same derivation holds without the nonlinearity a. Thus, we have shown equation 9 and we are done.
C.5 FROM COMPUTATION GRAPHS TO PARAMETER GRAPHS
We conjecture that Proposition I can be extended to automorphisms over parameter graphs instead of computation graphs. Here discuss the formal intuition behind this conjecture.
To prove that the network function remains unchanged after an automorphism of the parameter graph, could show that any automorphism of the parameter graph corresponds phism of the computation graph. One way to do this is through a surjective homomorphism from the computation graph to the parameter graph that is locally-injective (Fiala & Kratochvil) 2008 This requires that each I-neighbourhood in the computation graph is mapped injectively to neighbourhood in the parameter graph.
automor-
Intuitively, this allows us to define an automorphism on the parameter graph and lift the automor- phism the computation graph through the preimage of the homomorphism. Concretely, apply the parameter graph automorphism to the computation graph by permuting the preimage sets for each node in the parameter graph under the homomorphism.
Using this machinery, we expect that the proof of Proposition I can be extended to parameter graphs One point of difficulty is that our parameter graph construction allows for multigraphs which are valuable for succinct representations and also make it easier to stitch together subgraphs.
C.6 EQUIVARIANCE OF GRAPH METANETWORKS TO NEURAL DAG AUTOMORPHISMS
Here, we show that graph metanetworks are equivariant to architecture graph symmetries. In partic- ular, consider an architecture graph G (V, E). We also allow modifications the graph, that G does not need to be the exact DAG of the computation graph; for instance, G can be the undirected version of the computation graph. Consider fixed ordering of the nodes, and let X Raxd, be the input node features, and A E Raxaxd, be the input edge features. For any permutation y V V. we define the standard group action: x)( I [ is defined by (x(X) X
and (AA) paxxxd. defined by 7(A)) = A.j.k.
Consider a graph function that updates node features and edge features, say GNN Raxd, x Rexxxd. - Ignd x xxxx For inputs X and A, let GNN(X, A)x E Raxd, denote the updated node features, and GNN(X, A) E Raxxxxd denote the updated edge features. We say such function is permutation equivariant for all permutations 77 : [r]] - [ni], holds that
(NN()((), No. CNN(X,AA)))) (CN((((X) 200.0))) = GNN(X, A)
(20)
(21)
20

---

Published as a conference paper at ICLR 2024
Graph neural networks have exactly this permutation equivariance (Maron et al. 2019). Thus, since any architecture graph symmetry is a permutation [n] [n]. holds that graph metanetworks are equivariant to architecture graph symmetries. Thus, restate result: Proposition 8, Graph metanets are equivariant to parameter permutations induced by neural DAG automorphisms.
However, this would also seem to imply that graph metanetworks are equivariant to more permuta- tions than just architecture graph symmetries, which would be undesirable because this would mean reduced expressivity. The key to fixing this problem is via adding the additional node and edge features described in Appendix B Recall that choose additional node and edge features, say X and Á, that are invariant to neural DAG automorphisms. That is, 6(X)
X and x(A) Á for any neural DAG automorphism ø. Also, we add these features in a fixed manner, so they do not depend the choice of node indices. Thus, the output node features of the Graph Metanet on the original graph GNN((X,, [A, whereas after applying permutation they GNN((((()) ], (n),), A])x, These outputs are in general not equal up permutation. How- ever, when ry is neural DAG automorphism, they are equal, because
GNN(( XI, ()(,), A))
(22)
= -GNN( (x)), (o(4), o (A))))
X Å invariant to
(23)
-G X8), (4, A))
GNN equivariance.
(24)
The same argument holds for output edge features.
D THEORY: EXPRESSIVE POWER OF GRAPH METANETWORKS
Representation vs. Approximation. In this Appendix section, we prove results on the expressive power of Graph Metanets. For these results, there is some continuous target function ftarget that we wish to approximate (either StatNN, NP-NFN, or forward passes through neural network). In our proofs, for ease of exposition we show that if we assume that the MLPs of our Graph Metanets can be arbitrary continuous functions, then there is some Graph Metanet that is exactly ftarget-
However, MIPs cannot generally express every continuous function. Instead, ML.s are known to be universal approximators of continuous functions (Hornik) 199T), meaning that for any continuous fiarget with inputs from a compact domain, for every precision e there an MLP that is within (in some function-space metric) to ftarget- We will call any function class with this property universal function class. If our metanets were just single MLP, then it would be trivial to go from an exact representation statement an approximation statement. However, in general our metanets are compositions of several MLPs and other operations.
appeal to standard techniques that have been used in recent works in the study of expressive power of equivariant neural networks to show approximation results (Lim al.) 2023 Navon al. 2023). In partieular, if neural network architecture can be written composition of functions f ....o o fi, and each of the functions comes from a universal function class, then the composition also forms a universal function class. Thus, one need only show that each component can be approximated to show that an entire target function can be approximated. This is captured in the following lemma, which is proved in [im al. (2023).
Despite this discrepancy,
Lemma 1. Let X C ddo be compact, and let Fi,..., Fr be families of continuous functions, where F. consists of functions from Rd::1 - R.d. Let F be the set of compositions of these functions, so F= Ffff o... o X + d2 F.).
For each i, let F; be family of continuous functions that universally approximates F,, Then the family of compositions F = {fi o...o f1 ft F.) universally approximates F.
D.1 PROOF: SIMULATING EXISTING METANETWORKS
D.1.1 SIMULATING STATNN
In its most expressive form, StatNN (Unterthiner et al. 2020) computes statistics from each weight and bias of ML.P CNN input, and then applies ML. top of these statistics. The
21

---

Published as a conference paper at ICLR 2024
statistics considered by Unterthiner al. (2020) are the mean, variance, and qth percentiles for q E {0, 25, 50, 75, 100). These are all continuous permutation invariant functions of the input, i.c. for any of these statistics s and flattened weight vector w (c.g. w = vec(W(()) for a weight matrix w(l or w = bt4) for a bias vector b(d),, holds that s(w) = s s(Pw) for any permutation matrix P. This is the key to showing that Graph Metanets can express StatNN, since the global graph fea- ture update can compute these permutation invariant functions, and then the final invariant layer can apply an MI. the global graph feature.
Proposition 9. Graph Metanetworks can express StatNN (Unterthiner et al. 2020 on MLP input networks.
B()).... Proof. Let W ....., (L)) be the weights and B(1) ..... B(L) the biases of an input MLP. Let be continuous permutation invariant function with output in R and let MLPSSttNN 2Ld, + R be the MIP of the StatNN, so
,....
StatNN(ww0) ....., W(L) (b()) (y(L),(() B,.),...,,BB4 B(4) ML.
(s(W))).., S(())).
(25)
We will show that a Graph Metanetwork can express this function. Let the edge features ((i.)) E R be three dimensional, with the first dimension holding the parameter value, the second dimension denoting the layer index, and the third dimension denoting whether the edge is associated to a weight or bias parameter. We will not require node features for this proof. Thus, our Graph Metanetwork will take the form
((,)) < (e(.))
(26)
<- MLP
Σe EEE
(27)
and the result will be that u E R will be the desired StatNN output.
Since is a continuous permutation invariant function, it can be expressed in the form s(w)
p{E, (ww)) for continuous functions p Rdd - Red and db IR + Rd (Zaheer et al. 2017). Let
MLP R g2Lxd be given by
)(e())) 1 [[e [,))) = weight] ((LP(((()))) the =(((a))) 1 [|((,))) -
(MLP"(ee..))
(28)
(e(r)) = bias]
(29)
for 1 = 1....,,.. L. Here, [a] is the ith entry of the vector a, [A] denotes the ith entry of the matrix A, and 1 is one if the condition is true and zero otherwise. Note that we write the output of MLP as a matrix to make indexing simpler, but it functionally is just vector.
Then note that the input to MT.P* is EeEr e, which is matrix A R22xds that satisfies:
[A] Σ ob ((vec(W(()))
(30)
[A] L.: Σ (vvec(b())))
(31)
Hence, we have that
p[[A]24:)) = s(W(1))
(32)
p( (j21+1::) a(6(1).
(33)
Thus, we can let MLP" R2xxd - R. be given by
MLP"(A) SL (p((A] ((A]())))....p(())))))))) ((AA]2):)
(34)
and we are done.
D.1.2 SIMULATING NP-NEN
Proposition 10. Graph Metanetworks can express NP-NEN (hou et al. 2023a on MLP input networks.
22

---

Published as a conference paper at ICLR 2024
Proof. GNN Form. Here is the form of our GNN, where for simplicity of the proof, move the global graph feature update before the other updates and remove one of the node-update MLPs. Note that the MLP" is inside the sum - this is important for our construction.
Σ
MLP" (e(s)) u)
(35)
(IJ)EE
0; +-
Σ
MLP (vi, Uj G(.)),
(36)
(i,5) f MLP" MMLP"(U, U, e(i)) u).
(37)
Graph Form. We make the graph undirected by including the reverse edge for each edge in the DAG and labelling these backward edges. We assume the nodes and edges are endowed with certain features at the start (the global graph feature u initialized to 0). Each node has the features: layer number and node type (input neuron number, output neuron number, hidden neuron, or bias neuron number). Each edge has the features: weight value, layer number, weight type, whether is backward or not, and whether it is weight or bias edge. Notably, each of these features is invariant to neural graph automorphisms. The node and edge features belong to the following sets:
{0, ... L} out bias),... biast, hidden]
E
{ini,..., ind out,,
(38)
IR
e(i E (l....,L) D} (forward, backward] (weight, bias}
(39)
Without biases. For illustration, we will first prove this result for the case when we are only pro- cessing weights and biases. First, recall the form of the NP-NFN linear layer (Zhou al., 2023a). It takes MLP weights W E W = Rdd2xd=- x x... x R4 xdo and outputs new weight representations H(W) W. We let W denote the weight matrix of the Ith layer, and let w:() be the (i, j)th entry of this matrix. Then, the NP-NEN linear layer takes the form:
Σ af d''w!!?
H(W
+ ab w + aas 11-1 -
(40)
( in + +5 4. ol w:
(41)
where we define (0) = 0 and W(L+1) = 0 for the boundary cases of | = and = L. Here, the R learnable scalars, and the denotes a summation over all relevant indices. For instance, for W()) E R we have ('D)) +,k E (i(l) and W = Σn EØ w:()
First, we use the global graph feature to compute the first summands. Let MLP" have an output space of R-, and let
MLP" w))' ()) = [[e] 1[[e]2 =],
(42)
where [y]' R is the ith entry of the vector y. and 1 is if the condition inside true, and zero otherwise. Since the first column of e holds the weight and the second column holds the layer number, the global graph feature is updated
E TR.
(43)
Next, consider the node feature updates. We define MLP" have output space of IR? such that
[MPP Uj, ((.)), 4)]² = [e [[( = backward]
(44)
(ML""(r,, e(.)) 2)}² [enn]l [[(e(s))) = forward].
(45)
23

---

Published as a conference paper at ICLR 2024
For node v that is the kth node in layer I (where the input nodes at layer 0), the node update is then given by:
w +,6
w R.*
R
(46)
Finally, we define the edge update. We let MLP have an output space of I, such that for any forward edge (1,j) the update is (writing = (e 1 for brevity):
MT.P (os; Mj ((.s)) u)
(47)
Σ [4] + (d [0y) + az a + do. + a5 [m.] + (
(48)
By plugging in the above forms of u, Us and vy (and noting that [e(r,))] new edge representation is exactly the NP-NEN update, so we are done.
W,3 see that this
With biases. Recall the form of the NP-NEN linear layer where the input MLP also has biases:
H(W b) = (W,
(49)
L
W9
Σ
ai
w(()
+ az
df' + ag 1-- (
(50)
jtt
a4
Wil)
+
L (httti + w.
(51)
+
a
+ dg abb!
(52)
810
=
+ c2 w''
+ + 1+ (
(53)
+
Σ dinbbe))
+5bb
(54)
Note that define ((0) 0, and once again define (0)) 0 and (([-1) 0 for the boundary cases.
We define the global model MLP" to have output in R such that for l E {I....,LI.
MLPP" u))² [e] 1 [[e] = 1, [e] weight]
(55)
(e, [e]l 1 [|e]2 L, le] bias
(56)
The global graph feature then takes the form
2 (
W(L)) 60)
050]
E R2
(57)
Next, we consider the node feature updates. We define MLP" to have an output space of IR3 such that
MLP" Dj ((d)): w)]² = [ew,)) 12 [[ew,))] 13 = backward, (e = weight]
(58)
MMPP" Uy ((d)) u²)² = (ea.)] 12 1 [ee,,))3 forward, (e(s)) = weight]
(59)
[MLP" (v; Bj G(i)); u))3 [e(aa) [[e(,,))3 forward, [e(r)] bias]
(60)
For a non-bias node v that is the ktth node in layer the node update is then given by:
(I+1)
w(())
(00) E R
(61)
k
24

---

Published as a conference paper at ICLR 2024
Finally, define the edge update. We let MLP have output space of R, such that for any forward edge (i, the update is (writing = [e(i,)) for brevity):
MLP Uy, C(t.j), u)
(62)
L
Il [e(h,a)] = weight
Σ
a [u]* +
+ dz - 10,}²
(63)
1l aj
[o:]
( d'l dele(),)))
(64)
L
+
Σ
[ rb/0;13 + obم
(65)
+1 [[( bias]
Σ
1
L
[u]
+ 2 [o,}] + c3 (1,}
(66)
cf
[2]=t + dfle
(67)
By plugging in the above forms of u, Uú. and Vy. we see that this new edge representation is exactly the NP-NEN update, so we are done.
D.1.3 EXPRESSIVE POWER WITH BIAS NODES
Here, we explain how encoding biases in linear layers as bias nodes can aid expressive power (in the simple case of MLPs). Intuitively, this is because the bias node allows for different bias parameters of the same layer to communicate with each other in each layer of a message passing graph metanet, whereas this not as simple and requires at least two layers if biases are encoded as self-loops or node features.
For a specific example, when the graph metanet does not use global feature (as in our experiments), then the bias node allow message passing between biases within the same layer. In the NP-NFN case covered in Appendix [.I.2] we see that the global feature allows message passing between biases of any layer. When there is no global feature, letting 0yar denote the bias node of the Ith layer, then our graph metanet can nonetheless update
(68)
((()) < CU80
cbl
(69)
for some scalar c E R. This is one of the equivariant linear maps included NP-NEN. On the other hand, graph metanet with no global feature cannot compute this in one layer if the biases encoded as self-loops or node features this would require at least two layers.
D.2 PROOF: SIMULATING FORWARD PASSES
Proposition 11. On computation graph inputs, graph metanets can express the forward pass of any input feedfooward neural network as defined in Section 2.2
Proof. We will prove this statement by showing that a graph metanet can compute the activations (x) for an input network fe and input æ. In particular, we will show that 1 layers of graph metanet can compute the activation 1g(()) (a) for any node e with layer number at most I, where the layer number is defined as the maximum path distance from any input node to We do this by induction on For the base case, (x) = x. for any input node v, so at depth zero of the graph metanet the induction hypothesis is satisfied. Also, the activation for bias node is always which graph metanet can compute using a node-update MLP that outputs for any node with the bias-node node label.
25

---

Published as a conference paper at ICLR 2024
For hidden node i with layer number larly):
O, recall that we have (output nodes are handled simi-
(x)
Σ
(70)
(Hi)EE
Let u!- (1-1) denote the node feature computed for node at depth 1 - 1 of the graph metanet. By induction, we have v ((-1) 12,(0 for any node j in layer I or earlier. Recall the form of the node update in graph metanet layer:
8, + MLP
Ui
Σ
(u, Un ((.)),) u)
(71)
(IJ)E
Let MLP1 for depth of the GNN be given by
(LP vy ((.)). u) le
(72)
and let MLP? for depth be given by
(L, o(a).
(73)
Then the node update equation for the GNN is
0,
(
Σ
(74)
(LJ))R
Σ
(as)
(75)
(IJEE
h (x))
(76)
The same derivation holds for an output node i, replacing a with the identity map. So, we are done
by induction.
E ADDITIONAL EXPERIMENTAL DETAILS
E. PREDICTING ACCURACY
E.1.1 CLASSIFIER TRAINING DETAILS
For the predicting accuracy experiments, we train two datasets of about 30 000 CIFAR-10 im- age classification neural networks each. We use the Fast Forward Computer Vision (FFCV) li- brary (Leclere al. 2023) for fast training and build off their template training setup for CIFAR-IO. We use random horizontal flips, random translations, and Cutout (Deesies & Taylor 2017) as data augmentation. The learning rate schedule has a linear warmup and then a linear decay zero. All the other hyperparameters are given in Table 5 we sample many of the hyperparameters for each neural network. For the Diverse Architectures dataset, train about 6000 of each of the types of networks, while for the Varying CNNs dataset we train about 30 000 of the 2D-CNNs.
1.1 METANETWORK TRAINING DETAILS
Data splits. For the 50% training data experiment, we take random split of 15 000 training networks, and for the 5% training data experiment we take random split of 1500 training networks. Then 2000 random networks are selected for validation, and the rest are for testing.
For the OOD experiments, we only train and validate networks of the lowest hidden dimension of each architecture. This means that we only train and validate on CNNs (both 2D and 1D) of hidden dimension 24, DeepSets of hidden dimension 32, ResNets of hidden dimension 16, and ViTs of hidden dimension 32. 2000 of these shallower networks are chosen for validation, and the rest for training. Then test on all networks with higher hidden dimension.
26

---

Published as a conference paper at ICLR 2024
Table 5: Hyperparameters for the CIFAR-10 image classifiers that trained for the predicting
accuracy experiments in Section
Model
Hyperparameter
Distribution Choice
All
Learning rate
Uniform (1,3)
Weight decay Label smoothing Optimizer
10 Uniform((.5)
Uniform .2) Uniform((SGD, Adam, RMSprop))
CNN (2D)
Hidden dimension
Uniform (24, 28, 32)
# Convolution layers
Uniformu {1, 2,
# Linear layers
Uniform ( {1,
Normalization type
Uniform (artormm(BachNoomae
Norm)
Dropout
Uniform .25)
Kernel size
CNN (ID)
Hidden dimension
Uniform 28, 32})
# Convolution layers
Uniform {1, 2,
# Linear layers
Uniform [1, 2
Normalization type
Uniform (Batem(nBBannn GroupNorm]
Dropout
Uniform (0, 25)
Kernel size
DeepSets
Hidden dimension
Uniform ( (32, 64))
# Equivariant layers # Linear layers Normalization type
Uniform [2, 3, 4))
Uniformu 21,2)
(aatouu(BaahNoran Group.oron
Dropout
Uniform 20,)
ResNet
Hidden dimension
Uniform/ (16, 32))
# Blocks
Uniform/ (2,3,4)
# Linear layers
Normalization type
BatchNorm
Dropout
VIT
Hidden dimension
Uniform( (302 (32, 48)) Uniform (2, 3})
# Transformer blocks
# Linear layers
Normalization type
LayerNorm Uniform (0, 25)
Dropout
27

---

Published as a conference paper at ICLR 2024
Table 6: Additional results for predicting accuracy the dataset of small CNNs from Unterthiner et al. 2020). We see that GMNs outperform all metanetworks. Here, we use the Graph Transformer GRIT Ma et al. 2023) as our graph learning model for GMN, we find it performs better than the message passing models that we tried.
Metanetwork
Test T
NEN-HNP
.944 .001
NEN-NP
.922 .001
StatNN
NFT
.926 99261000
GMN (ours)
9938=0000
DocaSets, Vol Netwo ks
DMC, Vol Netwo
GNN, Val Networds
Predicted Accuracy
Prodicted Accuracy
Predicc Prediccedd Accuracy
bo
Actual Accuracy
Actual Accuracy
Actual Accuracy
DeepSets, Networks
Tast
GNN,
La bn 4 bs
Predicted Accuracy
Prodicted Accuraty
Predicced Accuraty ha La
3
6 Notuel Accursty
7
4
7
2 11
5
.8
Actuall Accuracy
Actual Accuracy
Figure 8: Scatterplots of actual test accuracies (x-axis) and crineettwor--reddiice test accuracies (y-axis) for image classifiers from the Diverse Architectures dataset with the OOD train/test split. Top row shows validation split input networks, and bottom row shows test split input networks.
Metanetworks. For each metanet, we first choose hyperparameters such that the total num- ber of trainable parameters is around 750 000. Then we search learning rates 00005, .0,0005,.00 and choose the learning rate that achieves the best validation R². Finally, we train the model using this best learning rate with random seeds and report the mean and standard deviation of the test set R² and Kendall 7 values.
We train all metanets with the Adam optimizer (Kingma Ba 2014). The training loss is binary cross entropy loss between the predicted and true accuracy; each metanet has sigmoid nonlinearity the end to ensure that its input within [0, 1].
Small CNN Experiments To compare against previously proposed metanetworks, we also run experiments on predicting accuracy of the small CNNs trained by Unterthiner
(020). following the experimental setup of Zhou et al. (023a). Results are given in Table E. We see that our GMNs outperform all other metanetworks. We use GRIT (Ma et al., (023) as our graph learning model, which is a type of Graph Transformer. Our GRIT model has 3 652 865 parameters, uses hidden dimension of 256, has 4 layers, is trained with a learning rate of 3e-4, and uses heads in its attention modules.
E.2 EDITING 2D INRS DETAILS
We closely follow the experimental setup of Zhou et al. (202aa) for the experiments in Section 4.2 As in their work, train the metanets for 500 iterations with batch size of 32, using the Adam optimizer with .001 learning rate. We follow their setup to parameterize the updated weights 0 as
28

---

Published as a conference paper at ICLR 2024
0 = + Meeanne(O),, where 0 are the parameters of the INR representing the unedited image, and y are learned scalars for each parameter that are initialized to .01. Our message passing GNN metanetworks have 8 945 671 parameters, hidden dimension of 512, and 4 message passing layers.
E.3 SELE-SUPERVISED LEARNING WITH INRS DETAILS
We closely follow the experimental setup of Navon et al. (023) for the self-supervised learning task. In particular, as they do in their experiments, choose hyperparameters for our graph metanet such that the number of trainable metanetwork parameters is about 100 000. Thus, we use graph metanetwork with GNN layers and hidden dimension of 76.
We follow the same training procedure as Navon et al. 2023), which uses a contrastive objective with data augmentations given by adding Gaussian noise to the weights and masking weights to zero. As in Navon et al. (2023), we evaluate metanet on three random seeds, and report the mean and standard deviation across these three runs.
29


Reviews:
[{"review_id": "q1oKs2yegE", "reviewer": "Reviewer_ZWAE", "summary": "The idea is to construct directed acyclic graphs (DAGs) that are also ‘parameter graphs’ (which represent parameters as weighted edges) that represent neural networks and feed them through a simple message-passing graph neural network ‘metanetwork’. According to them, they are distinguished from other works that do this by a few design choices of the graph construction. They show that the way they construct neural network DAGs is invariant to the order in which neurons on the same layer (ones between which order should not matter) are embedded, which is stated to be an issue with the way some other metanets represent networks. They state how they represent different kinds of layers in their ‘parameter graphs’, and show that it can represent layers others cannot (such as normalization layers, according to them) and have less scaling issues with parameter-sharing layers such as convolution and attention layers. They evaluate by attempting to learn the prediction accuracy of datasets of image classifier networks on CIFAR-10, one is of 2d CNNS, and one is of varying models (CNNs, ViTs, ResNets, etc) with competing methods (those which can represent the inputs).", "questions": "Please see weakness.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "It should be noted that I am not familiar with other papers regarding metanets and am basing my assessment largely on information from this paper. That being said, given that they represent the current state of this area fairly, this seems to be an impressive paper. They appear to address scaling issues with similarly expressive network representations and expand the variety of representable networks, which seem to be great contributions.", "weaknesses": "They address most of my concerns I had while reading, including some tests to compare against the newer state of the art metanets mentioned that were excluded from most of the result due to not being able to represent certain types of layers. Notably a cited competing method ‘NFN’ was left out from this, though seems to be addressed in the appendix, and as it apparently deals exclusively with MLPs I believe it is fair not to compare with the proposed method.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "IuoTvq8GNl", "reviewer": "Reviewer_ffRA", "summary": "This paper introduces a novel approach called Graph Metanetworks (GMNs) that utilizes graph neural networks to process diverse neural architectures. The authors address the challenge of generalising metanetwork architectures to different types of networks by building graphs representing the input neural networks and processing them using graph neural networks. The GMNs are proven to be expressive and equivariant to parameter permutation symmetries, and they demonstrate superior performance on various metanetwork tasks across different neural network architectures.", "questions": "- Since graph can encode a variety of architectures, I wonder if it is possible to generalise and especially extrapolate prediction to unseen architecture, e.g. the test MLP/CNN is wider or deeper than all training examples. \n\n- I do not see pooling layers discusses. Are there pooling layers involved in the networks in the dataset?\n\n- How to encode atrous/dilated convolutions?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "The proposed GMNs can generalise to different types of neural architectures. Unlike previous works e.g. Navon et al. (2023) that were tailored to specific networks, the GMNs can handle a wide range of architectures, including those with complex modules such as attention blocks. Upon the main idea of processing weights with graph networks in Zhang et al. (2023), this paper extends the method to a variety of neural layers that are common in modern neural architectures.\n\nThe authors also provide a proof of the expressiveness and equivariance of GMNs to parameter permutation symmetries, which is an important property for metanetworks.", "weaknesses": "- Although different neural layers and architectures can be encoded as the proposed parameter graph representation, however, the information in the spatial domain is missing (e.g. translation equivariance and receptive field in ConvNets). This could be an inherent and general limitation of the proposed method as well as other related work on weight domain.\n\n- There are also a variety of non-parametric operations in feed-forward neural networks that are not addressed by the paper. E.g. pooling layers (especially max pooling), atrous convolutions, padding and so on. I can expect these operations can be encoded as some additional indicating features but I wonder if there are better and less artificial solutions.\n\n- Comparison with Navon et al. (2023) only on the 1D sine curve toy dataset, but not on the more challenging 2D INR image tasks in the original paper. Other experiments take DeepSets and DMC as baselines which are 2017 and 2020 papers.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Y1qkhkiTxM", "reviewer": "Reviewer_fJUi", "summary": "The paper proposed a new approach to represent neural networks and their weights using a so called parameter graph in which nodes represent neurons and edges represent weights/parameters of the neural network. The paper demonstrates how to build this graph for different kinds of architectures including non-trivial and very practical cases such as transformers. The paper formally shows useful properties of the approach such as permutation equivariance corresponding to neuron equivariance. Finally, the results in two experiments show better performance than the baselines.", "questions": "1. Given that the same graph metanet can process diverse architectures, can the model trained in 5.1 be directly applied to large realistic models like ResNet-50 or large ViTs? For example, in the paper \"Zero-Cost Proxies for Lightweight NAS\" there are PyTorchCV networks with recorded accuracies. It would be very interesting to see how the results would look like for such challenging cases. Moreover, in Section 5.1 it would be interesting to see correlation performance of the methods from \"Zero-Cost Proxies for Lightweight NAS\" like gradnorm, which are very easy to compute. Another interesting experiment could be to track the predicted accuracy during training some network and see if it correlates well with the actual accuracy. Given that evaluation of large networks is very expensive, this approach could be an alternative way to track performance.\n\n2. Is it correct that the trained metanets cannot generalize to some architectures, for example to the kernels of larger size than seen during training because the edges would have more features? If yes, this should be clearly described in Limitations or another appropriate section.", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper has several good contributions:\n\n1. The paper addresses an interesting and promising topic of representing neural networks and their weights.\n2. The proposed parameter graph is a reasonable approach and looks much simpler than previous works such as DWSNets and NFN/NFT. It is similar to the concurrent work of Zhang et al. (2023), which is properly credited. \n3. Description of how to build graphs for different layer types such as convolution, self-attention and residual connections is very informative.\n4. The experiments show reasonably good results of the proposed approach.\n\nThe paper is also well written and organized.", "weaknesses": "The paper have several weaknesses. I'm willing to revise the rating based on authors' response.\n\n1. The paper says that \"While our graphs are DAGs, we are free to use undirected edges\". Would not the direction of edges be a useful feature in some cases? For example, sometimes networks take multiple inputs and have multiple outputs so there is no way to differentiate input vs output unless edge direction is used.\n\n2. The computational complexity vs other approaches is not analyzed. Can the model encode large models in a feasible way? \n\n3. There are very few experiments. The authors could have more experiments following previous works from DWSNets, NFN, etc. or be creative in designing more novel experiments. For example, generation of new networks mentioned in the intro (Erkoc¸ et al., 2023) could be a very appealing experiment.\n\n4. The experiments lack ablations to understand how the model behaves under different settings. These could be the number of layers/params in the graph metanet, different GNN architectures, different approaches to treat the bias term (e.g. comparing to Zhang et al. (2023)), ablating different components of the GNN in 2.3, etc.\n\n5. Source code is not available in the submission, which would be very helpful at least for how to build a graph given a neural network. Do the authors intend to open source the code?\n\n6. The purpose of Section 3 and Proposition 2 is a bit unclear, these could be replaced by more experiments or computational complexity analysis that are lacking.\n\n7. Section 5.2. lacks details. How the vector representation is obtained? Details of training the GNN are missing. More difficult INR tasks could be added.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (spotlight)