Universal Neural Functionals

Summary

This paper develops an algorithm for constructing functions that take neural network parameters / parameter-derived values as inputs but are equivariant / invariant to the inherent permutation symmetries in neural network weights. More specifically, this work is applicable to a larger class of neural network architectures than previous approaches (including RNNs), due to the less restrictive assumptions made. The authors evaluate their method on two types of task: 1) as a model for predicting generalisation from neural network weights, 2) as a learnable optimiser for training neural networks more effectively than traditional methods like SGD. They find moderate improvements with their method over previous approaches.

Strengths

- The approach is well-motivated from a mathematical perspective on permutation equivariant / invariant functions, and the experiments appear to remain faithful to the proposed methodology rather than resorting to an assortment of ad-hoc hacks / tricks. - The paper is on the whole very well written. In particular, the authors do a good job of explaining the mathematics that, although not conceptually too difficult, is inherently cumbersome / fiddly. - The work generalises previous work to a wider class of neural networks, rather than being specific to just a couple of architectures, and is therefore a good contribution to the area (I am not very familiar with the related work so I am relying on the authors' own summary of the related work to make this judgement).

Weaknesses

- The RNN generalisation prediction experiment only compares against one baseline, which is not far behind the proposed method. It is therefore difficult to judge the significance of these numbers. - The authors state "Navon et al. [2023] showed that permutation equivariance significantly improves performance on weight-space tasks, but their models only apply to the weight spaces of simple feedforward multilayer perceptrons (MLPs)." but do not appear to test any methods without permutation symmetries when performing their experiments. In particular, it would have been nice to see whether a learned optimiser that doesn't account for permutations symmetries (e.g. a plain-old neural network?) performs significantly worse than the proposed method or DeepSet, which also is designed to take permutation symmetries into account. - There are some small changes I would make to the paper for readability, please see my suggestions below.

Questions

## Questions - In Example 3.1 we seem to be assuming that weights from 2 separate layers ($m$ and $\ell$) will be permuted with the same permutation because they have the same size. Surely we can permute two layers on neurons in different ways? Please clarify as I may be missing something here. - On line 182, could you expand on what is meant by "Entries that are not explicitly assigned by the left-hand side are 0."? - In Eq. 15 are we manually computing the momentum term still, and only using the learned optimiser to produce the new direction / update at each step? - You mention that your assumption of permutation symmetry over the spatial dimension in CNNs is technically incorrect, but that this can be addressed somehow by use of positional encodings. Did you use this positional encoding fix in your experiments? - My understanding of the UNF is that it is technically not a neural network, as might be common in other work on learned optimisers, but rather we learn coefficients for a particular set of basis functions that you have derived, however the computations we end up needing are quite similar to neural networks and can therefore be implemented using standard deep learning libraries. Is this correct? ## Suggestions - In the second paragraph of the Preliminaries section, circa line 64, it would be useful to explicitly say "where $S_n$ is the symmetric group of $n$ elements", "$\sigma$ is a particular permutation in this group", "the usual / canonical action of $\sigma$ on the set of indices permutes them according to the permutation $\sigma$". Clarifying explanations like these do not take up much space but vastly improve the readability of the paper. - On line 72, it does not appear that $s$ is defined anywhere, but appears to be the non-linearity of the neural network. A short note would improve clarity. - Typo on line 246: Remove the word "is" from "STATNN is operates on basic statistical features of the weights"

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

- The authors discuss some general limitations of the current algorithm, such as computational cost and applicability to more general problems, though these are generally beyond the scope of the current paper.

Summary

The paper proposes Universal Neural Functionals that are models operating on neural network weights in a permutation equivariant way. Compared to previous works, UNFs are more general and can be more easily applied to architectures such as RNNs and Transformers. The UNFs show good results in generalization prediction for RNNs and also improve learned optimizers in several image and language tasks.

Strengths

1. The paper addresses the problem of learning from network weights which is challenging and has a lot of potential. 2. Interesting learning to optimize experiments, which is a very relevant application of this kind of method. 3. A method is inspired by the symmetry group theory, which is an interesting alternative to message passing and graph networks and potentially could be very powerful in this task. 4. The paper is well written and organized.

Weaknesses

1. The algorithm 1 that constructs equivariant weight-space layers need "specifications" of a network that indicate which dimensions of which layers are permuted simultaneously. While for simple architectures like MLPs providing them is easy, for architectures like RNN and Transformer it seems tricky and essentially requires digging into the exact implementation of each layer (e.g. how the weights in the multi-head attention layers are split in heads and q, k, v weights). It would be useful to see the specifications for RNNs and Transformers similarly to the ones for MLP in Appendix A. Given that the results for Transformers in Fig. 2 are not very strong, it may indicate that the specifications are incorrect. So the paper is a bit misleading in a sense that it claims to "automatically construct permutation equivariant models for any weight space", but it does not automatically constructs the specifications which is actually a quite tricky part for some complicated models with complicated implementations and "analyzing its computation graph" may not be enough since it usually reveals the connections between layers, but not the connections between dimensions. 2. The theory of the paper is a bit disconnected from the experiments making it harder to understand the theory part and its strengths. How Algorithm 1 works in practice for the networks in the experiments? How the valid partition looks like for those networks? How the basis looks like and are there any insights about the architecture given the basis (e.g. how the basis will look like for mlps, cnns, rnns and transformers)? 3. RNN generalization prediction - More baselines are needed. For example, Deep Sets could be one simple baseline, also since the architecture of the RNN models is the same, a simple MLP on concatenated parameters can be another baseline. - it's questionable that 0.8968 is significantly better than 0.8839. - Given those two weaknesses, I believe the setup should be changed to include more diverse architectures to highlight UNF's strength. For example, see the Graph Metanetworks paper for the inspiration on how to construct "Diverse Architectures" experiments. 4. Learned optimizers - L289: it's questionable that Deep Sets is "the default architecture choice for f". For example, in the paper "Practical tradeoffs between memory, compute, and performance in learned optimizers" (Metz et al, 2022) a simple per-parameter MLP (small_fc_lopt) is used that does not include deep set layers. So directly using that or similar optimizer as one of the baselines would be interesting. - Why Adam is not used instead of SGDM as the backbone of learned optimizers in Eq. 15 given that Adam is a standard choice for optimizing Transformers used in the experiments? - In CNN on CIFAR-10, why UNF is better than NFN-NP given that UNF assumes spatial dimensions can be permuted? Is positional encoding added to spatial dimensions in UNF? 5. Related work is missing some relevant papers like "Hyper-Representations as Generative Models: Sampling Unseen Neural Network Weights" and "NeRN - Learning Neural Representations for Neural Networks". Overall, given that there are very few experiments, lack of empirical analysis and the results are often similar to the baselines, I'm inclined towards rejection.

Questions

For SGDM do the authors actually learn α and γ0 or it's simply tuned?

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

2

Limitations

the limitations are discussed

Summary

The paper proposes Universal Neural Functionals (UNFs), which are models that process the weights of other neural networks. UNFs are equivariant to permutation symmetries of neurons and applicable to general architectures. The authors formalize neural network weights as sets of tensors, and develop an algorithm that constructs maximally expressive equivariant linear layers for processing any collection of tensors given a description of their permutation symmetries. The resulting UNFs outperforms StatNN, a baseline that uses statistical features, on an RNN generalization prediction task. Additionally, UNFs improves over several other architectures on learning optimizers for various image classifiers and language models.

Strengths

- UNF has the appealing property of being applicable to general architectures. Plus, the code appears easy to use on custom neural networks. This work is thus likely to be useful for the deep learning community. - The algorithm that computes maximally expressive equivariant linear layers for a set of tensors is a significant theoretical contribution. - Connection to past work is discussed well throughout the paper. - The paper is generally clearly written and well organized. The examples in Section 2 and 3 are particularly helpful in showing the wide range of architectures UNFs are applicable on and in helping readers understand the concept of valid partitions. - Both experiments uses datasets with decent size compared with related works and demonstrates promising performance of UNFs.

Weaknesses

- Despite the claims in the abstract and conclusion, there has been other work, most notably Lim et al. 2023, that developed permutational equivariant weight-space models applicable to general architectures. This unfortunately weakens the novelty of this paper. The narrative might need to be modified to emphasize other contributions, which are still significant, such as constructing maximally expressive equivariant linear layers and promising results for learned optimizers. - I find Section 3 a bit difficult to follow. Part of the reason could be tensor indices are inherently complicated. However, it would be helpful if the authors could provide intuition along with stating results. For example, why do valid partitions produce a basis and other partitions do not? - It is not clear whether the proposed architecture scales well, since the number of valid partitions can be very large especially when many indices permute simultaneously or there are many layers. - In both experiments, the proposed method is not compared to the most recent permutation equivariant architectures (Zhang et al. 2023, Kofinas et al. 2024, Lim et al. 2023). These papers are highly relevant, as they solve the same category of problems (processing or generating neural network weights) and feature similar advantages (being permutational equivariant, and Lim et al. 2023 also has the ability to process arbitrary weight spaces). Additionally, they have been published on line more than two months before NeurIPS’s deadline so are not generally considered contemporaneous.

Questions

- Among the many applications of weight-space networks, how did the authors decide on the two tasks to conduct experiment on? - In line 316, it is stated that UNF makes the stronger assumption than NFN that all tensor dimensions are permutable, but line 273-276 seem to suggest the UNF respect fewer symmetries. Could the authors clarify? **Minor issues / suggestions** - $\sigma_{d^l_i}$ and $\sigma$ in line 67 are not defined. - In Equation (1) and the two lines after it, $\sigma$ depends on $l$, unless one assumes all weights have the same dimensions. This dependency should be made more explicit. - In Example 2.1, is the set of weight matrices $\{W_1, …, W_{L}\}$? If so, shouldn’t the indices of $S$ and $\sigma$ range from 1 to $L$ instead of $L+1$? - Since Theorem 3.1 uses basis for linear equivariant maps, $\mathbb{L}$ defined in line 119 should be the space of linear equivariant maps, instead of the space of equivariant maps. - Line 165 typo “a a character” - What is $W$ on the right side of the equation in line 172? - Line 538 “next section” -> previous section

Rating

5

Confidence

4

Soundness

3

Presentation

3

Contribution

2

Limitations

The authors include an informative limitation section that points out challenges and future directions.

Summary

Extending from recent works, the authors propose a new neural network layer architecture that enforces permutation equivariance in the weight space. The proposed architecture is used for learned optimizers and compared with other existing methods in a series of tasks, archiving improvement over the state-of-the-art methods.

Strengths

1. Methodology is well-developed. 2. Tests on various architectures show improvement over state-of-the art methods.

Weaknesses

1. The presentation of the methodology is not very clear (specifically around Eq. 9 and 10). 2. The presentation of the results is also not very clear and the results are not sufficiently reported. 3. Whilst the training times for the test cases are reported, the inference speed is not reported.

Questions

See weaknesses. Please improve the presentation of methodology and results.

Rating

5

Confidence

2

Soundness

3

Presentation

2

Contribution

2

Limitations

The authors have adequately addressed the limitations.

Summary

The study develops permutation equivariant models, developing an algorithm for defining a permutation eqivariant map of tensors with arbitrary rank, applying them in training learned optimizations and generalization prediction and finds that these class of models have improved performance on weight space tasks. This algorithm can be adapted to residual networks and transformers, a substantive selling point of this work.

Strengths

The paper is pretty strong, showing good results for tasks which are traditionally hard to work on. The method is novel, and applicable for various weight space tasks. The methodology is well introduced and explained well.

Weaknesses

line 219, operate is mentioned twice. There is a limited number of datasets evaluated upon. I would like to see a table showing the computational efficiency/time of this method vs other methods.

Questions

Can you produce a table of computational speed + resource comparisons?

Rating

7

Confidence

2

Soundness

3

Presentation

4

Contribution

4

Limitations

NA

Response to rebuttal

Thank you for your response. > As suggested by a reviewer, we also ran a Deep Set comparison > for completeness, we are currently running an RNN generalization prediction experiment with a non-equivariant baseline (a simple MLP ingesting all the weights) and will update when the results are available. These extra baselines will be very helpful for contextualising the results of the paper. This was my primary criticism of the paper, so it's nice to see it addressed, even if all it took was adding a couple more baselines. I also thank the authors for clarifying my points of confusion in the paper. Hopefully small changes / additions to the paper can be made so that other people will not run into these same comprehension issues. Other reviewers have raised some interesting points, but my evaluation of the paper remains positive, and so I shall be maintaining my score of 7 (accept).