Equivariant Polynomial Functional Networks

Authors: Thieu Vo, Hoang V. Tran, Tho Tran Huu, An Nguyen The, Thanh Tran, Minh-Khoi Nguyen-Nhat, Duy-Tung Pham, Tan Minh Nguyen
Primary area: other topics in machine learning (i.e., none of the above)
Keywords: neural functional network, equivariant model, polynomial layer, monomial matrix group

Abstract:
Neural Functional Networks (NFNs) have gained increasing interest due to their wide range of applications, including extracting information from implicit representations of data, editing network weights, and evaluating policies. A key design principle of NFNs is their adherence to the permutation and scaling symmetries inherent in the connectionist structure of the input neural networks.  Recent NFNs have been proposed with permutation and scaling equivariance based on either graph-based message-passing mechanisms or parameter-sharing mechanisms. However, graph-based equivariant NFNs suffer from high memory consumption and long running times. On the other hand, parameter-sharing-based NFNs built upon equivariant linear layers exhibit lower memory consumption and faster running time, yet their expressivity is limited due to the large size of the symmetric group of the input neural networks. The challenge of designing a permutation and scaling equivariant NFN that maintains low memory consumption and running time while preserving expressivity remains unresolved. In this paper, we propose a novel solution with the development of MAGEP-NFN (**M**onomial m**A**trix **G**roup **E**quivariant **P**olynomial **NFN**). Our approach follows the parameter-sharing mechanism but differs from previous works by constructing a nonlinear equivariant layer represented as a polynomial in the input weights. This polynomial formulation enables us to incorporate additional relationships between weights from different input hidden layers, enhancing the model's expressivity while keeping memory consumption and running time low, thereby addressing the aforementioned challenge. We provide empirical evidence demonstrating that MAGEP-NFN achieves competitive performance and efficiency compared to existing baselines.

Decision: reject
Reviewer scores: 6, 8, 5, 3

Reviews and discussion:
<thread>
creation date:
1730429453090

thread_type:
official_review

<thread_object>
creator: reviewer_zvnr

<content>
summary:
This work proposes a novel Neural Functional Network (NFN), i.e. a neural network that takes the weights of neural networks as input data. The method, termed MAGEP-NFN, belongs to the family of parameter-sharing NFNs and incorporates both the permutation symmetries of neural networks and the scaling symmetries of common non-linearities, such as ReLU and tanh. The proposed method alleviates the reduced expressivity of previous works by introducing non-linear equivariant layers through polynomials of the input weights and biases. MAGEP-NFN achieves competitive results in INR tasks and CNN generalization tasks, while having far fewer parameters than the compared baselines.

strengths:
The manuscript is fairly clear and easy to follow, despite being math heavy.

NFNs are a very promising research direction, and making them faster and more efficient is key to scaling them up to more real-world applications.
Hence, the fact that the proposed method can reach good performance with a reduced number of parameters
compared to previous works is very promising, and a significant step forward.

weaknesses:
The numbers reported for NFN in the INR classification experiments are considerably worse than
the official ones. Namely, NFN scores $46.6 \pm 0.072$ in CIFAR10, $92.9 \pm 0.218$ on
MNIST and $75.6 \pm 1.07$ 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?

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, and it is very hard for the reader to understand the usefulness of all individual components.
Further explanation of each module could provide more intuition about the method. Further, both formulas could benefit from some visualizations that show the connections between layers and the flow of information. 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.

questions:
[1] 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$?

[2] 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.

[3] What is the computational complexity of the method compared to having linear layers?

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

**Minor questions**:

[5] In Eq. 15, shouldn't it be $\pi_{i}(j)$ instead of $\pi_{i}^{-1}(j)$?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730474981653

thread_type:
official_review

<thread_object>
creator: reviewer_4jxm

<content>
summary:
This paper follows a recent line of work on *neural functional networks (NFNs)*, i.e. neural networks that are specifically designed to process other neural networks. In specific, the first NFNs are composed of linear layers (interleaved with non-linearities) that are equivariant to the *permutation* symmetries of the input NNs, while a recent follow-up called Monomial-NFN follows the same strategy, but additionally accounting for the *scaling* symmetries. However, the latter turned out to suffer from limited expressivity, due to the extensive parameter sharing that resulted from the existing symmetries. 

