Scale Equivariant Graph Metanetworks

Weakness 1. Writing.

Please refer to our global response regarding the structure of our paper and the writing style.

Weakness 3. Equivariant tasks.

Please refer to our global response, where we provide details regarding the INR editing task as well as to Table 1 of the attached PDF.

Weakness 4. Runtime and memory consumption.

We thank the reviewer for suggesting a comparison of runtime and memory consumption. We select the F-MNIST dataset and make two comparisons; at first we report the number of parameters of all the reported models. Regarding the runtime, we fix the number of parameters for fairness of comparison and report the training time, inference time and GPU memory consumption. Please refer to Table 2 of the attached PDF. We can see that ScaleGMN-B does not introduce performance degradation regarding the runtime, while both our methods are quiet slower than the baselines. Since we do not use computationally much heavier operations, this last result indicates that our implementation could be further optimized w.r.t. runtime.

Weakness 5. Larger and diverse input architectures.

Please refer to our global response regarding the datasets and experiments on larger and more complex architectures.

Weakness 7. DWS and INR2VEC on CIFAR-10

Please refer to our response to reviewer x2pG regarding these two experiments.

Question 3. Runtime performance of ScaleGMN and ScaleGMN-B.

Please refer to Table 2 of the attached PDF. As discussed in Weakness 4, ScaleGMN-B does not introduce performance degradation regarding the runtime.

For weaknesses: 1,3,4,5,7 and question 3 please refer to Comment.

Weaknesses

Weakness 2. Scaling data augmentations

An effective method to get such augmentations is to sample scaling matrices for every training datapoint - diagonal sign matrices for sign and diagonal positive scaling matrices for positive scale symmetries, respectively. Subsequently, we transform the input network parameters by applying the matrices to the hidden layers. (like in Eq.3 in our paper, but without the permutation matrices)

The way to perform this random sampling, however, is not straightforward and we have to deal the two cases separately.

Sign symmetries: Since the diagonal sign matrices are of the form $\mathbf{Q} = \text{diag}(q_1 , \dots , q_d ), q_i = ±1$, we sample each element of the matrix independently and uniformly at random with probability 0.5. We observe that augmenting the training set leads consistently to better results when compared to the original baselines. None of these methods however achieved results on par with ScaleGMN and ScaleGMN-B. (Table 3 in the PDF)

Positive-scale symmetries: The diagonal positive scaling matrices are of the form: $\mathbf{Q} = \text{diag}(q_1 , \dots , q_d ), q_i \in (0, +\infty)$. Hence, we have to sample from a continuous and unbounded distribution making the augmentation strategy much more difficult to design. Consulting the plots in the appendix A.5, we opted for an exponential distribution and experiment with the coefficient $\lambda$. Nevertheless, regardless of the distribution choice we cannot guarantee that the augmentations will be sufficient, due to the lack of upper bound. We observe that we were not able to surpass the original baselines, which indicates that designing an effective baseline of this type is not straightforward. (Table 3 in the PDF)

We thank the reviewer for proposing this important baseline.

Questions

Question 1. Canon/symm mappings for all activation functions.

In general, we cannot guarantee that these mappings are easy to construct for any activation function, since it is currently unknown if we can provide a general characterisation of the possible symmetries that may arise. Our results extend to all positively homogeneous activations, $\sigma(\lambda x) = \lambda \sigma(x), \lambda > 0$, and all odd ones, $\sigma(-x) = - \sigma(x)$. We refer the reviewer to Table 1 of [1], where we can see that LeakyReLU also falls in the first category. Regarding polynomial activations, which demonstrate non-zero scaling symmetries, one option would be: (1) norm division to canonicalise scale, and (2) sign symm/canon as discussed in the main paper (and the rebuttal). The above, cover a fairly broad spectrum of common activation functions.

Question 2. Forward vs Bidirectional and Augmented CIFAR-10.

Indeed, in most invariant tasks, the bidirectional variant does not provide significant advantages. This is inline with what has been observed in [2], where a forward variant is also sufficient. We speculate that this might be related to the fact that the tasks we considered may be solved by simulating the forward pass of the input NN alone. Given that the forward pass can be expressed by the forward variant alone, this might provide an explanation for its efficacy on the invariant tasks we considered. The significance of the bidirectional variant is highlighted mostly on equivariant tasks, where ScaleGMN-B significantly outperforms the forward variant. Please refer to the global response, where we report our results on INR-editing and discuss the limitations of the forward variant on the equivariant tasks.

Please refer to our response to Question 4 of reviewer x2pG for ScaleGMN-B on the Augmented CIFAR-10.

Question 4. Trainable parameters: GMN vs ScaleGMN(-B)?

