Graph Neural Networks for Learning Equivariant Representations of Neural Networks
Authors: Miltiadis Kofinas
Boris Knyazev
Yan Zhang
Yunlu Chen
Gertjan J. Burghouts
Efstratios Gavves
Cees G. M. Snoek
David W. Zhang
Venue: ICLR 2024 oral (2024)
Keywords: Deep weight space
Graph neural networks
Transformers
Permutation equivariance
Implicit neural representations
Networks for networks
Neural graphs
Primary area: learning on graphs and other geometries & topologies
TL;DR: We propose graph neural networks that learn permutation equivariant representations of other neural networks

Abstract: Neural networks that process the parameters of other neural networks find applications in domains as diverse as classifying implicit neural representations, generating neural network weights, and predicting generalization errors. However, existing approaches either overlook the inherent permutation symmetry in the neural network or rely on intricate weight-sharing patterns to achieve equivariance, while ignoring the impact of the network architecture itself. In this work, we propose to represent neural networks as computational graphs of parameters, which allows us to harness powerful graph neural networks and transformers that preserve permutation symmetry. Consequently, our approach enables a single model to encode neural computational graphs with diverse architectures. We showcase the effectiveness of our method on a wide range of tasks, including classification and editing of implicit neural representations, predicting generalization performance, and learning to optimize, while consistently outperforming state-of-the-art methods. The source code is open-sourced at https://github.com/mkofinas/neural-graphs.

Published as a conference paper at ICLR 2024
GRAPH NEURAL NETWORKS FOR LEARNING EQUIV- ARIANT REPRESENTATIONS OF NEURAL NETWORKS
Miltiadis Kofinas** Boris Knyazev²
Yan Zhang²
Yunlu Chen
Gertjan J. Burghouts* Efstratios Gavves'
Cees G. M. Snoek
David W. Zhang"
Universitty of Amsterdam
Samsung SAIT Al Lab, Montreal
Cameeii Mellon University 4TNO
ABSTRACT
Neural networks that process the parameters of other neural networks find applica- tions in domains as diverse as classifying implicit neural representations, generating neural network weights, and predicting generalization errors. However, existing ap- proaches either overlook the inherent permutation symmetry in the neural network rely on intricate weight-sharing patterns achieve equivariance, while ignoring the impact of the network architecture itself. In this work, we propose to represent neural networks as computational graphs of parameters, which allows us to har- ness powerful graph neural networks and transformers that preserve permutation symmetry. Consequently, our approach enables single model to learn from neural graphs with diverse architectures. We showcase the effectiveness of our method on a wide range of tasks, including classification and editing of implicit neural representations, predicting generalization performance, and learning to optimize, while consistently outperforming statc-of-the-as methods. The source code is open-sourced at hitps: /github. cwm/ckoc///ooooo finas/ /aeas//eeue /nuual graphs.
INTRODUCTION
How design neural networks that themselves take neural network parameters input? This would allow us to make inferences about neural networks, such as predicting their generalization error (Unterthiner et al., 2020), generating neural network weights (Schürholt al., 2022a), and classifying or generating implicit neural representations (Dupont et al., 2022) without having to evaluate them on many different inputs. For simplicity, let us consider deep neural network with multiple hidden layers. As a naive approach, we can simply concatenate all flattened weights and biases into one large feature vector, from which we can then make predictions as usual. However, this overlooks an important structure in the parameters: neurons in layer can be reordered while maintaining exactly the same function (Hecht-Nielsen, 1990). Reordering neurons of neural network means permuting the preceding and following weight matrices accordingly. Ignoring the permutation symmetry will typically cause this model to make different predictions for different orderings of the neurons in the input neural network, even though they represent exactly the same function.
In general, accounting for symmetrics in the input data improves the learning efficiency and underpins the field of geometric deep learning (Bronstein et al., 2021). Recent studies (Navon et al., 2023; Zhou et al., 2023a) confirm the effectiveness of equivariant layers for parameter spaces (the space of neural network parameters) with specially designed weight-sharing patterns. These weight-sharing patterns, however, require manual adaptation to each new architectural design. Importantly, a single model can only process neural network parameters for a single fixed architecture. Motivated by these observations, we take an alternative approach to address the permutation symmetry in neural networks: we present the neural graph representation that connects neural network parameters similar to the computation graph (see Figure 1). By explicitly integrating the graph structure in our neural network, a single model can process heterogeneous architectures, i.e. architectures with varying computational graphs, including architectures with different number of layers, number of hidden dimensions, non-linearities, and different network connectivities such as residual connections.
"Joint first and last authors. Correspondence to: mkoof nas@ggaa 1, noom, davi han.dd0nzngagggmaaaa 1.com

---

