Universal Neural Functionals

Authors: allan  zhou
chelsea  finn
james  harrison
google  deepmind
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 19:09:00

Abstract: A challenging problem in many modern machine learning tasks is to process weight-space features, i.e., to transform or extract information from the weights and gradients of a neural network. Recent works have developed promising weight-space models that are equivariant to the permutation symmetries of simple feedforward networks. However, they are not applicable to general architectures, since the permutation symmetries of a weight space can be complicated by recurrence or residual connections. This work proposes an algorithm that automatically constructs permutation equivariant models, which we refer to as universal neural functionals (UNFs), for any weight space. Among other applications, we demonstrate how UNFs can be substituted into existing learned optimizer designs, and find promising improvements over prior methods when optimizing small image classifiers and language models. Our results suggest that learned optimizers can benefit from considering the (symmetry) structure of the weight space they optimize.

Main paper: 
 introduction: 
 Many problems in machine learning require handling weight-space features, such as the weights, gradients, or sparsity masks of neural networks. For example, optimizers iteratively map the current weights and gradient history to updated weights. Taking this perspective, researchers have proposed a variety of data-driven methods that train a neural network to process these weight-space features. Examples applications of these neural functionals [Zhou et al., 2023a] include training neural networks to predict classifier generalization from weights [Eilertsen et al., 2020] , to optimize other networks [Metz et al., 2022] , and to classify or edit implicit neural representations (INRs) [De Luigi et al., 2023] .Until recently, researchers lacked a unifying and principled framework for designing neural functionals, and would implement a custom model for their particular weight-space task. A significant recent advance was the development of weight-space models that are permutation equivariant [Navon et al., 2023 , Zhou et al., 2023a] . Neuron permutation symmetries arise in a neural network's weight space because re-ordering hidden neurons has no effect on the network's function [Hecht-Nielsen, 1990] . A permutation equivariant neural functional can guarantee that under a neuron permutation of its input, its output permutes accordingly. 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). Permutation equivariant neural functionals [Zhou et al., 2023a] added the ability to process weights from simple feedforward convolutional networks (CNNs). However, in practice we may deal with the weight spaces of complex networks that have residual connections, recurrence, normalization layers, and so on. Extending existing approaches to each possible weight space would be tedious and challenging.Figure 1 : Illustration of the permutation symmetries in the weight space of a recurrent neural network (Example 2.2). Left: Each layer contains feedforward (ff) weights mapping between different layer's activations, and recurrent (rec) weights transforming activations over time. We can permute the hidden activations as illustrated without changing the final outputs h L t . Right: Permuting the hidden activations induces a permutation on the weights. Here, the rows and columns of the feedforward weights are permuted by (σ ℓ+1 , σ ℓ ), while the recurrent weights are permuted by (σ ℓ , σ ℓ ). Our algorithm automatically constructs permutation equivariant models for any collection of weight tensors given a description of its symmetries (Appendix A).We propose an approach that automatically constructs permutation equivariant models for any collection of tensors whose dimensions can permute according to a shared set of permutations. This naturally encompasses the permutation equivariance we might desire for any given weight space. We show that our algorithm constructs the most general linear layer that operates on a given weight space while guaranteeing equivariance to the specified permutation symmetries. Stacking multiple such layers with pointwise nonlinearities produces a deep permutation equivariant model, which we refer to as a universal neural functional.To evaluate the empirical effectiveness of UNFs, we apply them to tasks that require processing networks with complex architectures containing recurrence, layer normalization, residual connections, and more. We use UNFs to implement learned optimizers and then optimize small image classifiers, RNNs, and Transformer language models, observing promising improvements over prior methods. In a generalization prediction task, we use UNF to predict the performance of sequence-to-sequence RNN models from their weights. Our experiments show that universal neural functionals are flexible, can be easily applied to different weight spaces, and improve upon prior weight-space methods.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: The abstract contains exactly the description of the algorithm we developed and experiments we ran.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
