Neural Functional Transformers

Summary

Designing architectures for processing the weights of neural networks is an exciting and relatively new area with great implications. This paper proposes a novel, attention-based architecture for processing NN weights. The proposed architecture is equivariant to the permutation symmetries of neural weight spaces. Furthermore, the authors propose INR2Array, a method for learning an invariant latent space for INR weights.

Strengths

1. The paper is generally well-structured and well-written. 1. The authors propose a novel architecture for processing the weights of neural networks. The new attention-based architecture achieves comparable or superior results compared to baseline methods. 1. This is a natural and necessary extension of previous work on linear equivariant weight space layers.

Weaknesses

My main concerns are the limited novelty w.r.t prior works and the missing/unfair evaluation results, which makes it difficult to assess the performance of the proposed approach w.r.t to prior work, and the effect of the different part of the proposed method (e.g., INR2Array^NFT vs. NFT). 1. While attention-based equivariant layers are a natural extension, the technical novelty of the paper is somewhat limited with respect to prior works ([1], [2], [3]). 1. Missing results for NFT without INR2Array for the INR classification task in Table 1. This is a necessary comparison for evaluating the effect of the main contributions of the paper, attention-based base layers for weight spaces, and INR2Array. 1. The evaluation against DWS in Table 1 is flawed and unfair: The results for DWS are taken from [1], which uses a single INR copy for each image. The results for other baselines and the proposed method use 10 INR copies. As discussed in Section 7 of [1], additional INR copies can significantly improve the performance. Also, a more direct evaluation should use the same INR dataset, as choices of the INR arch., hyperparameters, and training procedure can affect the results of methods trained using this INR dataset. Using different INR datasets can potentially greatly vary the results. Finally, the augmentation scheme used in this paper is different than the one used in [1]. 1. Line 120: Need to provide an explanation/intuition for the provided claim. 1. The equivariance proof provided for Theorem 3.1 is insufficient. While I believe the results are valid, the claim (“To show this, we can expand the left hand side term-by-term using the definition of SA.”) should be extended, and more details are needed. Also, it would be beneficial to provide a concrete example (using a specific layer) to show equivariance. 1. As stated by the authors, the proposed layers are computationally expensive. Please provide runtime (or complexity) comparisons w.r.t linear equivariant weight-space layers. 1. Additional comments on experiments: - Please report the number of parameters for each method (In Table 1). - What is the NFN (and INR2Array^NFN) variant used in Table 1 (NP or HNP)? - It would be beneficial to include ablation on design choices, specifically the use of Random Fourier features and the use of Layer norm. - Please provide more details on the use of mixup augmentations: are you performing the mixup in weight space? - Also, it would be beneficial to compare the results of NFT with/without mixup augmentation. Minor: 1. In Sec. 2, it would be clearer if the authors explicitly state that elements from the NP group, generally do not preserve the function defined by the set of weights. [1] Navon et al., Equivariant architectures for learning in deep weight spaces, ICML 2023. [2] Zhou et al., Permutation Equivariant Neural Functionals, 2023. [3] Luigi et al., Deep Learning on Implicit Neural Representations of Shapes, ICLR 2023.

Questions

1. Can the proposed Arch. be easily extended to the HNP design and setup? From Table 3 it appears this could be very beneficial for some tasks. 1. The authors stated the difficulty of training large NFT architectures stably as a limitation of the method (Section 7). Since learning over weight spaces is a relatively novel task, it would be very beneficial if the authors could elaborate on any limitations/findings/failure cases from their experiments in the Appendix.

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

2 fair

Presentation

3 good

Contribution

2 fair

Limitations

Addressed, although more details would be beneficial, see Questions.

Summary

The paper studies an interesting problem of encoding weights of neural network using another network (Neural Functional Transformer or NFT). NFTs are built based on previous works, mainly DWS [31] and NFN [46], proposing permutation equivariant linear models to encode neural network weights. NFT's main contribution compared to these is a more expressive model based on Transformers that is permutation equivariant only to those weight permutations that result in the same function. While DWS already proposed such a minimal permutation equivariance, this paper achieves better performance due to leveraging Transformers and other tricks such as spatial INRs (INR2ARRAY). Overall, NFTs achieve better performance when encoding the weights of INRs as well as competitive performance in other tasks.

Strengths

1. The topic of encoding neural network weights is very interesting with a lot of potential. 2. The idea of combining Transformers with minimal equivariance seems as a powerful and promising way to approach this problem. 3. The improvements in most experiments are significant. 4. The paper is written and presented relatively well.

Weaknesses

