Universal Neural Functionals

We thank the reviewer for their thoughtful comments.

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.

We believe this scale of improvement is typical in the literature for generalization prediction; despite its simplicity StatNN is a very strong baseline. For example, in “GNNs for learning equivariant representations” (Kofinas et al (ICLR 2024 oral)), the best method NG-T outperformed StatNN by only ~0.02 on rank correlation. As suggested by a reviewer, we also ran a Deep Set comparison, which significantly underperforms both our method and StatNN.

Test rank correlation between predicted and actual performance (higher is better):

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.

This would be an interesting experiment to run, but to our knowledge a non-equivariant method (such as a simple neural network ingesting all weights) would be extremely computationally costly to run in learned optimization due to the high dimensionality of weights. For this reason, we are not aware of learned optimizers in the literature that use this type of (non-equivariant) architecture. Instead, prior works prefer to use per-parameter MLPs, as you mentioned.

However, 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.

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

Good point, two weights having the same size does not necessarily mean their dimensions are permuted the same way. This is just a limitation of our notation--$\mathbb{R}^{n_1 \times n_2}$ is our way of saying that permutation $\sigma_1$ affects dimension 1 and permutation $\sigma_2$ affects dimension 2. For a matrix with the same size but with different permutations, we might write it as belonging to $\mathbb{R}^{n_3 \times n_4}$, even if $n_3=n_1$ and $n_4 = n_2$. Unfortunately, we are not aware of a better notation for conveying this point, but we will clarify it in the text.

On line 182, could you expand on what is meant by "Entries that are not explicitly assigned by the left-hand side are 0."?

Equation 11 is assigning values to entries of the tensor $E(W^{(m)})$, but in general the indices generated by the characters on the LHS would not cover all the possible indices of the tensor. So any indices not specified by Equation 11 are assumed to be 0.

In Eq. 15 are we manually computing the momentum term

The momentum terms fed as input are computed manually / in the standard way.

Did you use this positional encoding fix in your experiments?

We did not find positional encoding necessary in this case, no. We will clarify the text on this

My understanding of the UNF is that it is technically not a neural network [...] but rather we learn coefficients for a particular set of basis functions that you have derived [...]

Your statement "we learn coefficients for a particular set of basis functions that you have derived" is correct, and moreover the learned linear combination of basis functions forms a single "layer" that we stack with nonlinearities that we stack and optimize much like a neural network, which is why we refer to it as one. One can say that its layers are different from that of common neural networks, but also the layers of a convolutional network are different from that of an RNN.

There are some small changes I would make to the paper for readability, please see my suggestions below.

Thank you for catching these issues and providing suggestions, we will include them in the updated manuscript.

We thank the reviewer for their detailed feedback. We will also include the suggested relevant works in our discussion.

it does not automatically constructs the specifications which is actually a quite tricky part for some complicated models

We agree that constructing specifications can be tricky, and we will include the specifications for our RNNs and Transformers. Analyzing the computation graph to automatically deduce the specifications is an interesting idea and may be possible if we first specify how each of the primitive operations use the dimensions of their parameter tensors. For example, a Linear layer implemented as $Y=XW$ always connects the first dimension of its parameter W to the second dimension of the previous layer's parameter.

How Algorithm 1 works in practice for the networks in the experiments? How the valid partition looks like for those networks?

Thanks for raising this very interesting point. Although the bases in general can look quite complex, we can already analyze specific differences in the bases generated for processing MLPs and RNNs. For example, valid partitions of the kind we give in Example 3.2 ($\mathcal{W}^{(m)}=\mathcal{W}^{(\ell)}=\mathbb{R}^{n_1 \times n_1}$) appear for RNNs but not for MLPs, because in RNNs the outputs at one timestep are used as the inputs for the next timestep. Hence we have the same permutation action on the input space and output space of recurrent weights. We will expand our analysis of these differences in the paper to better connect theory and practice.

RNN generalization prediction: It's questionable that 0.8968 is significantly better than 0.8839.

This scale of improvement is typical in the literature for generalization prediction; despite its simplicity StatNN is a very strong baseline (see the additional comparison we just ran again Deep Sets). For example, in “GNNs for learning equivariant representations” (Kofinas et al (ICLR 2024 oral)), the best method NG-T outperformed StatNN by only ~0.02 on rank correlation.

RNN generalization prediction: More baselines are needed.

As suggested, we have added a Deep Sets (Zaheer et al, 2017) comparison. As the results below show, our method (UNF) and StatNN (Unterthiner et al, 2020) outperform the method based on Deep Sets by a wide margin.

Test rank correlation between predicted and actual performance (higher is better):

it's questionable that Deep Sets is "the default architecture choice for f".

To clarify, what we call the "Deep Set" baseline in our learned optimization experiments is actually the per-parameter MLP used in Metz et al, 2022. This is because if you exclude the $\gamma$ term in Deep Sets (Zaheer et al, Eq 4) and stack multiple layers, the result is equivalent to a per-parameter MLP. We apologize for the confusion and will clarify the text. Regardless, the actual experiments we ran are in fact comparing with the standard architecture choice for learned optimizers, such as the one used by Metz et al, 2022.

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