preliminaries: 
 We largely follow or extend the notation and naming of Zhou et al. [2023a] . Given a fixed neural network architecture, there is a weight space W of possible parameters (weights, biases, normalization scalings, etc.) . We refer to all such parameters as "weights". A particular set of weights W = W (1) , • • • , W (L) contains multiple "tensors", or multidimensional arrays. Depending on the architecture, W contains numerous symmetries [Hecht-Nielsen, 1990 , Godfrey et al., 2022] , i.e., transformations on the weight space that do not affect the network's behavior. Following prior work [Navon et al., 2023 , Zhou et al., 2023a] , this work focuses only on the permutation symmetries, which are called neuron permutations.Neuron permutations correspond to re-arranging the neurons within (hidden) layers, which have no canonical ordering. We make the simplifying assumption that all layers can be re-arranged-this assumption can be later corrected using positional encodings [Zhou et al., 2023a] . Assuming there are N independently permutable layers of neurons, the neuron permutation group is the direct product S = S n1 × • • • × S n N , where n i is the number of neurons being permuted in each layer.In general, each weight is a "tensor" (multi-dimensional array) of real numbers. Using M (a, b,• • • ) to denote arrays R a×b×••• , consider a rank-D ℓ tensor W (ℓ) ∈ M n d ℓ 1 , • • • , n d ℓ D ℓ. Each dimensiond ℓ i is permuted by σ d ℓ i .That is, the action of σ on the indices of the weight tensor is:EQUATIONDefining the multi-index ⃗ i := (i 1 , • • • , i D ℓ ), the action on the weight tensor is to permute the entries: i) , and the action on W is σW := σW (1) , • • • , σW (L) . We now elaborate on the definition of the group and action in several common cases. Example 2.1 (Multilayer perceptron). A multilayer perceptron (MLP) with L + 1 layers has activations h ℓ+1 = s W (ℓ) h ℓ + b (ℓ+1) , with h 1 being the first (input) layer and h L+1 the output. If each h ℓ is a vector of length n ℓ , then the weights are matrices W (ℓ) ∈ M (n ℓ+1 , n ℓ ) and the biases are vectors b (ℓ) ∈ M (n ℓ ). Then we have a neuron permutation group S = S n1 × • • • × S n L+1 , and σ ∈ S can be written σ = (σ ℓ ) L+1 ℓ=1 . The action on the weights and biases is:σW (ℓ) ⃗ i := W (ℓ) σ -1 ( ⃗EQUATIONwhere P (σ ℓ ) is the n ℓ × n ℓ permutation matrix corresponding to σ ℓ . This corresponds exactly to the "NP" setting in Zhou et al. [2023a] .Example 2.2 (Recurrent neural network). Consider a deep recurrent neural network (RNN) [Elman, 1990] without biases. We follow the presentation of Wang et al. [2023] :EQUATIONwhere h 1 t are the inputs and h L+1 t are the outputs at each timestep, with h ℓ 0 initialized to 0. The weight space consists of feedforward (ff) weights W ℓ ff ∈ M (n ℓ+1 , n ℓ ) and recurrent (rec) weights W ℓ rec ∈ M (n ℓ , n ℓ ). We again define the neuron permutation group S := S n1 × • • • × S n L+1 , but the action of the group on the weight space is now different. Here, re-arranging the neurons corresponds to transforming the weights:W ℓ ff → P (σ ℓ+1 ) W ℓ ff P (σ ℓ ) ⊤ and W ℓ rec → P (σ ℓ ) W ℓ rec P (σ ℓ ) ⊤ .As illustrated by Figure 1 , the feedforward weights transform just as in the MLP case (Eq. 2), but the recurrent weights' rows and columns must be transformed by the same permutation.Example 2.3 (Convolutional neural network). Consider a 1D convolutional neural network (CNN) without biases. Using ⋆ to denote cross-correlation, we have activations h ℓ+1 = s W (ℓ) ⋆ h ℓ , where the input is h 1 and the output is h L+1 . If each filter has spatial dimension k ℓ and each h ℓ has n ℓ channels, then we have rank-3 weight tensors W (ℓ) ∈ M (n ℓ+1 , n ℓ , k ℓ ) and neuron permutation group S = L ℓ=1 S n ℓ × S k ℓ . Looking at how each dimension of W (ℓ) permutes, we would have σ n ℓ+1 ∈ S n ℓ+1 permute the first dimension (output channels), σ n ℓ ∈ S n ℓ permute the second dimension (input channels), and σ k ℓ ∈ S k ℓ permute the third dimension (spatial).We note that permutating the spatial dimensions of a convolution filter would change the CNN's behavior and is not a true symmetry of the weight space. This is a notable difference between how our framework handles convolutional weight spaces compared to NFNs [Zhou et al., 2023a] , where the action of the neuron permutation group does not affect the spatial dimensions at all. Assuming that all dimensions of each weight tensor can permute simplifies the development of our framework, and undesired symmetry can be broken (if desired) by positional encodings of the input [Zhou et al., 2023a , Lim et al., 2023] .Equivariance and invariance. We are interested in functions T : W → W that are equivariant, meaning that it doesn't matter whether we apply a neuron permutation to the input or the output. We define L S (W, W) as the space of equivariant linear maps, i.e., those T satisfying:EQUATION)Our goal is to design a layer (i.e., a parameterized space of functions) that is equivalent to L S (W, W).In some applications, we may instead desire invariance, that is a function P satisfyingEQUATIONFollowing prior work [Navon et al., 2023 , Zhou et al., 2023a] , we can build invariant neural functionals by composing several equivariant layers with an invariant pooling layer, e.g., one that sums over every dimension of each weight tensor and concatenates the results.Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes] Justification: This is discussed at various points of the paper, including in the Conclusion (final section of the main paper).Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