While I really like the problem the paper tackles and using Transformers with minimal equivariance, unfortunately there are many weak points described below. I will be willing to raise my score if these are addressed. 1. Intuitive explanation of the difference between Eq. 4 and 5 is missing (why the union of sets is used and how is it implemented, by concatenation?) Perhaps, some illustration would help. Moreover, it would be very useful to see ablation results with each of the terms (4,5,6) ablated in the model. 2. The paper mentioned that the weight space can include different features such as weights and gradients, but the details of what features are actually used in this work are missing. It's a bit confusing to mention gradients if they are never used, because it creates mental links with completely different literature. Why self-attention between different weight space features (dimensionalities from 1 to c) are needed? 3. L186 mentions the limitation of Inr2Vec "requires every INR share the same initialization", but in the Experiments this advantage of the proposed method is not analyzed w.r.t. Inr2Vec, so it's unclear if this advantage is really useful. Also, Inr2Vec is not compared in Table 1 for some reason. 4. In Table 1, why there are no results for INR2ARRAY with DWS and no DWS results for CIFAR-10? Since, DWS has source code available I believe this should be straightforward to evaluate and would strengthen the paper if the proposed NFT still dominates. Why INR2ARRAY only improves when combined with NFT and not with NFN? 5. Section 5.1 is confusing, because it mixes methods, experiments and related work together. 6. In Table 3, why NFN can be combined with HNP, while NFT cannot? Using HNP seems important in this experiment, so combining it with NFT could improve the paper. 7. Editing INRs is an interesting experiment, but it is not very clear what would be the realistic use case. The experiment in DWS "Adapting a network to a new domain" seems more useful in this sense. 8. Overall, the experiments seem quite esoteric with unclear real world use cases. For example, the motivation for encoding the weights of INR is not very convincing. This is a general problem with the topic of encoding of neural network weights as researchers seem to have not found really useful and realistic use cases of these models. The experiment "Predicting CNN classifier generalization" tackles a more concrete and useful application, but these also seem tiny networks with a few thousand parameters (Small CNN Zoo in [40]). So overall, this paper would be stronger if some other more practical applications are considered and/or the model is scaled to encoding more realistic networks (like ResNets in the "NeRN - Learning Neural Representations for Neural Networks. ICLR 2023." paper). 9. Some relevant papers/methods are not mentioned/compared in this paper. These include "Self-supervised representation learning on neural network weights for model characteristic prediction. NeurIPs 2021.", which discussed and compared in DWS, and "NeRN - Learning Neural Representations for Neural Networks. ICLR 2023." **Minor** The paper needs some polishing to improve the reading experience. Some abbreviations like NFN and SIRENS were used before they are introduced/cited, e.g. L201: " Given a SIREN ...". There are some occasional typos (e.g. L266 "anount") and the hyperrefs are broken (always jump to the first page).

Questions

Please see Questions in Weaknesses above.

Rating

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

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

2 fair

Limitations

Limitations are briefly discussed, which is appreciated. The authors could empirically compare time/memory complexity of their method to NFN, DWS and naive self-attention, which could strengthen the paper.

Summary

This paper introduces a new architecture called Neural Functional Transformer (NFT), that uses an attention-based layer that is equivariant to neuron permutations. NFT is an architecture that is used for tasks where the model is taking an input another neural network's weights (in their case, either an MLP or CNN), which may be e.g. an INR of a some signal or the weights of another trained model. They benchmark their architecture on two INR tasks (image classification, image editing) and predicting classifier generalization from weights.

Strengths

The paper presents the problem of building architectures that operate on other network's weights, which is significant and has useful applications to e.g. working with INRs. The architecture they propose based on attention is original, interesting and clearly presented. The writing of the paper is clear throughout. Their experiments are relevant and explained well.

Weaknesses

Although there are no major weaknesses in this work, I found the empirical performance of NFT slightly underwhelming. - For CIFAR-10 classification, why isn't there a comparison to Spatial Functa (or usual functa) (which was the inspiration for using spatially structured latent representations)? As shown in Spatial Functa's table 2, they achieve 90% top-1 accuracy, which is significantly higher than 63% achieved by NFT. Both approaches are for image classification using INRs for the images, hence a comparison seems appropriate. Regarding using the same initialization for different INRs, rather than being a weakness, wouldn't that be the go-to approach anyway because it cuts the INR training time by using techniques like MAML to pick the initialization? - What about direct image classification with NFTs instead of first encoding the INR as an array and then performing classification with that array instead of the original image array? - For generalization prediction, why isn't there a method corresponding to NFT_HNP in Table 3?

Questions

- I think there's a typo in the caption of Figure 1 in the last sentence. - The weight spaces handled by the proposed layer are those of CNNs and MLPs. Do you think there are ways to extend this to other commonly used architectures like transformers? - Neuron permutations come from the fact that MLPs have layers composed like $W_1\sigma(W_2x)$ where $\sigma$ is a point-wise nonlinearity. From this, given a permutation matrix $P$, we get $W_1P^T\sigma(PW_2x) = W_1\sigma(W_2x)$ because $\sigma(PW) = P\sigma(W)$ for any matrix $W$. Is this property satisfied for matrices $P$ that are not permutation matrices? If so, we would also want equivariance to those transformations of the weight space. - I found Section 3.3 slightly unclear to read, particularly lines 155-158. - Where is the layer encoding in section 3.4? - In Section 5.2, do you train two INRs with the same initialization on the the original and edited images and then take their weight difference as the target for the NFT? - Typo in "amount" in line 266. - In section 5.3, define the abbreviation HNP for clarify (I think only NP was defined in the beginning of the paper).

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

