Permutation Equivariant Neural Functionals

Summary

The paper introduces and evaluates permutation equivariant neural functional networks (NFNs). Neural functionals are models which take weights of other neural networks, or more general weight-space features, like gradients or sparsity masks, as inputs. Their permutation equivariance addresses the issue that any particular vector of weight-space features is just an arbitrary representative of an equivalence class of weight-space features that correspond to exactly the same network. Specifically, the weights and biases in all hidden layers may by permuted arbitrarily, which corresponds to simultaneous permutations of rows and columns of the previous and following weight matrix, respectively. Permutation equivariant NFNs ensure predictions which are independent from the particular permutation. As usual for equivariant networks, this reduces the number of model parameters (of the NFN) significantly, and makes learning feasible in the first place. The authors consider two specific group actions, which are the first order permutation actions and trivial actions on invariant scalars. Equivariance refers throughout the paper to mappings that commute with the first order action in the input and output, while invariance refers to mappings from first order actions to scalars. Equation 3 presents a linear NFN layer, which computes features as usual by taking learned linear combinations of weight-space features that are summed over different combinations of rows and columns. Proposition 1 claims that this layer is 1) sufficient and 2) necessary for equivariance, i.e. spans the space of the most general linear equivariant maps (I have some concerns here, discussed in the weakness section). As nonlinearities, NFNs use pointwise nonlinearities, which commute with permutation actions. A straightforward extension to permutations of CNN channels is discussed towards the end of section 2.2. Besides the HNP model, which considers the "correct" hidden neuron permutation symmetry, an NP model, which additionally permutes input and output neurons and has a further reduced number of parameters, is introduced. In order for the network to be able to break this excess equivariance, positional encodings are added. Extensive experiments evaluate HPN-NFNs and NP-NFNs and compare them to baseline models. The first two applications are permutation invariant predictions of CNN test accuracies from their weights and invariant classifications of implicit neural representations. In both cases the NFNs perform significantly better than baselines, and the HNP variant wins over the NP task in the first but loses in the second one (they are still quite close). The other two applications are permutation equivariant predictions of a lottery ticket hypothesis winning ticket sparsity mask and the editing of implicit neural representations. In the former, NP-NFN performs close to a task-specific baseline, while HNP-NFN is prohibitively expensive, thus supporting the need for the NP-NFN model. In the latter, both NP-NFN and HNP-NFN outperform their baselines.

Strengths

The paper is well motivated and very clearly written - it was really a joy to read! After introducing the mathematical framework, the authors derive the most general linear equivariant mappings for the considered group actions (I have some concerns here, see the weakness section). Explicit equations for the required number of parameters of the different models are given. The experiments are very extensive and show the utility and superiority of NFNs in various applications.

Weaknesses

I have two technical concerns regarding the claims about the generality of the proposed linear equivariant map in proposition 1. Firstly, it should be mentioned that the proof assumes a specific choice of group action w.r.t. which the layer is equivariant, namely first order permutation actions. Permutation equivariant networks have also been built for other actions, like irreducible representations or higher order tensor product actions - these would lead to other layers. This point needs to be discussed and the statement of proposition 1 needs to be made precise, e.g. by saying that the map $T$ is equivariant w.r.t. the specific first order action $\sigma$ in the input and output. [edit: this claim was wrong] ~Secondly, I believe that the mapping in equation 3 is only then the most general $S_{n_1}\times\dots\times S_{n_{L-1}}$ equivariant intertwiner when $n_1\neq\dots\neq n_{L-1}$ is assumed. If, however, $n_i=n_j$ for some layers $i\neq j$, the same permutation group $S_{n_i} = S_{n_j}$ acts on these neurons, and it should be possible to linearly accumulate the corresponding weight-space features after $\star$-summing over the other indices. A similar summing over features is already present in equation 3, namely the first term, which sums invariant scalars that are accumulated from all layers. Note that this contribution can in practice become very large, since networks are in practice often constructed such that they have the same number of hidden neurons or channels throughout many different layers. The more general equivariant mappings should be benchmarked against the current NFNs in the experimental section.~ The sufficiency and necessity of the invariant NFN layers in section 2.3. is not supported by a similar theorem as proposition 1. Due to these weaknesses I decided to only give a weak accept to the current state. However, the paper has a huge potential and I would be happy to switch to a strong accept when the issues are addressed.

