Scale Equivariant Graph Metanetworks

Summary

The paper considers the emergent and fascinating field of learning over weight spaces, that is, neural nets that process other neural networks. The processing NN is referred to as a metanetwork. Previous approaches showcased the importance of accounting for the input NN’s symmetries by designing equivariant architectures. However, they were mainly focused on permutation symmetries. This paper proposes a GNN-based metanetwork, which is permutation and scale equivariant. The paper studies the expressiveness (in terms of simulating the forward and backward pass of the input NN) of the proposed arch. The proposed method is evaluated using several INRs and classifiers datasets.

Strengths

1. The paper deals with an important and timely problem of learning in deep weight spaces, and presents a novel architecture by incorporating scale and permutation equivariance. 2. The paper provides theoretical analysis and results regarding the expressive power of the proposed approach. 3. Empirical results show significant improvement over baseline methods.

Weaknesses

My main concern is the limited empirical evaluation and missing natural baselines. While the presented empirical results show significant improvement over baseline methods, at the current state of the learning over weight spaces literature, I would expect a more diverse, challenging, and comprehensive empirical evaluation. 1. The writing and formatting require enhancement and refinement—specifically, long sentences, many slashes, many footnotes, etc. Also, in my view, the proposed method is introduced too late in the paper (page 6). 2. A missing natural baseline is to use a permutation equivariant baseline like DWS/NFN or NG-GNN together with scaling data augmentations as in [1]. 3. The experimental section only considers invariant tasks. Additional experiments using equivariant tasks (e.g. INR editing, domain adaptation, etc.) would significantly strengthen the empirical study. 4. Some evaluation and comparison of runtime and memory consumption w.r.t. baselines would be beneficial. 5. Furthermore, adding experiments with larger input networks and diverse input architectures (like a varying number of layers) would again significantly strengthen the empirical study. 6. Also, adding some ablation regarding design choices would be beneficial. 7. Why are the results for DWS and INR2Vec missing in Table 1? 8. Minor: - Line 198: should be a^k. - Line 336: “transformations We” -> “transformations. We” References: [1] Improved Generalization of Weight Space Networks via Augmentations, ICML 2024.

Questions

1. Is it always feasible and relatively easy to design and implement either the canon or symm mappings for all activation functions? 2. The bidirectional version of the method achieves on-par performance as the non-bidirectional one, except for the Augmented CIFAR-10 dataset, where the performance is much worse. Could you provide some insights regarding this result? 3. Additionally, how does the ScaleGMNB and ScaleGMN compare in terms of runtime? 4. How do ScaleGMN and ScaleGMNB compare to GMN in terms of the number of trainable parameters and runtime? 5. Did you use any augmentation on the input NNs? 6. Are there any limitations on the choice of activations of the ScaleGMN network?

Rating

7

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

Yes.

Summary

This paper addresses the emerging and fascinating field of deep-weight space learning where neural nets used to process weights and biases of another deep model. The authors have introduced new methods based on GNN architecture called ScaleGMN and ScaleGMNB for Scale Equivariant Graph MetaNetworks. The latter is a bi-directional variant. Both approaches tackle the scale symmetries presented by the input neural model's activation functions. The authors claim the following contributions: - Extending the scope of metanetwork design from permutation to scaling symmetries. - Designing networks that are invariant or equivariant to scalar multiplication from arbitrary scaling groups. - Theoretical analysis of the expressive power of ScaleGMN. - Extensive empirical comparison with recent work on various datasets in the field of weight space learning.

Strengths

- With the growing number of works on Implicit Neural Representation (INR) and the increasing need to process neural networks, this paper tackles a crucial field and problem, advancing the area in a way that can benefit many practitioners. - The authors introduce a new structure that incorporates both permutation and scale symmetries by ensuring it is equivariant to scale and permutation. - The authors provide a theoretical analysis of the expressive power of the proposed approaches. Additionally, they show that ScaleGMN can simulate forward and backward passes of arbitrary inputs in the processed NN. - The empirical results show a significant improvement compared to recent works in the field of weight space learning.

Weaknesses

- The writing can be significantly improved, particularly by breaking down long sentences that are hard to follow. Additionally, the writing pace is somewhat slow. While this might be beneficial for the average reader, the authors allocate too much space to exposition and problem formulation. As a result, the presentation of the proposed methods begins relatively late in the paper. - The experimental section focuses only on invariant tasks, i.e. INR classification and NN generalization prediction. It would be interesting to see how well ScaleGMN and ScaleGMNB deal with equivariant tasks like the ones presented in [1,2] which are considered harder for weight space architectures. - The processed architectures are not diverse dealing only with small-sized feed-forward and CNN architectures. It would be interesting to see more diversity in the processed architectures like deeper nets, attention-based methods, etc. ------- [1] Equivariant Architectures for Learning in Deep Weight Spaces, Navon et al. [2] Permutation Equivariant Neural Functionals, Zhou et al.

