Equivariant Polynomial Functional Networks

Reply [2/2]

W6: Finally, the theoretical contributions need some improvements. For example, it is unclear what is the role of Proposition 4.2. and Theorem 4.3. (why are they important?). Additionally, as far as I understand Theorem 4.4. characterises all G-invariant polynomials (by the way, of what degree?). Is there a similar argument for G-equivariant polynomials? Q1. What is the degree of the resulting polynomial? Shouldn’t this be chosen as a hyperparameter by the user?

Answer to W6+Q1: Following the reviewers' suggestion, we have added more intuition and explanations to explicitly clarify the roles of Proposition 4.2 and Theorem 4.3. Specifically, Proposition 4.2 says that stable polynomial terms can be viewed as a generalization of the weights themselves, while Theorem 4.3 provides the reasoning behind the stability of these polynomials. Additionally, we have included another theorem from appendix (Theorem 4.4 in the revised version) to further discuss the linear dependence of stable polynomial terms, offering essential computational steps for calculating the equivariant and invariant layers through the weight-sharing mechanism.

It is important to note that Theorem 4.4 (renumbered as Theorem 4.5 in the revised version) does not characterize all $G$-invariant polynomials. Instead, it characterizes all $G$-invariant polynomials of the specific form presented in Equation (20) (now is Equation (12) in the revised version), which is a linear combination of stable polynomial terms. The degree of these polynomials is $L$, which is the maximal degree of the stable polynomial terms. A similar theorem for $G$-equivariant polynomial layers is provided in Appendix C.5.

W7: Only certain types of activation functions can be used. W8: It is hard to extend to other activation functions W9: It cannot operate on diverse architectures. I advise the authors to mention those clearly and discuss potential options to address them.

Answer: MAGEP-NFN is designed to handle any MLPs or CNNs architectures of any size. In general, our method is applicable to other architectures with different activation functions, provided that the symmetric group of the weight network is known. The idea is to use the weight-sharing mechanism to redetermine the stable polynomials and then calculate the constraints of the learnable parameters. Although the calculation may require adaptation, it remains feasible. These discussions were added to the limitation part of the paper.

Q2: Apart from the memory/runtime comparisons that this paper is missing, I think that the authors should include another baseline: “Scale Equivariant Meta Networks”, Kalogeropoulos et al., NeurIPS’24. This is going to be a fairer competitor as opposed to Monomial-NFN and it would be interesting to observe the performance vs runtime tradeoffs

Answer: We have compared our model with GNN[2] and ScaleGMN[3] for predicting CNN generalization, using HNP[1] as a reference. GNN shows a performance drop with separate activation subsets. While ScaleGMN greatly enhances performance on the Tanh subset, improvements on the ReLU subset are less substantial. In contrast, our model achieves an overall significant improvements on both datasets. We refer the reviewer to Table 15 in Appendix G of our revision for the results, and present them in Table 3 below as well for convenience.

Table 3: Performance comparison with Graph-based NFNs

References

[1] Zhou et al. Permutation Equivariant Neural Functionals. NeurIPS 2023

[2] Kofinas et al. Graph Neural Networks for Learning Equivariant Representations of Neural Networks. ICLR 2024

[3] Kalogeropoulos et al. Scale Equivariant Graph Metanetworks. NeurIPS 2024

Q3: I am wondering why in the editing experiment we don’t observe any performance improvement.

Answer: In the editing INRs task, the activation function used is the trigonometric sine function (sin), which inherently induces a flipping-sign group action $\mathcal{M}^{\pm 1}$. This group action is less diverse compared to the positive scaling group action $\mathcal{M}_n^{>0}$ of ReLU activation funcion. As a result, our model has access to less inductive information, which could explain the lack of observable performance improvement over the baselines. Nonetheless, our model achieves performance comparable to all the baselines and consistently matches or outperforms them, particularly sharing the best results with the NP model.

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

Reply [1/2]

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

W1: The methodological part (Section 4) is quite hard to follow and the information provided is highly technical. [...]

