Paper: Graph Metanetworks for Processing Diverse Neural Architectures (ICLR-2024)

Reviewer comment (Theory):
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.

Author response:
We add layer number (of both edges and nodes) as features for the input; this is mentioned in Appendix A. This allows the network to differentiate input and output nodes. Also, we can add edge direction as a feature (forward or backward), which is mentioned in Appendix A and found to be useful in the expressive power results in Appendix C.1.2.

Reviewer comment (Experiments):
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.

Author response:
As explained in the general comment, we have added an additional set of experiments in our revised manuscript, including a 2D-INR editing task. This includes a comparison with NFN and the more recent NFT method [Zhou et al. 2023b].  

We agree that there are many exciting applications for this approach, especially given our ability to handle heterogeneous architectures within a single trained model. While we hope to explore these, we emphasize that our current submission contains significant contributions in terms of the novel method that we proposed, our theoretical analysis of it, and the empirical experiments validating the method in practice.

Reviewer comment (Experiments):
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.

Author response:
We are disappointed that this contribution of our work was not clearer. We emphasize that these results are extensions/generalizations of established theorems presented in prior published work. Additionally, these results theoretically justify our application of graph neural networks as metanets.  

While we certainly agree that more experiments could strengthen our paper, and we have added more since submission, we stand by the inclusion of these theoretical results as a major contribution.  

See our general comment for more on the value of our theory.

Reviewer comment (Experiments):
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.

Author response:
Given that in Section 5.1 we show that the same graph metanet can predict the accuracy of small ResNets and ViTs, we think our models can likely be useful for larger models. We agree with the reviewer that these larger scale accuracy prediction tasks could be interesting and useful. In some sense, future versions of our metanetwork may allow predictions of new types of “scaling laws” for input neural nets.

Reviewer comment (Writing):
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.

Author response:
While we don’t establish a concrete answer to the specific question of kernel size, we do not have evidence that our method cannot generalize to novel architectures. We rely on GNNs, which can be applied to any input graph. Consequently, we can train and evaluate our metanets on vastly different parameter graphs, including graphs with kernel sizes never seen during training.  

To handle the specific case of kernel size, we use positional encoding of the kernel location as edge features. One could have the center kernel position encoded as (0,0) and then encode arbitrary kernel sizes by moving radially outwards from the center when computing positional encoding. This would allow our model to handle different kernels at training and testing time. Note that we do not establish empirical generalization to this particular distribution shift. However, we do show that our metanets generalize to a similar distribution shift: increasing network width.

Reviewer comment (Theory):
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.

Author response:
> “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.”

We are not sure what the reviewer means here, a clarification may be helpful. For instance, the receptive field is readily available in the parameter graph representation, since the kernel sizes are encoded, so any metanetwork can trivially determine the receptive field. Notably, we include node and edge type features, so for instance a node corresponding to a convolutional layer channel has a node feature indicating that it is a convolution node.

Thank you for quickly clarifying! We see what you mean now.

However, we do not think this is necessarily much of an issue. To recap, for each weight in a convolution filter, we include an edge between the nodes representing the input and output channel. This edge is given a feature that is a positional encoding of the (x,y) spatial location of the weight in the original filter. Then we use a sinusoidal positional encoding for this (x,y) feature. Note that sinusoidal positional encodings like this are used in Vision Transformers, as they also take 2D grids, flatten them, and then add in 2D information again via (x,y) positional encodings. As Vision Transformers and similar approaches are of course very successful, this approach of encoding spatial information via positional encodings is probably very reasonable.

If we want to include 2D spatial information more explicitly, then we can take another approach using 2D convolutions in our metanetwork. For instance, if we have a 5x5 convolution filter, this would give 25 edges between each input and output channel. The edge features of these edges can be themselves obtained by a standard 2D convolution network applied on the 5x5 filter, which more explicitly accounts for 2D spatial information. After this, our standard graph metanet framework can be followed, with further processing done by a GNN.

Reviewer comment (Experiments):
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.

Author response:
> “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.”

Our method is able to represent a wide variety of network architectures while retaining expressivity. **We already handle many more operations than previous work**, e.g. NFN and DWSNets only handle linear or convolutional layers (and not normalization, residuals, attention, group equivariant linear maps, or spatial parameter grids as we handle).  This is achieved through a combination of parameter graph design and feature encoding. 

For instance, with regards to spatial pooling layers, these are essentially a no-op in the parameter graph. In the sense that the same parameter graph accurately captures the equivariances that we’re seeking whether pooling is included or not. In fact, the image classifiers that we process in Section 5.1 have a global mean pooling (we have added this note to the revision).

Following our current design philosophy, adding features is exactly the right solution to indicate to the GNN that a pooling layer has been applied. It is also effectively free in terms of computation cost. We are not sure what the reviewer means when they say that adding features is “artificial”, but we note that adding features is very efficient (essentially zero cost), general, and simple.

Reviewer comment (Experiments):
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.

Author response:
> “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.”

As noted in the general comment, we now include experiments on 2D INR editing tasks from [Zhou et al. 2023a]. Our method outperforms the NFN baseline (NFN-HNP is essentially the same approach as DWSNets) but is beaten by the more recent specialized NFT architecture on only one of the tasks. We expect this gap can be closed by using more powerful/sophisticated graph neural network architectures, but emphasize at this time that our more general approach is applicable across a wide range of input networks while NFN/NFT are designed only for simple MLPs and simple CNNs (without e.g. normalization layers or residual connections).

We have also included additional results using NFN and NFT to predict generalization on a fixed CNN architecture. Our more general approach outperforms or matches the performance of these specialized baselines.

Reviewer comment (Experiments):
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.

Author response:
> “Since graph can encode a variety of architectures, I wonder if it is possible to generalize and especially extrapolate prediction to unseen architecture, e.g. the test MLP/CNN is wider or deeper than all training examples.”

**Our initial submission already had experiments** that evaluate the capability of our approach to generalize to unseen architectures. Please see the “OOD” setting in Section 5.1, and Figure 5. Notably, our GNNs generalize well to wider networks than those seen in training, for all of the diverse architectures: 2D CNN, 1D CNN, DeepSets, ResNet, ViT.

Reviewer comment (Reproducibility):
I do not see pooling layers discusses. Are there pooling layers involved in the networks in the dataset?

Author response:
> “I do not see pooling layers discusses. Are there pooling layers involved in the networks in the dataset?”

Yes, we use mean pooling across the spatial dimensions before the fully connected prediction layers for the image classifiers in Section 5.1. We have noted this in the revision.

Reviewer comment (Reproducibility):
How to encode atrous/dilated convolutions?

Author response:
> “How to encode atrous/dilated convolutions?”

One way to do this is to view the atrous/dilated convolution as a larger sparse kernel in the standard way, with zeros in many entries. E.g. a 3x3 kernel with dilation 2 can be viewed as a 5x5 kernel with zeros in all but 9 entries.

Another, more efficient way to encode these is to add another edge feature indicating the dilation factor for the convolutional weights. E.g. standard convolutions get an additional edge feature of 1, while convolution with dilation factor $k$ gives an edge feature of $k$.