Questions

- Why DWS and INR2VEC are missing in Table 1? (CIFAR-10 | Augmented CIFAR-10 experiments). - In common bi-directional architectures (e.g. LSTM) we see performance (runtime) degradation compared to non-bi-directional design is that the same for ScaleGMN and ScaleGMNB? What is the computational complexity of both methods? - Can ScaleGMN and ScaleGMNB handle input data with heterogeneous activation functions, i.e. one network with ReLU activations and another with only tanh activations? (as long as they respect Prop. 4.1) - The results for the Augmented CIFAR-10 experiments are odd. ScaleGMNB performs worse than all baselines except MLP, while in other experiments, it outperforms them. Do the authors have an explanation for this observation? - Adding a figure that illustrates ScaleGMN and ScaleGMNB architectures and their design w.r.t permutation and scale symmetries, would be beneficial.

Rating

6

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

See above.

Summary

This work develops new GNN-based metanetworks that are equivariant to scaling symmetries induced by nonlinearities in input neural networks. Their ScaleGMNs extend metanetworks, which are typically only permutation equivariant (if at all equivariant), to also account for other symmetries in input neural networks' parameters. The architecture is proved to be equivariant to the desired symmetries and also expressive in that it can simulate forward and backward passes of the input. Experiments show improvements over merely-permutation-equivariant metanetworks.

Strengths

1. Great writing. Nice introduction and related work, as well as good setup and notation in Section 3. 2. Nice theoretical results. That ScaleGMN is can express the forward and backward pass is a good way to check that its expressive power is not overly limited when adding the additional scaling equivariances. Also, there is an interesting discussion in Appendix A.2 on equivariant for bidirectional ScaleGMNs. 3. Large empirical improvements, especially on INR classification (without many of the unfair or expensive tricks that others use!), with ScaleGMN. I say some previous tricks are "unfair" because, for instance, random probe features of Kofinas et al. as used on INRs can essentially be used to see the input image pixels, and hence the prediction task is also taking in the image as well as the INR representing it. That ScaleGMN can beat the prior methods without these tricks is very impressive. 4. Several other interesting empirical findings. These include the fact that ScaleGMN does not need random Fourier features or data augmentations, and that the bidirectional version can vary in performance (sometimes drastically as in augmented CIFAR-10).

Weaknesses

1. It would be interesting to see how ScaleGMNs perform on different tasks, especially an equivariant task (rather than just invariant tasks) such as INR editing. 2. ScaleInv is oddly described. The equation on Page 6 should probably have $\tilde x_i$ or something similar as its arguments, instead of $x_i$ (because if the $\rho^k$ are just general MLPs as you say right after the equation, then this is not scale invariant). 3. At the end of Page 6 and beginning of Page 7, you say that sign canonicalization can only be used for dimension 1, but this is not quite true. In Ma et al. 2023 and Ma et al. 2024, algorithms are given for canonicalizing with respect to the sign group, for use on inputs to a neural network. References * [Ma et al. 2023] Laplacian Canonization: A Minimalist Approach to Sign and Basis Invariant Spectral Embedding. https://arxiv.org/abs/2310.18716 * [Ma et al. 2024] A Canonization Perspective on Invariant and Equivariant Learning. https://arxiv.org/abs/2405.18378

Questions

1. Do you have a way to handle translation symmetries in nonlinearities?

Rating

8

Confidence

4

Soundness

4

Presentation

4

Contribution

4

Limitations

Limitations are discussed on Page 9

References for Author Rebuttal

[1] Zhou, Allan, et al. "Permutation equivariant neural functionals." Advances in neural information processing systems 36 (2024). [2] Kofinas, Miltiadis , et al. "Graph neural networks for learning equivariant representations of neural networks." ICLR, 2024. [3] Navon, Aviv, et al. "Equivariant architectures for learning in deep weight spaces." International Conference on Machine Learning. PMLR, 2023. [4] Lim, Derek, et al. "Graph Metanetworks for Processing Diverse Neural Architectures." The Twelfth International Conference on Learning Representations. [5] Zhou, Allan, Chelsea Finn, and James Harrison. "Universal neural functionals." arXiv preprint arXiv:2402.05232 (2024). [6] Unterthiner, Thomas, et al. "Predicting neural network accuracy from weights." arXiv preprint arXiv:2002.11448 (2020). [7] Schürholt, Konstantin, et al. "Model zoos: A dataset of diverse populations of neural network models." Advances in Neural Information Processing Systems 35 (2022): 38134-38148.

Reviewer response

I would like to express my gratitude to the authors for the time and effort they have dedicated to this rebuttal. Most of my concerns have been addressed and after reading the other reviews and comments I decided to raise my score.