Answer: Following the reviewers' suggestions, we have significantly revised Section 4 to include additional intuitions and insights, along with simpler explanations throughout the process of deriving invariant and equivariant layers. We believe that these updates have made the section more readable and informative.

W2: Therefore, I am afraid that this problem might be reflected in reproducibility/extensibility, as it will be hard for the interested reader to adopt/re-implement this method, and even harder to deeply understand the structure of the resulting layers.

Answer: We have included an Implementation section in the Appendix F, featuring pseudocode designed for better clarity. Our MAGEP-NFN's architecture incorporates several components expressed using concise and adaptable einsum operators. Moreover, the supplementary materials contain code for implementing MAGEP-NFN layers, along with comprehensive documentation to support accurate reproduction of the results.

W3: I would recommend the authors invest more in explanations and intuition about each part of their methodological section, most importantly in explaining the terms involved in the equivariant/invariant layer equations.

Answer: Following reviewers' suggestions, we have revised Section 4 by adding more explanations and intuition for every part of this section. We also explain the terms involved in the equivariant and invariant equations.

W4: First, they argue that their method improves expressivity compared to Monomial-NFN. Intuitively, probably this is indeed the case, but the authors have not provided any theoretical argument to support this claim.

Answer: We believe that the superior expressiveness of our equivariant polynomial layers compared to equivariant linear layers is a straightforward and self-evident fact. It is a similar situation to when we claim that polynomials are more general than linear functions. Consequently, we deemed it unnecessary to formalize this as a theorem. However, in response to the reviewers' suggestion, we have added a remark at the end of Section 4.2 to explicitly compare the equivariant and invariant polynomial layers proposed in our work with the equivariant and invariant linear layers presented in [4]. This remark emphasizes that our equivariant layers are more general and expressive than those in [4].

References

[4] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

W5: Additionally, the authors mention that GNNs have increased memory consumption and runtime. As far as I can see, this claim is also not supported with, at least, experimental evidence. I suggest the authors report the above metrics and compare those of their method to GNNs (metanetworks), NFN, Monomial-NFN, and ScaleGMN (see below), to improve our understanding of the trade-offs.

Answer: To compare computational and memory costs, we have included runtime and memory consumption for our model and previous ones in predicting CNN generalization task in the tables below. We have also include these results in Table 16 and 17 in Appendix H of our revision. For graph-based architectures, we compare with two recent works: GNN[2] and ScaleGMN[3]. Our model runs significantly faster and uses much less memory than these graph-based networks and NP/HNP[1]. Introducing additional polynomial terms slightly increases our model's runtime and memory usage compared to Monomial-NFN [4]. However, this trade-off results in considerably enhanced expressivity, which is evident across many tasks like Predict CNN Generalization or INRs Classification.

Table 1: Runtime of models

Table 2: Memory consumption

References

[1] Zhou et al. Permutation Equivariant Neural Functionals. NeurIPS 2023

[2] Kofinas et al. Graph Neural Networks for Learning Equivariant Representations of Neural Networks. ICLR 2024

[3] Kalogeropoulos et al. Scale Equivariant Graph Metanetworks. NeurIPS 2024

[4] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

Reply [2/2]

Q2: Related to question 1, in equations 20-21, is it necessary to sum over the whole range of $s$ and $t$ ? This seems like a very inefficient operation for very deep networks (e.g. ResNet 101). Isn't it possible to sum over a few adjacent layer indices? An ablation study on the number of hops would be useful here.

Answer: Yes, it is necessary to sum over s and t. It should be noted that many terms in Equation (21), which is Equation (15) in the revised version, vanish after applying the weight-sharing process. Ultimately, the additional polynomial terms in the final equivariant and invariant layers are not excessive. To validate this observation, we have included a comparison of the computational and memory costs of our models with those of linear methods (and graph-based approaches as well) in the response to your next question. Additionally, an ablation study examining the necessity of these components in the equivariant and invariant layers has been provided in the response to your earlier concern.

Q3: What is the computational complexity of the method compared to having linear layers?