Published as a conference paper at ICLR 2024
layer (-
layen
Figure 1: Representing a neural network and its weights as neural graph. We as- sign neural network parameters to graph features by treating biases b; corre- sponding node features Vi, and weights W ij as edge features Eiy connecting the nodes in adjacent layers.
Neural netwerk feodforwand activation in layer !)
b Σ W x
Neural network as neural graph:
feature: V. b.
feature: E0 W
Wi
We make the following contributions. First, we propose a simple and efficient representation of neural networks as neural graphs, which ensures invariance to neuron symmetries. The perspective of permuting neurons rather than parameters makes our model conceptually much simpler than prior work. We detail various choices on how the neural graph can encode a neural network, including the novel concept of "probe features" that represent the neuron activations of forward pass. Second, adapt existing graph neural networks and transformers to take neural graphs as input, and incorporate inductive biases from neural graphs. In the context of geometric deep learning, neural graphs constitute a new benchmark for graph neural networks. Finally, we empirically validate our proposed method on wide range of tasks, and outperform state-of ther approaches by large margin.
2 NEURAL NETWORKS AS NEURAL GRAPHS
Two aspects determine a neural network's s behavior: the parameters, and the architecture that defines how the parameters are used. Prior work has only modeled one of these, namely the parameters, while assuming fixed architecture. This limits the existing models changes (regardless of whether they affect the parameters or not, such as through skip- kiip-conectioons) cannot be accounted for. These changes can also alter the network's symmetry group, which requires a complete redesign for every new architecture. Our approach therefore considers both the parameters as well as the architecture that specifies the computational instructions of the neural network.
single input architecture; any
Our method is straightforward: we represent the input neural network as a graph, with its nodes corresponding to individual neurons in the neural network and its edges corresponding to connections between neurons. The weights of the input neural network determine the edge features and the biases determine the node features. We refer to this as the neural graph. Importantly, the natural symmetries in these graphs correspond exactly to the neuron permutation symmetries in neural networks: when we permute the nodes of the neural graph, the adjacency matrix is permuted in a way such that the connections between the same neurons remain the same. Therefore, different neuron permutations of the input neural network which correspond to the same function result in the same graph as well, which precisely what we desire'. Another benefit of the neural graph representation that there is an extensive body of prior research on graph neural networks that we can make use of to obtain powerful models. Graph neural networks and transformers naturally exhibit equivariance to the permutation symmetries of graphs, and therefore our neural graphs that represent neural networks.
In the following section, we explain how we convert Multi- Perceptrons (MLPs) into neural graphs. Afterwards, we show how to extend the conversion to Convolutional Neural Networks (CNNs). Finally, we show how our method can process heterogeneous architectures.
2.1 MLPS AS GRAPHS
We first outline the procedure for constructing a neural graph G = (V, with node features V E _axdv and edge features E E d Here, n denotes the total number of nodes in the graph, and dv; dg denote the number of node and edge features, respectively. Consider an MLP with L fully connected layers. The weight matrices for this MLP are (W(1) W }. and the biases {b(1) ...., b(() Each weight matrix ( has dimensions d, d and each bias b(0) has dimensions d,. The total number of nodes is then given by n El= d,, where dy is the dimension of the input. We define the edge and node feature matrices containing the weights and biases as
Neurons in the input and output layers are typically an exception to this, they are not freely exchangeable without changing the underlying function. Section 3 addresses this through positional embeddings.
2

---

Published as a conference paper at ICLR 2024
follows:
de
4
de
W
0
W
b
b(2)
0
W
b(4)
where W(IT denotes the transposed weight matrix. The edge features form sparse block matrix where the sparsity pattern depends on the number of nodes per layer: the first diagonal block has size dg, the second dj, and so on. We can verify that putting W as the first off-diagonal blocks means that they have the expected size of d x dj-1. The first do node features in V are set to 0 to reflect the fact that there are no biases at the input nodes. We can show that our neural graph has one-to-one correspondence to the neural network's computation graph. This ensures that distinct MLPs get mapped to distinct neural graphs.
Similar to the example above, many applications of neural networks in parameter space only include the neural network parameters as input. For such applications, an MLP has scalar weights w ij E RR
and biases bl¹) R comprising the elements of E, V. resulting in -ie-dimesioonal features, i.e. dv dF 1. Depending on the task at hand, however, we have the flexibility to incorporate additional edge and node features. We explore some examples of this Section 2.4.
2.2 CNNS AS GRAPHS
So far, we have only described how to encode basic MLPs as graphs. We now address how to gener- alize the graph representation to alternative network architectures, namely convolutional networks. To make the exposition of our method clear, we will use the following interpretation of convolutional layers and CNNs. Convolutional layers take as input a multi-channel input image (e.g. an RGB image has C' = 3 channels) or multi-channel feature map, and process it with a filter bank, i.e. collection of convolutional kernels often termed filters. Each filter results single-channel feature map: applying all filters in the filter bank results in collection of feature maps, which we concatenate together in a multi-channel feature map. CNNs, in their simplest form, are a stack of convolutional layers mixed with non-linearities in between.
Permutation symmetries in a CNN work similarly to an MIP; permuting the filters layer while simultancously permuting the channels of cach filter in the subsequent layer effectively cancels out the permutations, shown visually in Figure 5 in Appendix C.2. Under the aforementioned interpretation, single-channel slices of multi-channel feature map (or single-channel slices of the input) correspond to nodes in our neural graph. Each node is connected via edges incoming from a particular convolutional kernel from the filter bank. By treating each channel node, CNN neural graph respects the permutation symmetries of the CNN.
We now describe the neural graph representation for each component in CNN (with more details in Appendix C.2). As a working example, let us consider a CNN with L convolu- tional layers. It consists of filters { [w()} and biases {b()) } for layers E (1...., L), where W E xxu (()) E Rd4, and wi, hy denote the width and the height of kernels at layer I.
Convolutional layers. Convolutional layers are the core operation in convolutional network. Since channels in CNN correspond to nodes in the neural graph, we can treat the biases the same way as in an MLP, namely as node features - see Equation (1). The kernels, however, cannot be treated identically due to their spatial dimensions, which do not exist for MLP weights. To resolve this, we represent the kernels by flattening their spatial dimensions to a vector. To ensure spatial self-consistency across kernels of different sizes, first vero-pad all kernels s = (Wmax: m and then flatten them. This operation allows for a unified representation across different kernel sizes; can process all kernels with the same network. The maximum kernel size is
maximum size
3

---

Published as a conference paper at ICLR 2024
chosen hyperparameter per experiment; this operation is visualized in Figure in Appendix C.2. After this operation, the kernels can be treated as a multi-dimensional equivalent of linear layer weights. We construct the edge features matrix similarly to Equation (1); the only difference is that this matrix no longer has scalar features. Instead, we have E E TRaxxxde with de = Wmax Amax-
Flattening layer. CNNS are often tasked with predicting single feature vector per image. In such cases, the feature maps have to be converted to a single feature vector. Modern CNN architectures (He al., 2016) perform adaptive pooling after the last convolutional layer. which pools the whole feature map in a single feature vector, while traditional CNNs (Simonyan & Zisserman, 2015) achieved that by flattening the spatial dimensions of the feature maps. The downside of the latter approach that CNNs are bound to specific input resolution and cannot process arbitrary images. Our neural graph not bound to any spatial resolution, and as such, its construction does not require any modifications to integrate adaptive pooling. While the CNNs in all our experiments employ adaptive pooling, we also propose two mechanisms to address traditional flattening, which we diseuss in Appendix C.2.
Linear layers. Linear layers are often applied after flattening adaptive pooling to produce the final feature vector representation for an image. The most straightforward way to treat linear layers in CNN is in the exact same fashion as in an MLP. The downside of this approach is that linear layers and convolutional layers have separate representations, as their dg will typically differ. An alternative is to treat linear layers 1 x 1 convolutions, which allows for unified representation between linear and convolutional layers. When treated as convolutions, the linear layers are padded to the maximum kernel size and flattened. In our experiments, explore both options, and choose the most suitable via hyperparameter search.
2.3 MODELLING HETEROGENEOUS ARCHITECTURES
One of the primary benefits of the neural graph representation is that it becomes straightforward to represent varying network architectures that can all be processed by the same graph neural network. Notably, we do not require any changes to accommodate a varying number of layers, number of dimensions per layer,
completely different architectures and connectivities between layers. When dealing with a single architecture, we can opt to ignore certain architectural components that shared across instances, as our method and related methods can learn to account for them during training. These include activation functions and residual connections. Thus, we will now describe how can incorporate varying non-linearities and residual connections in heterogeneous architectures of CNNs or MLPs.
Non-linearities. Non-linearities are functions applied elementwise to each neuron, and can thus be encoded as node features. We create embeddings for a list of common activation functions and add them to the node features.
Residual connections. Residual connections are an integral component of modern CNN architec- tures. A residual connection directly connects the input of layer to its output as y f(x) + x. Incorporating residual connections in our neural graph architecture is straightforward, since we can include edges from each sender node x to the respective receiving node in y. Since residual connections can be rewritten y f(x) + Ix, where I is the identity matrix, the edge features have a value of for each neuron connected.
Transformers. The feedforward component of transformers (Vaswani al., 2017) can be treated as an MLP and its neural graph follows a similar schema as previously outlined. We explain the conversion rules for normalization layers in Appendix C.3 and the conversion rules for multi-head self-attention in Appendix C.4. After individually converting the parts, we can compose them into a neural graph for transformers.
2.4 NODE AND EDGE REPRESENTATION
Our neural graph representation gives us the flexibility to choose what kinds of data serve node and edge features. Though we mainly focus on weights and biases, there are also other options that can use, for example we use the gradients in learning optimize setting.
Edge direction. Our basic encoding only considers the forward pass computations of the neural network, yielding directed acyclic graph. To facilitate information flow from later layers back
4

---

Published as a conference paper at ICLR 2024
to earlier layers we add reversed edges to the neural graph. Specifically, we include ET as additional edge features. Similarly, we can also include E + E' as extra features representing undirected features.
Probe features. Humans tend to interpret complicated functions by probing the function with a few input samples and inspecting the resulting output. We give the graph neural network similar ability by adding extra features to every node that correspond to specific inputs. In particular, we learn a set of sample input values that we pass through the input neural network and retain the values for all the intermediate activations and the output. For example, consider the simple input neural network f(x) = ( (W()) + b(2) + b(2) for which we acquire an extra node feature:
Vrobee (zx (I.a(w ...(wwl)) (w()) x +b()))
per E {Imm Im.......MM M where {Tm Im=1....MM is set of learned input values. For simplicity, we use the same set for all neural graphs. The features Equation are then included as additional node features. Notably, probe features are invariant to all augmentations on the input neural network's parameters as long as it maintains the exact same function for the output and hidden layers.
Normalization.
Existing works on parameter space networks (Navon et al., 2023; Zhou et al., 2023a) perform feature normalization by computing the mean and standard deviation separately for each neuron in the training set. This operation, however, violates the neuron symmetries; since neurons can be permuted, it becomes essentially arbitrary to normalize neurons across neural networks. We propose a simple alternative that respects permutation equivariance: we compute a single mean and standard deviation for each layer (separately for weights and biases), and use them to standardize our neural graph, i.e. W = (w()- - (w /ow 600) = (D - mb /a,"" , Ll.
Positional embeddings. Before processing the neural graph, we augment each node with learned positional embeddings. To maintain the permutation symmetry in the hidden layers, nodes corre- sponding the same intermediate layer share the same positional embedding. Although this layer information is implicitly available in the adjacency matrix, using positional embeddings allows immediate identification, eliminating the need for multiple local message-passingg steps. However, distinguish between the symmetries of input and output nodes and those of hidden nodes. In a neural network, rearranging the input or output nodes generally alters the network's 's underlying function. In contrast, the graph representation is indifferent to the order of input and output nodes. To address this discrepancy, we introduce unique positional embeddings for each input and output node, thereby breaking the symmetry between and enabling GNNs and transformers to differentiate between them.
3 LEARNING WITH NEURAL GRAPHS
Graph neural networks (GNNs) and transformers are equivariant with respect to the permutation symmetries of graphs. We present one variant of each and adapt them for processing neural graphs.
GNN. Graph neural networks (Scarselli et al., 2008; Kipf & Welling, 2017) in the form of message passing neural networks (Gilmer et al., 2017) apply the same local message passing function at every node. While various GNN variants exist, only few of them are designed to accommodate edge features, and even fewer update these edge features in hidden layers. Updating edge features important in our setting since our primary features reside on the edges; when per-weight outputs required, they become imperative. We choose PNA (Corso et al., 2020) our backbone, state-of-the-an graph network that incorporates edge features. However, PNA does not update its edge features. To address this gap, we apply common extension to it by updating the edge features at each layer using a lightweight neural network Φe:
(k+1) =
- (k+1)
(k)
(1))
where & is the layer index in our network. This ensures that the edge features are updated per layer based on incident node features and produce representation that depends the graph structure.
Navon et al. (2023) suggest that the ability to approximate the forward pass of the input neural network can be indicative of expressive power. Part of the forward pass consists of the same operation for each edge: multiply the weight with the incoming activation. Motivated by the insights on algorithmic alignment by Xu al. (2020), adapt the message-passing step include this mulriplicative
5

---

Published as a conference paper at ICLR 2024
interaction between the node and edge features. In particular, we apply FiLM to the message passing step (Perez et al., 2018; Brockschmidt, 2020):
m Oscale (ei)) @ @m (lv,,, vy]) + (e)))
Note that this differs from the FiLM-GNN (Brockschmidt, 2020) in that we compute the scaling factors based on the edge features and not based on the adjacent node's features.
Transformer. The transformer encoder (Vaswani et al., 2017) can be seen as graph neural network that operates on the fully connected graph. Similar to GNNs, the original transformer encoder and common variants do not accommodate edge features. We use the transformer variant with relational attention (Diao & Loynd, 2023) that adds edge features to the self-attention computation. The relational transformer is already equipped with updates to the edge features. As per GNNs, we further augment the transformer with modulation to enable multiplicative interactions between the node and edge features. In particular, we change the update the value matrix in the self-attentin module:
(Wralue scale Ciy)
( + W Cij- shift
Vij
EXPERIMENTS
We assess the effectiveness of our approach on a diverse set of tasks requiring neural network processing, with either per-parameter outputs (equivariant) or a global output (invariant). We refer to the graph network variant of our method as NG-GNN (Neural Graph Graph Neural Network), and the transformer variant as NG-T (Neural Graph Transformer). We refer to the appendix for more details. The code for the experiments is open-soureed to facilitate reproduction of the results.
4.1 INR CLASSIFICATION AND STYLE EDITING
First, we evaluate our method on implicit neural representation (INR) classification and style editing, comparing against DWSNet (Navon et al., 2023) and NEN (Zhou et al., 2023a). These tasks involve exclusively the same MLP architecture per dataset.
Setup. For the INR classification task, we use two datasets. The first dataset contains a single INR (Sitzmann et al., 2020) for each image from the MNIST dataset (LeCun et al., 1998). The second contains one INR per image from Fashion MNIST (Xiao et al., 2017). An INR is modeled by a small MLP that learns the mapping from an input coordinate the grayscale value of the image (or to the RGB value for colored images). Each INR is separately optimized to reconstruct its corresponding image. We use the open-source INR datasets provided by Navon et al. (2023).
In the style editing task, we assess the model's 's ability to predict weight updates for the INRs using the MNIST INR dataset. The objective to enlarge the represented digit through dilation. We follow the same training objective Zhou et al. (2023a).
Results. In Figures 2a and 2b, we observe that our approach outperforms the equivariant baseline by up to +11.6% on MNIST and up to +7.8% on Fashion MNIST in INR classification. In Figure 2c, we also observe performance gains over both DWSNet and NFN in the style editing task. Interestingly, the baseline can perform equally well in terms of training loss, but our graph-based approach exhibits better generalization performance. We do not report any non equivariian baselines as they all perform considerably worse. The results highlight that increasing the number of probe features improve the performance. Furthermore, the probe features are effective even in a setting where we require an output per parameter as opposed
further
global prediction. Despite the advancements. performance gap persists when compared to the performance of neural networks applied directly to the original data. We hypothesize that key factor contributing to this gap overfitting, compounded by the absence of appropriate data augmentation techniques to mitigate it. This hypothesis is supported by the findings of Navon et al. (2023), who observed that the addition of INR views as a data augmentation strategy increased accuracy by +8%.
Importance of positional embeddings. We ablate the significance of positional embeddings in the task of MNIST INR classification. Without positional embeddings, NG-GNN achieves an accuracy of a3.9+..a%. and NG-T attains 77. This is a decrease of 7.5 and 14.5 points respectively, which highlights the importance of positional embeddings.
httpps: htt.ggttto........eeeoeoeeoeoo wtm.comm////////
-graps

---

Published as a conference paper at ICLR 2024
100
80
2.5
DWSSNtt
Accuracy (%6
(Ouns)
80
75
NG-T (Ours)
Accuracy
2.2
NG-T (Ours)
(Ours)
NG-GNN (Ours)
NG-GNN (Ours)
DWSNet
DWSNet
85
16
64
65
0
16
64
0.4
16
64
Number of probe fectures
Number of probe features
Number of probe features
(a) MNIST classification
(b) Fashion MNIST classification
(c) MNIST dilation
Figure 2: INR classification and style editing. Results are averaged over 4 seeds. Classification results are in Figures 2a and 2b, and dilation results are in Figure 2c. Both NG-GNN and NG-T outperform DWSNet and NEN already with probe features and significantly improve with more probe features.
Table I: Predicting CNN generalization from weights. Kendall's on CIFARIO-GE (Unterthiner et al., 2020) and CNN Wild Park. Higher is better. We report the mean and standard deviation for 3 seeds.
Method
CIFARIO-GS
CIFARIO Wild Park
NFNHNP (Zhou et al., 2023a) StatNN (Unterthiner et al., 2020)
0.934
0.915
0.719
NG-GNN (Ours)
0.930+o.001
0.804 009
NG-T (Ours)
0.935
0.817
4.2 PREDICTING CNN GENERALIZATION FROM PARAMETERS
Next, we consider the task of predicting the generalization performance of a CNN image classifier based on its parameters (Unterthiner et al., 2020; Eilertsen et al., 2020; Schürholt et al., 2022b).
Setup. We consider two datasets with increasing complexity. First, we use the CIFARIO-GS dataset from the Small CNN Zoo (Unterthiner et al., 2020), where GS denotes that images are converted to grayscale. We closely follow the CNN generalization setting from Zhou et al. (2023a); we use Kendall's correlation r to measure the performance, and we explore both MSE and binary cross entropy as loss functions. We compare our method with NFN (Zhou et al., 2023a) and StatNN (Unterthiner et al., 2020), method that computes statistics (mean, standard deviation, quantiles) for weights and biases of each layer, concatenates them, and processes them with an MLP.
The Small CNN Zoo comprises identical networks, i.e., all CNNs have the exact same convolutional layers, non- mineaiinasssies and linear layers, and only vary in their weights and biases, resulting from different initialization and training hyperparameters. Our method can in principle be applied to datasets that contain neural networks with varying architectures, but we have yet to see whether single model can learn a joint representation. To test this, introduce new dataset of CNNs, which we term CNN Wild Park. The dataset consists of 117,241 checkpoints from 2,800 CNNs, trained for up 1,000 epochs on CIFARIO. The CNNs vary in the number of layers, kernel sizes, activation functions, and residual connections between arbitrary layers. We describe the dataset in detail in Appendix E. We compare our method against StatNN (Unterthiner et al., 2020). NFN is inapplicable in this setting, due to the heterogeneous architectures present in the dataset. Similar to the CIFARIO-GS we use Kendall's T to as an evaluation metric.
Results. We report results in Table 1. In both settings, our NG-T outperforms the baselines, while NG-N performs slightly worse than NG-T. Although the performance gap in CIFAR10-GS, narrow, on the Wild Park dataset we significant performance gap between StatNN and our neural graph based approaches. Here the architecture of the CNN classifier offers crucial information and the parameters alone appear to be insufficient.
Importance of non-linearity features. We ablate the significance of non-linearit features on CNN Wild Park. Without non-linearity features, NG-GNN achieves O. as,ss and NG-T 0.728+ resulting in a performance decrease by 2.6 and 8.9, respectively. This highlights that including the non-linearities as features is crucial, especially for the Transformer.
7

---

Published as a conference paper at ICLR 2024
Table 2: Learning to optimize. Test image classification accuracy (%) after 1,000 steps. We repeat the evaluation for 5 random seeds and report the mean and standard deviation.
Optimizer
FashionMNIST (validation task)
CIFAR-10 (test task)
Adam (Kingma & Ba, 2015) FF (Metz et al., 2019)
80.97 +0.00
54.76=22.2
85.08
57.55=1.on
LSTM (Metz et al., 2020)
85.69 a..23
59. 10..0m
NEN (Zhou et al., 2023a)
83. 33.78
57.93=0.64
NG-GNN (Ours)
85.91 20.37
64.37
NG-T (Ours)
86.52:t.11
60.79
4.3 LEARNING TO OPTIMIZE
novel application where leveraging the graph structure of a neural network may be important, yet underexplored, is "learming to optimize" (L20) (Chen et al., 2022; Amos, 2022). L20's task is to train a neural network (optimizer) that can optimize the weights of other neural networks (optimizee) with the potential to outperform hand-designed gradient-descent optimization algorithms (SGD, Adam (Kingma & Ba, 2015), etc.). Original L,20 models were based on recurrent neural networks (Andrychowicz et al., 2016; Wichrowska al., 2017; Chen et al., 2020). Recent L20 models are based on simple and more efficient MLPs, but with stronger momentum-ben based features (Metz al., 2019; 2020; Harrison et al., 2022: Metz et al., 2022a).
Setup. We implement two strong learnable optimizer baselines: per-parameter feed forward neural network (FF) (Metz et al., 2019) and its stronger variant (FF layerwise LSTM, or just LSTM) (Metz et al., 2020). FF predicts parameter updates for each parameter ignoring the structure of an optimizee neural network. LSTTM adds layer features (average features of all the weights within the layer) with the LSTM propagating these features between layers. We also combine FF with NFN (Zhou et al., 2023a) as another baseline, and add Adam as a non-learnable optimizer baseline. We compare these baselines to FF combined with our NG-GNN or Relational transformer (NG-T). The combinations of FF with the NFN, NG-GNN, and NG-T models are obtained analogously to LSTM, i.e., the weight features are first transformed by these models and then passed to FF to predict parameter updates.
In all learnable optimizers, given an optimizee neural network with weights (W()) )==1...... the features include weight gradients (V) JI= (1=0)))11............. and momentums at five scales: [0.5, 0.9, 0.99, 0.999, 0.9999] following Metz et al. (2019). These features are further preprocessed with the log and sign functions (Andrychowicz et al., 2016), so the total number of node and edge features is dy = dg = 14( (4xmx mmaax) in our experiments, where Wmax Amax comes from flattening the convolutional kernels according Section 2.2.
We follow a standard L20 setup and train each optimization method on FashionMNIST (Xiao et al., 2017) followed by evaluation on FashionMNIST For FashionMNIST, we use a small three layer CNN with 16, 32, and 32 channels, and 3 x 3 (10ax 1mmx) ) kernels in each layer followed by global average pooling and classification layer. For CIFAR-10, we use a larger CNN with 32, 64, and 64 channels.
well on CIFAR-10 (Krizhevsky et al., 2009).
In order perform well on CIFAR- 10, the learned optimizer has to generalize to both a larger architecture and a different task. During the training of the optimizers, we recursively unroll the optimizer for 100 inner steps in each outer step and train for up to 1,000 outer steps in total. Once the optimizer is trained, we apply it for 1,000 steps on a given task. We train all the learnable optimizers with different hyperparameters (hidden size, number of layers, learning rate, weight decay, number of outer steps) and choose the configuration that performs the best on NaShionMNIST. This way, CIFAR-10 is test task that is not used for the model/hyyprpprraeeer selection in our experiments.
Results. On the validation task (FaShionMNISTT), the best performing method is NG-T followed by NG-GNN (Table 2). The best performing baseline is LSTM followed by FF. On the test task (CIFAR-10), NG-GNN outperforms NG-T indicating its strong generalization capabilities. LSTM is again the best performing baseline followed by NFN. The training loss and test accuracy curves on CIFAR- 10 also reveal that NG-GNN's 's gain over other methods gradually grows demonstrating its great potential in the L20 area (Figure 3). We believe that NG-T may also generalize better to new tasks once it is trained more tasks, following recent trends in 1.20 (Metz al., 2022b).
8

---

Published as a conference paper at ICLR 2024
NFN
2.25
FF
2.60
training 125
Adam
NFN
1.50
1.25
NGTT
LSTM
MS-GNN
0.75
620 training Eoration
1000
training Roration
Figure 3: Training (left) and testing (right) curves on CIFAR-10 for the baseline optimizers (Adam, FF, LSTM, NFN) and the optimizers trained with our NG-GNN and NG-T. NG-GNN and NG-T outperform the baselines with the NG-GNN performing the best on this task.
5 RELATED WORK
Networks for networks. A recent line of work studies how learn representations for trained classifiers (Baker et al., 2017: Eilertsen et al., 2020; Unterthiner et al., 2020; Schürholt et al., 2021; 2022b;a) predict their generalization performance, or other properties that provide insight into neural networks. Dupont et al. (2022); De Luigi et al. (2023) learn low-dimensional encodings for INRs for generation other downstream tasks. These works either flatten the network parameters or compute parameter statistics and process them with standard MLP. Using statistics respects the symmetries of the neural network, but discards fine-grained details that can be important for the task. On the other hand, architectures that flatten the network parameters are not equivariant with respect to the symmetries of the neural network. This is undesirable because models that are functionally identical obtained by permuting the weights and biases appropriately can receive vastly different predictions. Schürholt et al. (2021) proposed neuron permutation augmentations to address the symmetry problem, however such an approach can require vastly due to the size of the permutation group: II, d, Other related works (Peebles et al., 2022; Ashkenazi al., 2023; Knyazev et al., 2021; Erkoç al., 2023) encode and/or decode neural network parameters mainly for reconstruction and generation purposes. However, similarly to the aforementioned works, these do not address the symmetry problem principled way.
training compute
The symmetry problem. Three recent studies address these shortcomings and propose equivariant linear layers that achieve equivariance through intricate weight-sharing patterns (Navon al., 2023; Zhou et al., 2023a;b). They compose these layers to create both invariant and equivariant networks for networks. So far, their performance on INR classification shows that there still a gap when compared to image classification using CNNs. Similar to these studies, our proposed method maintains the symmetry of the neural network. contrast to these works, by integrating the graph structure in neural network, process heterogeneous architectures for a much wider gamut of applications.
no longer limited to homogeneous architectures and can
CONCLUSION AND FUTURE WORK
We have presented an effective method for processing neural networks with neural networks by representing the input neural networks as neural graphs. Our experiments showcase the breadth of applications that this method can be applied to. The general framework is flexible enough that domein- adaptations, egg., including gradients as inputs, are simple to add. Furthermore, by directly using the graph, open the door for variety of applications that require processing varying architectures, as well as new benchmarks for graph neural networks.
Limitations. While our method is versatile enough to handle neural networks with a diverse range of architectural designs, the scope of our investigation was limited to two different architecture families (MLPs and CNNs), and only theoretically demonstrated how to represent transformers as neural graphs. Another limitation is that method's strong performance INRs is confined to 2D images, which restricts its applicability. Extending our approach to handle neural radiance fields (Mildenhall al., 2020) would substantially broaden its utility.

---

Published as a conference paper at ICLR 2024
ACKNOWLEDGMENTS
We sincerely thank Siamak Ravanbakhsh for the invaluable guidance on the equivariance discussion in Appendix A. Furthermore, we thank SURF (www.. surf. n1) for the support in using the National Supercomputer Snellius. The experiments were also enabled in part by computational resources provided by Calcul Quebec and Compute Canada.
REFERENCES
Brandon Amos. Tutorial on amortized optimization for learning to optimize over continuous domains. arXiv preprint MxXi:2202200666 2022. 8
Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W Hoffmmn, David Pfau, Tom Schaul, Brendan Shillingford, and Nando De Freitas. Learning to learn by gradient descent by gradient descent. In Advances in Neural Information Processing Systems 29 (NIPS), 2016. 8
Maor Ashkenazi, Zohar Rimon, Ron Vainshtein, Shir Levi, Elad Richardson, Pinchas Mintz, and Eran Treister. NeRN Learning Neural Representations for Neural Networks. In 1th International
Conference Learning Representations (ICLR), 2023. 9
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. Layer Normalization. arXiv preprint arXlv:1607.064450, 2016. 16
Bowen Baker, Otkrist Gupta, Ramesh Raskar, and Nikhil Naik. Accelerating Neural Architecture Search using Performance Prediction. arXiv preprint arXiv: 10823, 2017. 9
Mare Brockschmidt. GNN-FiLM: Graph Neural Networks with Feature-wise Linear Modulation. In Proceedings of the 37th International Conference on Machine Learning (ICML), 2020.
Michael M Bronstein, Joan Bruna, Taco Cohen, and Petar Veličković. Geometric Deep Learning: Grids, Groups, Graphs, Geodesics, and Gauges. arXiv preprint arXiv:2104. 1347888..44. 2021. 1, 13
Tianlong Chen, Weiyi Zhang, Zhou Jingyang, Shiyu Chang, Sijia Liu, Lisa Amini, and Zhangyang Wang. Training Stronger Baselines for Learning to Optimize. In Advances in Neural Information Processing Systems 33 (NeurIPS), 2020.
Tianlong Chen, Xiaohan Chen, Wuyang Chen, Zhangyang Wang, Howard Heaton, Jialin Liu, and Wotao Yin. Learning to Optimize: A Primer and A Benchmark. The Journal of Machine Learning Research, 23(1):8562-8622, 2022. 8
Gabriele Corso, Luca Cavalleri, Dominique Beaini, Pietro Liò, and Petar Veličković. Principal Neighbourhood Aggregation for Graph Nets. In Advances in Neural Information Processing Systems 33 (NeurIPS), 2020.
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 11th International Conference on Learning Representations (ICLR), 2023. 9
Cameron Diao and Ricky Loynd. Relational Attention: Generalizing Transformers for Graph- Structured Tasks, In 11th International Conference on Learning Representations (ICLR), 2023.
Emilien Dupont, Hyunjik Kim, SM Eslami, Danilo Rezende, and Dan Rosenbaum. From data to functa: Your data point is function and you can treat it like one. In Proceedings of the 39th International Conference on Machine Learning (ICML), 2022. 1,9
Gabriel Eilertsen, Daniel Jönsson, Timo Ropinski, Jonas Unger, and Anders Ynnerman. Classifying the classifier: dissecting the weight space of neural networks. In 24th European Conference on Artificial Intelligence (ECAI), 2020. 7.9
Ziya Erkoç. Fangchang Ma, Qi Shan, Matthias Nießner, and Angela Dai. HyperDiifussin: Generating Implicit Neural Fields with Seight-Space Diffusion. In International Conference on Computer Vision (ICCV), 2023.
10

---

Published as a conference paper at ICLR 2024
Justin Gilmer, Samuel S Schoenholz, Patrick F Riley, Oriol Vinyals, and George E Dahl. Neural Message Passing for Quantum Chemistry. In Proceedings of the 34th International Conference on Machine Learning (ICML), 2017. 5
James Harrison, Luke Metz, and Jascha Sohl-Dickstein. Closer Look at Learned Optimization: Stability, Robustness, and Inductive Biases. In Advances in Neural Information Processing Systems 35 (NeurIPS), 2022.
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep Residual Learning for Image Recognition. In Proceedings of the IEEE Conference on Computer Vision and Puttern Recognition (CVPR), 2016.
Robert Hecht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Advanced Neural Computers, pp. 129-135. Elsevier, 1990.
Sergey loffe and Christian Szegedy. Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. In Proceedings of the 32nd International Conference
Machine Learning (ICML), 2015. 16
Diederik P Kingma and Jimmy Ba. Adam: Method for Stochastic Optimization. In 3rd International Conference for Learning Represenaations (ICLR), 2015. 8
Thomas N. Kipf and Max Welling. Semi-Supervised Classification with Graph Convolutional Networks. In 5th International Conference on Learning Representations (ICLR), 2017. 5
Boris Knyazev, Michal Drozdzal, Graham W Taylor, and Adriana Romero Soriano. Parameter Prediction for Unseen Deep Architectures, In Advances in Neural Information Processing Systems 34 (NeurIPS), 2021.
Alex Krizhevsky, Geoffrey Hinton, et al. Learning Multiple Layers of Features from Tiny Images. 2009. 8
Yann LeCun, Léon Bottou, Yoshua Bengio, and Patrick Haffner. Gradient- learning applied to document recognition. Proceedings of the IEEE, 86(11):227822324 1998. 6
Luke Metz, Niru Maheswaranathan, Jeremy Nixon, Daniel Freeman, and Jascha Sohl-Dickstein. Understanding and correcting pathologies in the training of learned optimizers. In Proceedings of the 36th International Conference on Machine Learning (ICML), 2019. 8
Luke Metz, Niru Maheswaranathan, C Daniel Freeman, Ben Poole, and Jascha Sohl-Dickstein. Tasks, stability, architecture, and compute: Training more effective learned optimizers, and using them to train themselves. arXiv preprint arXiv:2009. 11433, 2020. 8
Luke Metz, C Daniel Freeman, James Harrison, Niru Maheswaranathan, and Jascha Sohl-Dickstein. Practical tradeoffs between memory, compute, and performance in learned optimizers. In Ist Conference Lifelong Learning Agents (CoLLAs), 2022a. 8
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 preprini (tt.:221110996600 2022b. 8
Ben Mildenhall, Pratul P Srinivasan, Matthew Tancik, Jonathan T Barron, Ravi Ramamoorthi, and Ren Ng. NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis. In European Conference Computer Vision (ECCV), 2020. 9
Aviv Navon, Aviv Shamsian, Idan Achituve, Ethan Fetaya, Gal Chechik, and Haggai Maron. Equiv- ariant Architectures for Learning in Deep Weight Spaces. In Proceedings of the 40th International Conference on Machine Learning (ICML), 2023. 1, 5, 6, 9, 14, 17, 18
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: 12892, 2022. 9
11

---

Published as a conference paper at ICLR 2024
Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville. FILM: Visual Reasoning with a General Conditioning Layer. In Proceedings of the AAAI Conference on Artificial Intelligence (AAAI), 2018. 6
Franco Scarselli, Marco Gori, Ah Chung Tsoi, Markus Hagenbuchner, and Gabriele Monfardini. The graph neural network model. IEEE transactions on neural networks, 20(1), 2008. 5
Konstantin Schürholt, Dimche Kostadinov, and Damian Borth. Self-Supervised Representation Learning on Neural Network Weights for Model Characteristic Prediction. In Advances Neural Information Processing Systems 34 (NeurIPS), 2021. 9
Konstantin Schürholt, Boris Knyazev, Xavier Giró-i Nieto, and Damian Borth, Sepeen-eerrssentationss as Generative Models: Sampling Unseen Neural Network Weights. In Advances in Neural Information Processing Systems 35 (NeurIPS), 2022a. 1,9
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. In Advances in Neural Information Processing Systems 35 (NeurIPS), 2022b. 7,
Karen Simonyan and Andrew Zisserman. Very Deep Convolutional Networks for Large-Scale Image Recognition. In 3rd International Conference on Learning Representations (ICLR), 2015. 4
Vincent Sitzmann, Julien Martel, Alexander Bergman, David Lindell, and Gordon Wetzstein. Implicit Neural Representations with Periodic Activation Functions. In Advances in Neural Information Processing Systems 33 (NeurIPS), 2020. 6
Thomas Unterthiner, Daniel Keysers, Sylvain Gelly, Olivier Bousquet, and Ilya Tolstikhin. Predicting Neural Network Accuracy from Weights. arXiv preprint arXiv:2002. 1448, 2020. 1, 7, 9, 19
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention Is All You Need. In Advances in Neural Information Processing Systems 30 (NeurIPS), 2017. 4, 6
Olga Wichrowska, Niru Maheswaranathan, Matthew W Hoffman, Sergio Gomez Colmenarejo, Misha Denil, Nando Freitas, and Jascha Sohl-Dickstein. Learned Optimizers that Scale and Generalize. In Proceedings of the 34th International Conference Machine Learning (ICML), 2017. 8
Han Xiao, Kashif Rasul, and Roland Vollgraf. Fashion- a Novel Image Dataset for Bench- 2017. 6,8
marking Machine Learning Algorithms. arXiv preprint
Keyulu Xu, Jingling Li. Mozhi Zhang, Simon S Du, Ken-ichi Kawarabayashi, and Stefanie Jegelka. What Can Neural Networks Reason About? In 8th International Conference for Learning Representations (ICLR), 2020. 5, 15
Allan Zhou, Kaien Yang, Kaylee Burns, Yiding Jiang, Samuel Sokota, Zico Kolter, and Chelsea Finn. Permutation Equivariant Neural Functionals. In Advances in Neural Information Processing Systems 36 (NeurIPS), 2023a. 6, 7, 13, 14, 18
Allan Zhou, Kaien Yang, Yiding Jiang, Kaylee Burns, Winnie Xu, Samuel Sokota, Zico Kolter, and Chelsea Finn. Neural Functional Transformers. In Advances in Neural Information Processing Systems 36 (NeurIPS), 2023b. 9,
12

---

Published as a conference paper at ICLR 2024
A EQUIVARIANCE PROPERTIES
In this section, we will address three points. First, we state the symmetry group of the neural graph explicitly. Then, we restate the symmetry group prior works and establish that any model that is equivariant to neural graphs will automatically be equivariant to the symmetry of the prior work. Lastly, we discuss why we believe that our choice of symmetry is usually more appropriate for processing neural network parameters as input.
A.1 NEURAL GRAPH SYMMETRIES
Since we represent the input as a graph, we can use many existing properties of graphs. Namely, is well-established that the symmetries of graphs are given by the symmetric group S, of order Σt= di (the total number of nodes). Let p denote the group representation that maps permutation E S, to the corresponding permutation matrix. A group action a on the graph = (V, E) can be defined as:
at (V. E)) = (p(+)V, p(T) Ep(x)).
We want to process the neural graph with a neural network that maintains the equivariance with respect the specified symmetries. One such example is the standard message passing graph neural network (MPNN), which is equivariant (ef. Bronstein et al. (2021)) since all individual operations are designed to be equivariant.
A.2 NEURON PERMUTATION GROUP
Similar Zhou et al. (2023a), focus the neuron permutation (NP) group to simplify the exposition. The main difference to the symmetry usually present in neural network parameters' is that input and output neurons can be freely permuted well. As describe in Section 2.4, positional embeddings rather than modifying the architecture to disambiguate the ordering of input and output neurons, so this difference is not directly relevant the following discussion.
We now restate the definitions of Zhou et al. (2023a) in our notation to allow for discussion of the similarities and differences. First, the NP group defined as S = Sáo ... x Sáz The elements of the group are then n' = (To..... ). For a specific layer I, the actions / are defined as
B 8( (V",bb0) (p(T)))O),
We can already see the similarity between Equations (6) and (7). To make this connection more concrete, notice that S is a subgroup of Sn, This is because is defined as direct product of symmetric groups whose degrees sum up to by definition. We can also verify that any action of corresponds an action of S,. To see this, define the group representation p' of S as the block permutation matrix:
[((m))
0
((m))
p((T')
(8)
p(TL)
Then, we can rewrite 3 as 3 in exactly the same form as a::
8' (V. E)) (d'(T'VV, E
(9)
Since S is subgroup of S,, any model that is S,, equiiaaria must also be S-equivariant. In other words, choosing models like MPNNs with our neural graph ensures that we satisfy the permutation symmetries of neural networks established in the literature.
Defined by Zhou al. (2023a) the hidden neuron permutation (HNP) group
13

---

Published as a conference paper at ICLR 2024
A.3 DISCUSSION
Why would prefer symmetry group the other? The main factor that Sn encompasses many choices of S at once. Using single S, -q-equivariant model, our method is automatically equivariant to any choice of .o, dz that sums 7ı; the same model can process many different types of architectures. In contrast, prior works (Navon et al., 2023; Zhou et al., 2023a) propose models that are solely equivariant S, which restrict the models to specific choice for dg..... dg at time; hence, it is not possible to process different network architectures with DWSNet (Navon et al., 2023) and NFN (Zhou et al., 2023b). Thus, when there a need to process different architectures with the same model, the neural graph representation is more suitable.
Another aspect to consider that DWSNet and NFN learn different parameters for each layer the input neural network. In contrast, neural network operating the neural graph shares parameters for the update function that it applies at every node, irrespective the layer belongs to. This difference is particularly relevant in scenarios where performance the primary focus. Under such conditions, the inherent equivariance of our model may be overly constraining, failing to meet the criteria for minimal equivariance as outlined by Zhou et al. (2023b). This has two potential effects: the reduced number of parameters could force the model to learn more generalizable features that work across layers, it could simply result in not having enough parameters to fit the data. We remedy this shortcoming by integrating layer embeddings that identify the layer each node belongs to. These features are added the node representations and enable graph neural networks to differentiate each node depending on the layer.
single, fixed architecture is
Experimentally, we observe that our model not only outperforms the baselines but also generalizes better. In Figure 4, we plot the train and test losses for MNIST INR classification. We observe that even points where the train losses match, the test loss of NG-T lower. The results confirm that the parameter sharing from S,, equivariance not only does not harm the performance but might even benefit generalization.
1.5
DWSNett troin loss
NG-T train loss
DWSNtet loss
0.5
NG.T
Epach
20
40
60
80
Figure 4: Generalization comparison between DWSNet and NG-T on MNIST INR classification. We observe better generalization for NG-T compared to DWSNet
B EXPRESSIVITY OF GNNS FOR NEURAL GRAPHS
The ability to approximate the forward pass of an input neural network can indicate the expressiveness of neural network operating other neural networks as inputs (Navon et al., 2023). We demonstrate that a general passsagepaasinggggg neural network with L layers can simulate the forward pass of any input MLP with L layers. We first place the input for which we intend to compute the forward pass in the node features the input nodes in the neural graph. Then we simulate the I-th layer's 's forward pass with the 1-th layer in the MPNN. To illustrate how this can be achieved, we start by decomposing the first layer of an MLP with activation function or into scalar operations:
f(x); ( +b)) =a b, + Σ W W
(10)
14

---

Published as a conference paper at ICLR 2024
We compare this to the update operation for the node features v, + in an MPNN:
v
Qu
Via Om Von
(11)
We can draw parallels between the operations in the MPNN and how the input MLP computes the subsequent layer's activations. For improved clarity, we assume in the following that the indices j match between the parameters and the graph's 's nodes and edges. A more precise version would need to add an offset do to the index i the graph. First, we recall that the edge features eiy contain the weights of the MLP Waj. the node features V, contain the input features x, (for which we compute the forward pass), and the node features V; contain the biases b,. The first layer of the MPNN can approximate (x)), by approximating the scalar product W N with the neural network Óm. Given these inputs, &m merely needs to approximate multiplication between the last two inputs. The addition of the bias, followed by the activation function easily be approximated with the node update MLP Φu.
For the remaining layers in the MLP, the &h merely needs ensure that it preserves the input node features until it requires them for approximating the (-th MLP layer in the I-th MPNN layer. For this, simply include layer embedding as layer position indicator in the input node features. Finally, the output of the simulated forward pass coincides with the node features of the final MPNN layer.
This construction approximates the forward pass of an input neural network by using general MPNN that operates on the neural graph. The algorithmic alignment between the operations of the MPNN and the computational steps of the forward pass indicates that the forward pass not only possible to express, but it can also be learned efficiently (Xu et al., 2020).
C NEURAL GRAPH REPRESENTATIONS
C.I MLPS AS GRAPHS
The neural graph specifies the computation of the neural network. Edges apply the function Vij Wiys; where Wij is the weight associated with the edge eij and the input comes from the tail node ws. The node Uj applies the sum operation z; bj En Miy all incoming edges' Yij-
C.2 CNNS AS GRAPHS
Figure 5: Permutation symmetries in CNN. shown in a 2-layer CNN sub-network. The first convolutional layer has a filter bank of C2 filters, each with C1 channels, and the second layer has C3 filters, each with C2 channels. Permuting the filters in the first layer results in permuting the channels of the feature map. Applying permutation in the channels of the filters in the second layer cancels out the per- mutations, resulting in permutation invariance.
C2
C2 liber
Os Mito
Flattening layer We propose two options for dealing with the flattening layer of traditional CNN architectures: adding one extra (virtual) layer, or repeating nodes in the last convolutional layer. For either option, we are binding the neural graph to a specific spatial resolution, and thus, we need to encode that in the graph structure well.
In the first option, we repeat every node in the last convolutional layer by the number of spatial locations. As example, if the CNN produces feature map that has spatial dimensions of 7 x at that layer, then we would need to copy the nodes 49 times. When we copy a node, we also need to copy the edges that point it. The weights connecting these nodes to the subsequent linear layer do not require any changes, since they already account for the spatial structure. Finally, we add positional embeddings the copied nodes indicate their different locations in space.
15

---

Published as a conference paper at ICLR 2024
0
D
Mc
Pall
0
Ayan
22
Flatzon
Maa
Figure 6: Representation of a convolutional kernel. We zero-pad to predetermined maximum kernel size (5 x in this example) and flatten it.
alternative option is to add an extra (virtual) layer, between the last convolutional layer and the subsequent linear layer. In that layer, we generate as many nodes as the spatial resolution of the feature map (e.g. in the example above), for each node in the last convolutional layer. Each node in the last convolutional layer is only connected to its own set of virtual nodes, with weight of 1. The new virtual nodes don't t have any node features.
C.3 NORMALIZATION LAYERS AS GRAPHS
Normalization layers, c.g. BatchNorm (loffe & Szegedy, 2015) and LayerNorm (Ba et al., 2016) are formulated as y @ x + B, which can be rewritten as a linear layer with diagonal weights y diag(>)) + B. As such, represent the d-dimensional input to the normalization layer x, we add an additional d nodes that correspond to the d-dimensional output y. The additional node features capture the additive terms 9, while the edge features capture the multiplicative terms We only add edges from x; to the corresponding y, to model the dlement-wiss multiplication.
treat normalization layers like linear layers: given d nodes that
C.4 TRANSFORMERS AS GRAPHS
Head self-attention layers initially apply linear projections to the inputs X. In total, assuming heads, we have 3// linear layers applied independently to inputs.
Q - A
(12)
K₁ XWK
(13)
VA = XWW,
(14)
with h E (1,..., H}. Each head is followed by dot-product attention layer Y, (Q,,
where is the softmax function. Finally, we concatenate all heads and perform final linear projection:
MHSA(X) Concat Y Y....., Y WO
(15)
A multi-head self attention layer takes d-dimensional vectors inputs and produces dig-dimensional vectors as output of each head, which are then concatenated and linearly projected to d dimensions. In the neural graph construction, each dimension of each head, and d nodes for each dimension of the output.
add d nodes for each dimension of the input, H dH nodes for
We model the different types of linear projections with multidimensionaal edge features. More
specifically, for each edge feature we have e/² ((w.) (w/),yr (WK))..) Since dot-
product attention a parameter- free operation, we don't 't model explicitly and let the neural graph network approximate it. The concatenation of all heads is automatically handled by the neural graph itself by connecting the appropriate nodes from each head through the output weights Wo to the corresponding output node. The final projection Wo is treated standard linear layer.
16

---

Published as a conference paper at ICLR 2024
D EXPERIMENTAL SETUP DETAILS
We provide the code for our experiments in the supplementary material. Here, we list the hyperpa- rameters used in experiments.
D.1 INR CLASSIFICATION AND STYLE EDITING
The NG-T T model consists of 4 layers with 64 hidden dimensions, 4 attention heads, and dropout at 0.2 probability. In the classification setting, concatenate the node representations that belong to the final layer of the INR and insert this into a classification head. Together with the classification head, this results in 378,090 trainable parameters.
The NG-GNN model consists layers, 64 hidden dimensions, and applies dropout at 0.2 probability, In the classification setting, concatenate the node representations that belong to the final layer of the INR and insert this into a classification head. Together with the classification head, this results in 348,746 trainable parameters.
We use the same architectures without the final aggregation and the classification head in the style editing experiment. The NG-GNN additionally uses the reversed edge features.
We train for 100 epochs on MNIST and for 150 epochs on FMNIST. We apply the same early stopping protocol Navon et al. (2023).
D.2 PREDICTING CNN GENERALIZATION FROM PARAMETERS
We do a small hyperparameter sweep for all models. For the StatNN baseline, we use 1000 hidden
dimensions resulting in a model resulting in 1,087,001 trainable parameters.
For the NG-T we use 4 layers with 8 hidden dimensions for the edges and 16 for the nodes, 1 attention head, and no dropout. We concatenate the node representations of the final layer and insert them into the classification head.
For the NG-GNN model, we use 64 hidden dimensions, resulting in 369,025 trainable parameters. We concatenate the node representations of the final layer and insert them into the classification head.
We train all models for epochs.
D.3 PROBE FEATURES
We use probe features in the INR classification and INR style editing experiments. We do not use probe features in the task of predicting CNN generalization, or the learning to optimize task.
E DATASET DETAILS
E. SMALL CNN WILD PARK
We construct the CNN Wild Park dataset by training 2800 small CNNs with different architectures for 200 to 1000 epochs CIFAR10. We retain checkpoint of its parameters every 10 steps and also record the test accuracy. The CNNs vary by:
Number of layers [2, 5] (note that this does not count the input layer).
Number of channels per layer dr E [4, 8, 16,32].
Kernel size of each convolution k [3, 5, 7].
Activation functions at each layer are one of ReLU, GeLU, tanh, sigmoid, leaky ReLU, or the identity function.
Skip connections between two layers with at least one layer in between. Each layer can have at most one incoming skip connection. We allow for skip connections even in the case when the number of channels differ, increase the variety of architectures and ensure independence between different architectural choices. We enable this by adding the skip connection only the min(Cn, Cm) nodes.
17

---

Published as a conference paper at ICLR 2024
We divide the dataset into train/val/test splits such that checkpoints from the same not contained in both the train and test splits.
F EXPERIMENT RESULTS
We provide the exact results of our experiments for reference and easier comparison.
Table 3: Classification of MNIST INRs. All graph-based models outperform the baselines.
Model
probe features
Accuracy in %
MLP (Navon et al., 2023)
17.6+a.0
Set (Navon al., 2023)
23.7+o.
DWSNet (Navon al., 2023)
85.7
NG-GNN (Ours)
91. 11.4
NG-GNN (Ours)
4
91,8+>.5
NG-GNN (Ours)
16
92.8+
NG-GNN (Ours)
64
94.7
NG-T T (Ours)
92. 92.4
NG-T (Ours)
4
03.3+
NG-T (Ours)
16
94.929
NG-T (Ours)
64
97.320.
Table Classification of Fashion MNIST INRs. All graph-based models outperform the baselines.
Model
probe features
Accuracy in %
MLP (Navon al., 2023)
19.9
Set NN (Navon al., 2023)
22.3
DWSNet (Navon al., 2023)
67.1
NG-GNN (Ours)
68.0g9
NG-GNN (Ours)
71.929.
NG-GNN (Ours)
16
73.1+a.5
NG-GNN (Ours)
64
74.2
NG-T (Ours)
72.7
NG-T (Ours)
73.829
NG-T (Ours)
16
74.1
NG-T (Ours)
64
74.810.0
Table 5: Dilating MNIST INRs. Mean-squared error (MSE) computed between the reconstructed image and dilated ground-truth image. Lower is better.
Model
# probe features
MSE in 10-2
DWSNet (Navon et al., 2023)
2.58
NEN (Zhou et al., 2023a)
2.55+a.00
NG-GNN (Ours)
2.3810.0
NG-GNN (Ours)
2.2G220.1
NG-GNN (Ours)
16
2.174mm.
NG-GNN (Ours)
64
2.06
NG-T (Ours)
1.96
NG-T (Ours)
1.88
NG-T (Ours)
16
1.82
NG-T (Ours)
641
1.75.0.01
18

---

Published as a conference paper at ICLR 2024
Table 6: Ablation study on the importance of non-linearity embeddings in the task of predicting CNN generalization from weights. Performance is measured as Kendall's CNN Wild Park. Higher is better. Removing the non- non-linearity embeddings results in performance decrease by 2.6 for NG-GNN and 8.9 for NG-T. The ablation highlights that including the non-linearities as features in the neural graph is crucial, especially for the Transformer variant.
Method
Kendall's 7
StatNN (Unterthiner et al., 2020)
0.719
NG-GNN (Ours)
0.80440.00
NG-GNN w/o activation embedding
0.77810.018
NG-T (Ours)
0.817
NG-T w/o activation embedding
0.72820-010
Table 7: Ablation study on the importance of positional embeddings in the task of MNIST INR classification. Removing positional embeddings results in decrease of 7.5 points for NG-GNN and 14.5 for NG-T, which highlights the importance of positional embeddings.
Method
Accuracy in %
NG-GNN (Ours)
91.410.6
NG-GNN w/o positional embeddings
83.9
NG-T (Ours)
92.410.a
NG-T w/o positional embeddings
77.9+0.7
19


Reviews:
[{"review_id": "mBYsQFmJ0A", "reviewer": "Reviewer_6cHp", "summary": "This work builds new neural networks that process other neural network parameters, by processing the computation graphs of the data neural networks. They do this processing via graph neural networks and graph transformers. Their models can then process htereogeneous architectures, as opposed to previous equivariant models. Experiments are conducted on several tasks involving processing neural networks.", "questions": "1. Why do the MNIST dilation numerical results differ so much from those in the Zhou et al. paper? They achieve about .070, whereas you achieve about .02.\n2. Could you report how many probe features are used in each of your experiments?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "1. The method can handle different nonlinearities, residual connections, and sizes of neural network.\n2. The graph framework is flexible, as it allows different types of base model, e.g. the GNN and Transformer that they consider in this work.\n3. Many types of empirical evidence, which shows the benefits of the method. The learned optimization experiments are particularly interesting.", "weaknesses": "1. While probe features improve performance, they are giving privileged information that is not quite in the same learning regime as other related works, which only take in parameters. With enough probe features, you are essentially inputting the original MNIST image into your neural network.\n2. Many claims of invariance or equivariance to permutation symmetries, without any proofs.", "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": "mip0SL435o", "reviewer": "Reviewer_rtFz", "summary": "This paper introduces a new approach to designing neural networks that process the parameters of other neural networks. This is achieved by representing the input neural network as a graph. In doing so, a graph neural network can operate on the input graph while respecting parameter permutation symmetries.\n\nThe neural network is converted to a graph by the introduction of a neural graph representation. This neural graph representation is related to the computation graph, but is more compact in some cases. Neural graph representations are designed for MLPs, CNNs, and residual networks. As part of the neural graph representation, the papers introduce novel node and edge features that improve the expressivity of the graph neural networks.\n\nThe proposed method is evaluated over a diverse set of experiments: classifying implicit neural representations, predicting network generalization, and learning to optimize. Across all settings, the proposed method significantly outperforms the evaluated baselines.", "questions": "- It is stated that the proposed neural graphs \"ensure invariance to neuron symmetries\". Are you able to outline how this might be proved?\n- Related to the previous question, the authors write that \"natural symmetries in the neural graphs correspond exactly to neuron permutation symmetries\", is this a statement that can be formalized? It is not clear to me that this is a 1:1 correspondence for all graphs considered. However, it is stated that this can be shown (Sec 2.1).\n- You observe that the baseline methods are able to perform equally well on the training loss, but fail to generalize as well (Sec 4.1). Why do you think this is? Did you explore adding regularization or similar to the baseline methods to help with generalization? I wonder if the probe features or other modifications are providing some of this benefit for the proposed method.\n\n\nMinor comments:\n\n- I thought the probe features is a very neat idea that, intuitively, adds a lot of expressive power to the proposed method.\n- At the end of the introduction, it is written that the proposed method \"outperforms state-of-the-art approaches by a large margin\". I think perhaps this would be better quantified with some specific values.\n- In Section 2.3, non-linearities are described as being added to the node features. Is this done via concatenation? And what is done when there is no activation?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "The proposed method is novel. I also agree with the authors that it is an improvement over baseline methods in that this approach can handle multiple architectures with the same model and does not require bespoke layer design.\n\nThe paper is very well written and easy to understand. I was able to follow completely and feel that adequate detail was provided for me to reproduce the results. \n\nThe empirical evaluation is thorough and conclusive. The proposed method gives a significant boost to performance over the baselines. The paper evaluates performance over several established tasks. There is a good amount of variety and error bars are included for all experiments, further establishing the consistent benefits. The authors also included full source code for their method and experiments.\n\nI feel that the contributions of this work are significant overall. This is a nice approach to designing neural networks that process other neural networks, which alleviates some of the complexity in prior work.", "weaknesses": "The CNN graph construction is effective but feels quite hacky. To use the method, the user must first specify a maximum kernel size. While this is unlikely to be a problem in practice, because the kernel size of modern CNNs does not vary much, it is not obvious how to extend this to other network layers that exhibit parameter sharing. For example, attention layers share parameters over sequence length which may vary significantly more than kernel size. Moreover, other standard building blocks like normalization layers are not covered in this work and there is no clear recipe provided for designing the corresponding neural graph representations.\n\nThere is little theoretical justification for the proposed approach. Prior work proves the expressivity of their methods alongside their group equivariance properties. Neither of these are explored formally in this work.\n\nI consider the empirical results to be quite complete, but an ablation of the various design decisions introduced would be valuable. For example, I'd like to better understand how much value the probe features, non-linearity identification, positional encoding, and other components contribute to the overall performance.", "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": "hOdBUXE5QH", "reviewer": "Reviewer_qkCh", "summary": "The paper considers the task of creating neural networks that take other neural networks as input. This problem has been considered in the machine learning literature which often does not take into account the different symmetries of neural network parameters, or only considers simple and fixed architectures such as multi-layer perceptrons or convolutional neural networks (Navon et al., 2023; Zhou et al., 2023a). The paper extends those works and proposes a more flexible approach where different architectures of neural networks can be given as input their neural network (i.e. different number of layers, widths, non-linearities). The paper also considers other types of layers such as residual connections and normalizations that take respect permutation equivariance across the different layers.", "questions": "Assuming that the architecture of neural networks is fixed during training and test, how scalable is the proposed approach compared to the baselines (Navon et al., 2023; Zhou et al., 2023a)? In particular, can the proposed approach consider the same size of neural networks as the baselines for a fixed maximum memory allocation? Or can the proposed approach consider even larger architectures than the baselines?\n\nIs the expressive power of the proposed approach the same as the baselines?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "In general, the paper is well motivated and easy to understand. The paper also shows strong improvement over baselines in three different tasks which are: INR classification of 2d images and style editing, predicting CNN performance and learning optimizers. In the supplementary material, an ablation study of the proposed probe features (that consider representations at different layers of the neural networks) is provided. I think that the paper is relevant to the machine learning community. It is a stepping stone to more general neural networks that take other neural networks as input.", "weaknesses": "Despite considering new types of layers (e.g. residual connections) compared to the literature, other types of layers such as multi-head attention layers are still missing. Moreover, only the MLP and convolutional layers have been tested in the experiments. The paper could be improved a lot if it experimentally showed that the other types of proposed layers can also help improve the performance of their model. \nThe paper also does not discuss the expressive power of their approach compared to the baselines, nor how difficult it is to scale their approach to very large neural networks/graphs.", "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 (oral)