To overcome this limitation, the authors in the present manuscript propose a Polynomial-NFN, i.e. a combination of polynomial layers (potentially interleaved with non-linearities) that are equivariant to permutation + scaling symmetries. The paper mainly focuses on the characterisation of polynomial equivariant layers, using a construction named stable polynomials. Experimentally, the method is tested on a variety of common NFN benchmarks showing improved performance in the majority of the cases.

strengths:
- **Significance and Motivation**. Similarly to many equivariant machine learning tasks, designing models that are equivariant to parameter symmetries, typically requires dealing with a computational complexity vs expressivity dilemma. The authors correctly identify this dilemma, promptly (since NFNs/metanetworks are increasingly becoming more popular), and with their approach, they attempt to find a better trade-off.
- **Novelty**. This is the first work that proposes polynomials equivariant to parameter symmetries and may potentially mark a first step towards improving this paradigm as an alternative to linear layers + non-linearities, as well as GNNs.
- **Empirical performance**. In practice, the improvement in expressivity is shown to be reflected in the downstream performance in most of the tested cases.

weaknesses:
- **Presentation and Exposition**. My major concern with this paper is that the methodological/technical contributions are not well-presented. In particular,
   - The methodological part (Section 4) is quite hard to follow and the information provided is highly technical. This is not necessarily an issue, but currently, the authors do not provide any intuition/insights or simpler explanations of their thought process (why did they choose this proof technique, i.e. via stable polynomials), derivations (e.g. a proof sketch) and resulting equivariant/invariant layers (why do they look like that). 
   - 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. 
   - 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.


- **Claims and Theory**. Another major concern is that the authors have made a few claims that are not adequately supported with evidence:
   - 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.
   - 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. Additionally, a computational complexity table with comparisons would help a lot.
   - 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?


- **Method**. Additionally, examining the method per se, even if it improves expressivity, I believe it inherits several of the limitations of Monomial-NFN:
   - Only certain types of activation functions can be used.
   - It is hard to extend to other activation functions
   - It cannot operate on diverse architectures.
I advise the authors to mention those clearly and discuss potential options to address them.

- **Experiments**.
   - 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.
   - I am wondering why in the editing experiment we don’t observe any performance improvement.

questions:
- Due to the lack of insights in the methodological section, several questions arise. For instance,
   - What is the degree of the resulting polynomial? Shouldn’t this be chosen as a hyperparameter by the user?
   - I assume that for different activation functions (or more precisely for different symmetries induced by the activation function) the resulting layers should be different. However, this is not clear to me from the Equations 20 and 21. Could the authors explain what are the differences?
   - What are the matrices $\Phi$ and $\Psi$? They are kind of abruptly introduced.


**Minor:**
- **Related work and background**:  I have the impression that these two sections (2 and 3) are quite similar to those of the Monomial-NFN paper. Especially, regarding the background section, I felt it was too long, which might not be necessary, since the reader can be pointed to Monomial-NFN or ScaleGMN to find a detailed explanation of the background.
- The paper misses an important citation: “Equivariant Polynomials for Graph Neural Networks”, Puny et al., ICML’23. This work discusses equivariant polynomials for permutation symmetries present in graphs.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730592598912

thread_type:
official_review

<thread_object>
creator: reviewer_c8vp

<content>
summary:
This paper proposes monomial matrix group equivariant polynomial neural functional networks (MAGEP-NFN), a new class of neural functionals whose input and output are the weights of neural networks. MAGEP-NFN builds on a prior work Monomial NFN which adds scale equivariance and sign equivariance of the original NFN which is only equivariant to neuron permutations. MAGEP-NDF introduces *stable polynomial terms* which are the product of a series of adjacent weights and biases of a neural networks that are also equivariant under different neuron permutations. These stable polynomial terms can be thought of as some "features" of the original weights, similar to how polynomials are used as features in kernels. MAGEP-NFN layers can process these stable polynomials in addition to the original weights. Experimental results demonstrate that MAGEP-NFNs outperform previous NFNs that use weight-sharing schemes in terms of generalization prediction and perform similarly on INR editing tasks.