Answer: To compare computational and memory costs, we've included runtime and memory consumption data for our model and previous ones in predicting CNN generalization task in the tables below. We have also include these results in Table 16 and 17 in Appendix H of our revision. For graph-based architectures, we compare with two recent works: GNN[4] and ScaleGMN[5]. Our model runs significantly faster and uses much less memory than these graph-based networks and NP/HNP[2]. Introducing additional polynomial terms slightly increases our model's runtime and memory usage compared to Monomial-NFN [4]. However, this trade-off results in considerably enhanced expressivity, which is evident across many tasks like Predict CNN Generalization or INRs Classification.

Table 2: Runtime of models

Table 3: Memory consumption

References

[2] Zhou, Allan, et al. "Permutation equivariant neural functionals." NeurIPS 2024.

[3] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

[4] Kofinas et al. Graph Neural Networks for Learning Equivariant Representations of Neural Networks. ICLR 2024.

[5] Kalogeropoulos et al. Scale Equivariant Graph Metanetworks. NeurIPS 2024.

Q4: Can the method work in a combined Small CNN Zoo dataset with both ReLU and tanh activations (and potentially other similar groups)?

Answer: No, our method is specifically designed to work with architectures that share the same weight space and symmetry group. CNNs with ReLU activations and CNNs with tanh activations have different symmetry groups, making them incompatible within the same framework.

Q5: In Eq. 15, shouldn't it be $\pi_{i}(j)$ instead of $\pi^{-1}_{i}(j)$

Answer: We believe that it should be $\pi^{-1}_{i}(j)$.

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

Reply [1/2]

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

Regarding the Ethics Concerns: Thank you for pointing out the similarity in the writing style between the introduction, related works, and background sections of our paper and those in [3]. While we have used different wording, we acknowledge that, as our work addresses the same problem as [3], the related works, definitions, terminologies, and formulas in the background sections inevitably overlap. These are fundamental to the problem and therefore remain consistent across both works.

To address your concern and reduce the perceived similarity, we have revised the writing style of these sections, removed some of the repeated content, particularly in the background, and cited to [3]. This ensures that our paper remains distinct while still providing all necessary technical details for reproducibility and clarity.

References

[3] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

W1: The numbers reported for NFN in the INR classification experiments are considerably worse than the official ones. Namely, NFN scores in CIFAR10, on MNIST and on FashionMNIST. This also means that the performance of MAGEP-NFN is quite poor. According to the appendix E, the authors use the same datasets as Zhou et al. Is there a different setup used in the experiments?

Answer: For the INRs classification task, we follow the same settings as in Monomial-NFN[3], which utilized the dataset from [2]. The key difference between [2] and [3] is that [2] augments the data tenfold, resulting in 450,000 training instances, while [3] does not use augmentation. By not augmenting the dataset, we enable a fairer comparison, highlighting architectures that can inherently handle transformations in weight space without relying on additional data.

References

[2] Zhou, Allan, et al. Permutation equivariant neural functionals. NeurIPS 2024.

[3] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

W2.1: Almost no explanation is given for equations 20 and 21, that comprise the core functionality of the proposed method. These equations are very math-heavy, [...]

Answer: Following the reviewers' suggestion, we have revised the entirety of Section 4, incorporating additional intuition and explanations to make the stable polynomial terms and their applications in constructing equivariant and invariant polynomial layers in Equations (20)-(21), which are now Equations (12) and (15) in the revised version, easier to understand.

W2.2: Finally, it is unclear if all the individual components contribute to the effectiveness of the method. An ablation study on those components could make the method more efficient, without sacrificing accuracy

Answer: We have conducted an ablation study to evaluate the significance of the components introduced in our work. Specifically, we categorize the terms as follows:

• Non Inter-Layer Terms: These are terms that involve only the mapping of Non-inter-layer weights and biases, $(W^l, b^l)^{l=1,…,L}$, to the output weight space, consistent with prior works (DWSNet [1], NP[2], HNP [2], Monomial-NFN[3]) on neural functional network layers.

• Inter-Layer Terms: These are the novel terms introduced in our paper, $[W],[WW],[bW],[Wb]$, designed to capture relationships between weights and biases across multiple layers.