Questions

The paper is very clearly written and I don't have any remaining questions.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Some limitations and potential future work are discussed in the conclusion. It should be added that one could consider more general permutation group actions.

Summary

This paper proposes the NF-Layer that maps the weight space of a deep neural network (DNN), including MLP and CNN, to another weight space, possibly with a different number of channels. Neural Functional Networks (NFNs) are then constructed using the NL-Layers to process the weight space of a DNN. NF-Layer is designed to be S-equivariant, that is, if the input weights of a DNN are permutated so as not to change the output of the DNN, the output of the NF-Layer, which is also in another weight space, is also permutated in the same order. Such a property enjoys the equivariant characteristics inherent in the DNN weights and enhances the efficiency in modeling NFNs, like the convolution structure enhances the efficiency of modeling image inputs. Experimental results are presented to demonstrate the usefulness of NF-Layers in applications of NFNs, including accuracy prediction, classification of implicit neural representations (INRs), winning ticket prediction, and style-editing through INRs.

Strengths

* An interesting application of the equivariant architecture based on parameter-sharing. * Introduction of the NP setting in the permutation definition. Although it requires a stronger assumption than the HNP setting, the resulting NF-Layer can be modeled more efficiently under the NP assumption. * Experiments suggesting the advantage of the proposed NF-Layers in interesting applications of NFNs.

Weaknesses

* A network architecture with equivariance based on parameter-sharing is not original and has been proposed in [51], as the authors suggest. * NF-Layers cannot directly handle some DNN models, including ResNet and Transformer. (Minor) * Figure 1 is difficult to understand at first glance.

Questions

I can understand the permutation equivariant property in DNN weights, but I cannot intuitively understand why the prediction accuracy of NFNs improves with NF-Layers with S-equivariance. Is it because the parameter to be trained is reduced by the S-equivariant structure, while retaining its expressiveness? It would be better to explain the multi-channel extension of permutations more carefully in ll.86-90. When I first read it, I mistakenly thought it was referring to the convolution channel. Also, it is better to describe that the proposed NF-Layer may change the channel size in this sense in the explanation of Figure 1. In Figure 2, $W_{\ast,j}^{(i+1)}$ (with yellow c) -> $W_{j,\ast}^{(i)}$? Equation in Figure 2 should be matched with Eq.(3). In ll.155-160, I could not understand the motivation to use position embedding with NFN_{NP}. In which experiments with NFN_{NP}, position embeddings are used?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

2: You are willing to defend your assessment, but it is quite likely that you did not understand the central parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

3 good

Presentation

3 good

Contribution

3 good

Limitations

As described by the authors, there is a concurrent work by Navon et al.[45] that has a very similar motivation and methodology. But I think this paper has its original contributions that the authors mention and deserves to be published.

Summary

This paper studies the problem of defining linear layers (and by extension, neural networks) that operate on neural network weight spaces. The core idea of this work is to take into account weight permutation symmetries, similar to Navon et al., ICML’23. In particular, the weights of certain feedforward architectures, such as MLPs and CNNs, can be permuted in several ways, without altering the function that the neural net represents. Therefore, the authors seek to define functions on the weights, that will be equivariant (or invariant) to these very symmetries. To achieve this, they characterise the space of linear equivariant/invariant functions, following the framework of Ravanbakhsh et al., ICML’17, i.e. by identifying the parameters of these functions that should be shared. The authors propose two variants, one that is derived from the detailed characterisation of the weight permutation symmetries that arise from hidden neurons ($\text{NFN}_{\text{HNP}}$), and another one that assumes extra permutation symmetries in the input/output layer (that can be broken with positional encodings), and is, therefore, more parameter efficient. Interestingly, these results easily carry over from MLPs to CNNs. The proposed layers are experimentally tested in a battery of tasks, such as predicting NN generalisation or sparsity masks for NN pruning, showing considerable improvement against baselines that do not take into account permutation symmetries, in addition to a heavily reduced parameter count.

Strengths

