Permutation Equivariant Neural Functionals
Authors: Allan Zhou
Kaien Yang
Kaylee Burns
Adriano Cardace
Yiding Jiang
Samuel Sokota
J Zico Kolter
Chelsea Finn
Venue: NeurIPS 2023 poster (2023)
Keywords: equivariance
permutation
implicit neural representation
generalization
TL;DR: This paper explores the design of permutation equivariant neural networks called "neural functional networks" (NFNs), designed to process the weights or gradients of other neural networks.

Abstract: This work studies the design of neural networks that can process the weights or gradients of other neural networks, which we refer to as *neural functional networks* (NFNs). Despite a wide range of potential applications, including learned optimization, processing implicit neural representations, network editing, and policy evaluation, there are few unifying principles for designing effective architectures that process the weights of other networks. We approach the design of neural functionals through the lens of symmetry, in particular by focusing on the permutation symmetries that arise in the weights of deep feedforward networks because hidden layer neurons have no inherent order. We introduce a framework for building *permutation equivariant* neural functionals, whose architectures encode these symmetries as an inductive bias. The key building blocks of this framework are *NF-Layers* (neural functional layers) that we constrain to be permutation equivariant through an appropriate parameter sharing scheme. In our experiments, we find that permutation equivariant neural functionals are effective on a diverse set of tasks that require processing the weights of MLPs and CNNs, such as predicting classifier generalization, producing "winning ticket" sparsity masks for initializations, and classifying or editing implicit neural representations (INRs). In addition, we provide code for our models and experiments at https://github.com/AllanYangZhou/nfn.