To assess their impact, we perform experiments on the invariant task of predicting CNN generalization on ReLU subset, using the architecture specified in Equation (21). The results of our experiments are presented in the table below. We have also include these results in Table 18 in Appendix I of our revision.

Table 1: Ablation study to evaluate the significant of each component

We can see that each newly introduced Inter-Layer terms provide additional information to the network, and when combined with the Non Inter-Layer terms, the performance is boosted considerably.

Q1: What does the outer sum $\sum_{L \geq s > t \geq 0}$ in equation 20 denote? Is it a double sum over s and t

Answer: Yes, it is a double sum over s and t with the constraint $L \geq s > t \geq 0$.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

We appreciate the reviewer’s helpful suggestion. Regarding the concern about plagiarism, while we acknowledge that certain background sections in the initial pages of our paper share some overlap with [Tran2024], we respectfully disagree that this constitutes plagiarism. Our paper spans 59 pages, introducing novel ideas supported by comprehensive mathematical proofs.

As our work addresses the same problem within a similar context as [Tran2024], a degree of overlap in sections such as the introduction (Section 1), related work (Section 2), and preliminaries (Section 3) is unavoidable. However, these overlaps are minor and form a proportionally insignificant part of the overall content, which is predominantly original.

Since academic integrity is just as important as technical integrity, as you mentioned, it is important to approach allegations of plagiarism with care. We respectfully request that the reviewer focuses on the paper's novel contributions and original ideas and considers the plagiarism concern to be a misunderstanding.

Reference

[Tran2024] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

W1: [...] While the paper emphasizes that MAGEP-NFN enjoys superior runtime and memory efficiency, there are no comparisons of these quantities for MAGEP-NFN and relevant baselines [...]. It would be good to add a runtime and memory efficiency comparison to the paper which in my opinion would make the paper stronger

Answer: To compare computational and memory costs, we have included runtime and memory consumption for our model and previous ones in predicting CNN generalization task in the tables below. We have also include these results in Table 16 and 17 in Appendix H of our revision. For graph-based architectures, we compare with two recent works: GNN[3] and ScaleGMN[4]. Our model runs significantly faster and uses much less memory than these graph-based networks and NP/HNP[2]. Introducing additional polynomial terms slightly increases our model's runtime and memory usage compared to Monomial-NFN [4]. However, this trade-off results in considerably enhanced expressivity, which is evident across many tasks like Predict CNN Generalization or INRs Classification.

Table 1: Runtime of models

Table 2: Memory consumption

References

[2] Zhou et al. Permutation Equivariant Neural Functionals. NeurIPS 2023

[3] Kofinas et al. Graph Neural Networks for Learning Equivariant Representations of Neural Networks. ICLR 2024

[4] Kalogeropoulos et al. Scale Equivariant Graph Metanetworks. NeurIPS 2024

[5] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

W2: Furthermore, while the theoretical contribution is nice and sound, MAGEP-NFN does not seem to bring substantial empirical improvement over relevant baselines.

Answer: Our model shows some significant improvement in the task Predict CNN Generalization and Classify INRs (See Table 1, 2). In the ReLU subset of Small CNN Zoo, while Monomial-NFN only observes the performance gain when augmenting the dataset, our method surpass the best baseline HNP by a large margin (0.933 versus 0.926). Moreover, since most recent graph-based baseline ScaleGMN - with much more training time and memory consumption - reach the number 0.928 in this task (please refer to Q2 in General Response), this further highlights the performance gain of our method. On the task Classify INRs, our model surpass all baselines, with the gap of 7.7% in MNIST, and 2.9% in CIFAR-10.

Q1: How would one apply this approach to methods that support arbitrary computation graphs such as [1]?

Answer: MAGEP-NFN is designed to handle any MLPs or CNNs architectures of any size. In general, our method is applicable to other architectures, provided that the symmetric group of the weight network is known. The idea is to use the weight-sharing mechanism to redetermine the stable polynomials and then calculate the constraints of the learnable parameters. We hope this addresses your question. If not, please let us know what we can provide in order to have a more comprehensive response.

References