Going from a GMN model to ScaleGMN requires adapting the MSG and UPD functions to be scale equivariant, leading to more learnable parameters, as opposed to using plain MLPs. Another design choice of ScaleGMN that introduces more parameters is using different MLPs for the I/O nodes (Appendix A.1.4).

Please refer to Table 2, where we report the training and inference runtimes.

Question 5. Did you use any augmentation on the input NNs? No, we do not use any augmentation on the input NNs.

Augmentations are only used for the dataset "Augmented CIFAR-10", where we follow the augmentation procedure from [3] for comparison reasons with the rest of the baselines. ScaleGMN and ScaleGMN-B rely solely on the original training dataset and on built-in equivariance/invariance.

Question 6. Are there any limitations on the choice of activations of the ScaleGMN network?

Importantly, our method does not impose any limitations on the choice of the activation functions.

We are able to select any activation, because these are only applied within the MLPs of the invariant modules. As discussed in Section 5 of our paper, the MLPs (equipped with non-linearities) are only applied after the canon/symm function. In case one chose to place activations in a different computational part of ScaleGMN, this would indeed limit their options so as not to compromise scale equivariance. However, this is not the case in our method. We thank the reviewer for noticing this important detail.

---

[1] Godfrey, Charles, et al. "On the symmetries of deep learning models and their internal representations." Advances in Neural Information Processing Systems 35 (2022): 11893-11905.

[2] Kofinas, Miltiadis, et al. "Graph Neural Networks for Learning Equivariant Representations of Neural Networks." The Twelfth International Conference on Learning Representations.

[3] Zhou, Allan, et al. "Permutation equivariant neural functionals." Advances in neural information processing systems 36 (2024).

[4] Shamsian, Aviv, et al. "Improved generalization of weight space networks via augmentations." arXiv preprint arXiv:2402.04081 (2024).

Weaknesses

Weakness 1. Writing Style

We thank the reviewer for their suggestions. Please refer to the global response about the writing pace and style.

Weakness 2. Equivariant tasks

To evaluate the performance of our method on tasks that require permutation and scale equivariance, we selected the task of INR editing. Please refer to our global response and to Table 1 of the attached PDF. In summary, the bidirectional variant, ScaleGMN-B, surpasses all baseline metanetworks, including GNNs using additional information (probe features)

Weakness 3. Processing diverse and complex architectures

Please refer to our global response regarding the experiments on larger and more complex architectures.

Questions

Question 1. DWS and INR2VEC on CIFAR-10

In Table 1 of our paper we include, together with our results, all the baselines from the literature. For the CIFAR-10 dataset we did not include any experiments with DWS and INR2VEC, as they were not included in [1], where the task was proposed. However, we have now run both experiments using the DWS model, which achieved $34.45 \pm 0.42$ on CIFAR-10 and $41.27 \pm 0.026$ on Augmented CIFAR-10. As expected, these results are on par with the rest of the permutation equivariant models.

Regarding INR2VEC [2], we have not reimplemented it and tested on the CIFAR-10 dataset - we only used results present in the literature. However, INR2VEC demonstrated significantly worse performance than DWS [3] and NFN [1] on the task of INR classification, as can be found in the literature, and therefore testing it on extra data will probably offer limited added value, as it. Nevertheless, we plan to also include this baseline in an updated version.

Question 2. Performance (runtime) degradation between ScaleGMN and ScaleGMNB and complexity

The reviewer here makes a correct comment regarding the complexity of the forward (ScaleGMN) and the bidirectional variant (ScaleGMN-B) of our method. In the latter case, we add backward edges to our graph and introduce a backward message function to discern the two message directions. Subsequently, we concatenate the outputs of the two message functions and apply the UPD function. Consequently, the additional complexity of ScaleGMN-B is introduced solely by the extra message function, with a complexity of $O(E)$. Given that ScaleGMN has the complexity of a standard GNN model $O(V+E)$, the final complexity of ScaleGMN-B is $O(V+2E)$.

To measure the differences regarding the runtime performance, we conduct a controlled experiment on the F-MNIST dataset. We observe that ScaleGMN-B only needs $6.83$ more seconds per epoch, when compared to ScaleGMN. Regarding the inference time this difference is $0,0132$ seconds per datapoint. Please refer to our response to reviewer ugM4 as well as to Table 2 in the PDF within our global response for the complete results.

Question 3. Heterogeneous activation functions.

This is an interesting question and of importance to ensure the generality of our method. In principle, our method does not impose any limitations regarding the homogeneity of the activation functions of the input neural networks. To see this, observe that the only part of the metanetwork that gets affected by the symmetry induced by the activation function, is the symmetrisation/canonicalisation function. In other words, all the modules of the metanetwork can be reused for any input NN with arbitrary activations, as long as the datapoints are symmetrised/canonicalised accordingly.