3 good

Presentation

4 excellent

Contribution

3 good

Limitations

Adequately addressed.

Summary

This paper aims to construct neural functional architectures that can handle high-dimensional weight-space objects. Specifically, this paper designs a novel class of weight-space models with neuron permutation symmetries using nonlinear self-attention and cross-attention layers. The proposed method is called neural functional transformers (NFTs). NFTs enable greater expressivity and higher classification accuracy on implicit neural representations (INRs) over existing methods.

Strengths

1. The overall written is good, and the authors provide code for reproduction. 2. The topic of this paper is cutting-edge and the proposed method is novel. 3. The proposed method is simple and its performance for INR classification is good.

Weaknesses

1. It seems straightforward to make self-attention permutation equivariant. What is the actual contribution of section 3.1? Do you compare the performance of original attention and your modified attention? 2. Equation (4) (5) (6) are difficult to read. The authors may represent them via a better way. 3. In table 3, the performance is worse than NFN which is inconsistent with table 1 and table 2. Can you explain the reason of this? 4. Do you evaluate your method on 3D dataset?

Questions

See "Weaknesses".

Rating

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

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some 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

2 fair

Limitations

This paper has no potential negative societal impact.

Summary

The authors consider the principled design of neural functionals, i.e. neural networks that take other neural networks as inputs. Example tasks include predicting the generalisation error of a trained network for a given set of hyperparameters, or trying to reconstruct the images or 3D object encoded in an implicit neural representation (INR, e.g. Mildenhall et al. arXiv:2003.08934). The authors propose a transformer-style architecture that consists of a stack of transformer blocks, each one consisting of a self-attention mechanism followed by an two-layer fully-connected neural net. The key challenge is to design a self-attention mechanism that respects the various symmetries of that preserve the network's behaviour. Naïve self-attention is equivariant to weight permutations that change network outputs. Previous work (Navon et al. [31] and Zhou et al. [46]) designed linear layers that are equivariant to neuron permutation symmetries. The key contribution of this paper is to propose a particular form of the self-attention mechanism for neural functionals. It really is a combination of three attention mechanisms (eq. 4-6), which has however to be approximated in a rather heavy way: from self-attention between weight-space features to self-attention between row-column *sums* of the weight matrices. The proposed mechanism works for both fully-connected layers and convolutional layers. Given that weight space tasks remain computationally challenging, the authors also propose an embedding method that maps implicit neural representations into arrays, an extension of the inr2vec method of De Luigi et al. [arXiv:2302.05438]. Its performance compares favourably to other methods when tested on weight-space INR classification of standard image data sets, where we seek to reconstruct the image encoded in the weights of an INR.

Strengths

- The authors propose a well-designed architecture for a neural functionals. - The authors show that the self-attention has several desirable properties, even in its computationally simplified form. - The proposed method works well on some tasks chosen by the authors. As I have no practical experience with neural functionals, I may have missed some relevant related work or may misjudge the novelty of the proposed solution, hence my low confidence.

Weaknesses

I believe the article could benefit from some discussion of key concepts regarding neural functionals, for example the type of tasks in standard benchmarks, to make the article easier to read on its own.

Questions

The authors introduce a neural functional architecture that is transformer-based and can deal with fully-connected and convolutional networks. An obvious question is whether this approach can be extended to transformers? I would appreciate some discussion of this issue from the authors.

Rating

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

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some 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

The authors discuss computational limitations of the vanilla form of their self-attention layer, and discuss computationally feasible alternatives.

Thank you for your reply

I have read the author's reply, as well as the other reviews, and I will keep my score.

Replying to Rebuttal

Thank you for your response and for providing additional results. I have some following questions and concerns: > Regarding the NFT without INR2Array results (reported in response to Reviewer 5PwT): It seems like the proposed architecture (NFT) achieves lower results compared to NFN on two out of the three tasks presented in the paper, namely "INR classification" and "predicting CNN classifier generalization". This result casts doubt on the effectiveness of the proposed architecture compared to previous methods. This is amplified since the proposed method is computationally expensive compared to, e.g., NFN as stated by the authors: "INR2Array with NFT takes 2.79x times as long per training step compared to with NFNs (NP)" > Regarding the results for DWS: The original DWS paper reported 85.7 for MNIST using a single INR per image. Could you please explain your results of 74.7 for a dataset with 10 INRs per image?

Update my rating from 5 to 6.

Thank you for your response. All my questions have been addressed. By combining the figure in pdf response and equations, I understand how minimal equivalence is achieved. I see the contributions in section 3.1. Moreover, the authors add ablation experiments which prove the importance of such minimal equivalence. So I update my rating from 5 to 6 and my confidence from 2 to 3. I suggest that the authors describe the weight-space more in Introduction or Preliminaries.