**Significance**: The proposed method has wide applicability, as correctly pointed out by the authors in the first paragraph of the introduction, and can potentially have a substantial impact on multiple neural-network-related problems (meta-learning, neural network editing, INR processing etc.). **Presentation, clarity and reproducibility**: Although the concepts studied in this paper and the neural network architectures resulting from the study of the relevant symmetries, have a fair amount of complexity and may require a good understanding of the study of symmetries, the authors have made a very good effort to make their paper as accessible as possible. In particular, the illustrations, the nice summarisation of the core layer in Eq. (3) as well as the parameter counts in Table (1) and the accompanying explanations provide a good overview of the method that I believe will be appreciated by future readers. Moreover, the authors have formulated their layers in such a way that makes them easy to implement, which is also apparent from the pseudocode provided in the appendix. **Experimental evidence**. The authors have tested their layers against a diverse set of problems and the reported results clearly motivate the need for equivariant/invariant layers and support all the design choices made. **Extended scope**: A few months ago, another work by Navon et al., ICML’23 characterised the same family of equivariant/invariant layers as well (providing additional theoretical support). This is obviously not a problem, since the two works were developed in parallel. It is therefore nice to see that the present paper has a bonus, i.e. an extended scope compared to the aforementioned one since (1) the formulation seamlessly allows to extend the method to CNNs, and (2) the NP formulation allows for more practical implementation, a significant reduction in the number of parameters, and empirically strong performance. (3) This is probably subjective and a matter of taste, but I appreciated the alternative way to derive the layers (see sec B.3. in the appendix), using the parameter-sharing framework of Ravanbakhsh et al., ICML’17. I found it easier to follow and slightly more intuitive, which might be of independent technical interest in the future. Therefore, although the layers have been rediscovered, I feel that the contributions are still valuable.

Weaknesses

**Clarity of the proofs (appendix)**: I found the explanations in section B.4. a bit hard to follow. Personally, I am familiar with the work of Navon et al., so I could easily understand the concepts, but I fear a reader not versed in the topic, might get lost in this part. Since it is crucial in order to gain a deep understanding of the paper, I would recommend that the authors try to simplify some concepts (e.g. by giving more details for a particular sub-case). **Related work**: I believe that the authors do not give enough credit to the work by Navon et al.. As I mentioned before, the two works are concurrent, and they apparently independently discovered the same core idea, but I think that the authors should be more upfront about this and mention this early on in the paper and cite Navon et al. more prominently.

Questions

**Experiments, baselines and implementation details**: Most of the following questions ask for clarifications and I do not consider them as weaknesses. However, I think a discussion is needed and it would be useful to include the answers to the below in the paper. In detail: - **NP case**. What are the positional encodings that are used in the experimental section? If I understand correctly they are fixed (handcrafted). Have the authors tried learnable ones instead? What is the number of added parameters when using learnable positional encodings? How do the authors explain the fact that using positional encodings in NFN_NP performs better than NFN_HNP? - **NFN_PT case**. Since this method performs well in the pruning mask prediction experiment, why didn’t the authors also include its performance in the rest of the experiments (e.g. generalisation and INR classification)? - **Pruning**: Why did the authors use a generative model (CVAE) for the winning ticket experiment? Why not a deterministic predictor (a 0-1 classifier)? Have the authors tested this? How far are the generated sparsity masks from the ones that IMP yields? Why isn’t the performance of MLP and MLP_aug included in this table? - In section 3.2, would it be possible for completeness to add some baselines that classify images and 3D shapes, using different representations instead of INRs, e.g. state-of-the-art results using image-based (CNNs) or point cloud-based classifiers? Could the authors comment on why the performance is still far from the one that the community has obtained by applying these baselines? This can be clearly observed in CIFAR-10. In addition, apart from the inr2vec baseline, it might be helpful to add the baseline from Dupont et al., ICML’22 (the “functa” representation), where INRs are represented as vectors of latent modulations. - **NN editing**. In section 3.4., regarding the CIFAR dataset, the differences between the competing methods are not clear qualitatively. Could the authors comment on this? - Is there any theoretically-backed advantage of using a combination of equivariant layers followed by an invariant one instead of using only invariant ones (when the ground-truth function is invariant)?

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