Experiments. Experimentally, we opted to split the datasets into two subsets (one with ReLU Nets and one with tanh Nets). This choice was made solely to evaluate our method separately for each type of symmetry and assess the different invariant methods that we employed for each case. Following, the reviewer's suggestion, we extend our evaluation to heterogeneous activation functions (a dataset containing both ReLU and tanh Nets). We conducted epxeriments on the CIFAR-10-GS dataset and report the results on Table 4 in the attached PDF of our global response. The baselines are reported as in [4]. Interestingly, we observe that ScaleGMN demonstrates superior performance compared to the previous baselines, significantly exceeding the performance of the next best model. We thank the reviewer for suggesting this experiment - we will include this in the updated version of the manuscript.

Question 4. The results of ScaleGMN-B on Augmented CIFAR-10.

This was indeed a confusing result, which was merely due to suboptimal hyperparameter search. Unfortunately, due to the large size of this dataset ($20$ times larger than "CIFAR-10") and limited computational resources, the hyperparameter search for the ScaleGMN-B on the Augmented CIFAR-10 dataset had not finished by the time of the submission. Hence, the reported result does not reflect the real performance of the model. We completed the hyperparameter search post-submission, and achieved accuracy equal to $56.95 \pm 0.57$ - this result follows the same pattern with the rest of the datasets.

Question 5. Figure of ScaleGMN and ScaleGMNB.

We thank the reviewer for suggesting to include a figure depicting our architectures. We will consider designing one and including it in an updated version of the manuscript.

---

[1] Zhou, Allan, et al. "Permutation equivariant neural functionals." Advances in neural information processing systems 36 (2024).

[2] De Luigi, Luca, et al. "Deep Learning on Implicit Neural Representations of Shapes." The Eleventh International Conference on Learning Representations.

[3] Navon, Aviv, et al. "Equivariant architectures for learning in deep weight spaces." International Conference on Machine Learning. PMLR, 2023.

[4] Kofinas, Miltiadis, et al. "Graph Neural Networks for Learning Equivariant Representations of Neural Networks." The Twelfth International Conference on Learning Representations.

Weaknesses

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

Metanetworks can indeed find interesting applications that require our model to be permutation and scale equivariant. Please refer to the global response for our experiments on INR editting and the corresponding results in Table 1 of the attached PDF. The bidirectional variant of our method, ScaleGMN-B, is able to surpass all the GNN-based metanetworks, even the ones using probe features.

---

Weakness 2. ScaleInv is oddly described. The equation on Page 6 should probably have or something similar as its arguments, instead of (because if the are just general MLPs as you say right after the equation, then this is not scale invariant).

Thank you for spotting this. Indeed, there is a typo in the definition of ScaleInv (below L265). The arguments of the function $\rho^k$ (universal approximators - MLPs) should have been $\tilde{\mathbf{x}}_i$ instead of $\mathbf{x}_i$, where $\tilde{\mathbf{x}}_i$ are explained later in the text (L277) and are the outputs of a canonicalisation or a symmetrisation function, i.e. $\tilde{\mathbf{x}}_i = \text{canon}(\mathbf{x}_i)$ or $\tilde{\mathbf{x}}_i = \text{symm}(\mathbf{x}_i)$, which ensures invariance. We will fix this in an updated version of the manuscript.

---

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

We are grateful to the reviewer for bringing up these two recent references that deal with sign canonicalisation - we were not aware of these works and, indeed, they are useful for our setup. Given the fact that symmetrisation introduces additional parameters (i.e. the internal MLP, see L276), we are interested in conducting experiments in the future with the proposed canonicalisation method, so as to examine if similar performance can be achieved with a reduced parameter count. Additionally, we will update our text accordingly to complement our discussion with this missing point.

Questions

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

Translation symmetries (such as those induced by the softmax activation) are an important next step in this research direction. Our method currently does not handle this case. A potentially straightforward modification might be to follow the same rationale with scale equivariant networks: first, define a translation invariant module via canonicalisation and second, use it to achieve equivariace (e.g. translate the input by the output of the invariant module). For example, for symmetries of the form $\mathbf{x}’ = \mathbf{x} + a$, where $a$ is a scalar, we can canonicalise as follows $\tilde{\mathbf{x}} = \mathbf{x} - \frac{1}{N} \sum_{i=1}^N x_i$.

We believe that characterizing even more nonlinearities (or families of them) and designing the respective invariant modules, is a prosperous future work towards implementing a unified framework able to handle various types of networks.