[1] Zhou et al. Universal neural functionals.

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.

Thank you for the added response and additional results. Please make sure they are properly incorporated in the paper. I will keep my current score.

Reply [1/4]

Thank you for your thoughtful review and valuable feedback. Below we address your concerns.

W1: The paper claims "low memory consumption and efficient running time", however these claims are not supported by empirical evidence. The previous paper (Tran et al., 2024) shows that it can reduce memory from 6GB to 0.5GB, however, the extra (stable) polynomial terms may increase memory/runtime of the method in this paper

Answer: To compare computational and memory costs, we have included runtime and memory consumption for our model and previous ones in predicting CNN generalization task in the tables below. We have also include these results in Table 16 and 17 in Appendix H of our revision. For graph-based architectures, we compare with two recent works: GNN[2] and ScaleGMN[3]. Our model runs significantly faster and uses much less memory than these graph-based networks and NP/HNP[1]. Introducing additional polynomial terms slightly increases our model's runtime and memory usage compared to Monomial-NFN [4]. However, this trade-off results in considerably enhanced expressivity, which is evident across many tasks like Predict CNN Generalization or INRs Classification.

Table 1: Runtime of models

Table 2: Memory consumption

References

[1] Zhou et al. Permutation Equivariant Neural Functionals. NeurIPS 2023

[2] Kofinas et al. Graph Neural Networks for Learning Equivariant Representations of Neural Networks. ICLR 2024

[3] Kalogeropoulos et al. Scale Equivariant Graph Metanetworks. NeurIPS 2024

[4] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

W2: No comparison to graph-based approaches (Lim et al.,2023; Kofinas et al., 2024) in Table 1-4 (ideally in terms of memory/efficiency vs results), so the significance of the results is hard to estimate

Answer: We compare our model with GNN[2] and ScaleGMN[3] for predicting CNN generalization, using HNP[1] as a reference. GNN shows a performance drop with separate activation subsets. While ScaleGMN greatly enhances performance on the Tanh subset, improvements on the ReLU subset are less substantial. In contrast, our model achieves an overall significant improvements on both datasets. We refer the reviewer to Table 15 in Appendix G of our revision for the results, and present them in Table 3 below as well for convenience.

Table 3: Performance comparison with Graph-based NFNs

References

[1] Zhou et al. Permutation Equivariant Neural Functionals. NeurIPS 2023

[2] Kofinas et al. Graph Neural Networks for Learning Equivariant Representations of Neural Networks. ICLR 2024

[3] Kalogeropoulos et al. Scale Equivariant Graph Metanetworks. NeurIPS 2024

W3: Some inconsistencies in evaluation, which is not explained, making comparison to previous literature difficult. [...]

Answer: For the INRs classification task, we follow the same settings as in Monomial-NFN[4], which utilized the dataset from [1]. The key difference between [1] and [4] is that [1] augments the data tenfold, resulting in 450,000 training instances, while [4] does not use augmentation. By not augmenting the dataset, we enable a fairer comparison, highlighting architectures that can inherently handle transformations in weight space without relying on additional data.

In the CNN generalization task, we evaluate all models on two subsets using ReLU and Tanh activations, while [2] uses the entire dataset for evaluation. Importantly, [2] cannot handle scaling symmetries, leading to a performance drop reported in [4]: 0.897 for the ReLU subset and 0.893 for the Tanh subset.

References

[1] Zhou et al. Permutation Equivariant Neural Functionals. NeurIPS 2023

[2] Kofinas et al. Graph Neural Networks for Learning Equivariant Representations of Neural Networks. ICLR 2024

[3] Kalogeropoulos et al. Scale Equivariant Graph Metanetworks. NeurIPS 2024

[4] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

Reply [4/4]

W12: The paper could be more accessible, e.g. by having illustrations of the proposed vs. previous method.

Answer: We have added a remark in Section 4.2 to illustrate a comparison of the proposed method vs. the previous method. In particular, our invariant polynomial layer is derived from the parameter-sharing mechanism. In contrast, the invariant equivariant layer proposed in [4] is an ad hoc formulation and does not result from a parameter-sharing mechanism. Consequently, there is no direct relationship between our invariant layer and the invariant layer in [4].

