Graph Metanetworks for Processing Diverse Neural Architectures

“1. ... 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. “

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.

“2. The computational complexity vs other approaches is not analyzed. Can the model encode large models in a feasible way?”

As message passing GNNs are generally efficient, and as much work has been done on GNNs, our methodology inherits these benefits. While we have not yet considered tasks on large networks, we note that GNNs are capable of training on graphs with 300 million edges (and hence networks with 300 million parameters) on a single GPU with 24 GB of RAM in 32-bit precision (e.g. https://arxiv.org/pdf/2110.14446.pdf). Thus, using multiple GPUs, GPUs with more RAM, and standard memory-saving tricks such as lower precision and gradient checkpointing, we believe that GMNs process input networks with billions of parameters. For inference, GMNs can ostensibly process input networks with tens of billions of parameters, due to lower memory cost than training.

“3. 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…”

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.

“4. The experiments lack ablations to understand how the model behaves under different settings…”

We agree that an ablation study would improve the empirical rigor of our work. We are working on this and hope to include it soon. Speaking informally, we found our metanets to be robust and effective across different choices of the GNN architecture, and choices of feature encoding.

“5. Source code is not available in the submission, which would be very helpful at least for how to build a graph given a neural network. Do the authors intend to open source the code?”

Unfortunately, we are unable to provide source code at this time. We are in the process of preparing an open-source release, and hope to have this ready in time for the camera-ready publication. In the meantime, we have added high level descriptions for our graph constructions in Appendix A.1.

“6. 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.”

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.

Thank you for your response. We greatly appreciate you taking the time to read our rebuttal and paper in further detail.

Additional 2D INR results are useful but indicating standard deviation would be beneficial

Apologies for this oversight. We will of course include standard deviations in the final version, as we have with other empirical results.

Thank you for your detailed questions about the feature constructions that we use. This has highlighted that this part of the paper deserves further clarification, and we will certainly do so in a later revision. In general, none of these features are strictly necessary and simply serve the purpose of improving network expressivity without sacrificing equivariance. Here are brief responses to the specific points that you raised:

having in the main text "we are free to use undirected edges" while actually using "edge direction" in some results sounds misleading

We will clarify this in the main text. Certainly, our proposed approach works with and without edge direction feature encoding. Some of our theoretical results lean on edge direction features (without violating equivariance), but empirically we found that this is unnecessary.

It sounds like the indices are fixed for all input networks so the metanet outputs would be different depending on the neuron permutation

(We will assume MLP input networks for ease of exposition, but the same argument can be generalized across all architectures.)

To clarify, neuron indices are not included for hidden neurons. We expose only the neuron indices for the input layer and output layer as these should not be included in the desired permutation equivariance. We assume that the input neurons come in some pre-fixed permutation, which is the case in typical learning settings. We (and prior work) are then only concerned with hidden neuron permutations. However, input/output permutations could also be trivially included if desired (by not exposing their neuron indices).

Wider networks (meaning more hidden units) are possible as the network cannot differentiate between any two neurons in the same hidden layer. This becomes a standard(-ish) application of GNNs generalizing to novel in-distribution graphs. We verify this empirically in our paper.

Using "layer number" might be problematic too

Including the layer number is not necessary, but is helpful for empirical performance. Residual edges can be encoded using the layer index of the outgoing nodes (each residual connection certainly has an output neuron generating it). Layer number is identical for all neurons within a single layer, and thus these cannot be distinguished by the GNN, and equivariance is retained. For what it's worth, this same general idea was independently proposed in the concurrent work of Zhang et al., where learned positional encodings are used as features to identify each layer. We found that the simpler layer index was sufficient.

We agree that generalization to deeper networks is an interesting question. There are several ways that this could be tackled. For example, one could try encoding the layer index as a continuous value from 0->1 with 0 for the input layer and 1 for the output layer. Or indeed, you could remove layer index entirely as a feature and rely on the GNN to learn this via message passing.

There are "layer type" features, however what kind of types are predefined is unclear.

We include layer type features for all unique layer types that we can represent (e.g. linear, convolution, multi-head attention, etc.). The features are simply categorical variables, with each layer type (in some predefined order) being given values in increments of 1. More sophisticated choices could be explored.

If a novel layer type was provided at inference time, then we would not generally expect our metanets to generalize well when these features are being used. But this is intuitive, as the network was never trained on networks with these layer types. However, one could forego the layer-type features entirely --- paying the cost of reduced expressiveness for improved generalization.

We agree that further empirical investigation would help to build an increased understanding of our metanets. But we would like to emphasize again that we have presented a novel and effective approach to designing equivariant metanetworks. While we will continue to strengthen our empirical evaluation by adding ablations and other improvements, the existing results show consistently high performance of our method on several tasks, alongside valuable new capabilities that other methods cannot achieve: generalization to different input architectures and handling diverse and heterogeneous input networks.

“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.

“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.

“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.

“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.

“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.

“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$.

“They address most of my concerns I had while reading, including some tests to compare against the newer state of the art metanets mentioned that were excluded from most of the result due to not being able to represent certain types of layers. Notably a cited competing method ‘NFN’ was left out from this, though seems to be addressed in the appendix, and as it apparently deals exclusively with MLPs I believe it is fair not to compare with the proposed method.”

We appreciate the author’s positive review. We have now run additional experiments on simple input networks (simple CNNs and MLPs) to compare our method against both NFN and the more recent NFT [Zhou et al. 2023b]; see the general comment for the results. In summary, our Graph Metanetworks generally outperform or match existing metanetworks on these tasks with simple input networks. Further, as you mention, our Graph Metanetworks can process significantly more diverse and complex input networks than these methods.

We hope that this addresses any remaining concerns, and would appreciate the opportunity to answer any other questions that you have.