We learn these hyperparameters using the same ES optimizer as we use for all the other learned optimizer results.

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?

The main purpose of the learned optimizer experiments is to demonstrate the impact due to architecture for learned optimizers with the same backbone, in this case SGDM. However, we are currently running experiments with Adam as the backbone as well.

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?

This is a good question--we did not find it necessary to add positional encoding for the spatial positions in UNF. One explanation is that spatial parameter sharing makes UNF more parameter-efficient, which can be helpful for ES optimizers like the ones used for learned optimization, since they can struggle at very large (high dimensional) parameter spaces. We will expand on this point in our experimental discussion.

We thank the reviewer for their careful analysis and feedback. We will incorporate the suggestions and fixes proposed here.

Despite the claims in the abstract and conclusion, there has been other work …

We agree that the novelty of this work lies in the construction of maximally expressive equivariant layers, and in empirical results on learned optimizers for complex architectures like RNNs and Transformers. We will update the abstract, intro, and conclusion to emphasize this aspect of our contribution and also address the contribution of related graph-network approaches.

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? We agree that the explanation can be difficult to follow due in part to dealing with indices for arbitrary-rank tensors. One justification for the definition of valid partitions is given in the proof in Appendix B.1: the valid partitions can be identified with equivalence classes of indices (a,b) that are used to define the basis matrices (Eq (16)).

For intuition, we may also work out parameter sharing patterns for simple permutation equivariant layers. Following Ravanbakhsh et al, we can find the parameter sharing for a layer by studying the orbits of the input and output index spaces. By studying a few examples (such as the one we give in Example 3.3) one observes that each orbit must be identified with a valid partition of the indices. We will expand our examples to include this intuition.

the number of valid partitions can be very large especially when many indices permute simultaneously or there are many layers

We agree that this can be a limitation, and characterize the exact growth of the number of basis functions in Eq 19. Our goal and theoretical contribution here was to characterize maximally expressive equivariant layers, i.e., including all possible basis terms. It is an interesting area of future research to consider whether one could select a good subset of the basis functions that perform well in practice, while also being computationally cheaper.

the proposed method is not compared to the most recent permutation equivariant architectures

We will attempt to include comparisons to these architectures for the updated manuscript, though full comparisons are challenging because, to our knowledge, Lim et al does not provide source code and Kofinas et al have not yet published their learned optimization code for direct comparison.

how did the authors decide on the two tasks to conduct experiment on?

Learned optimization is a natural and challenging task for methods that can learn to edit the weights of other networks, and would also potentially have plenty of downstream impact. Generalization prediction is also an interesting weight-space task studied in prior relevant work (Navon et al, 2023, Zhou et al 2023).

Although many past weight-space papers also included experiments on INRs, we omit INR experiments in this paper because all the INRs involved in those experiments were actually MLPs, whereas the focus of our work is to extend to processing more general architectures.

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

Good point, L273-276 was only meant when comparing UNF to Deep Sets (the standard architecture for learned optimizers until recently), not NFNs. We will clarify the text on this point.

What is $W$ on the right side of the equation in line 172?

Should be the input $W^{(m)}$, thank you for catching this

In Example 2.1, is the set of weight matrices $W_1, \cdots, W_L$? If so, shouldn’t the indices of $S$ and $\sigma$ range from $1$ to $L$ instead of $L+1$

Not quite, a feedforward network with $L$ weight matrices has $L+1$ layers of neurons, when we include the input and output layers. And neurons are the things that give rise to our permutation symmetries. For example, consider 2 weight matrices $W_1,W_2$. There are three layers of neurons: input, hidden, and output.

We thank the reviewer for their review and suggestions, which we have incorporated into the draft manuscript.

presentation of the methodology is not very clear (specifically around Eq. 9 and 10)

We apologize for any confusion, and have polished the presentation throughout, including more intuitive explanations for Eq. 9 and 10 and also expanded examples. If there are any specific points of confusion to address, we would also welcome more detailed comments.

the results are not sufficiently reported

We have run additional experiments in response to various requests from other reviewers--we welcome any additional specific feedback on the presentation or reporting of the results.

We thank the reviewer for their thoughtful assessment--we are glad that the review recognized the novelty of the work and the challenging nature of the problem space.

There is a limited number of datasets evaluated upon.

We will expand the evaluation in a few ways, based on suggestions from reviewers. First, we added a Deep Sets baseline for the RNN generalization experiment, and found that UNF continues to perform best (see table of test rank correlations below). Additionally, we are currently generating generalization prediction datasets for other architectures such as Transformers, and will include those results once available.

Rank correlation between predicted and actual performance (higher is better):

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

Appendix Section C.3 currently contains information about the computational costs of running our UNF methods, we will update it to also include computational cost numbers for the baselines.