Permutation Equivariant Neural Functionals
Allan Zhou'
Kaien Yang'
Kaylee Burns'
Adriano Cardace²
Yiding Jiang?
Samuel Sokota J. Zico Kolter
Chelsea Finn' CCannggie Mellon University
Stanford University
University of Bologna
ayz8cs.stanke stanford. edu
Abstract
This work studies the design of neural networks that can process the weights or gradients of other neural networks, which we refer neural functional networks (NFNs). Despite a wide range of potential applications, including learned opti- mization, processing implicit neural representations, network editing, and policy evaluation, there are few unifying principles for designing effective architectures that process the weights of other networks. We approach the design of neural functionals through the lens of symmetry, in particular by focusing on the permuta- tion symmetries that arise in the weights of deep feedforward networks because hidden layer neurons have no inherent order. We introduce a framework for build- ing permutation equivariant neural functionals, whose architectures encode these symmetries as an inductive bias. The key building blocks of this framework are Na-Laye (neural functional layers) that we constrain to be permutation equiv- ariant through an appropriate parameter sharing scheme. In our experiments, we find that permutation equivariant neural functionals are effective on diverse of tasks that require processing the weights of MI.ss and CNNs, such predicting classifier generalization, producing "winning ticket" sparsity masks for initializa- tions, and classifying or editing implicit neural representations (INRs). In addition, we provide code for our models and experiment
Introduction
As deep neural networks have become increasingly prevalent across various domains, there has been growing interest in techniques for processing their weights and gradients as data. Example applications include learnable optimizers for neural network training
56
2 44, extracting
information from implicit neural representations of data [61 58), corrective editing of network weights (57 I (6), policy evaluation [24), and Bayesian inference given networks as evidence (60) We refer to functions of neural network's 's weight-space (such as weights, gradients, or sparsity masks) as neural functionals; when these functions are themselves neural networks, we call them neural functional networks (NFNs).
In this work, we design neural functional networks by incorporating relevant symmetries directly into the architecture, following a general line of work in "geometric deep learning" X 54
34 5
For neural funetionals, the symmetries of interest are transformations of a network's weights that preserve the network's 's behavior. In particular, we focus on neuron permutation symmetries, which are those that arise from the fact that the neurons hidden layers have no inherent order.
Neuron permutation symmetries are simplest in feedforward networks, such as multilayer perceptrons (MLPs) and basic convolutional neural networks (CNNs). These symmetries are induced by the fact that the neurons in each hidden layer of feedforward network can be arbitrarily permuted without changing its behavior [27]. In MLPs, permuting the neurons in hidden layer corresponds to
http::://gthab..ccm/laaaaaaaaaaa
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

NF La-Layer
NF-Layer
Hy
H40H
W
H
Figure 1: The internal operation of our permutation equivariant neural functionals (NFNs). The NFN processes the input weights through series of equivariant Layers, with each one producing weight-space features with varying numbers of channels. In this example, neuron permutation symmetry simultancously permutes the rows of (( and the columns of W7(3), This permutation propagates through the NFN in an equivariant manner.
permuting the rows of the weight matrix W(0,, and the columns of the next weight matrix W(++1) shown the left-hand side of Figure Note that the same permutation must be applied the rows (+ and columns of W1++1), since applying different permutations generally changes network behavior and hence does not constitute a neuron permutation symmetry.
We introduce a new framework for constructing neural functional networks that are invariant or equivariant neuron permutation symmetries. Our framework extends long line of work on permutation equivariant architectures [52
that design equívariant layers for a particular permutation symmetry interest. Specifically, we introduce neural functional layers (NF-Layers) that operate on weight-space features (see Figure I while being equivariant to neuron permutation symmetries. Composing these NF-Layers with pointwise non- -inciiiassss produces equivariant neural functionals.
68 25 63
We propose different NF-Layers depending on the assumed symmetries of the input weight-space: either only the hidden neurons of the feedforward network can be permuted (hidden neuron permuta- tion, HNP), all neurons, including inputs and outputs, can be permuted (neuron permutation, NP). Although the HNP assumption is typically more appropriate, the corresponding NF-Layers can be parameter inefficient and computationally infeasible some settings. In contrast, NF-LLayes derived under NP assumptions often lead to much more efficient architectures, and, when combined with positional encoding scheme we design, can even be effective on tasks that require breaking input and output symmetry. For situations where invariance is required, also define invariant NF-Layers that can be applied on top of equivariant weight-space features.
Finally, we investigate the applications of permutation equivariant neural functionals on tasks involv- ing both feedforward MLPs and CNNs. Our first two tasks require (1) predicting the test accuracy of CNN image classifiers and (2) classifying implicit neural representations (INRs) of images and 3D shapes. We then evaluate NFNs on their ability to (3) predict good sparsity masks for initializations (also called winning tickets [19), and on (4) weight-space "style-editing" task where the goal is to modify the content an INR encodes by directly editing its weights. multiple experiments across these diverse settings, we find that permutation equivariant neural functionals consistently outperform non-equivariant methods and are effective for solving weight-space tasks.
Relation to DWSNets. The recent work of Navon et al. [45] recognized the potential of leveraging weight-space symmetries build equivariant architectures on deep weight-spaces; they characterize a weight-space layer which is mathematically equivalent to our NF-Layer in the HNP setting. Their work additionally studies interesting universality properties of the resulting equivariant architectures. and demonstrates strong empirical results for suite of tasks that require processing the weights of MLPs. Our framework additionally introduces the NP setting, where we make stronger symmetry assumptions to develop equivariant layers with improved parameter efficiency and practical scalability. We also extend our NFN variants to process convolutional neural networks (CNNs) as input, leading to applications such as predicting the generalization of CNN classifiers (Section 3.1..

---

Table 1: Permutation symmetries of L-layer feedforward networks with no. .g...., R, neurons at each layer. All feedforward networks are invariant under hidden neuron permutations (HNP), while NP assumes that input and output neurons can also be permuted. We show the corresponding equivariant Layers which process weight-space features from U, with Gj input channels and Co output channels.
Equivariant NF-Layer
Group
Abbry
Permutable layers
Signature
Parameter count
8
II Sp, I 1-1 Sn,
NP
All layers
u + Um
OlacaL22 (c,colL+ + no +oo+o n2))²)
5
HNP
Hidden layers
:: - um
None
T: + u
CCOo 2
Equivariant neural functionals
We begin by setting up basic concepts related to (hidden) neuron permutation symmetries, before defining the equivariant Na-Layers Sec. 2.2 and invariant NF-Layers in Sec. 2.3
2.1 Preliminaries
Consider an L-layer feedforward network having n, neurons at layer i, with no and nt be- ing the input and output dimensions, respectively. The network is parameterized by weights W W(t) Rm: x i E [1..A]] and biases
{v(+) E R | [I. We denote the combined collection U' == (W, v) belonging to weight-space, U == WxV.
Since the neurons in a hidden layer E (1,.... L. - 1) have no inherent ordering, the network is invariant to the symmetric group S., of permutations of the neurons in layer i. This reasoning applies to every hidden layer, so the network is invariant to S = Sn, ... X Sun, which we refer to as the hidden neuron permutation (HNP) group. Under the stronger assumption that the input and output neurons also unordered, the network is invariant to S == So ... x Sag, which refer to as the neuron permutation (NP) group. We focus on the NP setting throughout the main text, and treat the HNP case in Appendix [ See Table I for concise summary of the relevant notation for each symmetry group we consider.
Consider an MLP and permutation ø = (oo. .L) E S. The action of the neuron permutation group is to permute the rows of each weight matrix ( by dg, and the columns by 0;-1. Each bias vector (i) is also permuted by o; So the action is oU = (aW, oe), where:
[oW]; jk = W-( of (1)
[ow]]
()
Until now we have used U = (W,") u) to denote actual weights and biases, but the inputs to a neural functional layer could be any weight-space feature such gradient, sparsity mask, or the output of a previous NF-La (Figure[]. []. Moreover, we may consider inputs with 2 1 feature channels, belonging to 4t + c - U, the direct sum of copies of 44 Concretely, each U E Ut consists E [1..T.] and biases v = { v(0) E R i E [[....] } with the channels in the final dimension. The action defined in Eq. T exxendds to the multiple channel
of weights W = { W16) 1Rn,xt--xe
case if we define w-( jk
w-(0)
j.k.: E R and v/+ = 1 Rs.
=
The focus of this work is on making neural functionals that are equivariant (or invariant) to neuron permutation symmetries. Letting e; and c, be the number of input and output channels, we refer to a function f 2 + 4 as S-equivariant if o fL)) action of S on the input and output spaces is defined by Eq. Similarly, a function f .4 + RR is S-invariant if f(oU)= fU) for all and U.
f(aU') ) for all a E S and U E U where the
f.g are equivariant, then their composition f o g also equivariant; if g is equivariant and is invariant, then f y is invariant. Since pointwise nonlinearities are already permutation equivariant, our remaining task is to design a linear NF-Layer that is S-equivariant. We can then construct equivariant neural functionals by stacking these NF-Layers with pointwise nonlincarities.

---

H(W)
HOW
Wi+b
W
W
We-D
W(++1)
W
W
W(i
H(W)
Σ alowss all w o wi+0+ 6 wit 866-
wi-l of Wi
Figure 2: A permutation equivariant NF-Layer takes in weight-space features as input (bottom) and outputs transformed features (top), while respecting the neuron permutation symmetries of feedforward networks. This illustrates the computation of a single output element H (W) defined in Eq. Each output is a weighted combination of rows or column sums of the input weights, which preserves permutation symmetry. The first term contributes weighted combination of row-and- column sums from every input weight, though this is omitted for visual clarity.
2.2 Equivariant NF-Layers
We now construct linear S-equivariant layer that serves as a key building block for neural functional networks. In the single channel case, we begin with generic linear layers T-=:0) 0) vec(U) + Ovec(U'), where vec(U) E Rdim(()) is U flattened as a vector and 9 E x is matrix of parameters. We show Appendix B.2 that any S-equivariant T-: 0) must satisfy system of constraints on known as equivariant parameter sharing. We derive this parameter sharing by partitioning the entries of by the orbits of their indices under the action of S, with parameters shared each orbit (54) of the appendix describes the parameter sharing in detail.
Table of
Equivariant parameter sharing reduces the matrix-vector product @vec(U)) to the NF-Layer present. For simplicity we ignore V and assume here that Ut = W and defer the full form to Eq, 3 in the appendix. Then H W + W maps input (W(1) ... W() to (H(W)(") .... H(W) Recall that the inputs not necessarily weights, but could be arbitrary weight-space features including the output of previous NF-Layer. For Wi)) E R Xx- the corresponding output is H ((W)() R x with entries computed:
H(W)
Eatewll + ((0 *. + + wii-1) + ( + c (.-11(((43) (i+11) +
+d ( jk
(2)
Note that the terms involving (i-1) or W(++1) should be omitted for i = 0 and L. respectively, and * denotes summation or averaging over either the rows or columns. Recall that the multi- channel case, each w-(i) is vector in IR so each parameter is Co X O; matrix. We also provide concrete pseudocode description of H in Appendix A Figure 2 visually illustrates the NF-Layer in the single-channel case, showing how the row or column sums from each input contribute each output. To gain intuition for the operation of H, it is straightforward to check S-equivariance:
Proposition 1. The NF-Layer H U U (Eq. p and Eq. 3) is S-equivariant, where the group' action input and output spaces is defined by Eq. 7 Moreover, any linear S-equivariant map T: + U is equivalent to H for some choice of parameters a, b, e,d.
Proof (sketch). We can verify that H satisfies the equivariance condition (o (oW) for any i,j, & by expanding each side of the equation using the definitions of the layer and action

---

(Eq. []. Moreover, Appendix B.2 shows that any S-equivariant linear map T- 8) must have the same equivariant parameter sharing as H, meaning that it must be equivalent to H for some choice of parameter a, b, c,
I
Informally, the above proposition tells us that H can express any linear S-equivariant function of weight-space. Since S is subgroup of S, H is also S-equivariant. However, it does not express every possible linear S-equivariant function. We derive the full S-equivariant NF-Layer H u U in Appendix C
Table D summarizes the number of parameters (after parameter sharing) under different symmetry assumptions. While in general a linear layer 0)::0) ur - U has cr dim(24² parameters, the equivariant NE-Layers have significantly fewer free parameters due parameter sharing. The S-equivariant layer H has O (c) while the S-equivariant layer H has O (c:co + no n2)?) parameters. The latter's ' quadratic dependence on input and output dimensions can be prohibitive in some settings, such in classification where the number of outputs can be tens of thousands.
Extension to convolutional weight-spaces. In convolution layers, since neurons correspond to spatial channels, we let n, denote the number of channels at the 5th layer. Each bias v(() E R has the same dimensions as in the fully connected case, so only the convolution filter needs to be treated differently since has additional spatial dimension(s) that cannot be permuted. For example, consider ID CNN with filters W W E R X i [[...]] }. where n, R4-1 are the output
and input channel dimensions and is the filter width. We let W w(()) j.k. R denote the 8thh
filter in the jth output channel, then define the S-action the same way as in Eq. |
We immediately observe the similarities to multi-channel features: both add dimensions that are not permuted by the group action. In fact, suppose we have c-channel features U E LL where UI the weight-space of a ID CNN. Then we combine the filter and channel dimensions of the weights, with W E R. x ((w)), This allows us to use the multi multi-channel NF-Layer H 2 - ue Any further channel dimensions, such as those for 2D convolutions, can also be folded into the channel dimension.
It is common for CNNs in image classification to follow convolutional layers with pooling and fully connected (FC) layers, which opens the question of defining the S-action when layer & is FC and layer is convolutional. If global spatial pooling removes all spatial dimensions from the output of é - 1 (as in e.g., ResNets (26) and the Small CNN Zoo (64)), then we can verify that the existing action definitions work without modification. We leave more complicated situations (e.g., when nontrivial spatial dimensions are flattened as input to FC layers) to future work.
10-encoding. The S-equivariant layer H/ more parameter efficient than / (Table D. but its NP assumptions are typically too strong. To resolve this problem, we can add either learned or fixed (sinusoidal) position embeddings to the columns of (1) and the rows of W(L) and v)(4); this breaks symmetry at input and output neurons even when using S-equivariant layers. our experiments, we find that 10-encoding makes competitive or superior H. while using fraction of the parameters.
2.3 Invariant NE-Layers
Invariant neural functionals can be designed by composing multiple equivariant NF-Layers with an invariant NF-Layer, which can then be followed by an MLP. We define an S-invarian layer P u 2 by simply summing or averaging the weight matrices and bias vectors across any axis that has permutation symmetry, i.e., P(U)- (w()).. ... W(L) (1) ,.... ..44) We can then apply a fully connected layer to the output of to produce an invariant vector of arbitrary size. In fact. this combination expresses every possible linear invariant function on U:
Proposition 2. Any S-invariant linear function Ut + Rdd can be expressed in the form f a for some choice of linear f 2R + ddd, Proof: See section B.3
We define the analogous S-invariant layer P in Eq. 20 of the appendix.
5

---

Table 2: Test of generalization prediction methods on the Small CNN Zoo 64 weights and test accuracies of many small CNNs trained on different datasets, such as CIFAR-IO-GS or SVHN-GS. NEN outperforms other methods on both datasets. Uncertainties indicate max and min over two runs.
which contains the
HFN
NFNxP
STATNN
CIFAR- 10-GS
0.934 + 0.001
0.922 + 0.001
0.915 + 0.002
SVHN-GS
0.931 = 0.005
0.856 = 0.001
0.843 + 0.000
3
Experiments
Our experiments evaluate permutation equivariant neural functionals on a variety of tasks that require either invariance (predicting CNN generalization and extracting information from INRs) equivariance (predicting "winning ticket" sparsity masks and weight-space editing of INR content).
Throughout the experiments, we construct neural functional networks (NFNs) using the NF-Layers described in the previous section. Although the specific design varies depending on the task, we will broadly refer our permutation equivariant NFNs NFNNP and NFNHNP, depending NF-Layer variant they use (see Table We also evaluate a "pointwise" ablation of our equivariant NF-Layer that ignores interactions between weights by only using the last term of Eq. computing HWW);, == d' dw(() jk We refer to NENs that use this pointwise NF-Layer as NFNPT-
which
Using our benchmarks, we compare the performance of NFNs with other methods for processing weights, including standard MLPs that operate on the flattened weight inputs. To encourage per- mutation equivariance,
optionally augment the MLP's training with permutations using Eq. []] On relevant datasets we also compare with the recently developed DWSNets [47), an equivariant architecture similar to NENNNP, and with inr2vec [12] recent non-equivariant approach for learning useful representations from weights.
3.1 Predicting CNN generalization from weights
Why deep neural networks generalize despite being heavily overparameterized is a longstanding research problem in deep learning. One recent line of work has investigated the possibility of directly predicting the test accuracy of the models from the weights [64] [6]. The goal is to study generalization in data-driven fashion and ultimately identify useful patterns from the weights.
Prior methods develop various strategies for extracting potentially useful features from the weights before using them to predict the test accuracy (29 67 64 30 42). However, using hand-crafted features could fail to capture intricate correlations between the weights and test accuracy. Instead, we explore using neural functionals to predict test accuracy from the weights of feedforward convolutional neural networks (CNN) from the Small CNN Zoo dataset [64], which contains thousands of CNN weights trained on several datasets with shared architecture, but varied optimization hyperparameters. We compare the predictive power of NEN and NFN against a method of Unterthiner et al. [64] that trains predictors on statistical features extracted from each weight and bias, and refer to it as STATNN. To measure the predictive performance of each method, we use Kendall's s [31 popular rank correlation metric with values in [11
In Table we show the results on two challenging subsets of Small CNN Zoo corresponding to CNNs trained on CIFAR-10-GS and SVHN-GS (GS stands for grayscaled), We see that NFN consistently performs the best on both datasets by a significant margin, showing that having access to the full weights increase predictive power over hand-designed features in STATNN. Because the input and output dimensionalities are small on these datasets, NEN only uses moderately more (~ 1.4x) parameters than NFNNP with equivalent depth and channel dimensions, while having significantly better performance.
3.2 Classifying implicit neural representations of images and 3D shapes
Given the rise of implicit neural representations (INRs) that encode data such as images and 3D- it is natural to wonder how to extract information about the
scenes [61
43
7 49
58
45
44] [5].

---

Table 3: Classification test accuracies (%) for implicit neural representations of MNIST, FashionM- NIST, and CIFAR-I0. 10. Equivariant architectures such as NFNs and DWSNets 47 outperform the non-equivariant MLP baselines, even when the MLP has permutation augmentations. Our NFNNP variant consistently outperforms all other methods across each dataset. Uncertainties indicate standard error over three runs.
HNP
NFNNP
DWSNets
MLP
MLPAu Aug
MNIST
92.5 = 0.071
92.9 1 0.218
74.4 + 0.143
14.5 + 0.035
21.0 + 0.172
FaShionMNIST
72.7 = 1.53
75.6 1.07
64.8 0.685
12.5 - 0.111
15.9 + 0.181
CIFAR-10
44.1 = 0.471
46.6 = 0.072
41.5 + 0.431
16.9 = 0.250
18.9 + 0.432
Table 4: Classification test accuracies (%) for datasets of implicit neural representations (INRs) of either ShapeNet- 10 [6] or ScanNet- 10 [00. Our equivariant NFNs outperform the MLP baselines and recent non-equivariant methods such as inr2vec [12.. Uncertainties indicate standard error over three runs.
NFN HNP
NFNNP
MLP
MLPAue Aug
inr2vec [2
ShapeNet-10 ScanNet-10
86.9 = 0.860
88.7 - 0.461
25.4 + 0.121
33.8 = 0,126
39.1 - 0.385
64.1 + 0.572
65.9 + 1.10
32.9 + 0.351
45.5 = 0.126
38.2 + 0.409
original data directly from INR weights. Compared to discrete signal representations (pixels, voxels, etc.), the advantage of directly processing INRs is that one can easily be agnostic to varying signal sizes and resolutions.
In this task, our goal is to classify the contents of INRs given only the weights as input. We consider datasets of SIRENs (58) that encode images (MNIST [38), FashionMNIST [66]. and CIFAR (35)) and 3D shapes (ShapeNet-10 and ScanNet-10 (53) ). For image datasets each SIREN network represents the mapping from pixel coordinate to RGB (or grayscale) value for single image, while for 3D shapes each network is a signed (or unsigned) distance function encoding a single shape. Each dataset of SIREN weights is split into training, validation, and testing sets. We construct and train invariant neural functionals to classify the INRs, and compare their performance against the MLP and MLPAun Aug baselines, which three-layer MLPs with ReLU activations and 1,000 hidden units per layer.
As Navon et al. test the performance of DWSNets on their independently constructed 2D-image own 2D- 2D-image INR datasets. Table 3 show that NENNP consistently achieves the highest test accuracies of any method on the 2D 2D-image tasks. More broadly, equivariant architectures significantly outperform the non-equivariant MLP approaches, even with permutation data augmentations.
INR datasets, we also present the results of training DWSNets
For the 3D-shape datasets we also report the performance of inr2vec 12] method for classifying 3D shapes from INR weights. Note that inr2vec's original setting assumes that all INRs in dataset are trained from the same shared initialization, whereas our problem setting makes no such assumption and allows INRs to be trained from random and independent initializations. As expected, Table 4 shows that NFNINP and NFNNP achieve significantly higher test accuracies than inr2vec, as well as the non-equivariant MLP baselines.
recent equivaarvti
In addition to superior generalization, Appendix Table shows that the NFNs are also better at fitting the training data compared to non-equivariant architectures. The MLPs achieve low train accuracy, even with an equal number of parameters as the NFNs. Interestingly, NFNNP matches or exceeds NFN performance on both CIFAR-10 and the 3D-shape datasets while using fewer parameters (e.g., 35% as many parameters on CIFAR-10), 10). Finally, we note that no weight space methods (including NFN) match the performance of near state-of-the-art methods on discrete data representations such as 2D image arrays and point clouds (see Section E.2 for details). All weight space methods still lack the geometric inductive biases that, e.g., convolutional networks have in image tasks.
7

---

Table 5: Test accuracy (%) of training with winning tickets (95% sparsity masks) produced either by running IMP or predicted by an NFN. We also show the performance of Random ticket (random mask of equivalent sparsity level), and Dense training (no sparsity). We show results for MLPs (trained on MNIST) and CNNs (trained on CIFAR-10). Uncertainties show standard error over initializations.
Dense
IMP
Random
NFNNP
NFNPT 4 0.07 95.0 0.01
CIFAR-10
63.1 = 0.06
44.0 0.06
21.1 = 0.26
41.4 = 0.08
MNIST 97.8 + 0.0
96.2 + 0.04
89.6 + 0.36
94.8 + 0.01
(x..3) 3)
(s. v)
Wr1
Method
Contrast (CIFAR-10)
Dilate
NFN
WF200
(MNIST)
10'10
MLP
0.031
0.306
MLP
0.029
0.307
6
NFNPT 0.029
0.197
NFNHNP 0.021
0.070
Original INR
Edited INR
NFNNP 0.020
0.068
Figure 3:
In weight-space style editing, an NFN directly edits the weights of an INR to alter the content it encodes. In this example, the NFN edits the weights to dilate the encoded image.
Table 6: Test mean squared error (lower is better) between weight-space editing methods and ground-truth image-space transformations.
3.3 Predicting "winning ticket" masks from initialization
The Lottery Ticket Hypothesis 19 20
LTH] conjectures the existence of winning tickets, or sparse initializations that train to the same final performance as dense networks, and showed their existence in some settings through iterative magnitude pruning (IMP). IMP retroactively finds a winning ticket by pruning trained models by magnitude: however, finding the winning ticket from only the initialization without training remains challenging.
We demonstrate that permutation equivariant neural functionals are a promising approach for finding winning tickets at initialization by learning over datasets of initializations and their winning tickets. Let Up E l/ be an initialization and let the sparsity mask M E {0, 1) dim((()) be winning ticket for the initialization, with zeros indicating that the corresponding entries of Uo should be pruned. The goal is to predict a winning ticket M given a held out initialization Un. such that the MLP initialized with Un and sparsity pattern M will achieve a high test accuracy after training.
We construct conditional variational autoencoder 32 59 cVAE] that learns generative model of the winning tickets conditioned on initialization and train on datasets of (initialization, ticket) pairs found by one step of IMP with sparsity level of P., 0.95 for both MLPs trained MNIST and CNNs trained on CIFAR- 10. Table compares the performance of tickets predicted by equivariant neural functionals against IMP tickets and random tickets. We generate random tickets by randomly sampling sparsity mask entries from Bernoulli(1
P..). In this setting, NFN is prohibitively parameter inefficient, but NFN is able to recover test accuracies that close that of IMP pruned networks in CIFAR-10 10 and MNIST, respectively. Somewhat surprisingly, NENpr performs just as well as the other NFNs, indicating that one can approach IMP performance in these settings without considering interactions between weights or layers. Appendix [_4 further analyzes how NFNpT learns to prune.
3.4 Weight-space style editing
Another potentially useful application of neural functionals is to edit (i.e., transform) the weights of a given INR to alter the content that it encodes. In particular, the goal of this task is to edit the weights of a trained SIREN to alter its encoded image (Figure B). We evaluate two editing tasks:
8

---

Dilate (MNIST)
Contrast (CIFAR-10)
Ground truth
3
5
4
G
9
6
8
3
Original INR
3
5
4
6
9
6
8
3
ML.Pagg
3
5
4
6
9
6
8
3
NEN
3
5
4
6
9
6
8
3
NENsr
3
5
4
6
9
6
8
3
Figure 4: Random qualitative samples of INR editing behavior on the Dilate (MNIST) and Contrast (CIFAR-10) editing tasks. The first row shows the image produced by the original INR, while the rows below show the result of editing the INR weights with an NFN. The difference between MLP neural functionals and equivariant neural functionals is especially pronounced
the more challenging Dilate tasks, which require modifying the geometry of the image. In the Contrast tasks, the MLP baseline produces dimmer images compared the ground truth, which is especially evident in the second and third columns.
(1) making MNIST digits thicker via image dilation (Dilate), and (2) increasing image contrast on CIFAR-I0 (Contrast). Both of these tasks require neural functionals process the relationships between different pixels to successfully solve the task.
To produce training data for this task, to dilate or increase the contrast of the MNIST and CIFAR-10 images, respectively. The training objective is minimize the mean squared error between the image generated by the NFN-edited INR and the image produced by image processing. We construct equivariant neural functionals to edit the INR weights, and compare them against MI.P-based neural functionals with and without permutation augmentation.
standard image processing libraries (28] OpenCV]
6 that permutation equivariant neural functionals (NFNHHN and NFNNp) achieve signifi- cantly better test MSE when editing held out INRs compared to other methods, on both the Dilate (MNIST) and Contrast (CIFAR-10) 10) tasks. In other words, they produce results that are closest to the "ground truth" image-space processing operations for each task. The pointwise ablation NFNrr performs significantly worse, indicating that accounting for interactions between weights and layers is important to accomplishing these tasks. Figure 4 shows random qualitative samples of editing by different methods below the original (pre-edit) INR. We observe that NFNs are more effective than Aug at dilating MNIST digits and increasing the contrast in CIFAR-10 images.
MI.Paug
4 Related work
The permutation symmetries of neurons have been topic of interest in the context of loss landscapes and model merging 21 62 17 I. Other works have analyzed the degree of learned permutation symmetry in networks that process weights [64] and studied ways of accounting for symmetries when measuring or encouraging diversity in the weight-space
13 However, these symmetries have not been a key consideration in architecture design for processing weight-space objects [2] 39 23 36
69 13 [
Instead, existing approaches try to encourage permutation equivariance through data augmentation [51] 44 In contrast, this work directly encodes the equivariance of the weight-space into our architecture design, which can result in much higher data and computational efficiency, as evidenced by the success of convolutional neural networks [37].
Our work follows a long line of literature that incorporates structure and symmetry into neural [8] This includes works that design permutation equivariant layers, originally for processing sets and graphs [52] 68 25 41. More generally. equivariant layers have been developed for processing arbitrary rank-k & tensors (and correspondingly. hypergraphs) under higher-order permutation actions 40 63 48). In this context, we can view feedforward networks as graphs that factor into special layered structure where each (hidden) layer of neurons is an independently permutable set of nodes, and the weights are adjacency matrices specifying the edge weights between nodes of adjacent layers. Then our equivariant NF-Layers give the maximal set
network architectures [37 8 54 34
9

---

of linear functions on these special graph structures. As discussed in Section T Navon et al. 47 recently developed an equivariant weight-space layer that is equivalent to our Layer in the HNP setting. Our work introduces the NP setting to improve parameter efficiency and scalability over the HNP setting, and extends beyond the fully connected case to handle convolutional weight-space inputs.
5 Conclusion
This paper proposes novel symmetry-inspired framework for the design of neural functional networks (NFNs), which process weight-space features such as weights, gradients, and sparsity masks. Our framework focuses the permutation symmetries that arise in weight-spaces due to the particular structure of neural networks. We introduce two equivariant NE-Layers as building blocks for NFNs, which differ in their underlying symmetry assumptions and parameter efficiency, then use them to construct a variety of permutation equivariant neural functionals. Experimental results across diverse settings demonstrate that permutation equivariant neural functionals outperform prior methods and are effective for solving weight-space tasks.
Limitations and future work. Although we believe this framework is step toward the principled design of effective neural functionals, there remain multiple directions for improvement. One such direction would concern extending the NF-Layers beyond feedforward weight spaces, in order to process the weights of more complex architectures such as ResNets [26] and Transformers [65]. Another useful direction would involve reducing the activation sizes produced by NF-Layers, in order to scale neural functionals to process the weights of very large networks. Finally, improvements to NF-Layers could account for the other symmetries of neural network weight spaces, such as scaling symmetries in ReLU networks 22
Acknowledgements
We thank Andy Jiang and Will Dorrell for helpful early discussions about theoretical concepts, and Yoonho Lee and Eric Mitchell for reviewing early draft of this paper. AZ and KB are supported by the NSF Graduate Research Fellowship Program. This work was also supported by Apple.
References
[1] S. K. Ainsworth, J. Hayase, and S. Srinivasa. Git re- basin: Merging models modulo permutation symmetries. arXiv preprint XrX:::2209 04836 2022.
[2] M. Andrychowicz, M. Denil, S. Gomez, M. W. Hoffman, D. Pfau, T. Schaul, B. Shillingford, and N. De Freitas. Learning learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
[3] Bengio, Y. Bengio, Cloutier, and Gescei. On the optimization synaptic learning rule. In Optimality in Biological and Artificial Networks?, pages 281-300 281-303. Routledge, 2013.
[4] Brea, B. Simsek, B. Illing, and W. Gerstner. Weight-space symmetry in deep networks gives rise to permutation saddles, connected by equal-loss valleys across the loss landscape. arXiv preprint arXiv: /907,0291 1, 2019.
[5] M. M. Bronstein, J. Bruna. T. Cohen, and P. Veličković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. urXiv preprint arXiv:2 13478, 2021.
[6] A. X. Chang, T. Funkhouser, L. Guibas, P. Hanrahan, Q. Huang, Z. Li, S. Savarese, M. Savva, Song, H. Su, J. Xiao, L. Yi, and F. Yu. ShapeNet: An Information 3D Model Repository. Technical Report 151 [cs.GR], Stanford University - Princeton University Toyota Technological Institute Chicago, 2015.
[7] Z. Chen and H. Zhang. Learning implicit fields for generative shape modeling, In Proceedings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition, pages 5939-5948. 2019.
10

---

[8] T. Cohen and M. Welling. Group equivariant convolutional networks. In International confer- ence on machine learning, pages 1990-2999, PMLR, 2016.
[9] S. Cohen, M. Geiger, Köhler, and M. Welling. Spherical CNNs. arXiv preprint 1801.00800 2018.
[10] A. Dai, A. X. Chang, M. Savva, M. Halber, T. Funkhouser, and M. Nießner. Scannet: Richly- annotated 3d reconstructions of indoor scenes. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 5828-588 5839, 2017.
[11] N. De Cao, W. Aziz, and I. Titov. Editing factual knowledge in language models. arXiv preprint 1xXiv2204.08866 2021.
[12] L. De Luigi, A. Cardace, R. Spezialetti, P. Zama Ramirez, S. Salti, and L. Di Stefano. Deep learning on implicit neural representations of shapes. International Conference on Learning Representations (ICLR), 2023.
[13] L. Deutsch, E. Nijkamp, and Y. Yang. A generative model for sampling high-performance and diverse weights for neural networks. arXiv preprint UrXiv:1905.0220 2019.
[14] E. Dupont, Y. W. Teh, and A. Doucet. Generative models as distributions of functions. arXiv preprint wrXiv:22022.47770 2021.
[15] E. Dupont, H. Kim, S. Eslami, D. Rezende, and D. Rosenbaum. From data functa: Your data point is a function and you should treat like one. arXiv preprint arXiv:2201. 12204 2022.
[16] G. Eilertsen, D. Jönsson, T. Ropinski, J. Unger, and A. Ynnerman. Classifying the classifier: dissecting the weight space of neural networks. arXiv preprint arXiv:2002.05668 2020.
[17] R. Entezari, Sedghi, O. Saukh, and B. Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. arXiv preprint arXiv:21 2021.
[18] M. Finzi, M. Welling, and A. G. Wilson. A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups. International Conference on Machine Learning, pages 33 18-3228. PMLR, 2021.
[19] J. Frankle and M. Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint /
2018.
[20] J. Frankle, G. K. Dziugaite, D. M. Roy, and M. Carbin. Stabilizing the lottery ticket hypothesis. arXiv preprint arXiv: 1,
[21] T. Garipov, P. Izmailov, D. Podoprikhin, D. Vetrov, and A. G. Wilson. Loss surfaces, mode connectivity, and fast ensembling of DNNs. Advances in neural information processing systems, 31, 2018.
[22] C. Godfrey, Brown, T. Emerson, and H. Kvinge. On the symmetries of deep learning models and their internal representations. Advances in Neural Information Processing Systems, 35: 11893-11905 2022.
[23] D. Ha, A. Dai, and Q V. Le. Hypernetworks. arXiv preprint arXiv: 2016.
[24] J. Harb, T. Schaul, D. Precup, and P. Bacon. Policy evaluation networks. CoRR, abs/2002.11833,
2020. URL URL httpps://axxii. www.vv.rrg/babaa hangg/abs/2000 11833.
[25] J. Hartford, D. Graham, K. Leyton-Brown, and S. Ravanbakhsh. Deep models of interactions across sets. In International Conference on Machine Learning, pages 1909-1918. PMLR, 2018.
[26] K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. CoRR, abs/1512, 3385:2, 2015.
[27] R. Hecht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Advanced Neural Computers, pages 129-135. Elsevier, 1990.
[28] Itseez. Open source computer vision library. attpps://github..o///ttseees 2015.

---

[29] Y. Jiang, D. Krishnan, H. Mobahi, and S. Bengio. Predicting the generalization gap in deep networks with margin distributions. International Conference on Learning Representations. 2019. URL httph::tt/p/erreeeeeeeee http:p///.ppeeeeeereeeeeeeeeeeep
[30] Y. Jiang, P. Natekar, M. Sharma, S. K. Aithal, D. Kashyap, Subramanyam, C. Lassance, D. M. Roy, G. K. Dziugaite, S. Gunasekar, et al. Methods and analysis of the first competition in predicting generalization of deep learning. In NeurIPS 2020 Competition and Demonstration Track, pages 170-190. PMLR, 2021.
[31] M. G. Kendall. A new measure of rank correlation. Biometrika, 30(1/2):81-93, 1938.
[32] D. P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint arXiv: 1312.6114,
2013.
[33] B. Knyazev, M. Drozdzal, G. W. Taylor, and A. Romero Soriano. Parameter prediction for unseen deep architectures. Advances in Neural Information Processing Systems, 34:29433- 29448, 2021.
[34] 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.
[35] A. Krizhevsky, G. Hinton, et Learning multiple layers of features from tiny images. 2009.
[36] D. Krueger, C.-W. Huang, R. Islam, R. Turner, A. Lacoste, and A. Courville. Bayesian hypernetworks. arXiv preprint 17X.v:1710004559, 2017.
[37] Y. LeCun, Y. Bengio, al. Convolutional networks for images, speech, and time series. The handbook of brain theory and neural networks, 3361(10) 1995.
[38] Y. LeCun, C. Cortes, and C. Burges. Mnist handwritten digit database. ATT Labs [Online]. Available:
2, 2010.
[39] K. Li and Malik. Learning optimize. arXiv preprint wrXiv:1606.01880 2016.
[40] H. Maron, H. Ben-Hamu, N. Shamir, and Y. Lipman. Invariant and equivariant graph networks. arXiv preprint arXiv: 1812.09929 2018
[41] H. Maron, O. Litany, G. Chechik, and E. Fetaya. On learning sets of symmetric elements. In International conference on machine learning, pages 6734 6744. PMLR, 2020.
[42] C. H. Martin and M. W. Mahoney, Implicit self-regularizztiion in deep neural networks: Evidence from random matrix theory and implications for learning. The Journal of Machine Learning Research, 22(1):7499-751) 2021.
[43] L. Mescheder, M. Oechsle, M. Niemeyer, S. Nowozin, and A. Geiger. Occupancy networks: Learning 3d reconstruction in function space. In Proceedings of the IEEE/CVF conference computer vision and pattern recognition, pages 4460-4470, 2019.
[44] L. Metz, J. Harrison, C. D. Freeman, A. Merchant, L. Beyer, J. Bradbury, N. Agrawal, B. Poole, I. Mordatch, A. Roberts, et al. Velo: Training versatile learned optimizers by scaling up. arXiv preprint arXiv:221 1.09600 2022.
[45] B. Mildenhall, P. P. Srinivasan, M. Tancik, J. T. Barron, R. Ramamoorthi, and R. Ng. Nerf: representing scenes as neural radiance fields for view synthesis (2020). arXiv preprint arXiv:2003.08934, 2020.
[46] E. Mitchell, C. Lin, A. Bosselut, C. Finn, and C. Manning. Fast model editing at scale. arXiv preprint 10. 2021.
[47] A. Navon, A. Shamsian, Achituve, E. Fetaya, G. Chechik, and Maron. Equivariant architectures for learning in deep weight spaces. arXiv preprint 2rX 2023.
12

---

[48] 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.
[49] J. Park, P. Florence, J. Straub, R. Newcombe, and S. Lovegrove. Deepsdf: Learning continuous signed distance functions for shape representation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 165-174, 2019.
[50] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury. G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, et al. Pytorch: An imperative style, -erhhreeformance deep learning library. Advances in neural information processing systems, 32, 2019.
[51] W. Peebles, I. Radosavovic, T. Brooks, A. A. Efros, and J. Malik. Learning to learn with generative models of neural network checkpoints. arXiv preprint arXiv:2209.1289) 2022.
[52] C. Qi, H. Su, K. Mo, and L. Guibas. Pointnet: deep learning on point sets for 3d classification and segmentation. evpr (2017). arXiv preprint 1612.0006 2016.
[53] C. Qin, H. You, L. Wang, C.-C. J. Kuo, and Y. Fu. Pointdan: A multi-scale 3d domain adaption network for point cloud representation. Advances in Neural Information Processing Systems, 32, 2019.
[54] S. Ravanbakhsh, J. Schneider, and B. Poczos. Equivariance through parameter-sharing. In International conference on machine learning, pages 2892-2900 PMLR, 2017.
[55] A. Rogozhnikov. Einops: Clear and reliable tensor manipulations with einstein- like notation. In International Conference on Learning Representations, 2022.
[56] T. P. Runarsson and M. T. Jonsson. Evolution and design of distributed learning rules. In 2000 IEEE Symposium on Combinations of Evolutionary Computation and Neural Networks. Proceedings of the First IEEE Symposium on Combinations of Evolutionary Computation and Neural Networks (Cat. No. 00, pages 59-63, IEEE, 2000.
[57] A. Sinitsin, V. Plokhotnyuk, D. Pyrkin, S. Popov, and A. Babenko. Editable neural networks. arXiv preprint arXiv:2004.00345, 2020.
[58] V. Sitzmann, J. Martel, A. Bergman, Lindell, and G. Wetzstein, Implicit neural representations with periodic activation functions. Advances in Neural Information Processing Systems, 33: 7462-7473, 2020.
[59] K. Sohn, H. Lee, and X. Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems, 28, 2015.
[60] S. Sokota, H. Hu, D. J. Wu, Z. Kolter, J. N. Foerster, and N. Brown. A fine-tuning approach to belief state modeling. In International Conference on Learning Representations, 2022. URL
net: htt.:e//oooommmmieccc.....aa http:://pppaeeeiiee..er///ootttp
[61] K. O. Stanley. Compositional pattern producing networks: A novel abstraction of development. Genetic programming and evolvable machines, 8: 1:1--1162, 2007.
[62] N. Tatro, P.-Y. Chen, P. Das, I. Melnyk, P. Sattigeri, and R. Lai. Optimizing mode connectivity via neuron alignment. Advances in Neural Information Processing Systems, 33:1 13:00000-1331, 2020.
[63] E. H. Thiede, T. Hy, and R. Kondor. The general theory of permutation equivarant neural networks and higher order graph variational encoders. arXiv preprint arXiv:2004.03399, 2020.
[64] T. Unterthiner, D. Keysers, S. Gelly, O. Bousquet, and I. Tolstikhin. Predicting neural network accuracy from weights. arXiv preprint arXiv:2002. 1448, 2020.
[65] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, F.. Kaiser, and Polosukhin. Attention is all you need. Advances neural information processing systems, 30, 2017.
13

---

[66] H. Xiao, K. Rasul, and R. Vollgraf. Fashion-mnist: novel image dataset for benchmarking machine learning algorithms, 2017.
[67] Yak, J. Gonzalvo, and Mazzawi. Towards task and architecture-inddeeenennt generalizatonooon gap predictors. arXiv preprint urXiv:1906.00555 2019.
[68] M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. Salakhutdinov, and A. Smola. Deep sets. doi: 10.48550. arXiv preprint preprint ARXIV.1003.06/14 2017.
[69] C. Zhang, M. Ren, and R. Urtasun. Graph hypernetworks for neural architecture search. arXiv preprint
2018.
14

---

Appendix
Equivariant NE-LLayer pseudocode
15
B NF-Layers for the NP Setting
16
B. Full definition
16
B.2 Proof of Proposition
17
B.2.1
Parameter orbits and sharing
18
B.2.2
Equivalence equivariant NF- Layer definition
19
B.3 Proof of Proposition
19
C NF-Layers for the HNP setting
20
Equivariant NE-Layer
20
Invariant NF-Layer
20
Additional experimental details
20
D.I
Predicting generalization
20
D.2
Predicting "winning ticket" masks from initialization
22
D.3
Classifying INRs
23
D.4
Weight-space style editing
23
E
Additional experiments and analysis
24
Analyzing train vs accuracies INR classification
................................
24
-en-weight-space baselines for image classification
24
E.3
Ablating positional encoding
24
E,4
Interpreting learned lottery ticket masks
24
Predicting MLP generalization from weights
27
Equivariant NF-Layer pseudocode
Here we present a multi-channel implementation of the S-equivariant NF-Layer presented in Eq (which ignores biases), using PyTorch (50] and Einops-like (55) pseudocode. That is, it implements linear layer Wo > W where c; and c, are the number input and output channels.
Note that our actual implementation differs from this pseudocode in a few ways: (1) supports the full weight-space U which includes biases, (2) it supports convolution weights as well as fully connected weights, and (3) it initializes parameters based on the fan-in of the NF-Layer, instead of from N (00 1).
class NPLayer (Mo.uuee):: (self, L, co, ci):
def
3
super (a__o)))__i(i
énitialize weights. comoutput channels, ciminput channels.
5
self. A nn. Parameeer (torch.rtorrh....oo randn (L, L, co, ci))
self mn. Paraaeeeee (torch. randn (L, co, ci))
7
self B_prev nn Parameter (torch. randn co, ci))
self.C nn. www.marameeeeeeeeemmm sarach randn (L, co, ci))
self next nn Parameter (torch randn co, ci))
15

---

10
self D nn. n.. Parameter (torch. (L, co, ci))
def (self, w): 13
Input is list of I, weight-space tensors with shapes: no) ((B, "L,
We return a list of L tensors with shapes: 15
(1, ((B, Cor ((,
17
where Co are the input and output channels.
W! Ki " W!! Each has shape (B,d-
Compute
20
row_means
w.meaan (-2) for in W]
Compute
W W Each has shape (B,g,
22
col_means
[w.meam mean (-1) for w in W] Compute (!! W!! with shape (B,
24
rowcol_means torch stack [w. mean (dim= -1)) for w in W], -1)
out_ []
26
for i, (w, b, b_prev, c_next, d) in (zipp
W, A,lf.A, self.B, self. beleeB,prev, self.C, self C_next, self
28
)):
Calculates E. Eat
29
h1 einsum (a, rowcol_means, 'co ci 8, B ci B co () ())
Calculates (-((())
+,k h2 einsum(b, row_means[i], 'co ci, ci iim
n_im1')
33
if i 0:
# Calculates (0--
h
h2 += einsum b_ree col col_means [i-1], 'co ci, B ci n_iml co () n_im1')
Calculates cw
36
j+ h3 einsum(c, col_means[i], 'co ci, ci n -> b co n ())
37
if i self 1: # Calculates cl
39
+J
40
h3 += einsum(c oinsum(c_next, (c_next, row_col[i+1], 'co ci, C n -> B co n_i 0)
Calculates d'w(l)
41
.k
42
h4 = einsum(d, w, 'co ci, B ci n niim -> B co n m_im1')
Calculates H(W(0) with shape (B, R4, n4-1).
44
out_k W. append (h1 + h2 + h3 + h4)
return out_w
B NF-Layers for the NP Setting
This section presents futher details related to the S-equivariant layer H : 14 + Ul and the S-invariant layer First, Sec. B.] gives the full definition of as a function of both weight-space and bias-
space features. Sec. B.2 provides proof that H(-) expresses any S-equivariant linear function on 2
(Proposition []. Sec. B.3 proves that combining our invariant layer P with fully connected layer can express any S-invariant function from /I to Rd (Proposition
B.1 Full definition
present the full definition of the S-equivariant NF-Layer H 2244 - U4 which completes Eq. 2q.2lp by also including the biases. The layer processes set of weights U (W, v) and outputs arrays (Y,²)
16

---

Table 7: S-equivariant parameter sharing for linear maps vee(U) + Ovec(U). Parameter sharing is a system of constraints on the entries of (0, o. s, tb) = B. Each table is organized by the layer indices s). For example, the first table says that for any s) where s - 1, constrain Oish = = s.p'.g' = ag 2 for any j, k. q and j', ",,'dd where k + p and k' * p'- After parameter sharing, we observe that there are only a constant number of free parameters for each (i, 8) pair, adding up to O (1.2) parameters total.
s
==i
8=
-=++1 1
other 8
j k.k k #g
Disk spg
1-11 k p a
j= p, # q
a
1,4+1 1
i q
0 1.5
bg b
k = P
8 k. = q
j == q
d'g
j - p.
8 i
8 s=1
other 8
1 P
0 1
j+1
Øtjkk sp
6 b-
= P
b
&
j ==1 p
8 s=i i
8 8=i i+1 1
other
1
j
0 b
i
fo
atp²
109 [b---1 j=p p
j= q
s=i
other
a i i p
b'
j j=1 P
with the same dimensions. The learnable parameters are in blue:
H(U) (())= WW X
(3)
jk = Σ Cg w + og. + by w *. + bbg b
k.t
+ ct w
j. + + 1 (1+1)
+ c (f-1) + d'g w/(
jk
1 = Σ w! + ol) w + Ui w((+2) + b 6, (1)
where denotes summation averaging over dimension.
B.2 Proof of Proposition
Here we show that the layer defined in Eq. 3 is not only S-equivariant, it can express every linear equivariant function 24 + U. Our strategy is to first consider general linear layers o(: U + U parameterized by 0 E θ and then constrain to l such that T(: () is S-equivariant. Applying this constraint will directly yield Eq. B
For our purposes, it is sometimes convenient to distinguish layer, row, and column indices of entries in U, so we split the parameters 0 = (0, sp, v) and write T : u - 24 in the form:
T(U:0) = (Y()),
W V = u
(4)
Y( = ΣΣΣΣ a=1
Jūsk (
jijk p
(5)
I
x(U),--
ΣΣΣ Σ
((())
Ge (s)
(6)
Papn
pq
Ppp p
q=1
p=1
17

---

We introduce the notation Ua to identify individual entries of U. Here a is tuple of length two or three, for indexing into either a weight or bias. We denote the space of valid index tuples of W and V by W and V, respectively, and define U == WU U as the combined index space of U. For example, if (i k) W. then Ua w jk
We can then define the index space Il == U x U for parameters 0 E O. We use 03 to index entry of I with upper and lower indices (a:, [] I. For example, if a (i.j. k) W and / (s,P. 9) W, we have 0% = &ifk
By flattening U into a vector, we can think of as a matrix with indices a and 3. Eq. can be rewritten in the flattened form:
T(U:0)a = Σ 0zzus
In the flattened representation, we can write the action of a E S on vec(U)) by a permutation matrix Po. Equivariance requires that PU)) = OP, vee(U) for any E S. Since the input U can be anything, we get the following constraint on A:
P.0 = OPo, Vo e.
(8)
We can re-express the action of 8 on U (Eq. write (a) which is defined:
I as an action on the index space U. For U,
= -(t.j.b))((t.))))
o,1(() (L.j,k) E w (i. (1,o,0)) (V,))
(9)
(10)
for any E S. We extend this definition into action of S on T:
a [a, 8] = [aa, oB), [x
(11)
We can now rewrite the equivariance constraint in Eq. using indices a, B::
[P,0] = 0g ((a) = -0 = [0Pa]
(12)
By relabeling +- (a)) we can rewrite this condition 03 = 00(a) Hence, equivariance is equivalent to requiring that 9 share parameters within orbits under the action of S on its indices a [54] Prop 3.1]. We characterize these orbits in Sec. B.2.1 yielding Table 7 which gives the parameter sharing as system of constraints the indices of #. These constraints characterize the space of S-equivariant T(,, 0). In Sec. B.2.2 we conclude the proof by showing that under parameter sharing, T(, 0) is equivalent Eq. 3
B.2.1 Parameter orbits and sharing
We now partition the parameters of 0 into orbits under the S-action on its index space, in order to generate the parameter sharing constraints given in Table 7 The index space of 0 is I U. There are four subsets of T:
1. TWW W x W: Contains (a, B3] = [(i, k), (s, (s 4)],, indexing parameters Giyk
apq
2. [WV W. V: Contains (a, 3 [(i,], k)), (s, indexing parameters diik
3. V = V X W: Contains (a, 3 = [(i.j), (s.P. g)). indexing parameters 2
4. TVV = x Contains [a, 3] = [ii,j), (s.p)). indexing parameters v/5
Consider the block of indices [WW = w x W, containing (a, 3] = [(i, j. k), (s, 4) indexing parameters dg. Since the S-action never changes the layer indices s), we can independently
consider orbits within sub-blocks of indices n.W [(i. k), (s,p. 4) Vj, kz p. }. The number of orbits within each sub-block TWW depends the relationship between the layer indices i and 3 s: they are either the same layer (s = i). they are adjacent (s = i or s = i i 1), or they are non-adjacent (s {i- - L, 1}). We now analyze the orbits of sub-blocks for a few cases.
18

---

If s = i- 1, then choose any two indices [a , (a(2),3(7))) E 1i, where the first satisfies
p o lc and the second satisfies p k. Then the orbits of each index are:
Orbit (00) = ((..j,)), (8.P.9) Vi
(13)
(((9) B(3)) {[((i,, (), (s, p. Vj.k,
Orbit
(14)
We that these two orbits actually partition the entire sub-block of indices 1 with each orbit characterized by whether or not k: We introduce the parameters d (for the first orbit) and byg- (for the second orbit). Under equivariant parameter sharing, all parameters of corresponding b,, depending on whether
nww are equal to either ag apr-1 or
+ k.
If s = i+1 1, we instead choose any two indices where the first satisfies * and the second satisfies
j q. Then the sub-block of indices I.W is again partitioned into two orbits:
{ [(i.j,,), (s. p. 4) kj.k, p. ) and { [(i,j, k), (s, Vj.k.p.:jj=q) (15)
depending on the condition j q- We name parameters a atftt and 1.41 for this sub-block, with one for each orbit.
We can repeat this process for sub-blocks of TWW where = and & 1,,1, 13. well as for the other three blocks of II. Table 7 shows the complete parameter sharing constraints on resulting from partitioning all possible sub-blocks into orbits.
Number of parameters. We also note that every layer pair (i, introduces only a constant number of parameters: the number of parameters in each cell of Table 7 has no dependence on the input, output, or hidden dimensions of 44. Hence the number of distinct parameters after parameter sharing simply grows with the number of layer pairs, i.e. O (L²).
B.2.2 Equivalence equivariant NF-Layer definition
All that remains is show that the map T(: 0) u with S-equivariant parameter sharing is equivalent to the NF-Layer H we defined in Eq. 3
(Table 7
Consider single term from Eq. 4 where simplify:
1. Substituting using the constraints of Table 7
Σ baik W(i-1) 1] a 1--11 1 1-1.p L.g
w(i-1) P.g
+ -
Σ Ewlt-)
(16)
kpP
a W (1-1) + -6 -w (1-1)
We can then reparameterize 6 can simplify every term of Eq. in a similar manner using the parameter sharing of Table 7 reducing the general layer to the S-equivariant Na-Lay
gii-1 - a ااد resulting in two terms that appear in Eq. We
B.3 Proof of Proposition
Analogous to Sec. B.2 we begin with general parameterized linear maps T(: 0) : 24 - Rd and characterize the parameter sharing implied by S-invariance. We then show that T(., under these constraints equivalent to f P. where f 12L RRd is parameterized linear map (i.e., fully connected layer) and P(U)- = [W()) w(L) (1 " o.4)). ...). We will assume d = 1 for the rest of this section, but it is straightforward to generalize our arguments to arbitrary
In general, we can parameterize linear maps 21 + by 0 E ddim(()) where:
T(U, 0) = Σ UaUa.
(17)
where a is an index for either weights or biases as in Sec. T(oU) implics:
B.2
The invariance constraint T(U)
0a 0 Or(a))
(18)
where af (a) is defined by Eq.
19

---

We can now characterize the orbits of the indices of & under S. In this case, the orbits (and resulting parameter sharing) are relatively simple. For a fixed i, all (i. k) E W form an orbit, and we introduce a shared parameter a'. Similarly, for each i all (i, E V form orbit, and we introduce shared parameter by for each. We can then rewrite the layer:
T(U = Σ a
< jk
+
Σ
=
= a a + +b
(19)
This is simply linear combination of the output of P(U) parameterized by {(a', b')}. and can be written the composition f where f is fully connected layer.
C NF-Laysrs for the HNP setting
C.1 Equivariant NF-Layer
Because expression for the S-equivariant NF-Layer analogous to Eq. instead define the layer in terms of its parameter sharing (Tables 8
would be unwieldy,
11 on
We can derive HNP-equivariant parameter sharing of 0 using a similar strategy to Sec. B.2.1 we partition the index spaces 1WW, IwV, IVW nIV into orbits under the action of S, and share parameters within each corresponding orbit of 0, Φ, (p, tb. The resulting parameter sharing is different from the NP-setting because while the action of S on L/ could permute the rows and columns of every weight and bias, the action of S on 24 does not affect the columns of W(1) or the rows of WLL,, ((L)) which correspond input and output dimensions (respectively).
The orbits are again analyzed within sub-blocks defined by the values of the layer indices (i, As with the NP setting, there are broadly four types of sub-blocks based on whether 8, i 8 i ++ 1, or $ {i 1, i, + }}. However, there are now additional considerations based on whether or s is an input or output layer. For example, consider the sub-block of 1WW where i s = 1, which we denote 1 The action on the indices in this sub-block can be written [a, 3] = (i,,o(() k), (1, a (p), 4)]. Importantly, the column indices k, g are never permuted since they correspond to the input layer. We see that IWW contains two orbits for each ls E [[L..no] and g E [1..mo], with the two orbits characterized by whether or not j = p. Hence we have 2n orbits and Table introduces 2ng parameters of parameters.
{ ag 1.1. ,6, 6111k.g k.q E [L..no] } for this sub-block
Now consider another sub-block of TWW where <i= 8 L. Now the action of Š on indices in this sub-block can be written (a, (] = [(i, (k)) (i,) (g))( Then have total of two orbits characterized by whether or not k: = p. rather than 2n orbits for the i = 1 case. Tables present the complete parameter sharing for each of 9. o, i,, t. resulting from analyzing every possible orbit within any sub-block of [WW [WV W IVV.
C.2 Invariant NF-Layer
While the NP-invariant NF-Layer sums over the rows and columns of every weight and bias, under HNP assumptions there is no need to sum over the columns of W11) (inputs) or the rows of W1)), (24) (outputs). So the HNP invariant NF-Layer P : + R mo is defined:
P(U)= P(U), (()) W(L) .((t))
(20)
where wi(!) and W:(L) denote summing over only the rows or only the columns of the matrix, respectively. Note that P satisfies S-invariance without satifying S-invariance.
D Additional experimental details
D.1 Predicting generalization
The model we use consists of three equivariant NF-Layers with 16, 16, and 5 channels respectively. We apply ReLU activations after each linear NF-Layer. The resulting weight-space features are passed
20

---

Table 8: HNP-equivariant parameter sharing on S 0, corresponding to the NF-Layer H :u - U.
i= 2
2<i<L L
i=
1
2.1
k * p
0g
k # p 1
1,-13
k k
12 k p
b
k k= p
bg
1---
k = p
i=1 1
1 1<ick L
i == L
j #9
8 s=i
ag 11kkg
j&P bg k1, j=p
b j=
a C0
L.L.P L.L.P
k k q
2 j
k =
d j= !
i=1
1 1<i 1-1
i= 1=
8= 3
0 1,2, j j=
1.2,
3 i79
122 i#9
4 L-LLLP j =
1-1,P i+1 q
j=9
i 1, <
i= 1,8= L
1<i < L. 8
1,ck
I.L.k.p 0
1 i L,s=1
i= L. =1
i= L,1< L
other 8
1,1,4
L,1.9 g
L.8.)
0
0
1 < i < I, < <
0
Table HNP-equivariant parameter sharing on & C 0, corresponding to the NE-Layer Ā u - U.
dikk ap
1<i<L
i= L
0
21-1
k
a L-b-11
k + p
b k k=p
b L.1-15
k p
i ==
1<i<L
i= L
8 8=i
a 1.1.k
j j$1 p
j#p p j=p
111. 0
j = p
i 1,1 s L
i 1,
1 L,S=L
1,sk ④
0 I.L.k.p
I.L.P
a
other
1 <
L, 1 < <<L L
i L, 1 < <<L
a L,
a 6
21

---

Table 10: equuequuiari parameter sharing on 4 C 8, corresponding to the La-Lr u - u.
Pssp
i=1
I<i<L
i= L
i
0 111 b111
a by
j&p
==p
j=p
1 < L -
i= L -
8 +
a2 j+p
1.op j p
bltht j=p
b 1.1 j ==p
1 I
< L,8=1
1si< L, 1 < op
1 < L,S = L
a.L.
other 8
I,, =
i <
Table 11: HNP-equivariant parameter sharing on t C 0, corresponding to the NF-Layer 4 - U.
V
1 Isich
i= T.
i
jøp
b L.L.i
j=p
1 1sic L,
i= L, S
1 SsisL1ss44
other
Liad
4
into an invariant NF-Layer with mean pooling. The output of the invariant NF-Layer is flattened and projected RR 1,000 The resulting vector then passed through an MLP with two hidden layers. each with 1,000 units and ReLU activations. The output is linearly projected to scalar and passed through sigmoid function. Since the output of the model be interpreted as a probability, we train the model with binary cross-entropy with hyperparameters outlined in Table 12 The model is trained for 50 epochs with early stopping based 7 on the validation set, which takes hour on Titan RTX GPU.
Table 12: Hyperparameters for predicting generalization on Small CNN Zoo.
Name
Values
Optimizer
Adam 0.001
Learning rate
Batch size
Loss
Binary cross-entropy
Epoch
50
D.2 Predicting "winning ticket" masks from initialization
Concretely, the encoder learns the posterior distribution (Z Uo, M) where Z E [ddim(4)xC is the latent variable for the winning tickets and the number latent channels. The decoder learns PM(M] Uo, Z). and both encoder and decoder are implemented using equivariant NF-Layers. For the prior p(Z) we choose the isometric Gaussian distribution, and train using the evidence lower bound (ELBO):
Co (M, Un) = (2~gp()1U. (M Uo, =)} DKI. (g)( Uo. M) 11 p()))
22

---

The initialization and spansity mask are concatenated so the input to the encoder qe is (U, e [ddim(())2 x2 After the bottleneck, we concatenate the latent variables and the original mask along the channels, i.e. the decoder input is (Uo, 2) E [Rdim(())((++1)
The first dataset uses three-layer MLPs with 128 hidden units trained on MNIST and the second uses CNNs with three convolution layers (128 channels) and fully-connected layers trained on CIFAR- 10. In each dataset,
include 400 pairs for training and hold out 50 for evaluation. The hyperparameter details are in Table [3 The encoder and decoder models contain equivariant Na-Layers with 64 hidden channels within each layer. The latent variable is 5 dimensions. Training takes 5H on Titan RTX GPU.
Table 13: Hyperparameters for predicting L.TH on MNIST and CIFAR-10.
Name
Values
Optimizer
Adam
Learning rate
10-3
Batch size
[4. 8]
Epoch
200
D.3 Classifying INRs
We use SIREN (58) for our INRs of CIFAR-10, FashionMNIST, and MNIST. For the SIREN models, we used three-layer architecture with 32 hidden neurons in each layer. We trained the SIRENs for 5,000 steps using Adam optimizer with a learning rate x 10 Datasets were split into 45,000 training images, 5,000 validation images, and 10,000 (MNIST, CIFAR- 10) 20,000 (FashionMNIST) test images. We trained 10 copies (MNIST, FaShionMNIIST) or 20 copies (CIFAR- 10) of SIRENs on each training image with different initializations, and single SIREN on each validation and test image. No additional data augmentation was applied. For 3D shape classification, we adopt the same protocol introduced in [1].. and train each SIREN to fit the Unsigned Distance Function (UDF) value of points sampled around a shape. Each SIREN is composed of a single hidden layer with 128 neurons. We use Adam
optimizer and we train for 1,000 steps.
We also trained neural functionals with three equivariant NE-Layers + ReLU activations, each with 512 channels, followed by invariant NF-Layers (mean pooling) and a three-layer MLP head with 1,000 hidden units and ReLU activation. Dropout was applied to the MLP head only. For the NFN 10 IO-encoiing, used sinusoidal position encoding with maximum frequency of 10 and 6 frequency bands (dimension 13). The training hyperparameters are shown in Table 14] and training took ~ 4H on Titan RTX GPU.
Table 14: Hyperparameters for classifying INRs on MNIST and CIFAR-10 using neural functionals.
Name
Values
Optimizer
Adam
Learning rate
1 x 10
Batch size
32 2 x 10y
Training steps
MLP dropout
0.5
We also experimented with larger MLPs (4,000 and 8,000 hidden units per layer) that have parameter counts comparable to those of the NFNs, but found that it did not significantly increase test accuracy, as shown in Table [6
D.4 Weight-space style editing
For weight-space editing, we use the same INRs as the ones used for classification but we do not augment the dataset with additional INRs. Let U, be the INR weights for the 7th image and SIREN(z, y; U') be the output of the INR parameterized by U at coordinates (x, y). We edit the INR
23

---

weights U. = U =
NF,( ) and is a learned scalar initialized to 0.01. Letting fi y) be the pixel values of the ground truth edited image (obtained from image-space processing), the objective is to minimize mean squared error:
N
d
(NFN) N
SIREN (z, fi(x,
(21)
1=1
Note that since the SIREN itself is differentiable, the loss can be directly backpropagated through U to the parameters of the NFN.
The neural functionals contain equivariant NE-Layers with 128 channels, one invariant NF-Layer (mean pooling) followed by linear layers with 1,000 hidden neurons. Every layer uses ReLU activation. The training hyperparameters can be found in Table 15 and training takes hour on Titan RTX GPU.
Table 15: Hyperparameters for weight-space style editing using neural functionals.
Name
Values
Optimizer
Adam
Learning rate
1 x 10-3
Batch size
32 5 x 10
Training steps
E
Additional experiments and analysis
E.1
Analyzing train vs test accuracies in INR classification
Table shows the training and test accuracies of our NFNs and other weight space methods across the benchmarks we consider in this paper. An interesting observation that addition to better test performance, NFN typically achieve much higher training accuracies compared to the MLPs. ML.PAug- the MI.P variant trained with permutation augmentations, especially struggles to achieve low training error. This indicates that training architectures to respect neuron permutation symmetries extremely challenging, perhaps because these symmetry groups are extremely large []] Therefore, encoding equivariance invariance via architecture design, as we do with NFNs, is much more effective.
E.2 Non- woi-heeght--pace baselines for image classification
For the image and 3D shape classification tasks discussed in Sec. 3.2 we compared the neural functionals to non- weight-space baselines, namely residual networks (26] for MNIST, FashionMNIST. and CIFAR-10, and PointNet [52] for ShapeNet-10 10 and ScanNet-10. The results, shown in Table [6] indicate that applying state of the art architectures pixels point clouds is more effective than any weight space method. This may be because these architectures can leverage geometric properties of image and shape classification tasks. While neural functionals handle weight space symmetries, it remains an open problem to incorporate the geometric symmetries of common machine learning problems into weight space methods.
E.3 Ablating positional encoding
We ablate the effect of IO-encoding on the NFNNP architectures. in the 2D-INR classification and style editing tasks. The results, reported in Tables [7 and [8] show that 10-encoding adds a very small (sometimes negligible) boost NEN performance, though never hurts. Since NFN often performs as well as or better than NEN this indicates that even the base NP variant can solve many weight-space tasks without needing to break that symmetry.
E.4 Interpreting learned lottery ticket masks
We further analyze the behavior of NENPT on lottery ticket mask prediction by plotting the mask score predicted for given initialization value at each layer. To make the visualization clear we train

---

Table 16: Classification train and test accuracies (%) on datasets of 2D INRS (MNIST, Fashion- MNIST, and CIFAR-10) and 3D INRs (ShapeNet-10 and ScanNet-10), Our equivariant NFNs outperform the MLP baselines, even when the MLPs have permutation augmentations to encourage invariance, and even when the MLPs are significantly larger to match NFN's s parameter counts (MLP-4000/8000) NFNs also outperform recent non-equivariant methods developed for 3D INRs such as inr2vec [12]. ResNet-18 18 and PointNet operate on gridded or pointcloud signals, instead of weight space representations, and their performance is given for reference. PointNet results were obtained directly from Qi et al. [52) runs.
Where given, uncertainties indicate standard error over three
MNIST
FashionMNIST
CIFAR-10
ShapeNet-10 10
ScanNet- 10
Train
94.9 + 0.579
82.3 2.78
75.5 + 0.810
100 0.0
100.0 0.0
NFNHNP
Test
92.5 - 0.071
72.7 + 1.53
44.1 + 0.471
86.9 0.860
64.1 +: 0.572
Train
95.0 0.115
81.8 + 0.868
66.0 0.694
100.0 0.0
100.0 + 0.0
NENSP
Test
92.9 + 0.218
75.6 1.07
46.6 0.072
88.7 + 0.461
65.9 + 1.10
Train
42.4 2.44
41.5 + 2.17
23.7 + 2.39
100.0 + 0.0
100.0 + 0.0
MLP
Test
14.5 + 0.035
12.5 + 0.111
16.9 + 0.250
25.4 0.121
32.9 + 01351
Train
20.5 0.401
14.9 1.45
19.1 1.75
34.0
42.7 22.7+00001
MLPAug Ang
Test
21.0 + 0.172
15.9 0.181
18.9 + 0.432
33.8 0.126
45.5 + 0.126
Train
72.6 + 1.39
30.4 + 0.521
MLP-4000
Test
15.5 + 0.090
17.1 + 0.120
Train
19.4 1.39
20.5 + 0.333
MLPP44000
Test
21.1 + 0.010
19.3 + 0.325
Train
77.8 = L.74
35.9 + 0.868
MLP-8000
Test
15.8 0.014
17.3 + 0..880
Train
20.3 + 3.30
18.1 + U.347
MI MI.P-8000
Test
21.3 + 0.075
19.6 + 0.060
Train
+ 0.0
93.8 0.090
inr2vec 12
Test
39.1 0.385
38.2 0.409
ResNet-18
Test
99.21 +: 0.28
95.48 + 0.13
93.95 + 0.08
PointNet
Test
94.3
72.7
Table 17: Classification test accuracies (%) when ablating the positional encoding for implicit neural representations (INRs) of MNIST, FashionMNIST, and CIFAR-IO. Uncertainties indicate standard еттог over three runs.
NFN
NFN (no positional encoding)
MNIST
92.9 + 0.218
92.9 + 0.077
CIFAR-10
46.6 + 0.072
46.5 + 0.160
FashionMNIST 75.6 + 1.07
73.4 +: 0.701
Table 18: Test mean squared error when ablating the positional encoding on the Dilate (MNIST) and Contrast (CIFAR-10) implicit neural representation (INR) editing task. Uncertainties indicate standard error over three runs.
NFNNP
NFNxp (no positional encoding)
Contrast (CIFAR-10)
0.020 + 0.000
0.020 + 0.000
Dilate (MNIST)
0.068 + 0.000
0.070 + 0.001
25

---

Layer Weights
Layer Blases
22.5
-28
-42
Layer Weights
Layer Blases
eane 1000
Predicted Mask Score
-15
39
-20
-40
-25
11-
-30
-0.075
-6.050
-0.025
3,025
-0.025
0.025
0.850
0.075
Layer 2 Weights
Layer 2 Biases
11
10
-3-
9-
-10
-0.075
050
-0.00s
0.000
3.025
8050
0.05
-0.075
-c.0s0
-0.025
0:000
n.025
0.75
Input Weight Value
Figure 5: Mask scores vs weight magnitude for modified NFNpT-
NFNPT on MLP mask prediction without layer norm, which can be viewed as a scalar function the initialization f() R R for each layer i. Figure for fixed latent value, the predicted mask score function of the initialization value (low mask scores are pruned, while high mask scores are not). These plots suggest that, in the MIP setting, neural functionals are learning something similar to magnitude pruning of the initialization. In our setting, this turns out to be strong baseline for lottery ticket mask prediction: the test accuracy of models pruned with the modified network 95.0%.
26

---

Table 19: Kendall's coefficient and R2 between predicted and actual accuracies of three- and five-layer MLPs trained on MNIST. Our equivariant neural functionals outperform the baseline from 644 which predicts generalization using only simple weight statistics features. Uncertainties indicate standard eтror over five runs.
NFN
NFN
STATNN
3-Layer
0.876 + 0.003
0.859 + 0.002
0.854 + 0.002
5-Layer
0.871 0.001
0.855 0.001
0.860 + 0.001
3-Layer
0.957 + 0.003
0.9424 + 0.003
0.937 + 0.002
R²
5-Layer
0.956 0.002
0.947 = 0.001
0.950 + 0.001
E.5 Predicting MLP generalization from weights
In addition to predicting generalization the Small CNN 7o0 benchmark (Section 3.1 also construct our own datasets to evaluate predicting generalization on MLPs. Specifically, we study three- and five-layer MLPs with 128 units in each hidden layer. For each of the two architectures, we train 2,000 MLPs on MNIST with varying optimization hyperparameters, and save 10 randomly-selectee checkpoints from each run construct dataset of 20,000 (weight, test accuracy) pairs. Runs are partitioned according to a 90% / 10% split for training and testing.
We evaluate NENNNN and NFN on this task and compare them to the STATNN haseline predicts test accuracy from hand-crafted features extracted from the weights. Table T9 shows that NFNNP and STATNN are broadly comparable, while NFN consistently outperform other methods across both datasets in measures of correlation: Kendall's and R² These results confirm that processing the raw weights with permutation equivariant neural functionals can lead to greater predictive power when assessing generalization from weights.
64 which
27


Reviews:
[{"review_id": "ZbTIcUHbf3", "reviewer": "Reviewer_VWST", "summary": "This paper considers the design of architectures whose inputs are the parameters of neural networks. They propose an equivariant weight-sharing scheme based on the permutational symmetries of neural networks: one can permute at least the internal neurons (the “HNP” case), and sometimes the input/output neurons as well (the “NP” case), of a network without changing the function represented by the end-to-end network. The parameter-sharing of the “NP” case is more advantageous but less often applicable, since the underlying problem must have additional symmetry structure to warrant input/output neuron permutational symmetries. To remedy this, the authors propose coupling the stronger “NP” case with positional encodings. They extend their framework to take as input convolutional layers, and test their proposed architecture on tasks including image classification from implicit neural representations, predicting sparsity masks, and weight-space editing. They compare to the use of MLPs with and without permutational augmentations, as well as problem-specific baselines, and obtain promising experimental results.", "questions": "1. Is there some advantage to classifying implicit neural representations, instead of the pixels directly? It would be helpful to motivate these problems more in the main body of the text.\n2. Can the authors comment on the expressivity of their proposed idea of using input encodings to adapt the NP setting to problems without input/output permutational symmetries? Is this universal, for example, or is there some loss of information from adding the positional encodings? (On a related note, how exactly are the positional encodings incorporated — are they literally added, or are they appended?) \n3. Could you be more specific in designing networks for a particular nonlinearity with different symmetries (e.g. scaling for ReLU)? See e.g. “On the Symmetries of Deep Learning Models and their Internal Representations” (Godfrey et al 2022), as well as other references in Navon et al 2023. \n4. How large are the models input to the neural functional network? Will all weight-space methods will fail for truly large weight space inputs, i.e. millions of parameters, even if you have very few learnable parameters thanks to symmetry?", "limitations": "There is not potential negative societal impact. The authors are upfront that the HNP case is less scalable than the NP case. One limitation not discussed is that the architecture described only takes into account permutational symmetries, and not other nonlinearity-dependent symmetries.", "rating": 6, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Originality: Outside of concurrent work, the idea of using permutational symmetries to train directly on weight and bias inputs is original and makes sense. This permutation subgroup is distinct from other groups that have been studied in the past (such as $S_n$ and $S_n \\times S_n$). Even beyond concurrent work, the idea of using positional encodings in conjunction with input/output permutational symmetries is a creative contribution. \n\nQuality: The experiments are varied and test against a reasonable set of baselines. \n\nClarity: The paper is generally quite clearly written and explained.\n\nSignificance: This paper enables the application of weight-space symmetry techniques to CNNs. The problem of learning with neural network weight inputs is topical, and such methods are likely to enjoy practical usage.\n", "weaknesses": "1. The primary weakness of this paper is its novelty in relation to other work by Navon et al (ICML 2023), which also articulates an architecture equivariant to the permutational symmetries of weight space and provides more thorough results on its universality. The authors state this work was concurrent (which I will take at face value). If its novelty is judged in relation to this work, then the main contributions are the NP setting with positional encodings, the extension to CNNs, and the distinct set of experiments, which are extensions of the key idea. This very relevant work is currently only mentioned near the end of the paper, but I would think it is important to have a more in-depth discussion of how the two fit together and the novelty of this work, and for this discussion to appear earlier in the paper as part of its framing.\n2. The paper only considers the permutational symmetries of neural networks, which are indeed perhaps the most general symmetries if one does not specify a certain nonlinearity. However, as noted in Godfrey et al 2022, a given nonlinearity may enjoy additional symmetries — for instance, ReLU enjoys a symmetry to positive scaling. This is not discussed in the paper, but would have been a more substantial contribution relative to Navon et al (2023). For instance, one straightforward way of incorporating this scaling symmetry could be to pick a positive scale based on the norm of the parameters of the first layer, and then use this scale to normalize the rest of the input weights and biases. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your insightful review, and for highlighting aspects of our contribution such as the NP setting and application to CNNs.\n\n> it is important to have a more in-depth discussion of how [DWS (Navon et al) and NFN] fit together and the novelty of this work, and for this discussion to appear earlier in the paper as part of its framing.\n\nWe will update the introduction to more clearly and prominently discuss DWS, and our contributions relative to that work–see the top-level reply for more details.\n\n> Is there some advantage to classifying implicit neural representations, instead of the pixels directly? It would be helpful to motivate these problems more in the main body of the text.\n\nThe results in Section 3.2 represent first steps towards better methods that operate directly on INRs, which we believe will eventually have multiple advantages over operating on discrete representations of data (pixels, point clouds, voxel grids, etc…):\n\n1. INRs are continuous and decouple the memory cost of the representation from the actual spatial resolution. This is important as we eventually move towards more complex and high resolution 3D signals such as entire 3D scenes [1], where for example working with gridded representations becomes less tractable. We believe the 3D shape SDF experiments are a promising first result for that direction (Table 4).\n1. Working with INRs directly opens the possibility for a single method that can elegantly classify different types of signals of different sizes and resolutions, or even different dimensions. Whereas, for example, changing the image resolution or size in pixel space can pose a problem for CNNs.\n\nWe will add this motivation to Section 3.2.\n\n> Can the authors comment on the expressivity of their proposed idea of using input encodings to adapt the NP setting to problems without input/output permutational symmetries? Is this universal, for example, or is there some loss of information from adding the positional encodings? \n\nThis is an interesting question–we don’t develop any such theoretical results for the NP architectures with positional encodings in this work. One could likely develop a universality result through a process analogous to Thm 3 in [2], which shows that positional encodings remove the permutation equivariance constraint and allow Transformers to approximate any function under certain conditions.\n\n> On a related note, how exactly are the positional encodings incorporated — are they literally added, or are they appended?)\n\nIn practice, we implement the positional encoding by concatenation (in the channel dimension), though adding the encoding should have a similar effect. For matrices other than the input/output, we simply append zeros to keep the channel dimension consistent.\n\n> The paper only considers the permutational symmetries of neural networks [...] One straightforward way of incorporating [ReLU] scaling symmetry could be to pick a positive scale based on the norm of the parameters of the first layer, and then use this scale to normalize the rest of the input weights and biases.\n\nIt is true that NFNs only consider permutation symmetries, while (depending on the activation) scaling symmetries may also exist in the weight space. We will discuss this limitation in the paper.\n\nSince NFNs are typically dealing with weights produced by an optimization process like SGD, some existing literature suggest that accounting for scale symmetry may be unnecessary in practice; to quote [3]\n\n> SGD’s implicit regularization balances weight norms and, therefore, scale invariance does not seem to play an important role in understanding symmetries of solutions found by SGD.\n\nThis could explain why we see decent performance on some tasks already. The suggested idea seems very interesting, though we likely do not have time to try it out in the limited discussion timeframe.\n\n> How large are the models input to the neural functional network? Will all weight-space methods fail for truly large weight space inputs, i.e. millions of parameters, even if you have very few learnable parameters thanks to symmetry?\n\nThe size of the inputs to the NFN depends on the task: for example in predicting winning tickets we have 3 layers and 128 neurons, while INR classification deals with 3-layer INRs having 32 hidden neurons each (see Appendix D for full details). It is possible that much larger weight spaces will be challenging to learn in for any NF-type architecture, though such weight spaces also pose practical problems due to increased memory and compute usage.\n\n## References\n[1] Mildenhall et al. Representing Scenes as Neural Radiance Fields for View Synthesis.\n\n[2] Yun et al. Are Transformers universal approximators of sequence-to-sequence functions?\n\n[3] Entezari et al. The Role of Permutation Invariance in Linear Mode Connectivity of Neural Networks."}, {"review_id": "DaPDc89h3R", "reviewer": "Reviewer_CEwM", "summary": "The authors proposed an architecture for processing other networks’ weights and implicit neural representations (INRs). The generalization abilities of this architecture are enhanced and the number of parameters is reduced by leveraging the symmetries of deep networks.\n\nThe authors claim the following contributions:\n- Proposing a new architecture named Neural Functional Networks (NFN) for processing weights of other networks including INRs. A key building block in NFNs are NFLayers which the authors constrained to be permutation equivariant.\n- Extensive experimental part that demonstrates the superiority of NFNs over various baselines.\n", "questions": "- Can NFNs handle heterogeneous networks in a single dataset? For example networks with varying input dimensions and hidden features.\n- What is the computational complexity of NFN / NFLayers?\n", "limitations": "A limitation section is included.", "rating": 3, "confidence": 5, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "- The paper is dealing with novel and interesting problems of processing networks’ weights and INRs, I believe that this intriguing domain can open the path to more interesting and valuable works for our community. \n- The paper is well-written and easy to follow. Specifically, the visualizations make the reading more accessible and straightforward.\n- The experimental part is extensive and includes multiple learning setups and datasets.\n", "weaknesses": "- My main concern is about the novelty of this work w.r.t [1]. In [1] the authors also proposed an architecture, coined DWSNets,  for processing weights and INRs. They also provide some theoretical guarantees about the expressive power of their method. It is worth mentioning that [1] is **not a concurrent work** according to the NeurIPS guidelines. Therefore, the authors should explain what is the novelty of the current work over [1].\n- Given the previous point, the authors should also include [1] as a baseline in the experimental part and cite it more explicitly in the paper.\n- There is a lack of available technical information regarding the process of generating INRs. Specifically, in sections 3.2, 3.3, and 3.4 the authors did not mention the number of samples used for train/val/test. I also wonder how long it took to generate the INRs and how many resources were used. The authors should also include information about the training paradigm of INRs (for example the number of optimization steps) to enable easier reproducibility. \n- Although performed in the weight space it will be interesting to see how well NFNs deal with more challenging style editing tasks like inpainting, deblurring, etc. Dilation and contrast editing are straightforward tasks in CV.\n- Given that it was the pioneering study demonstrating the permutation equivariance of feed-forward networks, it is advisable for the authors to include a citation to [2] (especially in lines 33-36). \n- In lines 155-158 the authors stated they have used positional encoding (PE) to boost NFNs performance (although it breaks the symmetry), what is the performance gain presented by using PE? I did not see such an ablation (although I might have missed it).\n----------\nCitations:\n\n[1] Equivariant Architectures for Learning in Deep Weight Spaces, Navon et al.\n\n[2] On the algebraic structure of feedforward network weight spaces, JHecht-Nielsen et al.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your suggestions and questions--we aim to clarify our contributions and strengthen the experiments with the suggested baselines.\n\n> explain what is the novelty of the current work over [DWS].\n\nWe agree that DWSNets are a very relevant recent work with significant overlaps and notable differences. We will update the introduction to more clearly and prominently discuss DWSNets, and our contributions relative to that work--see the top-level reply for details.\n\nWe would like to emphasize some contributions of  our work relative relative to DWS:\n\n1. We additionally focus on the NP setting, which leads to layers that are more scalable and much more parameter-efficient (see Table 1) than HNP/DWS layers while maintaining good performance in most of our experiments. NP layers are also easier to visualize, understand, and implement, as evidenced by Figure 2 and the pseudocode in Appendix A.\n1. We extend both HNP/DWS and NP layers to handle CNN weights, in addition to MLPs. This extension is not trivial as CNN filters have additional dimensions that are not permutable, so the symmetries of CNN weights and MLP weights are not the same.\n\nWith regards to technical foundations, as mentioned by Reviewer jZvx, independently developed approaches for deriving equivariant weight-processing layers can be a useful contribution to the community. DWSNets and our paper also contain many complementary experiments and benchmarks to demonstrate the applicability of these architectures.\n\n> Given the previous point, the authors should also include [DWS] as a baseline in the experimental part\n\nWe have run DWSNets on our 2D-INR classification datasets and present the results in our top-level reply. DWSNets perform somewhat worse than both NFN variants across the board, even when controlling the number of parameters and after sweeping learning rates. We will also include results for this baseline on the other applicable benchmarks in our revised paper.\n\n> The authors should also include information about the training paradigm of INRs (for example the number of optimization steps) to enable easier reproducibility.\n\nAppendix D.3 contains additional information about the INR datasets, including the train/val/test sizes and the optimization process for generating each INR. We will add more explicit references to this information in Sections 3.2-3.4 of the main paper. Producing each INR takes ~168 seconds with 2 CPU cores (no GPU required), and can be done in parallel over multiple machines/cores. For the easiest reproducibility, we will also publicly release the INR datasets after the anonymous period.\n\n> Although performed in the weight space it will be interesting to see how well NFNs deal with more challenging style editing tasks like inpainting, deblurring, etc.\n\nAs you note, although the style editing tasks are easy in image space the primary challenge is to be able to accomplish them in weight space, which is much more difficult. We agree that more challenging image editing tasks like inpainting and deblurring are a worthwhile goal, but ultimately become more of a test of geometric inductive biases (which convolutional networks have) rather than inductive biases related to weight space symmetries, which is the focus of this work.\n\n> It is advisable for the authors to include a citation to [2] (especially in lines 33-36).\n\nOur paper already cites [2] on Line 31, but we will also add a citation in lines 33-36.\n\n> used positional encoding (PE) to boost NFNs performance (although it breaks the symmetry)\n\nThe input/output positional encoding (PE) for the NP variant of NFNs _only breaks symmetry to permutations of the input and output neurons of the weight space_, but preserves symmetry to hidden layer permutations. We apologize for the confusing wording and will update the writing.\n\n> What is the performance gain presented by using PE?\n\nWe performed an ablation to answer this question--see the top-level reply for details. Results show that PE actually adds a very small (sometimes negligible) boost to NFN_NP performance. This suggests that the NFN_NP architecture by itself can solve many weight space tasks without needing to break permutation symmetry of the input and output neurons.\n\n> Can NFNs handle heterogeneous networks in a single dataset?\n\nIn principle, parameter sharing makes our NP-equivariant NF-layer agnostic to the widths of the input weights, i.e. the number of neurons at each layer can vary. Similarly, the HNP-equivariant NF-layer is agnostic to the number of neurons at the hidden layers, but the input and output dimensions must stay fixed. The depth (number of layers) of the input weights must be fixed in both cases.\n\nIn practice, handling heterogenous networks is challenging in modern ML frameworks since varying input sizes makes batch computation difficult on GPUs, so in our experiments the input networks have a fixed size. In principle, padding could mitigate this problem, though it may use GPU memory inefficiently.\n\n> What is the computational complexity of NFN / NFLayers?\n\nSuppose the input weights have constant width $n=n_0=\\cdots=n_L$, so that there are $Ln^2$ input weights in total. A naive linear layer operating on these weights would require $L^2n^4$ operations.\n\nFor the NP case, consider implementing Eq 3 without parallelization, by:\n1. First calculate all quantities of the form $W\\_{\\star,\\star}^{(i)}$, $W\\_{j,\\star}^{(i)}$, $W\\_{\\star,k}^{(i)}$. These require $O(Ln^2)$ operations.\n1. Calculate each term in Eq (3) separately, for each output index $(i,j,k)$. The first term requires $L^2$ operations, terms 2-5 require $Ln$ operations, and the final term requires $Ln^2$ operations.\n1. Add the terms to produce the final result for each $(i,j,k)$. This requires $O(Ln^2)$ operations.\n\nSo the layer can be implemented on $O(L^2 + Ln^2)$ operations. We will include this and the HNP result in the revised paper."}, {"review_id": "apRwskpRmW", "reviewer": "Reviewer_jZvx", "summary": "This paper studies the problem of defining linear layers (and by extension, neural networks) that operate on neural network weight spaces. The core idea of this work is to take into account weight permutation symmetries, similar to Navon et al., ICML’23. In particular, the weights of certain feedforward architectures, such as MLPs and CNNs, can be permuted in several ways, without altering the function that the neural net represents. Therefore, the authors seek to define functions on the weights, that will be equivariant (or invariant) to these very symmetries. \n\nTo achieve this, they characterise the space of linear equivariant/invariant functions, following the framework of Ravanbakhsh et al., ICML’17, i.e. by identifying the parameters of these functions that should be shared. The authors propose two variants, one that is derived from the detailed characterisation of the weight permutation symmetries that arise from hidden neurons ($\\text{NFN}_{\\text{HNP}}$), and another one that assumes extra permutation symmetries in the input/output layer (that can be broken with positional encodings), and is, therefore, more parameter efficient.\nInterestingly, these results easily carry over from MLPs to CNNs. The proposed layers are experimentally tested in a battery of tasks, such as predicting NN generalisation or sparsity masks for NN pruning, showing considerable improvement against baselines that do not take into account permutation symmetries, in addition to a heavily reduced parameter count.\n", "questions": "**Experiments, baselines and implementation details**: Most of the following questions ask for clarifications and I do not consider them as weaknesses. However, I think a discussion is needed and it would be useful to include the answers to the below in the paper. In detail:\n\n- **NP case**.  What are the positional encodings that are used in the experimental section? If I understand correctly they are fixed (handcrafted). Have the authors tried learnable ones instead? What is the number of added parameters when using learnable positional encodings? How do the authors explain the fact that using positional encodings in NFN_NP performs better than NFN_HNP?\n- **NFN_PT case**. Since this method performs well in the pruning mask prediction experiment, why didn’t the authors also include its performance in the rest of the experiments (e.g. generalisation and INR classification)?\n- **Pruning**: Why did the authors use a generative model (CVAE) for the winning ticket experiment? Why not a deterministic predictor (a 0-1 classifier)? Have the authors tested this? How far are the generated sparsity masks from the ones that IMP yields? Why isn’t the performance of MLP and MLP_aug included in this table?\n- In section 3.2, would it be possible for completeness to add some baselines that classify images and 3D shapes, using different representations instead of INRs, e.g. state-of-the-art results using image-based (CNNs) or point cloud-based classifiers? Could the authors comment on why the performance is still far from the one that the community has obtained by applying these baselines? This can be clearly observed in CIFAR-10. In addition, apart from the inr2vec baseline, it might be helpful to add the baseline from Dupont et al., ICML’22 (the “functa” representation), where INRs are represented as vectors of latent modulations.\n- **NN editing**. In section 3.4., regarding the CIFAR dataset, the differences between the competing methods are not clear qualitatively.  Could the authors comment on this?\n- Is there any theoretically-backed advantage of using a combination of equivariant layers followed by an invariant one instead of using only invariant ones (when the ground-truth function is invariant)?\n", "limitations": "- The authors have dedicated a paragraph to the limitations of their method. One thing that might be missing is a clarification (if I understand correctly) that NF layers can only be applied to the weights of neural nets that share the same architecture (fixed number of layers, fixed kernel size for CNNs). Note that L189-190 (“…which contains thousands of CNN weights trained on several datasets with varied hyperparameters”) seems contradictory to this. I think the phrase \"varied hyperparameters\" should be clarified to avoid confusion.\n\n- Moreover, given that classifying INRs with NF-layers seems to still have room for improvement in order to reach the performance achieved by traditional neural nets that operate on raw representations, it might be useful to discuss the possible reasons for this more thoroughly (perhaps stronger inductive biases, taking into account other properties of the geometry of weight spaces, might be missing?)\n \n- No foreseeable negative societal impact.", "rating": 7, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "**Significance**: The proposed method has wide applicability, as correctly pointed out by the authors in the first paragraph of the introduction, and can potentially have a substantial impact on multiple neural-network-related problems (meta-learning, neural network editing, INR processing etc.).\n\n**Presentation, clarity and reproducibility**: Although the concepts studied in this paper and the neural network architectures resulting from the study of the relevant symmetries, have a fair amount of complexity and may require a good understanding of the study of symmetries, the authors have made a very good effort to make their paper as accessible as possible. In particular, the illustrations, the nice summarisation of the core layer in Eq. (3) as well as the parameter counts in Table (1) and the accompanying explanations provide a good overview of the method that I believe will be appreciated by future readers. Moreover, the authors have formulated their layers in such a way that makes them easy to implement, which is also apparent from the pseudocode provided in the appendix.\n\n**Experimental evidence**. The authors have tested their layers against a diverse set of problems and the reported results clearly motivate the need for equivariant/invariant layers and support all the design choices made.\n\n**Extended scope**: A few months ago, another work by Navon et al., ICML’23 characterised the same family of equivariant/invariant layers as well (providing additional theoretical support). This is obviously not a problem, since the two works were developed in parallel. It is therefore nice to see that the present paper has a bonus, i.e. an extended scope compared to the aforementioned one since (1) the formulation seamlessly allows to extend the method to CNNs, and (2) the NP formulation allows for more practical implementation, a significant reduction in the number of parameters, and empirically strong performance. (3) This is probably subjective and a matter of taste, but I appreciated the alternative way to derive the layers (see sec B.3. in the appendix), using the parameter-sharing framework of Ravanbakhsh et al., ICML’17. I found it easier to follow and slightly more intuitive, which might be of independent technical interest in the future. Therefore, although the layers have been rediscovered, I feel that the contributions are still valuable. ", "weaknesses": "**Clarity of the proofs (appendix)**: I found the explanations in section B.4. a bit hard to follow. Personally, I am familiar with the work of Navon et al., so I could easily understand the concepts, but I fear a reader not versed in the topic, might get lost in this part. Since it is crucial in order to gain a deep understanding of the paper, I would recommend that the authors try to simplify some concepts (e.g. by giving more details for a particular sub-case).\n\n**Related work**: I believe that the authors do not give enough credit to the work by Navon et al.. As I mentioned before, the two works are concurrent, and they apparently independently discovered the same core idea, but I think that the authors should be more upfront about this and mention this early on in the paper and cite Navon et al. more prominently. \n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the detailed and insightful review, and agree that indepedently developed frameworks can be useful contributions to the community.\n\n> I found the explanations in section B.4. a bit hard to follow.\n\nWe will aim to improve the exposition in Section B.4. We also welcome any feedback on particular aspects that are confusing or could be improved.\n\n> I think that the authors should be more upfront about [Navon et al] and mention this early on [...]\n\nWe will update the introduction to more clearly and prominently discuss DWSNets, and our contributions relative to that work. See our top level reply for more detail.\n\n> Have the authors tried learnable [positional encodings] instead? What is the number of added parameters when using learnable positional encodings?\n\nWe will update the appendix to be clearer about this: the positional encodings (PE) are sinusoidal--learnable ones are also possible, though in Transformers it was observed that learned vs sinusoidal PE had essentially no impact on performance [1, Table 3 Row E].\n\nFor learned PE you would need a separate vector of length $c$ (feature channels) This would add $c(n_0 + n_L)$ parameters.\n\n> How do the authors explain the fact that using positional encodings in NFN_NP performs better than NFN_HNP?\n\nNote the effect is not consistent--the HNP variant performs better in some tasks, such as predicting generalization (Table 2). NFN_NP may perform better on some settings simply because it is more parameter-efficient and easier to train, or other optimization related details. Also, the PE ablation in our top-level reply suggests that differences in performance are largely due to differences in the HNP and NP architectures themselves, rather than PE.\n\n> NFN_PT case. Since this method performs well in the pruning mask prediction experiment, why didn’t the authors also include its performance in the rest of the experiments\n\nThis is a good question. We generally didn’t include the pointwise (PT) ablation because, in our early development, PT architectures performed very poorly on all problems _except_ for pruning. Since PT is much more restricted than the full (H)NP layers, poor performance on most settings was expected and we focused our attention on ablations in the setting where it did have signs of life (pruning).\n\n> Pruning: Why did the authors use a generative model (CVAE) for the winning ticket experiment? Why not a deterministic predictor (a 0-1 classifier)? Have the authors tested this? How far are the generated sparsity masks from the ones that IMP yields? Why isn’t the performance of MLP and MLP_aug included in this table?\n\nEven conditioned on a fixed initialization, winning tickets are _not_ deterministic due to the noise of SGD–one can quickly verify this experimentally. Therefore, it makes sense to learn a probabilistic model p(mask | init) rather than a deterministic mapping.\n\n> In section 3.2, would it be possible for completeness to add some baselines that classify images and 3D shapes, using different representations instead of INRs, e.g. state-of-the-art results using image-based (CNNs) or point cloud-based classifiers?\n\nWe will add these (non weight-space) baselines in Section 3.2. \n\n> Could the authors comment on why the performance is still far from the one that the community has obtained by applying these baselines? [...] it might be useful to discuss the possible reasons for this more thoroughly (perhaps stronger inductive biases, taking into account other properties of the geometry of weight spaces, might be missing?)\n\nAs suggested, we believe that this is primarily due to a difference in inductive biases. For example, it is common to use convolutional networks because of translation symmetry when solving CIFAR in image space. While NFNs are able to leverage weight space symmetries, it does _not_ have symmetry to translations of the CIFAR INRs. How to encode inductive biases related to the underlying natural signal, not just the weights, remains an interesting and challenging direction for future work. We will add this commentary to the paper.\n\n> In addition, apart from the inr2vec baseline, it might be helpful to add the baseline from Dupont et al., ICML’22 (the “functa” representation), where INRs are represented as vectors of latent modulations.\n\nFuncta cannot be run on our INR datasets since it requires a special (meta-learning) training process to produce the INRs, while the INRs in our datasets are produced by vanilla training methods. We would also interpret functa as operating in a different setting depending on how much control one has over the process that produces the input weights.\n\n> NN editing. In section 3.4., regarding the CIFAR dataset, the differences between the competing methods are not clear qualitatively. Could the authors comment on this?\n\nFor the CIFAR contrast task, it turns out that even the ground truth brightening operation can be relatively subtle to notice visually. Coupled with the fact that none of the methods are perfect at achieving the ground truth operation, this makes the qualitative differences difficult to observe (though some samples are more obvious than others, see the right column of 4). We will investigate to see if there is a better way to display the visual changes.\n \n> NF layers can only be applied to the weights of neural nets that share the same architecture [...] L189-190 (“…which contains thousands of CNN weights trained on several datasets with varied hyperparameters”) seems contradictory to this.\n\nThe CNNs in that CNN Zoo share the same architecture but vary in learning rate and other optimization hyperparameters. We will clarify this wording to avoid confusion.\n\n## References\n[1] Vaswani et al. Attention is all you need."}, {"review_id": "D1enReEp4G", "reviewer": "Reviewer_9eY1", "summary": "This paper proposes the NF-Layer that maps the weight space of a deep neural network (DNN), including MLP and CNN, to another weight space, possibly with a different number of channels. Neural Functional Networks (NFNs) are then constructed using the NL-Layers to process the weight space of a DNN. NF-Layer is designed to be S-equivariant, that is, if the input weights of a DNN are permutated so as not to change the output of the DNN, the output of the NF-Layer, which is also in another weight space, is also permutated in the same order. Such a property enjoys the equivariant characteristics inherent in the DNN weights and enhances the efficiency in modeling NFNs, like the convolution structure enhances the efficiency of modeling image inputs. Experimental results are presented to demonstrate the usefulness of NF-Layers in applications of NFNs, including accuracy prediction, classification of implicit neural representations (INRs), winning ticket prediction, and style-editing through INRs.\n", "questions": "I can understand the permutation equivariant property in DNN weights, but I cannot intuitively understand why the prediction accuracy of NFNs improves with NF-Layers with S-equivariance. Is it because the parameter to be trained is reduced by the S-equivariant structure, while retaining its expressiveness?\n\nIt would be better to explain the multi-channel extension of permutations more carefully in ll.86-90. When I first read it, I mistakenly thought it was referring to the convolution channel. Also, it is better to describe that the proposed NF-Layer may change the channel size in this sense in the explanation of Figure 1.\n\nIn Figure 2, $W_{\\ast,j}^{(i+1)}$ (with yellow c) -> $W_{j,\\ast}^{(i)}$?\n\nEquation in Figure 2 should be matched with Eq.(3).\n\nIn ll.155-160, I could not understand the motivation to use position embedding with NFN_{NP}. In which experiments with NFN_{NP}, position embeddings are used?", "limitations": "As described by the authors, there is a concurrent work by Navon et al.[45] that has a very similar motivation and methodology. But I think this paper has its original contributions that the authors mention and deserves to be published. \n", "rating": 7, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* An interesting application of the equivariant architecture based on parameter-sharing.\n* Introduction of the NP setting in the permutation definition. Although it requires a stronger assumption than the HNP setting, the resulting NF-Layer can be modeled more efficiently under the NP assumption.\n* Experiments suggesting the advantage of the proposed NF-Layers in interesting applications of NFNs.\n", "weaknesses": "* A network architecture with equivariance based on parameter-sharing is not original and has been proposed in [51], as the authors suggest.\n* NF-Layers cannot directly handle some DNN models, including ResNet and Transformer.\n\n(Minor)\n* Figure 1 is difficult to understand at first glance.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and interesting question.\n\n> A network architecture with equivariance based on parameter-sharing is not original and has been proposed in [51]\n\nTo clarify, [51] (Equivariance through parameter sharing) provides general strategies for developing layers equivariant to a given choice of group action. [51] does not study the permutation symmetries of neural networks' weights and does not develop ways of processing weights. Rather, we utilize strategies from [51] to develop our NF-layers.\n\n> NF-Layers cannot directly handle some DNN models, including ResNet and Transformer.\n\nThis is a good observation, architectures like ResNet have more complicated weight space permutation symmetries compared to feedforward networks so our definitions of the (hidden) neuron permutation groups do not directly apply. We believe that extending these layers to handle more general network topologies is an important (and nontrivial) direction of future work.\n\n> cannot intuitively understand why the prediction accuracy of NFNs improves with NF-Layers with S-equivariance. Is it because the parameter to be trained is reduced by the S-equivariant structure, while retaining its expressiveness?\n\nYou are correct that enforcing S-equivariance reduces the parameter space, and we know that the target function must be S-equivariant (or invariant). So ideally enforcing equivariance does not sacrifice our ability to express the target function, while reducing the space of parameters to search over, leading to better generalization.\n\n> It would be better to explain the multi-channel extension of permutations more carefully in ll.86-90. When I first read it, I mistakenly thought it was referring to the convolution channel.\n\nThanks for your suggestions–we will update our presentation of the NF-layers to more clearly explain that they can have arbitrary input and output channels, and avoid any confusion with channels in, e.g., the convolutional weight space being processed.\n\n> Possible error in figure 2?\n\nThanks for pointing out this error, we will update the figure to fix this. \n\n> In ll.155-160, I could not understand the motivation to use position embedding with NFN_{NP}. In which experiments with NFN_{NP}, position embeddings are used?\n\nWe use positional embeddings (PE) in all of our experiments. The reason is that our tasks (and most real world tasks involving weights) do not allow arbitrary permutations of the input and output layers of the networks being processed--only the hidden layers can be permuted. So the NP assumptions give an incorrect symmetry, which can be broken by PE.\n\nBy analogy, Transformers are equivariant to token permutations, but most sequence-to-sequence tasks are not actually permutation symmetric (order matters). So Transformers often use PE to break that symmetry.\n\nWe will update L155-160 to clarify this point.\n\n## References\n[1] Cohen and Welling. Group Equivariant Convolutional Networks."}, {"review_id": "MHlhZN0iuK", "reviewer": "Reviewer_7rfP", "summary": "The paper introduces and evaluates permutation equivariant neural functional networks (NFNs).\nNeural functionals are models which take weights of other neural networks, or more general weight-space features, like gradients or sparsity masks, as inputs.\nTheir permutation equivariance addresses the issue that any particular vector of weight-space features is just an arbitrary representative of an equivalence class of weight-space features that correspond to exactly the same network.\nSpecifically, the weights and biases in all hidden layers may by permuted arbitrarily, which corresponds to simultaneous permutations of rows and columns of the previous and following weight matrix, respectively.\nPermutation equivariant NFNs ensure predictions which are independent from the particular permutation.\nAs usual for equivariant networks, this reduces the number of model parameters (of the NFN) significantly, and makes learning feasible in the first place.\n\nThe authors consider two specific group actions, which are the first order permutation actions and trivial actions on invariant scalars.\nEquivariance refers throughout the paper to mappings that commute with the first order action in the input and output, while invariance refers to mappings from first order actions to scalars.\nEquation 3 presents a linear NFN layer, which computes features as usual by taking learned linear combinations of weight-space features that are summed over different combinations of rows and columns.\nProposition 1 claims that this layer is 1) sufficient and 2) necessary for equivariance, i.e. spans the space of the most general linear equivariant maps (I have some concerns here, discussed in the weakness section).\nAs nonlinearities, NFNs use pointwise nonlinearities, which commute with permutation actions.\nA straightforward extension to permutations of CNN channels is discussed towards the end of section 2.2.\n\nBesides the HNP model, which considers the \"correct\" hidden neuron permutation symmetry,\nan NP model, which additionally permutes input and output neurons and has a further reduced number of parameters, is introduced.\nIn order for the network to be able to break this excess equivariance, positional encodings are added.\n\nExtensive experiments evaluate HPN-NFNs and NP-NFNs and compare them to baseline models.\nThe first two applications are permutation invariant predictions of CNN test accuracies from their weights and invariant classifications of implicit neural representations.\nIn both cases the NFNs perform significantly better than baselines, and the HNP variant wins over the NP task in the first but loses in the second one (they are still quite close).\nThe other two applications are permutation equivariant predictions of a lottery ticket hypothesis winning ticket sparsity mask and the editing of implicit neural representations.\nIn the former, NP-NFN performs close to a task-specific baseline, while HNP-NFN is prohibitively expensive, thus supporting the need for the NP-NFN model.\nIn the latter, both NP-NFN and HNP-NFN outperform their baselines.", "questions": "The paper is very clearly written and I don't have any remaining questions.\n", "limitations": "Some limitations and potential future work are discussed in the conclusion.\nIt should be added that one could consider more general permutation group actions.", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper is well motivated and very clearly written - it was really a joy to read!\nAfter introducing the mathematical framework, the authors derive the most general linear equivariant mappings for the considered group actions (I have some concerns here, see the weakness section).\nExplicit equations for the required number of parameters of the different models are given.\nThe experiments are very extensive and show the utility and superiority of NFNs in various applications.", "weaknesses": "I have two technical concerns regarding the claims about the generality of the proposed linear equivariant map in proposition 1.\nFirstly, it should be mentioned that the proof assumes a specific choice of group action w.r.t. which the layer is equivariant, namely first order permutation actions.\nPermutation equivariant networks have also been built for other actions, like irreducible representations or higher order tensor product actions - these would lead to other layers.\nThis point needs to be discussed and the statement of proposition 1 needs to be made precise, e.g. by saying that the map $T$ is equivariant w.r.t. the specific first order action $\\sigma$ in the input and output.\n\n[edit: this claim was wrong]\n~Secondly, I believe that the mapping in equation 3 is only then the most general $S_{n_1}\\times\\dots\\times S_{n_{L-1}}$ equivariant intertwiner when $n_1\\neq\\dots\\neq n_{L-1}$ is assumed.\nIf, however, $n_i=n_j$ for some layers $i\\neq j$, the same permutation group $S_{n_i} = S_{n_j}$ acts on these neurons, and it should be possible to linearly accumulate the corresponding weight-space features after $\\star$-summing over the other indices.\nA similar summing over features is already present in equation 3, namely the first term, which sums invariant scalars that are accumulated from all layers.\nNote that this contribution can in practice become very large, since networks are in practice often constructed such that they have the same number of hidden neurons or channels throughout many different layers.\nThe more general equivariant mappings should be benchmarked against the current NFNs in the experimental section.~\n\nThe sufficiency and necessity of the invariant NFN layers in section 2.3. is not supported by a similar theorem as proposition 1.\n\nDue to these weaknesses I decided to only give a weak accept to the current state.\nHowever, the paper has a huge potential and I would be happy to switch to a strong accept when the issues are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review and detailed analysis of the technical aspects of our work.\n\n> It should be mentioned that the proof assumes a specific choice of group action w.r.t. which the layer is equivariant, namely first order permutation actions.\n\nWe will make it clear at the beginning of the proof that we are only concerned with equivariance to a particular choice of action (as defined in Equation 1).\n\n> Secondly, I believe that the mapping in equation 3 is only then the most general $S\\_{n\\_1} \\times \\cdots \\times S\\_{n\\_{L-1}}$ equivariant intertwiner when $n\\_1 \\neq \\cdots \\neq n\\_{L-1}$ is assumed.\n\nWe don't follow how two layers having the same number of neurons creates a special case. To aid us in understanding and responding to this point better, could you write out an example of any additional or missing terms explicitly, which would show that Eq 3 is not fully general?\n\nAdditionally, would this claim not contradict Proposition 1? Though it is possible that we made a mistake and dropped terms somewhere in the derivation of the layer.\n\n> The sufficiency and necessity of the invariant NFN layers in section 2.3. is not supported by a similar theorem as proposition 1.\n\nWe will add an analogous proposition as Prop 1, but for the sufficiency and necessity of the invariant layers in Section 2.3, using a similar strategy as in Appendix B.2."}]

Author rebuttal:
We thank the reviewers for their detailed and thoughtful comments and questions. Reviewer suggestions have helped us improve the writing and pointed us towards additional experiments that significantly strengthen the paper. A brief summary of changes and new experiments:

* We will update the introduction to more clearly and prominently discuss DWSNets, and our contributions relative to that work.
* We have run DWSNets on our benchmarks in order to provide a direct comparison.
* We performed ablation experiments on the effect of positional encoding (PE) on the NFN_NP architecture variant.

Additional changes and information are provided in the reviewer-specific replies.

## DWSNet discussion

In order to more clearly discuss DWSNets as a relevant work with significant overlap and notable differences, we will put the following discussion to the introduction:

> The recent work of Navon et al. [45] recognized the potential of leveraging weight space symmetries to build equivariant architectures on deep weight spaces; they characterize a weight-space layer which is mathematically equivalent to the equivariant NF-Layer we develop for the HNP setting. Their work additionally studies interesting universality properties of the resulting equivariant architectures, and demonstrates strong empirical results for a suite of tasks that require processing the weights of MLPs.


> Our independently developed framework additionally focuses on the NP setting, where making stronger symmetry assumptions enables us to develop equivariant layers with improved parameter efficiency and practical scalability. This work also extends both NFN variants to process convolutional neural networks (CNNs) as input, leading to applications such as predicting the generalization of CNN classifiers (Section 3.1).


## DWSNet Comparison

Here is an initial empirical evaluation of DWS on our INR classification benchmarks. We trained DWSNets both at width 32 (as in [1]) and width 512 (as in ours), with our data augmentation of creating multiple INR copies per image. We also tried [1]'s data augmentation scheme, but found that it tended to hurt performance. For each dataset and channel size we swept learning rates in $[1e-3, 5e-3, 1e-4, 5e-4]$. Initial results show somewhat lower test accuracies than NFNs, with comparable performance on CIFAR10. Although the DWS layers match our HNP variant in theory, real world performance depends heavily on many architectural hyperparameters, just as different CNN architectures can achieve different performance.

| Test accuracy  | MNIST | FashionMNIST | CIFAR |
|----------------|-------|--------------|-------|
| DWS (32 channel)| 74.7 | 67.5         | 42.3  |
| DWS (512 channel)| 61.6 | 62.0        | 42.9  |
| NFN_NP (512 channel)| 92.9 | 75.6      | 46.6  |
| NFN_HNP (512 channel)| 92.5 | 72.7     | 44.1  |

And here is the number of parameters for each method, for reference:
| No. params     | MNIST | FashionMNIST | CIFAR |
|----------------|-------|--------------|-------|
| DWS (32 channel)| 0.6M | 0.6M         | 1M  |
| DWS (512 channel)| 71M | 71M         | 134M  |
| NFN_NP (512 channel)| 45M | 45M       | 47M   |
| NFN_HNP (512 channel)| 69M | 69M      | 135M  |

## Positional Encoding (PE) ablation

We performed the PE ablation on both 2D INR classification and style editing tasks. The results show that PE actually adds a very small (sometimes negligible) boost to NFN_NP performance, though it never hurts. Since NFN_NP often performs as well as or better than NFN_HNP, this indicates that even the base NP variant can solve many weight space tasks without needing to break that symmetry.

**Table 3b**: Ablating the PE on the INR classification. Higher is better.
|             | NFN_NP         | NFN_NP (no PE) |
|-------------|----------------|----------------|
| CIFAR-10    | 46.6 ± 0.072   | 46.5 ± 0.160   |
| MNIST       | 92.9 ± 0.218   | 92.9 ± 0.077   |
| FashionMNIST| 75.6 ± 1.07    | 73.4 ± 0.701   |

**Table 6b**: Ablating the PE on style editing tasks. Lower is better.
|                | NFN_NP       | NFN_NP (no PE) |
|----------------|--------------|----------------|
| Contrast (CIFAR-10) | 0.020 ± 0.000 | 0.020 ± 0.000 |
| Dilate (MNIST) | 0.068 ± 0.000 | 0.070 ± 0.001 |

## References
1. Navon et al. Equivariant Architectures for Learning in Deep Weight Spaces.

Meta-review: The paper proposes an equivariant architecture for processing the weights of neural networks. Reviewers were generally positive with regard to the idea and the paper. A concern was raised regarding the limited contribution of this paper with respect to a recent ICML 2023 paper by Navon et al.
The authors submitted a rebuttal, acknowledged the similarity to this previous work, mentioned the differences, and promised to prominently discuss it in the revision.
In the discussion, most reviewers agreed that although these two papers have a lot in common, the current paper has some new contributions like the extension to CNN, the NP setup (permutation equivariance on input and output dimensions), and additional interesting experiments.
The AC has read the current paper and the paper by Navon et al. thoroughly. It is the AC's opinion that the extension to convents is a marginal contribution. The NP setting with the positional encoding is the more important contribution as it allows scaling the networks in an elegant way. In the discussion, the AC and reviewers agreed to accept the paper. The authors are asked to revise their paper in accordance with the comments above and their rebuttal.
Decision: Accept (poster)