universal neural functionals: 
 Since equivariance is preserved under composition, and pointwise non-linearities are already permutation equivariant, we can build deep equivariant models as long as we have an equivariant linear layer. Additionally, composing equivariant layers with an invariant pooling operation produces a deep invariant model. This section introduces a method for producing equivariant weight-space layers for any given weight space, which enables the flexible construction of universal neural functionals.The weight space is a direct sum of individual weight subspaces L) , so the problem of defining an equivariant layer on W can be decomposed into defining equivariant layers between each pair of weight subspaces W (m) and W (ℓ) , for all ℓ and m [Navon et al., 2023] .W = W (1) ⊕ • • • ⊕ W (We re-state this result in our own notation. For any ℓ, m pair we define L S W (m) , W (ℓ) as the space of equivariant maps between the two weight subspaces. It contains all T ℓm : W (m) → W (ℓ) satisfyingEQUATIONnoting that the action on the left and right hand sides of the equivariance condition are not, in general, the same.Assume that we already have a basis B sp for L S W (p) , W (s) . A basis function E ∈ B sp can be extended to Ē : W → W by defining:EQUATIONwhereĒ(W ) := Ē1 (W ), • • • , ĒL (W ) .Theorem 3.1 (Navon et al. [2023] ). Let { B ℓm } be bases for each L S W (m) , W (ℓ) . Then the union of these bases (extended by Eq. 7) is a basis for linear equivariant maps on W. That is, we have the basis B for L S (W, W) defined:EQUATION)This result tells us that we can construct an equivariant basis B for L S (W, W) by simply combining the equivariant bases { B ℓm } for each pair of weight subspaces.Algorithm 1 Basis for equivariant W (m) → W (ℓ) layer Require: W (m) , W (ℓ) 1: Initialize basis B ℓm ← { } 2: I ← { o 1 , • • • , o D ℓ , i 1 , • • • , iE P (X) c[o1],••• ,c o[D ℓ ] := R X c[i1],••• ,c[i Dm ]9:B ℓm ← B ℓm ∪ { E P } 10: end for 11: return B ℓm Since weights are tensors, our decomposed problem involves finding bases for permutation equivariant maps between tensors. Variants of this problem have been studied by numerous prior works-in particular, Maron et al. [2018] theoretically characterize a basis for equivariant maps between arbitrary-rank tensors, and provide a concrete implementation of the basis functions in the rank-2 case. Here, we describe a general algorithm that automatically constructs a basis for permutation equivariant maps between arbitrary-rank tensors. Concretely, it implements each basis function in terms of simple array operations that are amenable to efficient computation with modern deep learning frameworks.Functions in L S W (m) , W (ℓ) take input tensors indexed by { i 1 , • • • , i Dm } and produces output tensors indexed by { o 1 , • • • , o D ℓ }. We can construct a basis B ℓm for this space where each element is identified by a valid partition P of these indices. Recall that the indices(i 1 , i 2 , • • • ) of W (m) are permuted by σ d m 1 , σ d m 2 , • • • .We say that two indices i 1 and i 2 "permute simultaneously" if d m 1 = d m 2 . Definition 1. A valid partition is a partition P of the output and input indicesI = { o 1 , • • • , o D ℓ , i 1 , • • • , i Dm } intonon-empty subsets, such that each subset only contains indices that are permuted simultaneously.Example 3.1 (W (m) = W (ℓ) = R n1×n2 ).Here the output and input indices are{ o 1 , o 2 , i 1 , i 2 }. The partition { { o 1 , o 2 } , { i 1 , i 2 } } is not valid because o 1 ,o 2 are permuted by σ 1 , σ 2 , so they do not permute simultaneously. On the other hand,{ { o 1 , i 1 } , { o 2 , i 2 } } is a valid partition. Example 3.2 (W (m) = W (ℓ) = R n1×n1 ). This time, the partition { { o 1 , o 2 } , { i 1 , i 2 } } is valid because o 1 , o 2 are both permuted by σ 1 , as are i 1 , i 2 .To construct the equivariant basis, we enumerate all valid partitions and then map each partition P to a basis function E P . Concretely, we label each subset of P with a distinct character α, β, γ, • • • and then remap each of our original indices{ o 1 , • • • , o D ℓ , i 1 , • • • , i Dm } toa a character based on which subset the index was in. This mapping is best illustrated by continuing our previous example.Example 3.3 (W (m) = W (ℓ) = R n1×n2 ).Here input and output are both matrices, with combined indices { o 1 , o 2 , i 1 , i 2 }. We have two permutations (σ 1 , σ 2 ) ∈ S n1 × S n2 that can act on the rows and columns of the input and output matrices. There are four valid partitions:EQUATIONConsider P 2 -we assign a character to each subset:EQUATION)which tells us to remap the output indices (o 1 , o 2 ) → (α, β) and the input indices(i 1 , i 2 ) → (α, γ), producing the basis function E P2 W (m) αβ := γ W (m)αγ , where summation over γ can be inferred because it only contains an input index.Repeating this index-remapping process for each valid partition will generate a total of four basis functions E P1 , • • • , E P4 for L S W (m) , W (ℓ) . Our equivariant W (m) → W (ℓ) layer will be defined as the linear combination m) , which is the layer introduced in Hartford et al. [2018] .T ℓm W (m) ; λ := 4 k=1 λ k • E P k W (To generalize the previous example, for each valid partition of the indices P we label its subsets with characters α, β, γ, • • • and then construct a basis function:EQUATIONwhere c[•] maps each index to the subset of P that contains it. We sum over the characters in R, which is the (possibly empty) subset of characters that only contain input indices (i.e., only appear on the right-hand side). Entries that are not explicitly assigned by the left-hand side are 0. Algorithm 1 gives a formal description of the complete process for generating B ℓm .Theorem 3.2. Algorithm 1 produces a basis for the equivariant linear maps from W (m) to W (ℓ) .Proof. See Appendix B.1.Once Algorithm 1 has generated a basis of equivariant functions B ℓm , we can implement an equivariant layer using a vector λ ℓm ∈ R |B ℓm | of learned coefficients:EQUATION)Theorem 3.1 now tells us that we may now construct the equivariant weight-space layer by combining the bases { B ℓm } into a basis B of functions on W. The weight-space layer T (•, λ) can then be defined by a linear combination of the basis functions with learned coefficients λ. Explicitly, the full layer is defined:EQUATIONwhere m) , λ ℓm . Appendix A provides a concrete description of how we specify the weight space in code and how the algorithm is then used to automatically construct an equivariant weight space layer. Our open-source implementation is compatible with most JAX [Bradbury et al., 2018] neural network libraries. Theorem 3.3. The weight-space layer (Eq.-13) is S-equivariant, and can express any linear equivariant function on W.λ ℓ,: = { λ ℓm | ℓ = 1, • • • , L } and T ℓ W, λ ℓ,: = L m=1 T ℓm W (Proof. Each T ℓm is a linear combination of basis functions in B ℓm . Then, as described by Thm 3.1, Eq. 13 is a linear combination of functions that form a basis for L S (W, W).For an MLP weight space with neuron permutation group defined as in Example 2.1, this approach will generate the exact same layer as NFN NP [Zhou et al., 2023a] . This is because the layers each parameterize all possible linear maps equivariant to the same symmetry group, and hence can express the same set of functions.In practice, we may be interested in simultaneously processing multiple weight-space features, such as the weights and a history of gradients. These features can be stacked into a "channel" dimension analogous to the channels of convolutional networks. In that case, we must consider direct sums of weight spaces of the form W c = ⊕ c k=1 W, with elements that can be written as Extending equivariant layers to the multi-channel setting is quite common in the geometric deep learning literature and simply involves taking linear combinations along the channel dimension [Cohen and Welling, 2016, Ravanbakhsh et al., 2017] . That is, we modify the equivariant layer between subspaces as:1 W = (W [1], • • • , W [c]), for W [k] ∈ W. Then the action is σW := (σW [1], • • • , σW [c]) for σ ∈ S,EQUATIONwhere each λ ℓm b is now a learned c o × c i matrix instead of a scalar.The previous sections describes the construction of S-equivariant layers that operate operate on weight-space features in W c . We construct universal neural functionals by stacking multiple such layers (interleaved with pointwise non-linearities) into a deep, permutation equivariant model that can process weights. To construct a permutation invariant model, we can add an invariant pooling layer after the equivariant layers, as in prior work [Navon et al., 2023 , Zhou et al., 2023a] .Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [No] Justification: We provide code for implementing the proposed algorithms, but data and code for some of the experiments could not be released due to proprietary restrictions. However, we do include details for how to implement these experiments in the paper. Guidelines:• The answer NA means that paper does not include experiments requiring code. • If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 10. Broader Impacts Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed? Answer: [NA] Justification: There are no obvious relevant societal impacts. Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
experiments: 
 In this section, we refer to weight-space models constructed using our algorithm as universal neural functionals (UNFs). We compare their performance to prior methods on two types of weight-space tasks: predicting the generalization of recurrent sequence-to-sequence models, and training learned optimizers for a variety of architectures and datasets.One promising application of neural functionals is in predicting the generalization of neural network models from their weights Eilertsen et al. [2020] . We construct Tiny RNN Zoo2 , a dataset of recurrent neural networks trained to do arithmetic by completing given questions character-by-character. For example, given the input string "15+20=" the correct completion would be "35<EOS>". To construct the dataset, we train 10 4 sequence-to-sequence [Sutskever et al., 2014] models on example problems with input numbers up to five input digits. Both encoder and decoder RNNs contain a single GRU cell [Chung et al., 2014] The success rate of each RNN model is clearly invariant under permutation symmetries of its weights, so invariance is a natural inductive bias for any generalization predictor. We evaluate STATNN [Unterthiner et al., 2020 ] and a UNF-based predictor (note that NFNs are not applicable to the weights of recurrent networks). STATNN is operates on basic statistical features3 of the weights, and has been shown to be a very strong baseline on previous generalization prediction tasks [Unterthiner et al., 2020] . On the other hand, UNF operates on raw weight inputs and may be able to extract more nuanced signals than STATNN, as was shown (for CNN classifiers) in Zhou et al. [2023a] .In particular, STATNN computes the mean, variance, and (0, 25, 50, 75, 100)-percentiles of each weight tensor in the RNN and feeds them into a six-layer MLP with hidden width 600. UNF is a permutation invariant model, implemented using a three-layer equivariant backbone (16 hidden channels) followed by invariant pooling and a three-layer MLP (512 hidden neurons). We train each predictor with binary cross entropy loss (since the target SR is in [0, 1]), using the Adam optimizer with learning rate 0.001, batch size 10, and training for up to 10 epochs. We use the validation data only for early stopping, and assess the performance of each predictor on the test inputs using Kendall's τ , the rank correlation between predicted and actual success rate.Results. Table 1 shows the performance of each predictor on held out weight inputs. Our UNF-based predictor achieves significantly higher rank correlation between predicted and actual success rate, suggesting that the equivariant layers are able to extract more informative features from the raw weights compared to STATNN.Choosing the optimizer is a key step in training any modern neural network. Though most popular optimizers are variants of stochastic descent, the non-convexity of neural network training leaves few rigorous guidelines for ideal optimizer design. This has led some researchers to propose training good optimizers using some form of meta-learning [Bengio et al., 1990 , 2013 , Andrychowicz et al., 2016 , Wichrowska et al., 2017 , Metz et al., 2019] .Common optimizers today (including the learned ones) are equivariant to any permutation of the weights. This is because permuting the weights also permutes the gradients, so stochastic gradient descent and similar optimizers will produce permuted updates. However, equivariance to any permutation ignores the actual symmetry structure of the optimized neural network. Arguably the more appropriate constraint is to only require equivariance to the neuron permutation group, which enables more expressive optimizers while still respecting the symmetries of the weight space. As we will see, this can be achieved by using UNFs to implement a learned optimizer.Training learned optimizers that generalize well is extremely compute-intensive [Metz et al., 2022] , so we conduct our experiments in several smaller settings to analyze the impact of architecture Figure 2 : Training loss (negative log-likelihood) curves for different tasks and architectures using meta-learned optimizers. We implement learned optimizers with either universal neural functionals (UNFs), NFNs [Zhou et al., 2023a] , or Deep Sets [Zaheer et al., 2017] . Deep Sets are the current standard choice for implementing learned optimizers. Note that NFN is identical to UNF in the MLP case, different for CNN case, and not applicable to RNNs or Transformers. All loss curves are smoothed and averaged over 5 random initializations (3 for Transformer), with shaded regions showing standard error. choice on learned optimizer performance. In each setting, an optimizer is meta-trained to optimize an architecture type on a task from random initializations. Following Harrison et al. [2022] , our learned optimizers track momentum terms m γ t ← γm t-1 + ∇ t and produce updates of the form:EQUATIONHere αm γ0 t is a "nominal term" that biases the learned optimizer to behave like stochastic gradient descent with momentum coefficient γ 0 . The neural functional f (•) ingests weights W t , gradients ∇ t , momentum terms at several coefficients { m γi t } i , and the iteration t. During meta-training, we optimize network f and scalars α, β, γ 0 to minimize the task training loss after a fixed number of training steps T , the "inner training horizion." To avoid the issue of backpropagating through an optimization process, we estimate meta-gradients using persistent evolutionary strategies [Vicol et al., 2021] .Comparisons. The default architecture choice for f (•) in prior work is Deep Sets [Zaheer et al., 2017] , which offers equivariance to any permutation symmetry. We study the effect of replacing Deep Sets by UNFs. We also try the NFN NP architecture [Zhou et al., 2023a] where applicable, though it cannot be used on the RNN and Transformer experiments. Finally, we consider stochastic gradient descent with momentum (SGDM), which is equivalent to fixing β = 0 in Eq. 15. The SGDM baseline is also meta-trained to tune the learning rate α and momentum decay rate γ 0 . We compare the different learned optimizers in four tasks: MLP on FashionMNIST. Each optimizer trains an MLP classifier on a downsized and flattened version of the FashionMNIST dataset [Xiao et al., 2017] . We note that for MLP weight spaces, UNF are identical to NFN NP [Zhou et al., 2023a] .CNN on CIFAR-10. Each optimizer trains a convolutional classifier on a downsized 16 × 16 CIFAR-10. In this setting our algorithm produces a UNF that is different to NFN NP (see Example 2.3).. Each optimizer trains a character-level RNN-based language model (LM) on the One Billion Word Language Model Benchmark (LM1B) dataset [Chelba et al., 2013] .Transformer on LM1B. Each optimizer trains a Transformer LM on LM1B, this time predicting tokens instead of characters.We use an inner training horizon T = 2,000 for the first three tasks and T = 5,000 for the Transformer task, since it takes longer to train. When implementing f (•) for each method, we use a network with four layers, 32 hidden channels, and ReLU nonlinearities. The Deep Set optimizer uses exclusively Deep Set layers [Zaheer et al., 2017, Eq. 4] , while the UNF and NFN optimizers uses three Deep Set layers followed by a single UNF or NFN layer. See Appendix C.1-C.2 for full descriptions of the tasks and meta-training.Results. Figure 2 shows the training curves produced by each of the meta-trained optimizers in each experiment. Learned optimizers with deep architectures (UNF, Deep Set, or NFN) outperform SGDM, even after tuning SGDM's learning rate and momentum decay. UNF typically learns fastest and achieves the lowest training loss across all methods, though Deep Set and NFN can be comparable in some settings. One interesting observation is that UNF outperforms NFN in the CNN experiment. As noted in Example 2.3, UNFs make the stronger assumption that all tensor dimensions-including the spatial dimensions of the convolution filter-are permutable, while NFNs do not. Although the UNF assumption is technically incorrect, the stronger assumption leads to a lower parameter count (see Table 3 in the appendix) which may be easier for meta-optimization.Overall, our results show the promise of using UNFs to create more expressive learned optimizers that utilize the specific symmetry structure of the weight spaces they optimize. Further work could investigate their capacity for generalization to new tasks and architectures, for example by metatraining on diverse tasks [Metz et al., 2022] . Moreover, as Table 3 in the appendix shows, a necessary trade-off of UNFs being more expressive is that they require more parameters for an equivalent number of layers and hidden channels. Since learned optimizers are still much smaller than the networks they could optimize, this may not be a significant computational constraint in practice. Still, it could be a challenge to meta-optimization, since evolutionary strategies are known to struggle in higher dimensions. Hence, further work on efficient high-dimensional meta-gradient estimators would complement the development of expressive weight-space models like UNF. 
related work: 
 There is a long history of neural network architectures that are equivariant to various symmetry groups [LeCun et al., 1995 , Cohen and Welling, 2016 , Ravanbakhsh et al., 2017 , Kondor and Trivedi, 2018 , Cohen et al., 2018] . Existing frameworks for automatically constructing equivariant models [Finzi et al., 2021] produce equivariant matrices, which would be intractable for our task. Our work constructs efficient equivariant basis functions for a particular class of permutation symmetries that arise in the weight spaces of neural networks. Permutation equivariant networks have been developed for sets [Zaheer et al., 2017] , matrices whose rows and columns permute independently [Hartford et al., 2018] , and tensors under higher-order permutation actions [Thiede et al., 2020, Pan and Kondor, 2022] -the latter may also be viewed as equivariant models on graphs or polytopes [Maron et al., 2018 , Albooyeh et al., 2019] . This work observes that a weight space is a collection of tensors under higher-order permutation symmetries, and develops equivariant models for that setting.There has been significant interest in designing architectures that that either optimize or generate neural network weights [Schmidhuber, 1993 , Ha et al., 2016 , Krueger et al., 2017 , Kirsch and Schmidhuber, 2021 , Peebles et al., 2022 , Metz et al., 2022] . Some works have identified the importance of respecting the relevant symmetries when implementing black box meta-learners [Kirsch et al., 2022] . However, precise characterizations of equivariant models on neural weight spaces are relatively recent and were initially restricted to simple feedforward models [Navon et al., 2023 , Zhou et al., 2023a,b] .A recent alternative approach has been to leverage message passing neural networks (MPNNs) [Zhang et al., 2023] to process weights as edges of a graph. Concurrent to this work, Kofinas et al. [2024] demonstrated applications of MPNNs to learned optimization for MLPs and CNNs and Lim et al. [2023] extended MPNNs to process general weight-spaces. MPNN-based approaches benefit from more flexible adaptation to heterogenous inputs, and the computational cost of message passing does not grow as rapidly as our basis-this is because our approach guarantees each linear layer to be maximally expressive while MPNNs do not. We give a more detailed exposition of this trade-off in Appendix B.3 
conclusion: 
 We introduce a method for constructing permutation-equivariant neural functionals that operate on arbitrary weight spaces, removing a major limitation of previous frameworks that were only applicable to the weight spaces of simple MLPs and CNNs. Our algorithm constructs maximally expressive equivariant linear layers for processing any collection of tensors given a description of their permutation symmetries, and implements these layers in terms of efficient array operations in standard deep learning frameworks. We empirically validate that the resulting universal neural functionals (UNFs) are effective at tasks that involve processing the weights and gradients of convolutional image classifiers, recurrent sequence-to-sequence models, and Transformer language models. In particular, we find that UNFs show promising improvements over existing learned optimizer designs in small scale experiments.Limitations and future work. It remains to be demonstrated how UNFs can be applied to heterogenous weight-space inputs, e.g., to have a single UNF act as a learned optimizer for any input architecture. Moreover, our experimental results only validate the promise of UNF-based learned optimizers in relatively limited settings, and more work would needed to test generalization across arbitrary tasks. Finally, computational tractability may be a significant challenge for more complex architectures as the number of basis terms generated by Alg. 1 would grow rapidly for higher rank tensors with higher-order interactions (see Appendix B.2). Resolving these challenges would further improve the scalability and applicability of neural functionals to weight-space tasks. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: No relevant risks.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [NA] Justification: No third-party assets used.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: Here we discuss the concrete specification that precisely describes a weight space and must be provided as input to the algorithm before it can construct equivariant weight-space layers. Our implementation is compatible with most JAX [Bradbury et al., 2018] neural network libraries.Suppose we wish to process an MLP's weights that are stored in a (nested) Python dictionary: params = { "layer1": {"weight": Array [64, 32] , "bias": Array[64]}, "layer2": {"weight": Array [64, 64] , "bias": Array[64]}, } Then a specification should match the nested dictionary structure but provide a string or integer name for each dimension of each array. The name tells the algorithm which permutation affects which dimensions of each array.In this example, the specification closely follows the MLP description in Example 2.1, where W (1) ∈ M (n 2 , n 1 ) is permuted as W (1) → P (σ 2 ) W (1) P (σ 1 )⊤ . specification = { "layer1": {"weight": ("n2", "n1"), "bias": ("n2",)}, "layer2": {"weight": ("n3", "n2"), "bias": ("n3",)}, } Providing this specification object to our algorithm is sufficient for it to deduce the symmetry group, its action, and construct the corresponding equivariant layer.Since most neural networks consist of repeating layers or blocks, the process of constructing the specification can be semi-automated by first defining a function that creates the specification for a single layer or block and then re-using that function for each block. Although we did not find this necessary for our experiments, it may also be possible to automatically deduce the specifications for a network in common deep learning frameworks by analyzing its computation graph.B.1 Algorithm 1 generates a basis for L S W (m) , W (ℓ) Here we show that Algorithm 1 produces a basis B ℓm for L S W (m) , W (ℓ) , the space of linear equivariant maps between W (m) and W (ℓ) . Consider instantiating these linear maps as matrices multiplying flattened input vec W (m) . Maron et al. [2018] characterize a basis { B µ } µ for these matrices, where the entries of each basis matrix are defined:EQUATION)Here a ∈ I m and b ∈ I ℓ are multi-indexes for the input and output spaces, and µ ∈ I m × I ℓ / ∼ is an equivalence class of the combined input-output index space I m × I ℓ under the equivalence relation ∼ defined by a ∼ a ′ if and only if a i = a j ⇐⇒ a ′ i = a ′ j for all i, j, i.e. the two multi-indexes a, a ′ have the same equality pattern. Maron et al. [2018, Eq. 10b] , any equivariant linear map is defined:EQUATIONwhere I{•} is an indicator function for the given condition.Notice that each equivalence class µ is represented by what we call a valid partition of[D m + D ℓ ] := {1, • • • , D m + D ℓ },so this is already a sum over valid partitions as in Eq. 12. We can now observe that each term on the RHS is equivalent to one of our basis functions (Alg 1 Line 8). That is, for a given equivalence class µ represented by valid partition P:EQUATION)This is because for any I := (a, b) yielding a nonzero term on the LHS, if i, j ∈ [D m + D ℓ ] are grouped together by partition P then I i = I j , otherwise they would violate the equality pattern of µ. Therefore, we can replace all indices grouped together in a partition with a single shared symbol, i.e. the characters in Eq. 11.Hence, Algorithm 1 produces a basis that spans the same space of equivariant functions defined in Maron et al. [2018] , but constructs the basis functions in terms of efficient array operations instead of as matrices. Note that this is similar to the construction in Pan and Kondor [2022] , but generalized to multi-node sets (non-square tensors whose axes can potentially permute independently).Suppose we have a neuron permutation symmetry group ℓ) to be the number of indices that σ i ∈ S ni permutes in weight tensors W (ℓ) ∈ W (ℓ) (which could be 0). Finally, denote b(k) to be the k'th Bell number. Then the number of basis functions generated by Algorithm 1 is:S = S n1 × • • • × S n N , i.e., every neuron permutation σ is composed of N distinct permutations (σ 1 , • • • , σ N ). For each i = 1, • • • , N we define c i W (EQUATIONEach UNF layer can express any linear equivariant function on a given weight space (Thm 3.3).Compared to methods based on message-passing neural networks (MPNNs), this means UNFs can have very expressive individual layers, but may also be more computationally challenging due to the growth in the size of the basis (see next section).As an example, consider a simple "RNN" where h t+1 = W h t and h t ∈ R n has exchangeable entries, meaning that W → P W P T is a symmetry. Algorithm 1 would produce an equivariant basis with b(2 + 2) = 15 terms4 .On the other hand, we could construct a parameter graph [Lim et al., 2023] with n nodes and 2n 2 directed edges between them (allowing a forward and backward edge for each weight, equivalently n 2 undirected edges). Then using a similar construction to Lim et al. [2023, Appendix C.1.2] , we would get a linear GNN that computes:EQUATIONwhich is a linear combination of 6 equivariant basis functions, instead of 15. This leads to a potientially interesting trade-off between expressivity vs tractability. However, we also note that in practice MPNNs use non-linear MLPs in their message passing updates, and the comparison between UNF and MPNN-style approaches remains an open empirical question.Here we describe each of the experimental settings we evaluated the learned optimizers on. Across all experiments, the training loss is negative log-likelihood. [Xiao et al., 2017] . The MLP has a hidden size of 32 and ReLU activation function. We use a batch size of 128.CNN on CIFAR-10. Train a convolutional classifier on a downsized 16 × 16 CIFAR-10. The classifier has two convolutional layers (16 and 32 channels), followed by global average pooling and a linear classification head, and is trained with a batch size of 128.Trains a character-level RNN-based language model (LM) on LM1B [Chelba et al., 2013] . The RNN itself has one hidden layer with size 64, and uses identity-initialization [Le et al., 2015 ]. An embedding layer with dimension 32 maps tokens to embeddings before feeding into the RNN, and an output layer produces token predictions from the RNN output. The LM is trained to predict the next token with teacher forcing at batch size 64, on sequences of length 16.Transformer on LM1B. Train a Transformer LM on LM1B, this time predicting tokens instead of characters. The Transformer has two blocks with an embedding dimension of 32, and uses four self-attention heads. We train with a batch size of 8 on length-8 sequences.Call DS[c] a single equivariant Deep Set layer [Zaheer et al., 2017, Eq 4 ] with c output channels (similarly for UNF[c] and NFN[c]). Then f (•) in our learned optimizers (Eq. 15) is always implemented as a feedforward architecture:DeepSetOpt = DS[32] -> ReLU -> DS[32] -> ReLU -> DS[32] -> ReLU -> DS[1] UNFOpt = DS[32] -> ReLU -> DS[32] -> ReLU -> DS[32] -> ReLU -> UNF[1] NFNOpt = DS[32] -> ReLU -> DS[32] -> ReLU -> DS[32] -> ReLU -> NFN[1]For all methods, we initialize α = 0.1 and γ 0 = 0.9 before starting meta-training. For non-SGDM methods, we initialize β = 0.001, and provide six momentum values { m γi t } i with coefficients γ i = 0.1, 0.5, 0.9, 0.99, 0.999, 0.9999. The iteration number t is converted into an 11-dimensional sinusoidal encoding, and all inputs to f (•) are concatenated along the channel dimension. Concretely, this results in an input in W 19 . The output is in W 1 .We meta-train for 50,000 steps using Adam, estimating meta-gradients over 16 parallel training runs using persistent evolutionary strategies (PES) [Vicol et al., 2021] with a truncation length of 50 and a noise standard deviation of 0.01. The meta-training objective is training loss at the end of the inner training horizon (T = 5,000 for the Transformer setting, and T = 2,000 otherwise), and we apply a gradient clipping of 1.0.Size of each learned optimizer f (•). Since Deep Set layers are agnostic to the specific weight space being optimized, the Deep Set learned optimizer uses the same number of parameters in each task. The same is not true of UNF layers, where the number of parameters grows in proportion to the size of the bases generated by Algorithm 1. Table 3 lists the number of parameters in f (•) for each learned optimizer.Experiments were run on a mix of TPU v3 and v4 accelerators. On a TPU v3-8, training a UNF for our RNN generalization prediction task takes < 3 hours. Also on a TPU v3-8, meta-training a UNF for one of our learned optimizers takes ∼ 4 hours for the MLP task, ∼ 7 hours for the CNN task, and ∼ 20 hours for the RNN task.38th Conference on Neural Information Processing Systems (NeurIPS 2024).In the multichannel setting we overload notation and use W to refer to elements of W c , not W.Inspired by the Tiny CNN Zoo[Unterthiner et al., 2020].Notably, it computes statistics that are invariant to permutations of the weights.In this case, the full basis is also given byMaron et al. [2018, Appendix A]. 


Decision: The paper presents a novel approach to permutation equivariant models, developing an algorithm for creating permutation equivariant maps of tensors, which is applied to weight space tasks with promising results. The strengths of the study include its innovative methodology, which is well-explained and has been successfully applied to various neural network architectures, including RNNs and Transformers. The paper is well-organized and clearly written, with a strong theoretical foundation that is supported by experiments on a range of architectures, demonstrating improvements over state-of-the-art methods. Reviewers' concerns are properly addressed.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: Can you produce a table of computational speed + resource comparisons?
limitations: limitations: NA
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See weaknesses. Please improve the presentation of methodology and results.
limitations: limitations: The authors have adequately addressed the limitations.
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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
limitations: limitations: The authors include an informative limitation section that points out challenges and future directions.
paper_summary: 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.
strength_weakness: strength_weakness: 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: questions: For SGDM do the authors actually learn α and γ0 or it's simply tuned?
limitations: limitations: the limitations are discussed
paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: ## 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"
limitations: 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.

Comments: title: Response to rebuttal
comment: I appreciate reviewers' response, which address my concerns, therefore I raise the score.
Thanks for the reply. I will keep my score.
Thank you for the response, which clarifies most of my questions. My only remaining reservation is the lack of comparison, especially empirical ones, to several recent works that propose permutational equivariant weight-space models for general architectures (Zhang et al. 2023, Kofinas et al. 2024, Lim et al. 2023). I am maintaining my rating for now.
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).