strengths:
This paper brings both theoretical and empirical contributions to NFNs.
The techniques proposed in the paper are well-motivated and rigorously derived.
Experiments and mathematical details are quite thorough, especially given how mathematically involved these approaches tend to be (I was able to follow the derivation reasonably well).

Overall, I believe MAGEP-NFN provides a concrete improvement in the line of work on NFNs and points towards a direction to improve NFNs further.

weaknesses:
The paper makes some unsubstantiated claims about its advantages over prior methods and GNN-based methods in terms of runtime and memory consumption. Specifically, 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 (For example, it is not clear to me that MAGEP-NFN should be more efficient than vanilla NFN). It would be good to add a runtime and memory efficiency comparison to the paper which in my opinion would make the paper stronger. 

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

**Minor**
- Use `` and '' for quotations.

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



**Reference**

[1] Universal neural functionals. Zhou et al.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730833723667

thread_type:
official_review

<thread_object>
creator: reviewer_jcc7

<content>
summary:
The paper proposes a new method to model neural network weights based on Equivariant Neural Functional Networks. Compared to previous approaches, this method is claimed to be more expressive while still efficient. This is achieved by adding nonlinear (polynomial) equivariant layers acting on weights. The paper presents an extensive theoretical framework and shows improved empirical performance on the INR-based tasks and generalization prediction.

strengths:
1. The paper tackles a challenging problem of modeling neural network weights and their symmetries.
2. The proposed nonlinear equivariant layers appear to be original and interesting. 
3. The paper seems theoretically sound, however, I haven't checked the correctness.
4. The results are better than for the linear polynomials (Monomial-NFN).

weaknesses:
The paper's main weaknesses are experiments and writing/presentation.

### The experiments are weak in several aspects:
1. 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. 
2. 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.
3. Some inconsistencies in evaluation, which is not explained, making comparison to previous literature difficult. For example, in (Kofinas et al., 2024) for MNIST the INR classification performance was about 90-95% and for FashionMNIST 70-75%, while in this paper it's  77 and 62 respectively. On the Small CNN Zoo dataset (Unterthineretal.,2020) the method achieves 0.933, while graph-based approaches (Kofinas et al., 2024) achieve better results (0.935).
4. 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.
5. 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. 

### Writing and presentation
1. "These specialized networks are known as neural functional networks (NFNs)(Zhouetal.,2024b)" is a questionable statement.
2. "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.
2. Related works is not informative, because it just list existing papers without discussing how the current work is connected (improved, different, etc.) to them. 
3. 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.
4. 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.
5. 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. 
6. The paper could be more accessible, e.g. by having illustrations of the proposed vs previous method.

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

2. 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?
</content>
</thread_object>

</thread>
<thread>
creation date:
1732176739144

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General Response

comment:
Dear AC and Reviewers,

We greatly appreciate your thoughtful reviews and valuable feedback, which have significantly enhanced the quality of our paper. We are encouraged by your recognition of the following aspects: 1) Our work is acknowledged for contributing both theoretical insights and empirical advancements to neural functional networks (Reviewer C8VP), marking a significant step forward in this area (Reviewers 4jXM, C8VP, zVNR); 2) The paper addresses a novel and challenging problem of designing polynomials that are equivariant to parameter symmetries (Reviewers jcC7, 4jXM) and presents a well-motivated study (Reviewer C8VP). The theory is described as sound (Reviewer jcC7), with rigorous and clear derivations that are easy to follow (Reviewers zVNR, C8VP); 3) The experiments are noted for their thoroughness (Reviewer C8VP) and demonstrate promising outcomes (Reviewers jcC7, zVNR), including improved in expressivity (Reviewer 4jXM).

Below, we address some common points raised in the reviews:

**1. Regarding runtime and memory consumption:** 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*
||NP [1]| HNP [1]|GNN [2]|ScaleGMN[3]|Monomial-NFN[4]|MAGEP-NFN|
|-|-|-|-|-|-|-|
| Tanh subset | 35m34s | 29m37s  |4h25m17s |1h20m| **18m23s**|*28m12s*|
| ReLU subset | 36m40s | 30m06s |4h27m29s |1h20m| **23m47s**|*28m43s*|


*Table 2: Memory consumption*
|     | NP [1]| HNP[1]|GNN [2]|ScaleGMN[3]|Monomial-NFN[4]|MAGEP-NFN|
|-|-|-|-|-|-|-|
| Tanh subset | 838MB | 856MB |6390MB|2918MB| **582MB**|*584MB*|
| ReLU subset | 838MB | 856MB |6390MB|2918MB|**560MB**|*584MB*|

**2. Regarding comparison with more graph-based baselines:** 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*

|     |HNP[1]| GNN [2]|ScaleGMN[3]|MAGEP-NFN|
|-|:-:|:-:|:-:|:-:|
| ReLU subset |0.926|0.897|  *0.928* | **0.933**|
| Tanh subset |0.934|0.893 |**0.942** |*0.940*|

**3. Regarding experiments setup for INRs Classification and Predict CNN Generalization:** 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.

---

**Reference.**

[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 are glad to answer any further questions you have on our submission.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732186215779

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of Revisions

comment:
Incorporating feedback from reviewers, as well as some further informative empirical studies, we summarize the main changes in the revised version of our paper as follows:
1. We have rewritten the abstract and introduction for improved clarity, focusing particularly on the description of graph-based models and neural functional networks.
2. We have revised the writing style and wording of the abstract, related works, and introduction. Additionally, we removed some repeated content in the background section, as in [1], and included an extra citation for that work. These changes ensure our paper remains distinct while providing the necessary technical details for reproducibility and clarity.
3. 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
4. We have added a Memory and Runtime comparison in Appendix H. The results, presented in Tables 16 and 17, highlight the efficiency of our model compared to graph-based models ([2], [3]) and HNP [4]. 
5. We have conducted additional baseline experiments in Appendix G, comparing our model with recent graph-based baselines (GNN [2], ScaleGMN [3]) on the task of predicting CNN generalization on the ReLU subset. The results, presented in Table 15 (Appendix G), show that our model achieves considerable improvements across both datasets, slightly falling behind ScaleGMN on the Tanh subset but demonstrating a significant gap in the ReLU subset. 
6. We have conducted an additional ablation study in Appendix I to assess the contribution of our newly introduced components. The results are presented in Table 18, Appendix I. Each of our newly introduced Inter-Layer terms can provide additional performance gained for the network, and when combined together, the performance is boosted considerably.

**References**

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

[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] Zhou et al. Permutation Equivariant Neural Functionals. NeurIPS 2023
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper is the first to propose polynomials equivariant to parameter symmetries in neural field networks, allowing for a good balance between computational complexity vs expressivity dilemma. There were also concerns regarding comparisons. There were questions regarding the presentation that was too dense and could lead to challenging reproducibility, which still persisted by the reviewers, regarding deeper intuition and additional explanations for choosing stable polynomials and possible limitations of this approach with respect to the expressivity of the method, by not including all invariant/equivariant polynomials, the degree of the polynomial the limitations of monomials, and not discussing sufficiently the differences that arise for different activation functions.

By far the most controversial point was whether the paper committed plagiarism, as raised by one of the reviewers. After checking myself, it appears that even with the edits, there is quite some similarity. I do not want to go too much into the reasons why this similarity could be (for instance, maybe it's the same group of people co-authoring) but in general, plagiarism is a hard issue to address, especially when it comes to very large conferences with many unknown parties involved (at the time of the reviewing), since more often than not the truth is somewhere in the middle. The only thing that is certain, to avoid problems with other papers in the future (and how much is something or not a plagiarism), it is better to be strict to ascertain that there is no doubt.

As far as the content is concerned, I think the paper is good. However, considering that there are also some unclarities to improve, I suggest the authors to do one more round of serious rewriting, and resubmit.