However, the equivariant polynomial layer in our MAGEP-NFNs and the equivariant linear layer from [4] are related. Specifically, the equivariant layer in [4] is exactly the linear component of our equivariant polynomial layer. Due to the lengthy formulation and construction process, we have provided the details of the equivariant polynomial layers in Appendix B.

References

[4] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

Q1: How difficult is the implementation of the proposed nonlinear layers? Can the authors provide pseudo code?

Answer: We have included an Implementation section in the Appendix F, featuring pseudocode for better clarity. Our MAGEP-NFN's architecture incorporates several components expressed using concise and adaptable einsum operators. Moreover, the supplementary materials contain code for implementing MAGEP-NFN layers, along with comprehensive documentation to support accurate reproduction of the results.

Q2: Can the trained MAGEP-NFN be used on the wider networks (e.g. predict generalization for larger CNNs than in the Small CNN Zoo dataset or retraining the model would be necessary?

Answer: Yes, MAGEP-NFN is designed to handle any MLPs or CNNs architectures of any size. In general, our method is applicable to other architectures, provided that the symmetric group of the weight network is known. The idea is to use the weight-sharing mechanism to redetermine the stable polynomials and then calculate the constraints of the learnable parameters.

We hope we have cleared your concerns about our work. We have also revised our manuscript according to your comments, and we would appreciate it if we can get your further feedback at your earliest convenience.

Reply [3/4]

W10: In section 4 it's not described how the equations 20-21 are different from the previous work (Tran 2024) or section 3. It would be more clear to directly contrast nonlinear and linear layers in this context to understand the contribution better.

Answer: Equations (20)-(21), which are Equations (12) and (15) in the revised version, describe the invariant polynomial layer derived from the parameter-sharing mechanism of our MAGEP-NFNs. In contrast, the invariant equivariant layer proposed in [4] is an ad hoc formulation and does not result from a parameter-sharing mechanism. Consequently, there is no direct relationship between our invariant layer and the invariant layer in [4].

However, the equivariant polynomial layer in our MAGEP-NFNs and the equivariant linear layer from [4] are related. Specifically, the equivariant layer in [4] corresponds to the linear component of our equivariant polynomial layer. Due to the lengthy formulation and construction process, we have provided the details of the equivariant polynomial layers in the appendix. These discussions have been incorporated into the revised version of the paper.

References

[4] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

W11: Contribution 2 of the paper (linear independence) is overstated - there is no comprehensive study in the main text. In contrast, the graph-based approached handled diverse and complex architectures (such as transformers) more easily.

Answer: Regarding the linear indepedence of stable polynomials: The detailed analysis of the linear independence of stable polynomial terms, which is both extensive and technically intricate, is distributed between Section 4.1 in the main text and Appendix B. The primary result on linear independence is presented in Theorem B.6. In response to the reviewers' suggestions, we have included a simplified version of Theorem B.6 along with additional results related to stable polynomials in the main text.

Regarding comparison to graph-based approaches: It is worth noting that, although graph-based approaches can handle diverse and complex architectures with relative ease, they suffer from higher memory consumption and runtime. To provide envident for this fact, we have included runtime and memory consumption for our model and previous ones in predicting CNN generalization task in the tables below. We have also included these results in Table 16 and 17 in Appendix H of our revision. For graph-based architectures, we compare with two recent works: GNN[2] and ScaleGMN[3]. Our model runs significantly faster and uses much less memory than these graph-based networks and NP/HNP[1]. Introducing additional polynomial terms slightly increases our model's runtime and memory usage compared to Monomial-NFN [4]. However, this trade-off results in considerably enhanced expressivity, which is evident across many tasks like Predict CNN Generalization or INRs Classification.

Table 1: Runtime of models

Table 2: Memory consumption

References

[1] Zhou et al. Permutation Equivariant Neural Functionals. NeurIPS 2023

[2] Kofinas et al. Graph Neural Networks for Learning Equivariant Representations of Neural Networks. ICLR 2024

[3] Kalogeropoulos et al. Scale Equivariant Graph Metanetworks. NeurIPS 2024

[4] Tran et al. Monomial Matrix Group Equivariant Neural Functional Networks. NeurIPS 2024.

Reply [2/4]

W4: The shortcoming of the method is that it requires some manual design of the polynomials for each architecture, so it's unclear how to extend it to other architectures such as transformers or even if it's possible to use a trained MAGEP-NFN on the wider networks.

Answer: Our method is applicable to these architectures, provided that the symmetric group of the weight network is known. The main idea involves leveraging the weight-sharing mechanism to redefine the classes of "stable polynomials" and subsequently determine the constraints on the learnable parameters. While the calculations may need adaptation to maintain the new architecture's structure, they certainly remain feasible. As our paper focuses on enhancing the expressivity of Monomial-NFN while preserving its efficiency and accuracy within the architectures considered in the context of current NFNs in the literature, these mixed architectures fall outside the scope of this work.

W5: The experiments on INR and generalization prediction for simple small CNNs are not sufficient. Previous works (Lim et al.,2023; Kofinas et al., 2024) added more practical experiments with learning to optimize or diverse transformer architectures. Otherwise, the proposed approach appears to be overengineered for some specific toy cases.

Answer: For the task learning to optimize, the previous work [2] has not provided the official implementation in their github repository (https://github.com/mkofinas/neural-graphs) yet. We will include this benchmark when the code is available.

References

[2] Kofinas et al. Graph Neural Networks for Learning Equivariant Representations of Neural Networks. ICLR 2024

W6: "These specialized networks are known as neural functional networks (NFNs)(Zhouetal.,2024b)" is a questionable statement.

Answer: For clarity, we revised the sentence to "Neural functional networks (NFNs) [1] have recently gained prominence as specialized frameworks designed to process key aspects of DNNs, such as their weights, gradients, or sparsity masks, treating these as input data".

References

[1] Zhou et al. Permutation Equivariant Neural Functionals. NeurIPS 2023

W7: "graph-based equivariant NFNs suffer from high memory consumption and long running times" is an overstatement given that Tran et al., 2024 showed the memory consumption of those methods is at maximum 6GB and runtime at max 4.5 hours. So in practice, this does not seem to be a problem for any modern GPUs. Please rewrite or scale the experiments.

Answer: We agree with the reviewer that in terms of computational cost and runtime, modern GPUs is suffice to perform Graph-based NFNs experiments. What we want to mention in the paper is that when compared with other parameter sharing models, graph-based models fall behind in terms of runtime and memory consumption (Please refer to Table 1 and 2 in our manuscript). We will rewrite this statement to "Compared to graph-based models, parameter-sharing-based NFNs built upon equivariant linear layers exhibit lower memory consumption and faster running time.".

W8: Related works is not informative, because it just list existing papers without discussing how the current work is connected (improved, different, etc.) to them.

Answer: Thanks for your comment. We have updated the related works section in the revised version. The updated relates works section is now more informative with discussions on how MAGEP-NFN is connected, improved, and different to them.

W9: The whole Section 3, if I understand correctly, seems to be background (notation introduction, etc. from Tran 2024), and it's quite lengthy, which makes the amount of the original content in the paper quite small.

Answer: Section 3 indeed serves as the background section, providing the necessary notations, definitions, and propositions required for constructing stable polynomials and equivariant/invariant layers in Section 4. Due to the extensive number of required notations, definitions, and propositions, this section spans nearly three pages. Nevertherless, following the reviewers' suggestions, we have shortened this section to less than two pages and relocated some of this material to the appendix.

Also, we respectfully disagree with the reviewer that the amount of the original content in our paper is quite small. Note that our manuscript spans 58 pages, including detailed derivations, extensive theoretical discussions, and thorough empirical validations of our methods.

We would like to thank the reviewer again for your thoughtful reviews and valuable feedback.

We would appreciate it if you could let us know if our responses have addressed your concerns and whether you still have any other questions about our rebuttal.

We would be happy to do any follow-up discussion or address any additional comments.