Universal Neural Functionals
Authors: Allan Zhou
Chelsea Finn
James Harrison
Venue: NeurIPS 2024 poster (2024)
Keywords: architecture
permutation
weight-space
equivariance
learned optimization
Primary area: deep_learning_architectures

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.

Universal Neural Functionals
Allan Zhou Stanford University ayzêcs stanford, .effford.edu edu
Chelsea Finn Stanford University
James Harrison Google DeepMind j jamesharrisonggoogle.coo
Abstract
A challenging problem many modern machine learning tasks is to process weight-spact features, i.e., to transform or extract information from the weights and gradients of neural network. Recent works have developed promising weight-space models that equivariant to the permutation symmetries of sim- ple feedforward networks. However, they are not applicable to general architee- tures, 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 refer to universal neural functionals (UNFs), for any weight space. Among other ap- plications, demonstrate how UNFs 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. We open-source library for constructing UNFs https: bettp:///githbb com hamm/lllaaaagg/huuu wnwversalll...mmommm.mo..m.. seeiivereal_nuraa_ _eurralfuncttoona functional
Introduction
Many problems in machine learning require handling weight-space features, such the weights, gradients, or sparsity masks 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 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 [Filertsen et al. 2020). to optimize other networks [Metz al., 2022 and to classify edit implicit neural representations (INRs) [De Luigi
jet al. 2023).
Until recently, researchers lacked unifying and principled framework for designing neural function- als, and would implement a custom model for their particular weight-space task. significant recent advance was the development of weight-space models that are permutation equivarian! [Navon et al., 2023, Zhou et al., 2023a Neuron permutation symmetries arise in a neural network's 's weight space because re-ordering hidden neurons has effect the network's function Elecht-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 [hhou al.) 2023a] added the ability to process weights from simple feedforward convolutional networks (CNNs). However, in practice we may deal with the weight spaces 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.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

rec
( (e wi (se)
(oy) (se)
(levee) (ervs)
Figure 1: Illustration of the permutation symmetries in the weight space of a recurrent neural network (Example 2.2 Left: Each layer contains feedfowward (ff) weights mapping between different layer's activations, and recurrent (rec) weights transforming activations over time. We can permute the hidden activations illustrated without changing the final outputs h.. Right: Permuting the hidden activations induces a permutation on the weights. Here, the rows and columns of the feedforward weights are permuted by (0t+1, ot), while the recurrent weights are permuted by (ø Øt). Our algorithm automatically constructs permutation equivariant models for any collection of weight tensors given 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 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, apply them tasks that require processing networks with complex architectures containing recurrence, layer normalization, residual connections, and more. We use UNFs 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 suvence-to-sequence RNN models from their weights. Our experiments show that universal neural functionals are flexible, can be easily applied different weight spaces, and improve upon prior weight-space methods.
2 Preliminaries
We largely follow or extend the notation and naming of Zhou et al. (2023a]. Given a fixed neural network architecture, there is weight space W of possible parameters (weights, biases, normal- ization scalings, etc.). We refer to all such parameters as "weights"" A particular set of weights
W = (W())... .... w()) contains multiple "tensors", or multidimensional arrays. Depending on the architecture, W contains numerous symmetries (Hecht-Niclae
1990, Godfrey et al. 2022). ie.. 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-thi assumption can be later corrected using positional encodings (Zhou et al., [023a]] Assuming there are N independently permutable layers of neurons, the neuron permutation group is the direct product S Sn, X Sax- where n; is the number of neurons being permuted in each layer.
In general, each weight a "tensor" (multi-dimensional array) real numbers. Using (a, denote arrays Raxbx.... consider rank- Dr tensor (()) M (Maf."" ". .Rdy, ) Each dimension de is permuted by ozg. That is, the action of on the indices of the weight tensor is:
a (1..... ,in.) = (os (i1), odp (in))

---

Defining the multi-index /= (i..-.., the action on the weight tensor is to permute the entries:
[ow[]]; w (7) and the action on W is oW == (oW()) , ,,.,,,((,)
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 + layers has acti- vations h f++ (W(t) + b))+))), with h' being the first (input) layer and h the output. If each he is a vector of length ne, then the weights are matrices W(O) E (ne+11, and the biases are vectors b(0) E M (ne). Then we have neuron permutation group S Snn x SaLAY and a E S can be written a = (01)221 The action on the weights and biases is:
( + ((e41) - ((z) and (0)) + P(oe)blo
(2)
where P ((te) is the ny x ne permutation matrix corresponding to ag. This corresponds exactly to the "NP" setting in Zhou et al. 2023a
Example 2.2 (Recurrent neural network). Consider deep recurrent neural network (RNN) (Elman 1990] without biases. We follow the presentation Wang et al. (2023]:
h++1 = (Wit k + W
(3)
where h are the inputs and h are the outputs at each timestep, with hy initialized to 0. The weight space consists of feedforward (ff) weights W E M (me+1, ne) and recurrent (rec) weights we ree M (ne, nz).. We again define the neuron permutation group S == Sn, ... x Smita but the action of the group on the weight space is now different. Here, re-arranging the neurons corresponds
transforming the weights:
W + (( (PP)))
and W W P W
rec P(t))
As illustrated by Figure 7 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 permulation.
Example 2.3 (Convolutional neural network). Consider ID convolutional neural network (CNN) without biases. Using * to denote coss-correlation, we have activations h = s (W * h²), where the input is h2t and the output is h+++1 If each filter has spatial dimension ki and each le has ne channels, then we have rank-3 weight lensors (()) M (Ne+1, Thg ke) and neuron permutation group S = =I'' Sne x Skz. Looking at how each dimension of ( permutes, we would have Oneta E Smetz permute the first dimension (output channels), Tne E Saz permute the second dimension (input channels), and ok, E Sk, 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 notable difference between how our framework handles convolutional weight spuces compared to NFNs (Zhou et al. [223aa] 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 doesn't 't matter whether we apply neuron permutation to the input the output. We define Ls (W, W) as the space of equivariant linear maps, i.e., those T satisfying:
T ((o = o7 (W), Yo E S, W E W.
(4)
Our goal is to design a layer (i.e., a parameterized space of functions) that is equivalent to Ls (W, W). In some applications, we may instead desire invariance, that is a function P satisfying
P P(oW) P (W), Vo E S, W E W.
(5)
Following 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.
3

---

3 Universal neural functionals
Since equivariance is preserved under composition, and pointwise non-linearities are already permu- tation equivariant, we can build deep equivariant models as long we have an equivariant linear layer. Additionally, composing equivariant layers with an invariant pooling operation produces a deep invariant model. This section introduces method for producing equivariant weight-space layers for any given weight space, which enables the flexible construction universal neural functionals.
3.1 Decomposing equivariant weight-space maps
The weight space is a direct sum individual weight subspaces W (y()) @ ... W(7), so the problem of defining equivariant layer on W can be decomposed into defining equivariant layers between each pair of weight subspaces Wy(r) and W(O), for all & and m [Navon et al. 2023
We re-state this result our own notation. For any E, m pair we define Ls (W/m), W)()) as the space of equivariant maps between the two weight subspaces. It contains all T' (y(m) + W)(e satisfying
T (dW(m) = oT (w(m)
=
Ya, W(m),
(6)
noting 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 basis B for ILs (WW()) W(s)). A basis function E E B can be extended to E W + W by defining:
E (W(P))
f=s
B(WY
0
otherwise
where F(W): (F.(W): F. (W))
Theorem 3.1 (Navon et al. 2023 ). Let { BBm be bases for each Ls (Wi(), W(O). Then the
union of these bases (extended by Eq. 7 a basis for linear equivariant maps on W. That is, we
have the basis B for Lss (W, W) defined:
B
U EE EmE[L]
(8)
This result tells us that we can construct equivariant basis B for Tis (W, W)) by simply combining the equivariant bases { ( } for each pair of weight subspaces.
3.2 Equivariant layers between tensors
Since weights are tensors, our decomposed problem involves finding bases for permu- tation equivariant maps between tensors. Variants of this problem have been stud- ied by numerous prior works-in particular,
Algorithm Basis for equivariant Wi(m > WV(4) layer
Require: W(()), Vy())
1: Initialize basis 8.
(}
2: < (01, "..."...00 .D, oD,,,K "h. }
Maron et al. 2018 theoretically character-
3: for P in VALIDPARTITIONS (I) do
ize a basis for equivariant maps between arbitrary-rank tensors, and provide a con- crete implementation of the basis functions in the rank-2 case. Here, we describe a general algorithm that automatically con- structs basis for permutation equivariant maps between arbitrary- tensors. Con- cretely, implements each basis function in terms of simple array operations that are amenable to efficient computation with modern deep learning frameworks.
4: Label each subset 8p E P by unique character
CHIRR(,,)
5:
for E do
6: Map index cla) CHAR(ss) where 8p
7:
end for
Ex Kolpy!
Σn Xefil.-..tio.
9: 8em + Bem U EEP)
10: end for
11: return Bém

---

Functions in Ls (Wim), W( take input tensors indexed by { ii.... ... i.... } and produces output tensors indexed by (or oD J. We can construct a basis B for this space where each element is identified by valid partition P of these indices. Recall that the indices (2..) of ((m)) are permuted by (oars "arr,"y ".p. We say that two indices i and i2 "permute simultaneously"" if dT dz
Definition A valid partition a partition P of the ouppu and input indices I (o..... "0,,, .iDm into non-empty subsets, such that each subset only contains indices that are permuted simultaneously.
Example 3.1 (W/(m) W(() 1 x=2). Here the output and input indices are {01, 02, i1,
The partition {{0, 02 J. (1, i2 }
} is not valid because 01, 02 are permuted by 01, 02, so they do
not permute simultaneously. On the other hand, {{ (0, 11 }. O2, l2 } a valid partition.
Example 3.2 (W/m) = W(0 =
This time, the partition { 02 ,
is valid
because o 02 are both permuted by aı, are ,2,
To construct the equivariant basis, we enumerate all valid partitions and then map each partition P to a basis function Ep. Concretely, we label each subset of P with a distinet character ax, 3, .. and then remap each of our original indices { {o1.... :DD, 01, ... .iD... } to character based on which subset the index was in. This mapping is best illustrated by continuing our previous example.
Example 3.3 (Mp(m) = -(0 = IR x"*), Here input and output are both matrices, with combined indices {or, 02, 21, ig }. We have two permutations (01, 02) E Sn,, x Sny that can act on the rows and columns of the input and output matrices. There are four valid partitions:
P. ( (1, (2),222)
P2 = (1=](o))))(0)))(()))))
P3 (fo) (in (0),22))
P4 - (0-(oo)).(0,,,.i))
(9)
Consider P2-we assign character each subset:
P2 (1o 0z (2)) (2))
(10)
which tells to remap the output indices (01, 02) + (a, 3) and the input indices (i1, + (a, 7),
producing the basis function Ep, (W(m) ag = E, 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 total of four basis functions Ep,.... Ep, for Lg (W(m), W(()). Our equivariant Wi(m) > W)()) layer will be defined as the linear combination Im (W(m); = Eimr d Ep, (W(m),, which is the layer introduced in Hartford et al.) (2018).
To generalize the previous example, for each valid partition of the indices label its subsets with characters a, y,,,,.. and then construct basis function:
E(W(m)
)efommm)))))))r Defos) (eooo = Σ w((m) dfiil .cli.
(11)
where cf-] maps each index to the subset of P that contains it. We sum over the characters in R. which is the (possibly empty) subset 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 ] gives formal description of the complete process for generating B
Theorem 3.2. Algorithm produces a basis for the equivariant linear maps from W()) to W(O)
Proof. See Appendix B.1
Once Algorithm has generated a basis of equivariant functions Bem, we can implement an equivari- layer using a vector tm IR [glss" of learned coefficients:
Tel (wim); )0m)) = Σ X
Ep, (Wim)).
(12)
5

---

3.3 Equivariant layers on weight spaces
Theorem 3.1 now tells us that we may now construct the equivariant weight-space layer by combining the bases { } into basis B of functions on W. The weight-space layer T X) can then be
defined by linear combination of the basis functions with learned coefficients A. Explicitly, the full
layer is defined:
T (W, A) = (T' (W, ,7 (W,
(13)
where a = { |e= 1, }}} and T (W. Em T (W(m),
Appendix A provides a concrete description of how 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. 20188 neural network libraries.
Theorem 3.3. The weight-space layer (Eq. 13 is S-equivariant, and can express any linear equiv- ariant function on W.
Proof. Each T is a linear combination of basis functions in Ben Then, as described by Thm 3.1 Eq. 13is linear combination of functions that form basis for Ls (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 NENNp [Zhou et al.) 2023a] This because the layers each parameterize all possible linear maps equivariant to the same symmetry group, and hence can express the same of functions.
3.4 Multiple feature channels
In practice, we may be interested in simultaneously processing multiple weight-space features, such as the weights and 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 W, with elements that can be written a
W = (W[1).-
W[c),, for W[k] E W. Then the action is oW := [oW[1]. -......... oW [c]) for S, extending the (single channel) definition. The definition of equivariance can then be extended to layers of the form T-)) W + W",, where C₁, Co are the number of input and output channels.
Extending equivariant layers the multi-channe setting is quite common in the geometric deep learning literature and simply involves taking linear combinations along the channel dimension [Cohen and Welling 2016. RRavanbakhsh et al. (017). That is, modify the equivariant layer between subspaces as:
|
T (p(m));am) [KT] =
Em [K k.
Ep. (w(m)) [k)],
(14)
6=11 k=1
where each nm is now learned co o matrix instead of a scalar.
3.5
Deep models
The previous sections describes the construction of S-equivariant layers that operate operate on weight-space features in W" We construct universal neural functionals by stacking multiple such layers (interleaved with pointwise non- lieeeiiiiiiie)) into a deep, permutation equivariant model that can process weights. To construct permutation invariant model,
add an invariant pooling
layer after the equivariant layers, as in prior work [Navon et al. 2023 Zhou et al. 2023a
Experiments
In this section, we refer 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 requence-to-sequena models, and training learned optimizers for a variety of architectures and datasets.
In the multichannel setting we overload notation and use W to refer elements of W not W.

---

4.1 RNN generalization prediction
One promising application of neural functionals predicting the generalization of neural network models from their weights Eilertsen et al. (2020). We construct Tiny RNN Zod a dataset of recurrent neural networks trained to do arithmetic by completing given questions Character-by-chareeee example, given the input string "15+20=" the correct completion would be "35<EOS>". To construct the dataset, we train 10 ecquence-to-sequene
For
[Sutskever et al. 2014] models on example problems with input numbers up to five input digits. Both encoder and decoder RNNs contain single GRU cell [Chung et al.) 2014] with hidden size 128. Each model is trained with a distinct learning rate and batch size, and it's test success rate (SR) is recorded. The learning rate is sampled from log-uniform distribution over [10- 4, 10_2), 2I, and the batch size is sampled uniformly from { 64, 128, 256 J. With the goal of predicting test SR from weights, we split the Tiny RNN Zoo into 8000/ 1000 training, validation, and test examples.
The success rate of each RNN model is clearly invariant under permutation symmetries of its weights, so invariance is natural inductive bias for any generalization predic- tor. We evaluate STATNN [Unterthiner et al., [202] and a UNF-based predictor (note that NFNs are not applicable to the weights of recurrent networks). STATNN is operates on basic statistical feature of the weights, and has been shown to be a very strong baseline on previous general- ization 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, was shown (for CNN classifiers) in Zhou iet al. [2023a].
Method Deep Set
Test 7
0.8306 + 0.0006
STATNN UNF (Ours)
0.8839 0.0007
0.8968 + 0.0006
Table 1: Rank correlation between predicted and actual success rates of RNNs on arith- metic task. Predicting with UNF signif- icantly outperforms STATNN [Unterthiner eet al. 2020).
In particular, STATNN computes the mean, variance, and (0, 25, 50, 75, perepeeeessssles of each weight tensor in the RNN and feeds them into a six-layer MLP with hidden width 600. UNF is permutation invariant model, implemented using 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 7, the rank correlation between predicted and actual success rate.
Results. Table shows the performance of each predictor 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.
4.2 Learned optimizers
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
PInspired by the Tiny CNN Zoo [Unterthiner et al. 2020 Notably, it computes statisties that are invariant permutations of the weights.
7

---

MLP FashionMNIST
CNN CIFAR-10
RNN LM1B
Transformer LM1 B
Train NLL
2.2
3.0
10
0.8
2.0
9
1.8
2.5
0.6
8
1.6
2.0
7
1.4
1000
2000
2000
2000
4000
UNF SGDM
Deep Set
NFN
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 (UNFss, NFNs [Zhou et al. 2023a], or Deep Sets [Zaheer et al. [017]. 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 optimize an architecture type on a task from random initializations. Following Harrison et al. [2022). our learned optimizers track momentum terms m. <- -- 1 + V, and produce updates of the form:
W++1 < W. (m? + 35 (W, V., { {m? 4),
(15)
Here am," To is a "nominal term" " that biases the learned optimizer to behave like stochastic gradient descent with momentum coefficient %oo The neural functional f(-) ingests weights W. gradients V., momentum terms at several coefficients { m,? he and the iteration t.
During meta-training, we optimize network f and scalars a, 3, Yo to minimize the task training loss after a fixed number of training steps T, the "inner training horizion." " avoid the issue of backpropagating through an optimization process, we estimate meta-gradients using persistent evolutionary strategies (Vicol et al. 202T].
Comparisons. The default architecture choice for f(-) in prior work is Deep Sets [aheer et al. 2017]. which offers equivariance to any permutation symmetry. We study the effect of replacing Deep Sets by UNFs. We also try the NENNr 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 3 = 0 in Eq. [5] The SGDM baseline also meta-tran trained to tune the learning rate a and momentum decay rate %o- We compare the different learned optimizers in four tasks:
MIP Fashion Each optimizer trains MLP classifier a downsized and flattened version of the FashionMNIST dataset [Xiao et al. 2017). We note that for MLP weight spaces, UNF are identical NFNNP Zhou al. 2023a
CNN on CIEAR-10. Each optimizer trains a convolutional classifier on a downsized 16 x 16 CIFAR- 10. In this setting our algorithm produces UNF that is different to NFN (see Example 2.3).
RNN on LMIB. Each optimizer trains a character-level RNN- language model (LM) on the One Billion Word Language Model Benchmark (LMIB) dataset (Chelba et al. 2013).
Transformer on LMIB. Each optimizer trains Transformer LM on LMIB, this time predicting tokens instead of characters.
We use inner training horizon T' = 2,000 for the first three tasks and T = 5,000 for the Transformer task, since 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. [017]] Eq. 4], while the UNF and NFN optimizers uses three Deep Set layers followed by single UNF or NFN layer. See Appendix C.I C.2 for full descriptions of the tasks and meta- traiatt.aiinng..
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
8

---

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 the CNN experiment. As noted in Example 2.3, UNFs make the stronger assumption that all tensor dimensions-incloe 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 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 meta- training on diverse tasks [Metz al. 2022]. Moreover, as Table 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 significant computational constraint in practice. Still, could be challenge to meta-optimization, since evolutionary strategies are known struggle in higher dimensions. Hence, further work efficient high-dimensional meta- veea-gradient estimators would complement the development of expressive weight-space models like UNF.
5 Related Work
There is a long history of neural network architectures that equivariant to various symmetry groups [eCun et al.
Cohen and Welling, 2016 Ravanbakhsh et al. 2017 Kondor and Trivedi 2018 Existing frameworks for automatically constructing equivariant mod-
1995
2018 Cohen et al.
els Finzi et al. [002]] produce equivariant matrices, which would be intractable for our task. Our work constructs efficient equivariant basis functions for 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 al.
2018
Albooych et al. 2019]. This work observes that 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 al. 2017 Kirsch and
Schmidhuber [201]] [Peebles et al.) [202]] 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 simple feedforward models [Navon al. 2023 Zhou
ct al. 2023alb) b).
A recent alternative approach has been to leverage message passing neural networks (MPNNs) Zhang ct 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 more detailed exposition of this trade-off in Appendix B.3
Conclusion
We introduce a method for constructing permutation-equivarinn 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 efficiint array operations in standard deep learning frameworks. We empirically validate that the resulting universul neural functionals (UNFs) are effective at tasks that involve processing the weights and gradients of convolutional image classifiers, recurrent soguence-to-sequene models, and Transformer language models. In

---

particular, we find that UNFs show promising improvements over existing learned optimizer designs small scale experiments.
Limitations and future work. remains to be demonstrated how UNFs can be applied to het- erogenous weight-space inputs, e.g., to have single UNF act 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 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. I would grow rapidly for higher rank tensors with higher-order interactions (see Appendix (B.)). Resolving these challenges would further improve the scalability and applicability of neural functionals weight-space tasks.
7 Acknowledgements
We thank Jascha Sohl-Dicks Dickstein and Yiding Jiang for insightful general discussions about the project, and Louis Kirsch for helpful feedback on early drafts. AZ supported by the NSF Graduate Research Fellowship Program. We are grateful to the TPU Research Cloud (TRC) for providing compute for some of the experiments.
References
M. Albooyeh, D. Bertolini, and S. Ravanbakhsh. Incidence networks for geometric deep learning. arXiv preprint 1905. 11460, 2019.
M. Andrychowicz, M. Denil, S. Gomez, M. W. Hoffman, D. Pfau, Schaul, B. Shillingford, and N. De Freitas. Learning to leam by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
S. Bengio, Y. Bengio, J. Cloutier, and J. Gescei. On the optimization a synaptic learning rule. In Optimality in Biological and Artificial Networks?, pages 281-303. Routledge, 2013.
Y. Bengio, S. Bengio, and J. Cloutier. Learning synaptic learning rule. Université de Montréal, Département d' d'informatique de recherche .... 1990.
Bradbury, Frostig. P. Hawkins, M. Johnson, Leary, Maclaurin, Necula, Paszke, J. VanderPlas, S. Wanderman-Milne and Q. Zhang. JAX: composable transformations of Python+NumPy programs, 2018. URL http: //tp://gttttb.ttt com/ got//loogleaa
C. Chelba, Mikolov, M. Schuster, Q. Ge, T. Brants, P. Koehn, and T. Robinson. One bil- lion word benchmark for measuring progress in statistical language modeling. arXiv preprint 2013.
Chung. C. Gulcehre, K. Cho, and Y. Bengio. Empirical evaluation of gated recurrent neural networks on sequence modeling. arXiv preprint arXiv:1412.3555 2014.
T. Cohen and M. Welling, Group equivariant convolutional networks. In International conference on machine learning, pages 2990-2999, PMLR, 2016.
T. S.S.C, M. Geiger, J. Köhler, and M. Welling, Spherical CNNs. arXiv preprint arXiv:1801.10133 2018.
L. De Luigi, A. Cardace. R. Spezialetti, Zama Ramirez, S. Salti, and L. Di Stefano. Deep learning implicit neural representations of shapes. In International Conference Learning Representations (ICLR), 2023.
G. Eilertsen, D. Jönsson, T. Ropinski, Unger, and A. Ynnerman. Classifying the classifier: dissecting the weight space of neural networks, arXiv preprint arXh::2002.05688, 2020.
J. L. Elman. Finding structure in time. Cognitive science, 14(2):179-22 1990.
M. Finzi, M. Welling, and A. G. Wilson. A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups. In International Conference on Machine Learning, pages 3318-3328 PMLR. 2021.
10

---

C. Godfrey, D. Brown, T. Emerson, and H. Kvinge. On the symmetries of deep learning models and their internal representations. Advances Neural Information Processing Systems, 35:11893- 11905, 2022.
D. Ha, A. Dai, and Q. Le. Hypemetworks. arXiv preprint arXiv:1609.09106, 2016.
Harrison, L. Metz, and J. Sohl-Dickstein. A closer look learned optimization: Stability, robustness, and inductive biases. Advances in Neural Information Processing Systems, 01:3758-77733 2022.
J. Hartford, D. Graham, K. Leyton-Brown, Brown, and S. Ravanbakhsh. Deep models of interactions across sets. In International Conference on Machine Learning, pages 1909-1918. PMLR, 2018.
R. Hecht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Advanced Neural Computers, pages 129-135. Elsevier, 1990.
L. Kirsch and Schmidhuber. Meta learning backpropagation and improving it. Advances in Neural Information Processing Systems, 34: 14122-14134, 2021.
L. Kirsch, S. Flennerhag, H. van Hasselt, A. Friesen, Oh, and Y. Chen. Introducing symmetries to black box meta reinforcement learning. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 36, pages 7202-7210, 2022.
M. Kofinas, B. Knyazev, Y. Zhang, Y. Chen, J. Burghouts, E. Gavves, C. G. Snoek, and D. W. Zhang. Graph Neural Networks for Learning Equivariant Representations of Neural Networks. In 12th International Conference on Learning Representations (ICLR), 2024.
R. Kondor and S. Trivedi. On the generalization of equivariance and convolution in neural networks to the action of compact groups. In International Conference on Machine Learning, pages 2747-2755. PMLR, 2018.
D. Krueger, C.-W. Huang, R. Islam, R. Turner, A. Lacoste, and A. Courville. Bayesian hypernetworks. arXiv preprint arXiv: 1710.0479, 2017.
Q V. Le, N. Jaitly, and G. E. Hinton. A simple way to initialize recurrent networks of rectified linear units. arXiv preprint (tti::/50400044// 2015.
Y. LeCun, Y. Bengio, et al. Convolutional networks for images, speech, and time series. The handbook of brain theory and neural networks, 3361(10):1995, 1995.
D. Lim, H. Maron, M. T. Law, J. Lorraine, and Lucas. Graph metanetworks for processing diverse neural architectures. arXiv preprint arXiv:2312.04501, 2023.
H. Maron, H. Ben-Hamu, N. Shamir, and Y. Lipman. Invariant and equivariant graph networks. arXiv preprint 1rXiv:1812.099900 2018.
L. Metz, N. Maheswaranathan, J. Nixon, D. Freeman, and J. Sohl-Dickstein. Understanding and correcting pathologies in the training of learned optimizers. In International Conference on Machine Learning, pages 4556-4565. PMLR, 2019.
L. Metz, J. Harrison, C. D. Freeman, A. Merchant, L. Beyer, J. Bradbury, N. Agrawal, B. Poole, Mordatch, A. Roberts, et Velo: Training versatile learned optimizers by scaling up. arXiv preprint arXiv:2211.09777
2022.
A. Navon, A. Shamsian, I. Achituve, E, Fetaya, G. Chechik, and H. Maron. Equivariant architectures for learning in deep weight spaces, arXiv preprint arXiv:2301. 12780, 2023.
H. Pan and R. Kondor. Permutation equivariant layers for higher order interactions. In International Conference on Artificial Intelligence and Statistics, pages 5987-6001. PMLR, 2022.
W. Peebles, I. Radosavovic, T. Brooks, A. A. Efros, and J. Malik. Learning to learn with generative models of neural network checkpoints. arXiv preprint 2rxiv:2209.12892 2022.
Ravanbakhsh, Schneider, and B. Poczos. Equivariance through parameter-sharing. In Interna- tional conference machine learning, pages 2892-2901. PMLR, 2017.

---

J. Schmidhuber. A 'self-referential weight matrix. In ICANN'93: Proceedings of the International Conference on Artificial Neural Networks Amsterdam, The Netherlands 13-16 September 1993 pages 446-450. Springer, 1993.
1. Sutskever, O. Vinyals, and Q V. Le. Sequence to sequence learning with neural networks. Advances in neural information processing systems, 27, 2014.
E. H. Thiede, T. S. Hy, and R. Kondor. The general theory of permutation equivarant neural networks and higher order graph variational encoders. arXiv preprint arXiv:2004.0m X11:22044033990, 2020.
Unterthiner, D. Keysers, S. Gelly, O. Bousquet, and I. Tolstikhin. Predicting neural network accuracy from weights. arXiv preprint 1rXiv:2002.11448, 2020.
P. Vicol, L. Metz, and J. Sohl-Dickstein. Unbiased gradient estimation in unrolled computation graphs with persistent evolution strategies. In International Conference on Machine Learning, pages 10553-10563. PMLR, 2021.
L. Wang, K. Zhang, A. Zhou, M. Simchowitz, and R. Tedrake. Fleet policy learning via weight merging and an application to robotic tool-use. arXiv preprint arXiv:2310.01366 2023.
O. Wichrowska, N. Maheswaranathan, M. W. Hoffman, S. G. Colmenarejo, M. Denil, N. Freitas, and Sohl-Dickstein. Learned optimizers that scale and generalize. In International conference on machine learning, pages 3751-3760. PMLR, 2017.
H. Xiao, K. Rasul, and R. Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017.
M. Zaheer, Kottur, S. Ravanbakhsh, Poezos, R. Salakhutdinov, and A. Smola. Deep sets. doi: 10.48550. arXiv preprint ARXIV/7003.6// 14, 2017.
D. W. Zhang, M. Kofinas, Y. Zhang. Y. Chen, G. J. Burghouts, and C. G. Snock. Neural networks graphs! graph neural networks for equivariant processing neural networks. 2023.
A. Zhou, K. Yang, K. Burns, Y. Jiang, S. Sokota, J. Z. Kolter, and C. Finn. Permutation equivariant neural functionals. arXiv preprint arXiv:2302. (0t.:2332..4440.0 /4040, 2023a.
A. Zhou, K. Yang, Y. Jiang, K. Burns, W. Xu, Sokota, J. Z. Kolter, and C. Finn. Neural functional transformers. arXiv preprint 2355:22305.13546 2023b.
12

---

A Weight-space specifications
Here we discuss the concrete specification that precisely describes a weight space and must be provided input to the algorithm before it can construct equivariant weight-space layers. Our implementation compatible with most JAX [Bradbury et al. 2018 neural network libraries.
Suppose we wish to process an MLP's weights that are stored in (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 (12, is permuted W(1) + P (o2) W
specification =
"layer1": : {"weight"; ("n2", "ni"), "bias": ("n2" )}, "layer2": {"weight"; ("n3", "n2"), "bias": ("n3" })},
}
Providing this specification object our algorithm is sufficient for it deduce the symmetry group, its action, and construct the corresponding equivariant layer.
Since most neural networks consist of repeating layers 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 network in common deep learning frameworks by analyzing its computation graph.
B Further analysis of UNFs
B.1 Algorithm generates basis for Ls (Wim) mo(e))
Here we show that Algorithm 1 produces basis B for Ls (W(m), W(O), the space of linear cquivariant maps between ()(m) and MX(3), Consider instantiating these linear maps as matrices multiplying flattened input vee (w)))) Maron et al. (2018) characterize basis B" }, for these matrices, where the entries of each basis matrix are dehned:
{
(a, E
B
(16)
otherwise
Here E Im and b Ty multi-indexes for the input and output spaces, and ja Tm Tz/ is an equivalence class of the combined input-output index space Tm x Te under the equivalence relation ~ defined by a a' if and only if a; aj + af d, for all i, i.c. the multi-indexes a, a' have the same equality pattern.
Re-arranging Maron et al. 2018
Eq. 10b), any equivariant linear map is defined:
L(W(m)),
Σ Σ
Wp B" W(m)
Σ Wp Σ A(aa,) E ()W(
(17)
aEIn MEL. Xle/
PEL. Xle//
aEt.
where 7 I(.) is an indicator function for the given condition.
Notice that each equivalence class j is represented by what we call a valid partition of D De) = {1. Dm + Dr). so this is already sum over valid partitions as in Eq. 12 We can now observe
13

---

that each term on the RHS is equivalent to one of our basis functions (Alg) Line 8). That is, for given equivalence class μ represented by valid partition P:
Σ Illa.b) e p)W!") (p((m))).
(18)
This is because for any I= (a, b) yielding a nonzero term on the LHS, if E [Dm + Del are grouped together by partition then I - Zz. otherwise they would violate the equality patlern of j. Therefore, we can replace all indices grouped together in partition with a single shared symbol, i.e. the characters in Eq.
==
Hence, Algorithm produces 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 similar to the construction in Pan and Kondor [2022]. but generalized to multi-node sets (non-square tensors whose axes can potentially permute independently).
B.2 Size of basis produced by Algorithm
Suppose we have neuron permutation symmetry group S = Sm, ... X Sux- i.e., every neuron permutation a is composed of distinct permutations (aI, ,NN). For each = I,... .N we define &; (W(+) to be the number of indices that Ø; E Sn, permutes in weight tensors W E Wll (which could be 0). Finally, denote b(k) to be the k'th Bell number. Then the number basis functions generated by Algorithm is:
N ΣII
(a (Mile) (w(()) + (molm))).
1Btm
(19)
E,m 1=1
B.3 Comparison to MPNN-based approaches
Each UNF layer can express any linear equivariant function on 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, may also be more computationally challenging due the growth in the size of the basis (see next section).
As example, consider simple "RNN" where h++1 Wh, and E R has exchangeable entries, meaning that W + PW pT a symmetry. Algorithm would produce an equivariant basis with 6(2 2)
15 terms
2023 with n nodes and 2n² directed edges between them (allowing forward and backward edge for each weight, equivalently n² undirected edges). Then using a similar construction to Lim al. 2023 Appendix C.1.2], we would get linear GNN that computes:
On the other hand, we could construct a parameter graph [Lim et al.
f(W)= aW.,, bW,, + cW,, A + dW eW.,, fWjk.
(20)
which 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.
C Experimental details
C.1 Learned optimization tasks
Here we describe each of the experimental settings we evaluated the learned optimizers on. Across all experiments, the training loss is negative log-likelihood.
In this case, the full basis is also given by Maron et al. 2018
Appendix A].
14

---

Figure 3: Number of parameters used by f() in each learned optimizer, for each task. Note that NFN and UNF are identical for the MLP task. This count does not include the other meta-learned scalars Eq. 15
which are a, To. 3.
Task
UNF
Deep Set
NFN
MLP on FashionMNIST
3,783
2,788
3,783
CNN on CIFAR-10
7,369
2,788
41,603
RNN on LMIB
8,043
2,788
N/A
Transformer on LM1B
64,168
2,788
N/A
MLP on FashionMNIIST. Train three-layer MLP classifier downsized (8 x 8) and flattened version of the FashionMNIST dataset (Xiao et al. 2017]. The MLP has hidden size 32 and ReLU activation function. We use batch size of 128.
CNN on CIFAR-10. Train 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 trained with batch size of 128.
RNN on LMIB. Trains a character-level RNN-based language model (LM) on LMIB CChllba et al. 2013). The RNN itself has one hidden layer with size 64, and uses dlentity-iiitializttee [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 LMIB. Train Transformer LM LMIB, 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 batch size of 8 on length-8 sequences.
C.2 Learned optimization meta-training
Call DS [c] single equivariant Deep Set layer [Zaherr et al. 2017 Eq 4] with output channels (sim- ilarly for UNF [c] and NFN [<]). Then f()) in our learned optimizers (Eq. 15) is always implemented as feedforward architecture:
DeepSetOpt = DS [32] -> ReLU -> DS [32] [32] ReLU DS ReLU UNF [1]
-> ReLU -> DS [32] -> ReLU -> DS [1]
UNFŌpt DS [32] -> ReLU
NFNOpt = DS [32] -> ReLU -> DS [32] -> ReLU -> DS [32] -> ReLU -> NFN [1]
For all methods, we initialize 0.1 and 70 0.9 before starting meta-training. For non-SGDM
methods, we initialize / 0.001, and provide six momentum values (m, ); with coefficients 7, 0.1, 0.5, 0.9, 0.99, 0.999. 0.9999. The iteration number / is converted into an 11 11-dimensional sinusoidal encoding, and all inputs to f(-) are concatenated along the channel dimension. Concretely, this results in an input in 19 The output is in W
We meta- train for 50,000 steps using Adam, estimating meta- gradients over parallel training runs using persistent evolutionary strategies (PES) [Vicol et al. [021] with truncation length of 50 and noise standard deviation of 0.01. The meta-trainng objective 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 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 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 Table 3 lists the number of parameters in f(-) for each learned optimizer.
C.3 Compute
Experiments were run on mix of TPU v3 and v4 accelerators. On TPU v3-8, training a UNF for RNN generalization prediction task takes < hours. Also TPU v3-8, meta-training UNF
15

---

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.
16

---

NeurIPS Paper Checklist
1. Claims
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 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 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 long is clear that these goals are not attained by the paper.
2. Limitations
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 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 (c.g.. independence assumptions, noiseless settings, model sell-speciiiatiinn 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-tex, system might not be used reliably to provide closed captions for online lectures because it fails 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 faimesss.
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 't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
17

---

Answer: [Yes]
Justification:
We provide proofs for both Thm 3.2 and 3.3.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided not)?
Question:
Answer: [Yes]
Justification: We provide code implement the proposed method well details in the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, descríbing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., g., in the case of a large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is new model (e-g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing verifying the results.
5. Open access to data and code
18

---

Question: Does the paper provide open access to the data and code, with sufficient instruc- tions 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.
Please see the NeurlPS code and data submission guidelines (https://nips.ce
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
/
ssiololiccyy for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification:
See Appendix C.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented in the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Our results report error bars.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
19

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: See Appendix C,3
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't 't make it into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification:
We have reviewed the Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration 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 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.
20

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
No relevant risks.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses for existing assets
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?
Question:
Answer: [NA]
Justification: No third third-party -arttpprty assets used.
Guidelines:
The answer NA means that the paper does 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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.
21

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do not release assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: No human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
No crowdsourcing or human subjects.
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
22


Reviews:
[{"review_id": "WQkSWynfE3", "reviewer": "Reviewer_eHaW", "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.", "questions": "Can you produce a table of computational speed + resource comparisons?", "limitations": "NA", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 4, "strengths": "The paper is pretty strong, showing good results for tasks which are traditionally hard to work on. \n\nThe method is novel, and applicable for various weight space tasks.\n\nThe methodology is well introduced and explained well.", "weaknesses": "line 219, operate is mentioned twice. \n\nThere is a limited number of datasets evaluated upon.\n\nI would like to see a table showing the computational efficiency/time of this method vs other methods.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> There is a limited number of datasets evaluated upon.\n\nWe 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.\n\nRank correlation between predicted and actual performance (higher is better):\n\n| Method     | Kendall's tau       |\n|------------|---------------------|\n| Deep Set   | $0.8306 \\pm 0.0006$ |\n| STATNN     | $0.8839 \\pm 0.0007$ |\n| UNF (Ours) | $0.8968 \\pm 0.0006$ |\n\n> Can you produce a table of computational speed + resource comparisons?\n\nAppendix 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5pP2extXWI", "reviewer": "Reviewer_J2cE", "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.", "questions": "See weaknesses. Please improve the presentation of methodology and results.", "limitations": "The authors have adequately addressed the limitations.", "rating": 5, "confidence": 2, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. Methodology is well-developed.\n2. 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).\n2. The presentation of the results is also not very clear and the results are not sufficiently reported.\n3. Whilst the training times for the test cases are reported, the inference speed is not reported.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their review and suggestions, which we have incorporated into the draft manuscript.\n\n> presentation of the methodology is not very clear (specifically around Eq. 9 and 10)\n\nWe 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.\n\n> the results are not sufficiently reported\n\nWe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "tfS500mtko", "reviewer": "Reviewer_k71j", "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.", "questions": "- Among the many applications of weight-space networks, how did the authors decide on the two tasks to conduct experiment on?\n- 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?\n\n**Minor issues / suggestions**\n- $\\sigma_{d^l_i}$ and $\\sigma$ in line 67 are not defined.\n- 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.\n- 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$?\n- 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.\n- Line 165 typo “a a character”\n- What is $W$ on the right side of the equation in line 172?\n- Line 538 “next section” -> previous section", "limitations": "The authors include an informative limitation section that points out challenges and future directions.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "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. \n- The algorithm that computes maximally expressive equivariant linear layers for a set of tensors is a significant theoretical contribution. \n- Connection to past work is discussed well throughout the paper. \n- 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.\n- 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.\n- 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? \n- 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. \n- 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their careful analysis and feedback. We will incorporate the suggestions and fixes proposed here.\n\n> Despite the claims in the abstract and conclusion, there has been other work …\n\nWe 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.\n\n>  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?\nWe 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)).\n\nFor intuition, we may also work out parameter sharing patterns for simple permutation equivariant layers. Following [Ravanbakhsh et al](https://arxiv.org/abs/1702.08389), 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.\n\n>  the number of valid partitions can be very large especially when many indices permute simultaneously or there are many layers\n\nWe 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.\n\n> the proposed method is not compared to the most recent permutation equivariant architectures\n\nWe 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.\n\n> how did the authors decide on the two tasks to conduct experiment on?\n\nLearned 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).\n\nAlthough 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.\n\n>  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\n\nGood 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.\n\n> What is $W$ on the right side of the equation in line 172?\n\nShould be the input $W^{(m)}$, thank you for catching this\n\n> 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$\n\nNot 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "DgmksrzwbY", "reviewer": "Reviewer_7uT6", "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.", "questions": "For SGDM do the authors actually learn α and γ0 or it's simply tuned?", "limitations": "the limitations are discussed", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The paper addresses the problem of learning from network weights which is challenging and has a lot of potential.\n2. Interesting learning to optimize experiments, which is a very relevant application of this kind of method.\n3. 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.\n4. 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.\n\n2. 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)?  \n\n3. RNN generalization prediction\n\n- 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.\n- it's questionable that 0.8968 is significantly better than 0.8839.\n- 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.\n\n4. Learned optimizers\n\n- 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.\n- 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?  \n- 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?\n\n5. 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\".\n\nOverall, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their detailed feedback. We will also include the suggested relevant works in our discussion.\n\n> it does not automatically constructs the specifications which is actually a quite tricky part for some complicated models\n\nWe 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.\n\n> How Algorithm 1 works in practice for the networks in the experiments? How the valid partition looks like for those networks?\n\nThanks 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.\n\n> RNN generalization prediction: It's questionable that 0.8968 is significantly better than 0.8839.\n\nThis 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.\n\n> RNN generalization prediction:  More baselines are needed.\n\nAs 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.\n\nTest rank correlation between predicted and actual performance (higher is better):\n\n| Method     | Kendall's tau       |\n|------------|---------------------|\n| Deep Set   | $0.8306 \\pm 0.0006$ |\n| STATNN     | $0.8839 \\pm 0.0007$ |\n| UNF (Ours) | $0.8968 \\pm 0.0006$ |\n\n> it's questionable that Deep Sets is \"the default architecture choice for f\".\n\nTo 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](https://arxiv.org/abs/2203.11860). 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.\n\n> For SGDM do the authors actually learn α and γ0 or it's simply tuned?\n\nWe learn these hyperparameters using the same ES optimizer as we use for all the other learned optimizer results.\n\n> 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?\n\nThe 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.\n\n> 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?\n\nThis 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "Ii2f9KYxXp", "reviewer": "Reviewer_fP4N", "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.", "questions": "## Questions\n- 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.\n- On line 182, could you expand on what is meant by \"Entries that are not explicitly assigned by the left-hand side are 0.\"?\n- 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?\n- 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?\n- 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?\n\n## Suggestions\n- 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.\n- 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.\n- Typo on line 246: Remove the word \"is\" from \"STATNN is operates on basic statistical features of the weights\"", "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.", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "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.\n- 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.\n- 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.\n- 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.\n- There are some small changes I would make to the paper for readability, please see my suggestions below.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their thoughtful comments.\n\n> 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.\n\nWe 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.\n\nTest rank correlation between predicted and actual performance (higher is better):\n\n| Method     | Kendall's tau       |\n|------------|---------------------|\n| Deep Set   | $0.8306 \\pm 0.0006$ |\n| STATNN     | $0.8839 \\pm 0.0007$ |\n| UNF (Ours) | $0.8968 \\pm 0.0006$ |\n\n> 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.\n\nThis 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.\n\nHowever, 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.\n\n> 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\n\nGood 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.\n\n> On line 182, could you expand on what is meant by \"Entries that are not explicitly assigned by the left-hand side are 0.\"?\n\nEquation 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.\n\n> In Eq. 15 are we manually computing the momentum term\n\nThe momentum terms fed as input are computed manually / in the standard way.\n\n> Did you use this positional encoding fix in your experiments?\n\nWe did not find positional encoding necessary in this case, no. We will clarify the text on this\n\n> 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 [...]\n\nYour 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.\n\n> There are some small changes I would make to the paper for readability, please see my suggestions below.\n\nThank you for catching these issues and providing suggestions, we will include them in the updated manuscript.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: 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.
Decision: Accept (poster)