- The authors have dedicated a paragraph to the limitations of their method. One thing that might be missing is a clarification (if I understand correctly) that NF layers can only be applied to the weights of neural nets that share the same architecture (fixed number of layers, fixed kernel size for CNNs). Note that L189-190 (“…which contains thousands of CNN weights trained on several datasets with varied hyperparameters”) seems contradictory to this. I think the phrase "varied hyperparameters" should be clarified to avoid confusion. - Moreover, given that classifying INRs with NF-layers seems to still have room for improvement in order to reach the performance achieved by traditional neural nets that operate on raw representations, it might be useful to discuss the possible reasons for this more thoroughly (perhaps stronger inductive biases, taking into account other properties of the geometry of weight spaces, might be missing?) - No foreseeable negative societal impact.

Summary

The authors proposed an architecture for processing other networks’ weights and implicit neural representations (INRs). The generalization abilities of this architecture are enhanced and the number of parameters is reduced by leveraging the symmetries of deep networks. The authors claim the following contributions: - Proposing a new architecture named Neural Functional Networks (NFN) for processing weights of other networks including INRs. A key building block in NFNs are NFLayers which the authors constrained to be permutation equivariant. - Extensive experimental part that demonstrates the superiority of NFNs over various baselines.

Strengths

- The paper is dealing with novel and interesting problems of processing networks’ weights and INRs, I believe that this intriguing domain can open the path to more interesting and valuable works for our community. - The paper is well-written and easy to follow. Specifically, the visualizations make the reading more accessible and straightforward. - The experimental part is extensive and includes multiple learning setups and datasets.

Weaknesses

- My main concern is about the novelty of this work w.r.t [1]. In [1] the authors also proposed an architecture, coined DWSNets, for processing weights and INRs. They also provide some theoretical guarantees about the expressive power of their method. It is worth mentioning that [1] is **not a concurrent work** according to the NeurIPS guidelines. Therefore, the authors should explain what is the novelty of the current work over [1]. - Given the previous point, the authors should also include [1] as a baseline in the experimental part and cite it more explicitly in the paper. - There is a lack of available technical information regarding the process of generating INRs. Specifically, in sections 3.2, 3.3, and 3.4 the authors did not mention the number of samples used for train/val/test. I also wonder how long it took to generate the INRs and how many resources were used. The authors should also include information about the training paradigm of INRs (for example the number of optimization steps) to enable easier reproducibility. - Although performed in the weight space it will be interesting to see how well NFNs deal with more challenging style editing tasks like inpainting, deblurring, etc. Dilation and contrast editing are straightforward tasks in CV. - Given that it was the pioneering study demonstrating the permutation equivariance of feed-forward networks, it is advisable for the authors to include a citation to [2] (especially in lines 33-36). - In lines 155-158 the authors stated they have used positional encoding (PE) to boost NFNs performance (although it breaks the symmetry), what is the performance gain presented by using PE? I did not see such an ablation (although I might have missed it). ---------- Citations: [1] Equivariant Architectures for Learning in Deep Weight Spaces, Navon et al. [2] On the algebraic structure of feedforward network weight spaces, JHecht-Nielsen et al.

Questions

- Can NFNs handle heterogeneous networks in a single dataset? For example networks with varying input dimensions and hidden features. - What is the computational complexity of NFN / NFLayers?

Rating

3: Reject: For instance, a paper with technical flaws, weak evaluation, inadequate reproducibility and incompletely addressed ethical considerations.

Confidence

5: You are absolutely certain about your assessment. You are very familiar with the related work and checked the math/other details carefully.

Soundness

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

A limitation section is included.

Summary

This paper considers the design of architectures whose inputs are the parameters of neural networks. They propose an equivariant weight-sharing scheme based on the permutational symmetries of neural networks: one can permute at least the internal neurons (the “HNP” case), and sometimes the input/output neurons as well (the “NP” case), of a network without changing the function represented by the end-to-end network. The parameter-sharing of the “NP” case is more advantageous but less often applicable, since the underlying problem must have additional symmetry structure to warrant input/output neuron permutational symmetries. To remedy this, the authors propose coupling the stronger “NP” case with positional encodings. They extend their framework to take as input convolutional layers, and test their proposed architecture on tasks including image classification from implicit neural representations, predicting sparsity masks, and weight-space editing. They compare to the use of MLPs with and without permutational augmentations, as well as problem-specific baselines, and obtain promising experimental results.

Strengths

Originality: Outside of concurrent work, the idea of using permutational symmetries to train directly on weight and bias inputs is original and makes sense. This permutation subgroup is distinct from other groups that have been studied in the past (such as $S_n$ and $S_n \times S_n$). Even beyond concurrent work, the idea of using positional encodings in conjunction with input/output permutational symmetries is a creative contribution. Quality: The experiments are varied and test against a reasonable set of baselines. Clarity: The paper is generally quite clearly written and explained. Significance: This paper enables the application of weight-space symmetry techniques to CNNs. The problem of learning with neural network weight inputs is topical, and such methods are likely to enjoy practical usage.

Weaknesses

1. The primary weakness of this paper is its novelty in relation to other work by Navon et al (ICML 2023), which also articulates an architecture equivariant to the permutational symmetries of weight space and provides more thorough results on its universality. The authors state this work was concurrent (which I will take at face value). If its novelty is judged in relation to this work, then the main contributions are the NP setting with positional encodings, the extension to CNNs, and the distinct set of experiments, which are extensions of the key idea. This very relevant work is currently only mentioned near the end of the paper, but I would think it is important to have a more in-depth discussion of how the two fit together and the novelty of this work, and for this discussion to appear earlier in the paper as part of its framing. 2. The paper only considers the permutational symmetries of neural networks, which are indeed perhaps the most general symmetries if one does not specify a certain nonlinearity. However, as noted in Godfrey et al 2022, a given nonlinearity may enjoy additional symmetries — for instance, ReLU enjoys a symmetry to positive scaling. This is not discussed in the paper, but would have been a more substantial contribution relative to Navon et al (2023). For instance, one straightforward way of incorporating this scaling symmetry could be to pick a positive scale based on the norm of the parameters of the first layer, and then use this scale to normalize the rest of the input weights and biases.

Questions

1. Is there some advantage to classifying implicit neural representations, instead of the pixels directly? It would be helpful to motivate these problems more in the main body of the text. 2. Can the authors comment on the expressivity of their proposed idea of using input encodings to adapt the NP setting to problems without input/output permutational symmetries? Is this universal, for example, or is there some loss of information from adding the positional encodings? (On a related note, how exactly are the positional encodings incorporated — are they literally added, or are they appended?) 3. Could you be more specific in designing networks for a particular nonlinearity with different symmetries (e.g. scaling for ReLU)? See e.g. “On the Symmetries of Deep Learning Models and their Internal Representations” (Godfrey et al 2022), as well as other references in Navon et al 2023. 4. How large are the models input to the neural functional network? Will all weight-space methods will fail for truly large weight space inputs, i.e. millions of parameters, even if you have very few learnable parameters thanks to symmetry?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

3 good

Limitations

There is not potential negative societal impact. The authors are upfront that the HNP case is less scalable than the NP case. One limitation not discussed is that the architecture described only takes into account permutational symmetries, and not other nonlinearity-dependent symmetries.

Thank you for your replies

Thank you for your replies. I understand the relationship between [51] and this paper. I would like to keep my score.

Response to Authors

Dear authors, Thank you for your reply. Most of my concerns have been addressed. I would advise incorporating the explanations you give here in an updated version (e.g. prominently discuss and compare against Navon et al., details about positional encodings including the ablation study, adding non-weight space baselines in 3.2. and discussing the differences). I continue to support the acceptance of the paper and my initial score remains. *Minor*: The only thing I am sceptical about is the absence of comparison to Dupont et al., ICML'22. The authors correctly mention that the INRs in this paper are produced by specialised training, but since with this method one obtains a vector representation of each INR, then this representation can be fed to, e.g., an MLP to solve all the *invariant* tasks tested in this paper. This is not a major concern, but it would be nice to include an experimental comparison (missing from Navon et al. as well). I would also encourage the authors to discuss my last question in an updated version (*“Is there any theoretically-backed advantage of using a combination of equivariant layers followed by an invariant one instead of using only invariant ones (when the ground-truth function is invariant)?”*) – this was not answered in the rebuttal.

Thank you for the response

Thank you to the authors for their response. They clarified the use of operating on INR representations, made a valid point about why scale symmetry may be superfluous, and generally answered my other questions. My only remaining concern retains to the novelty of this work relative to Navon et al. I think the experimental comparison and added paragraph will help on this front, but because the contributions relative to Navon et al are still fairly minor, I retain my original rating of weak accept.