Title: Permutation Equivariant Neural Functionals
Venue: NeurIPS 2023
Accepted: yes
Keywords: equivariance

permutation

implicit neural representation

generalization
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

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

1 Introduction

As deep neural networks have become increasingly prevalent across various domains, there has been a growing interest in techniques for processing their weights and gradients as data. Example applications include learnable optimizers for neural network training [3, 56, 2, 44], extracting information from implicit neural representations of data [61, 45, 58], corrective editing of network weights [57, 11, 46], policy evaluation [24], and Bayesian inference given networks as evidence [60]. We refer to functions of a neural network'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' [8, 54, 34, 5]. For neural functionals, the symmetries of interest are transformations of a network's weights that preserve the network's behavior. In particular, we focus on neuron permutation symmetries , which are those that arise from the fact that the neurons of 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 a feedforward network can be arbitrarily permuted without changing its behavior [27]. In MLPs, permuting the neurons in hidden layer i corresponds to

W

(1)

W

(2)

W

(3)

W

(4)

H

1

(

W

)

(1)

H

1

(

W

)

(2)

H

1

(

W

)

(3)

H

(

W

)

(4)

H

2

(

H

1

(

W

))

(1)

H

2

(

H

1

(

W

))

(2)

H

2

(

H

1

(

W

))

(3)

H

(

H

(

W

))

(4)

NF-Layer

NF-Layer

1

2

1

permuting the rows of the weight matrix W ( i ) , and the columns of the next weight matrix W ( i +1) as shown on the left-hand side of Figure 1. Note that the same permutation must be applied to the rows W ( i ) and columns of W ( i +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 to neuron permutation symmetries. Our framework extends a long line of work on permutation equivariant architectures [52, 68, 25, 63, 41] that design equivariant layers for a particular permutation symmetry of interest. Specifically, we introduce neural functional layers (NF-Layers) that operate on weight-space features (see Figure 1) while being equivariant to neuron permutation symmetries. Composing these NF-Layers with pointwise non-linearities produces equivariant neural functionals.

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 permutation, HNP), or 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 in some settings. In contrast, NF-Layers derived under NP assumptions often lead to much more efficient architectures, and, when combined with a positional encoding scheme we design, can even be effective on tasks that require breaking input and output symmetry. For situations where invariance is required, we 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 involving 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) a weight-space 'style-editing' task where the goal is to modify the content an INR encodes by directly editing its weights. In 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 to 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 a 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).

2 Equivariant neural functionals

We begin by setting up basic concepts related to (hidden) neuron permutation symmetries, before defining the equivariant NF-Layers in Sec. 2.2 and invariant NF-Layers in Sec. 2.3.

2.1 Preliminaries

Consider an L -layer feedforward network having n i neurons at layer i , with n 0 and n L being the input and output dimensions, respectively. The network is parameterized by weights W = { W ( i ) ∈ R n i × n i -1 | i ∈ J 1 ..L K } and biases v = { v ( i ) ∈ R n i | i ∈ J 1 ..L K } . We denote the combined collection U := ( W,v ) belonging to weight-space, U := W×V .

Since the neurons in a hidden layer i ∈ { 1 , · · · , L -1 } have no inherent ordering, the network is invariant to the symmetric group S n i of permutations of the neurons in layer i . This reasoning applies to every hidden layer, so the network is invariant to ˜ S := S n 1 ×··· × S n L -1 , which we refer to as the hidden neuron permutation (HNP) group. Under the stronger assumption that the input and output neurons are also unordered, the network is invariant to S := S 0 ×··· × S n L , which we 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 B. See Table 1 for a concise summary of the relevant notation for each symmetry group we consider.

Consider an MLP and a permutation σ = ( σ 0 , · · · , σ L ) ∈ S . The action of the neuron permutation group is to permute the rows of each weight matrix W ( i ) by σ i , and the columns by σ i -1 . Each bias vector v ( i ) is also permuted by σ i . So the action is σU := ( σW,σv ) , where:

Until now we have used U = ( W,v ) to denote actual weights and biases, but the inputs to a neural functional layer could be any weight-space feature such as a gradient, sparsity mask, or the output of a previous NF-Layer (Figure 1). Moreover, we may consider inputs with c ≥ 1 feature channels, belonging to U c = ⊕ c i =1 U , the direct sum of c copies of U . Concretely, each U ∈ U c consists of weights W = { W ( i ) ∈ R n i × n i -1 × c | i ∈ J 1 ..L K } and biases v = { v ( i ) ∈ R n i × c | i ∈ J 1 ..L K } , with the channels in the final dimension. The action defined in Eq. 1 extends to the multiple channel case if we define W ( i ) jk := W ( i ) j,k, : ∈ R c and v ( i ) j := v ( i ) j, : ∈ R c .

The focus of this work is on making neural functionals that are equivariant (or invariant) to neuron permutation symmetries. Letting c i and c o be the number of input and output channels, we refer to a function f : U c i →U c o as S -equivariant if σf ( U ) = f ( σU ) for all σ ∈ S and U ∈ U c i , where the action of S on the input and output spaces is defined by Eq. 1. Similarly, a function f : U c → R is S -invariant if f ( σU ) = f ( U ) for all σ and U .

If f, g are equivariant, then their composition f ◦ g is also equivariant; if g is equivariant and f is invariant, then f ◦ g 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 nonlinearities.

(

H

(

W

)

(

i

)

jk

b i

,

i W

(

i

)

*,

k

+

b i

,

i

-1

W

(

i

-1)

k

,*

+

d i W

(

i

)

jk

c

i

,

i W

(

i

)

j

,*

+

c

i

,

i

+1

W

(

i

+1)

*,

j

+

=

(

∑

s

a i

,

s W

(

s

)

*,*

)

+

W

(

i

-1)

k

,*

W

(

i

)

*,

k

H

(

W

)

(

i

)

j

,

k

W

(

i

+1)

*,

j

W

(

i

)

j

,*

W

(

i

)

jk

W

(

i

+1)

W

(

i

)

W

(

i

-1)

H

W

)

(

i

)

2.2 Equivariant NF-Layers

We now construct a 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 ( · ; θ ) : vec ( U ) ↦→ θ vec ( U ) , where vec ( U ) ∈ R dim( U ) is U flattened as a vector and θ ∈ R dim( U ) × dim( U ) is a matrix of parameters. We show in Appendix B.2 that any S -equivariant T ( · ; θ ) must satisfy a 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 in each orbit [54]. Table 7 of the appendix describes the parameter sharing in detail.

Equivariant parameter sharing reduces the matrix-vector product θ vec ( U ) to the NF-Layer we now present. For simplicity we ignore V and assume here that U = W and defer the full form to Eq. 3 in the appendix. Then H : W c i →W c o maps input ( W (1) , · · · , W ( L ) ) to ( H ( W ) (1) , · · · , H ( W ) ( L ) ) . Recall that the inputs are not necessarily weights, but could be arbitrary weight-space features including the output of a previous NF-Layer. For W ( i ) ∈ R n i × n i -1 × c i , the corresponding output is H ( W ) ( i ) ∈ R n i × n i -1 × c o with entries computed:

Note that the terms involving W ( i -1) or W ( i +1) should be omitted for i = 0 and i = L , respectively, and ⋆ denotes summation or averaging over either the rows or columns. Recall that in the multichannel case, each W ( i ) jk is a vector in R c i so each parameter is a c o × c i matrix. We also provide a 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 to each output. To gain intuition for the operation of H , it is straightforward to check S -equivariance:

Proposition 1. The NF-Layer H : U c i →U c o (Eq. 2 and Eq. 3) is S -equivariant, where the group's action on input and output spaces is defined by Eq. 1. Moreover, any linear S -equivariant map T : U c i →U c o is equivalent to H for some choice of parameters a, b, c, d .

Proof (sketch). We can verify that H satisfies the equivariance condition [ σH ( W )] ( i ) jk = H ( σW ) ( i ) jk for any i, j, k by expanding each side of the equation using the definitions of the layer and action

(Eq. 1). Moreover, Appendix B.2 shows that any S -equivariant linear map T ( · , θ ) 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, d .

Informally, the above proposition tells us that H can express any linear S -equivariant function of a weight-space. Since ˜ S is a 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 1 summarizes the number of parameters (after parameter sharing) under different symmetry assumptions. While in general a linear layer T ( · ; θ ) : U c i → U c o has c i c o dim( U ) 2 parameters, the equivariant NF-Layers have significantly fewer free parameters due to parameter sharing. The S -equivariant layer H has O ( c i c o L 2 ) , while the ˜ S -equivariant layer ˜ H has O ( c i c o ( L + n 0 + n L ) 2 ) parameters. The latter's quadratic dependence on input and output dimensions can be prohibitive in some settings, such as 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 i denote the number of channels at the i th layer. Each bias v ( i ) ∈ R n i has the same dimensions as in the fully connected case, so only the convolution filter needs to be treated differently since it has additional spatial dimension(s) that cannot be permuted. For example, consider a 1D CNN with filters W = { W ( i ) ∈ R n i × n i -1 × w ∣ ∣ i ∈ J 1 ..L K } , where n i × n i -1 are the output and input channel dimensions and w is the filter width. We let W ( i ) jk := W ( i ) j,k, : ∈ R w denote the k th filter in the j th output channel, then define the S -action the same way as in Eq. 1.

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 ∈ U c where U is the weight-space of a 1D CNN. Then we combine the filter and channel dimensions of the weights, with W ( i ) ∈ R n i × n i -1 × ( cw ) . This allows us to use the multi-channel NF-Layer H : U wc i →U wc o . 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 ℓ -1 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.

IO-encoding. The S -equivariant layer H is more parameter efficient than ˜ H (Table 1), 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 W (1) and the rows of W ( L ) and v ( L ) ; this breaks the symmetry at input and output neurons even when using S -equivariant layers. In our experiments, we find that IO-encoding makes H competitive or superior to ˜ H , while using a fraction of the parameters.

2.3 Invariant NF-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 -invariant 2 layer P : U → R 2 L by simply summing or averaging the weight matrices and bias vectors across any axis that has permutation symmetry, i.e., P ( U ) = ( W (1) ⋆,⋆ , · · · , W ( L ) ⋆,⋆ , v (1) ⋆ , · · · , v ( L ) ⋆ ) . We can then apply a fully connected layer to the output of P 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 U → R d can be expressed in the form f ◦ P , for some choice of linear f : R 2 L → R d . Proof: See section B.3.

Table 2: Test τ of generalization prediction methods on the Small CNN Zoo [64], which contains the weights and test accuracies of many small CNNs trained on different datasets, such as CIFAR-10-GS or SVHN-GS. NFNHNP outperforms other methods on both datasets. Uncertainties indicate max and min over two runs.

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) or 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 to our permutation equivariant NFNs as NFNNP and NFNHNP, depending on which NF-Layer variant they use (see Table 1). We also evaluate a 'pointwise' ablation of our equivariant NF-Layer that ignores interactions between weights by only using the last term of Eq. 2, computing H ( W ) i jk := d i W ( i ) jk . We refer to NFNs that use this pointwise NF-Layer as NFNPT.

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 permutation equivariance, we optionally augment the MLP's training with permutations using Eq. 1. On relevant datasets we also compare with the recently developed DWSNets [47], an equivariant architecture similar to NFNHNP, and with inr2vec [12], a 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, 16]. The goal is to study generalization in a 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 raw 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 a shared architecture, but varied optimization hyperparameters. We compare the predictive power of NFNHNP and NFNNP 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 τ [31], a popular rank correlation metric with values in [ -1 , 1] .

In Table 2, 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 NFNHNP consistently performs the best on both datasets by a significant margin, showing that having access to the full weights can increase predictive power over hand-designed features as in STATNN. Because the input and output dimensionalities are small on these datasets, NFNHNP only uses moderately more ( ∼ 1 . 4 × ) 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 3Dscenes [61, 43, 7, 49, 58, 45, 14, 15], it is natural to wonder how to extract information about the

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 a 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 MLPAug baselines, which are three-layer MLPs with ReLU activations and 1,000 hidden units per layer.

As Navon et al. [47] test the performance of DWSNets on their independently constructed 2D-image INR datasets, we also present the results of training DWSNets on our own 2D-image INR datasets. Table 3 show that NFNNP consistently achieves the highest test accuracies of any method on the 2D-image tasks. More broadly, equivariant architectures significantly outperform the non-equivariant MLP approaches, even with permutation data augmentations.

For the 3D-shape datasets we also report the performance of inr2vec [12], a recent non-equivariant method for classifying 3D shapes from INR weights. Note that inr2vec's original setting assumes that all INRs in a 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 NFNHNP and NFNNP achieve significantly higher test accuracies than inr2vec, as well as the non-equivariant MLP baselines.

In addition to superior generalization, Appendix Table 16 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 NFNHNP performance on both CIFAR-10 and the 3D-shape datasets while using fewer parameters (e.g., 35% as many parameters on CIFAR-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.

NFN

Original INR

Edited INR

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 U 0 ∈ U be an initialization and let the sparsity mask M ∈ { 0 , 1 } dim ( U ) be a winning ticket for the initialization, with zeros indicating that the corresponding entries of U 0 should be pruned. The goal is to predict a winning ticket ˆ M given a held out initialization U 0 , such that the MLP initialized with U 0 and sparsity pattern ˆ M will achieve a high test accuracy after training.

We construct a conditional variational autoencoder [32, 59, cV AE] that learns a generative model of the winning tickets conditioned on initialization and train on datasets of (initialization, ticket) pairs found by one step of IMP with a sparsity level of P m = 0 . 95 for both MLPs trained on MNIST and CNNs trained on CIFAR-10. Table 5 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 m ) . In this setting, NFNHNP is prohibitively parameter inefficient, but NFNNP is able to recover test accuracies that are close to that of IMP pruned networks in CIFAR-10 and MNIST, respectively. Somewhat surprisingly, NFNPT 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 E.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 3). We evaluate two editing tasks:

Dilate (MNIST)

Contrast (CIFAR-10)

46

826

Ground truth

Original INR

8

3

MLPAug

8

m

NFNHNP

8

3

8

NFNNP

(1) making MNIST digits thicker via image dilation ( Dilate ), and (2) increasing image contrast on CIFAR-10 ( Contrast ). Both of these tasks require neural functionals to process the relationships between different pixels to successfully solve the task.

To produce training data for this task, we use standard image processing libraries [28, OpenCV] to dilate or increase the contrast of the MNIST and CIFAR-10 images, respectively. The training objective is to 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 MLP-based neural functionals with and without permutation augmentation.

Table 6 shows that permutation equivariant neural functionals (NFNHNP and NFNNP) achieve significantly better test MSE when editing held out INRs compared to other methods, on both the Dilate (MNIST) and Contrast (CIFAR-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 NFNPT 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 MLPAug at dilating MNIST digits and increasing the contrast in CIFAR-10 images.

4 Related work

The permutation symmetries of neurons have been a topic of interest in the context of loss landscapes and model merging [21, 4, 62, 17, 1]. 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, 33]. 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 network architectures [37, 8, 54, 34, 9, 18]. 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 rankk 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 a 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

of linear functions on these special graph structures. As discussed in Section 1, Navon et al. [47] recently developed an equivariant weight-space layer that is equivalent to our NF-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 a 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 on the permutation symmetries that arise in weight-spaces due to the particular structure of neural networks. We introduce two equivariant NF-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 a 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].

References

S. K. Ainsworth, J. Hayase, and S. Srinivasa. Git re-basin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209.04836 , 2022.

M. Andrychowicz, M. Denil, S. Gomez, M. W. Hoffman, D. Pfau, T. Schaul, B. Shillingford, and N. De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems , 29, 2016.

S. Bengio, Y. Bengio, J. Cloutier, and J. Gescei. On the optimization of a synaptic learning rule. In Optimality in Biological and Artificial Networks? , pages 281-303. Routledge, 2013.

J. 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:1907.02911 , 2019.

M. M. Bronstein, J. Bruna, T. Cohen, and P. Veliˇ ckovi´ c. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478 , 2021.

A. X. Chang, T. Funkhouser, L. Guibas, P. Hanrahan, Q. Huang, Z. Li, S. Savarese, M. Savva, S. Song, H. Su, J. Xiao, L. Yi, and F. Yu. ShapeNet: An Information-Rich 3D Model Repository. Technical Report arXiv:1512.03012 [cs.GR], Stanford University - Princeton University Toyota Technological Institute at Chicago, 2015.

Z. Chen and H. Zhang. Learning implicit fields for generative shape modeling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition , pages 5939-5948, 2019.

T. Cohen and M. Welling. Group equivariant convolutional networks. In International conference on machine learning , pages 2990-2999. PMLR, 2016.

T. S. Cohen, M. Geiger, J. Köhler, and M. Welling. Spherical CNNs. arXiv preprint arXiv:1801.10130 , 2018.

A. Dai, A. X. Chang, M. Savva, M. Halber, T. Funkhouser, and M. Nießner. Scannet: Richlyannotated 3d reconstructions of indoor scenes. In Proceedings of the IEEE conference on computer vision and pattern recognition , pages 5828-5839, 2017.

N. De Cao, W. Aziz, and I. Titov. Editing factual knowledge in language models. arXiv preprint arXiv:2104.08164 , 2021.

L. De Luigi, A. Cardace, R. Spezialetti, P. Zama Ramirez, S. Salti, and L. Di Stefano. Deep learning on implicit neural representations of shapes. In International Conference on Learning Representations (ICLR) , 2023.

L. Deutsch, E. Nijkamp, and Y. Yang. A generative model for sampling high-performance and diverse weights for neural networks. arXiv preprint arXiv:1905.02898 , 2019.

E. Dupont, Y. W. Teh, and A. Doucet. Generative models as distributions of functions. arXiv preprint arXiv:2102.04776 , 2021.

E. Dupont, H. Kim, S. Eslami, D. Rezende, and D. Rosenbaum. From data to functa: Your data point is a function and you should treat it like one. arXiv preprint arXiv:2201.12204 , 2022.

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.05688 , 2020.

R. Entezari, H. Sedghi, O. Saukh, and B. Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. arXiv preprint arXiv:2110.06296 , 2021.

M. Finzi, M. Welling, and A. G. Wilson. A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups. In International Conference on Machine Learning , pages 3318-3328. PMLR, 2021.

J. Frankle and M. Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv:1803.03635 , 2018.

J. Frankle, G. K. Dziugaite, D. M. Roy, and M. Carbin. Stabilizing the lottery ticket hypothesis. arXiv preprint arXiv:1903.01611 , 2019.

T. Garipov, P. Izmailov, D. Podoprikhin, D. P. Vetrov, and A. G. Wilson. Loss surfaces, mode connectivity, and fast ensembling of DNNs. Advances in neural information processing systems , 31, 2018.

C. Godfrey, D. 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.

D. Ha, A. Dai, and Q. V. Le. Hypernetworks. arXiv preprint arXiv:1609.09106 , 2016.

J. Harb, T. Schaul, D. Precup, and P. Bacon. Policy evaluation networks. CoRR , abs/2002.11833, 2020. URL https://arxiv.org/abs/2002.11833 .

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.

K. He, X. Zhang, S. Ren, and J. Sun. Deep residual learning for image recognition. CoRR, abs/1512 , 3385:2, 2015.

R. Hecht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Advanced Neural Computers , pages 129-135. Elsevier, 1990.

Itseez. Open source computer vision library. https://github.com/itseez/opencv , 2015.

Y. Jiang, D. Krishnan, H. Mobahi, and S. Bengio. Predicting the generalization gap in deep networks with margin distributions. In International Conference on Learning Representations , 2019. URL https://openreview.net/forum?id=HJlQfnCqKX .

Y. Jiang, P. Natekar, M. Sharma, S. K. Aithal, D. Kashyap, N. 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.

M. G. Kendall. A new measure of rank correlation. Biometrika , 30(1/2):81-93, 1938.

D. P. Kingma and M. Welling. Auto-encoding variational bayes. arXiv preprint arXiv:1312.6114 , 2013.

B. Knyazev, M. Drozdzal, G. W. Taylor, and A. Romero Soriano. Parameter prediction for unseen deep architectures. Advances in Neural Information Processing Systems , 34:2943329448, 2021.

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.

A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.

D. Krueger, C.-W. Huang, R. Islam, R. Turner, A. Lacoste, and A. Courville. Bayesian hypernetworks. arXiv preprint arXiv:1710.04759 , 2017.

Y. LeCun, Y. Bengio, et al. Convolutional networks for images, speech, and time series. The handbook of brain theory and neural networks , 3361(10):1995, 1995.

Y. LeCun, C. Cortes, and C. Burges. Mnist handwritten digit database. ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist , 2, 2010.

K. Li and J. Malik. Learning to optimize. arXiv preprint arXiv:1606.01885 , 2016.

H. Maron, H. Ben-Hamu, N. Shamir, and Y. Lipman. Invariant and equivariant graph networks. arXiv preprint arXiv:1812.09902 , 2018.

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.

C. H. Martin and M. W. Mahoney. Implicit self-regularization in deep neural networks: Evidence from random matrix theory and implications for learning. The Journal of Machine Learning Research , 22(1):7479-7551, 2021.

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 on computer vision and pattern recognition , pages 4460-4470, 2019.

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:2211.09760 , 2022.

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.

E. Mitchell, C. Lin, A. Bosselut, C. Finn, and C. D. Manning. Fast model editing at scale. arXiv preprint arXiv:2110.11309 , 2021.

A. Navon, A. Shamsian, I. Achituve, E. Fetaya, G. Chechik, and H. Maron. Equivariant architectures for learning in deep weight spaces. arXiv preprint arXiv:2301.12780 , 2023.

H. Pan and R. Kondor. Permutation equivariant layers for higher order interactions. In International Conference on Artificial Intelligence and Statistics , pages 5987-6001. PMLR, 2022.

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

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, high-performance deep learning library. Advances in neural information processing systems , 32, 2019.

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.12892 , 2022.

C. Qi, H. Su, K. Mo, and L. Guibas. Pointnet: deep learning on point sets for 3d classification and segmentation. cvpr (2017). arXiv preprint arXiv:1612.00593 , 2016.

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.

S. Ravanbakhsh, J. Schneider, and B. Poczos. Equivariance through parameter-sharing. In International conference on machine learning , pages 2892-2901. PMLR, 2017.

A. Rogozhnikov. Einops: Clear and reliable tensor manipulations with einstein-like notation. In International Conference on Learning Representations , 2022.

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.

A. Sinitsin, V. Plokhotnyuk, D. Pyrkin, S. Popov, and A. Babenko. Editable neural networks. arXiv preprint arXiv:2004.00345 , 2020.

V. Sitzmann, J. Martel, A. Bergman, D. Lindell, and G. Wetzstein. Implicit neural representations with periodic activation functions. Advances in Neural Information Processing Systems , 33: 7462-7473, 2020.

K. Sohn, H. Lee, and X. Yan. Learning structured output representation using deep conditional generative models. Advances in neural information processing systems , 28, 2015.

S. Sokota, H. Hu, D. J. Wu, J. Z. Kolter, J. N. Foerster, and N. Brown. A fine-tuning approach to belief state modeling. In International Conference on Learning Representations , 2022. URL https://openreview.net/forum?id=ckZY7DGa7FQ .

K. O. Stanley. Compositional pattern producing networks: A novel abstraction of development. Genetic programming and evolvable machines , 8:131-162, 2007.

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:15300-15311, 2020.

E. H. Thiede, T. S. Hy, and R. Kondor. The general theory of permutation equivarant neural networks and higher order graph variational encoders. arXiv preprint arXiv:2004.03990 , 2020.

T. Unterthiner, D. Keysers, S. Gelly, O. Bousquet, and I. Tolstikhin. Predicting neural network accuracy from weights. arXiv preprint arXiv:2002.11448 , 2020.

A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, Ł. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems , 30, 2017.

H. Xiao, K. Rasul, and R. Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017.

S. Yak, J. Gonzalvo, and H. Mazzawi. Towards task and architecture-independent generalization gap predictors. arXiv preprint arXiv:1906.01550 , 2019.

M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. Salakhutdinov, and A. Smola. Deep sets. doi: 10.48550. arXiv preprint ARXIV.1703.06114 , 2017.

C. Zhang, M. Ren, and R. Urtasun. Graph hypernetworks for neural architecture search. arXiv preprint arXiv:1810.05749 , 2018.

Appendix

A Equivariant NF-Layer pseudocode

Here we present a multi-channel implementation of the S -equivariant NF-Layer presented in Eq. 2 (which ignores biases), using PyTorch [50] and Einops-like [55] pseudocode. That is, it implements a linear layer H : W c i →W c o , where c i and c o are the number of input and output channels.

Note that our actual implementation differs from this pseudocode in a few ways: (1) it 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 (0 , 1) .

B NF-Layers for the NP Setting

This section presents futher details related to the S -equivariant layer H : U → U and the S -invariant layer P . First, Sec. B.1 gives the full definition of H as a function of both weight-space and biasspace features. Sec. B.2 provides a proof that H ( · ) expresses any S -equivariant linear function on U (Proposition 1). Sec. B.3 proves that combining our invariant layer P with a fully connected layer can express any S -invariant function from U to R d (Proposition 2).

B.1 Full definition

We present the full definition of the S -equivariant NF-Layer H : U → U , which completes Eq. 2 by also including the biases. The layer processes a set of weights U = ( W,v ) and outputs arrays ( Y, z )

̸

̸

̸

̸

̸

Table 7: S -equivariant parameter sharing for linear maps vec ( U ) ↦→ θ vec ( U ) . Parameter sharing is a system of constraints on the entries of ( ϑ, ϕ, φ, ψ ) = θ . Each table is organized by the layer indices ( i, s ) . For example, the first table says that for any ( i, s ) where s = i -1 , we constrain ϑ ijk spq = ϑ i,j ′ ,k ′ s,p ′ ,q ′ = a i,i -1 ϑ for any j, k, p, q and j ′ , k ′ , p ′ , q ′ where k = p and k ′ = p ′ . After parameter sharing, we observe that there are only a constant number of free parameters for each ( i, s ) pair, adding up to O ( L 2 ) parameters total.

̸

̸

̸

̸

with the same dimensions. The learnable parameters are in blue:

where ⋆ denotes summation or averaging over a dimension.

B.2 Proof of Proposition 1

Here we show that the layer defined in Eq. 3 is not only S -equivariant, it can express every linear equivariant function U → U . Our strategy is to first consider general linear layers T ( · ; θ ) : U → U parameterized by θ ∈ Θ , and then constrain to θ such that T ( · ; θ ) is S -equivariant. Applying this constraint will directly yield Eq. 3.

For our purposes, it is sometimes convenient to distinguish layer, row, and column indices of entries in U , so we split the parameters θ = ( ϑ, ϕ, φ, ψ ) and write T : U → U in the form:

̸

̸

̸

̸

We introduce the notation U α to identify individual entries of U . Here α is a 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 := W ∪ V as the combined index space of U . For example, if α = ( i, j, k ) ∈ W , then U α = W ( i ) jk .

We can then define the index space I := U × U for parameters θ ∈ Θ . We use θ α β to index an entry of θ with upper and lower indices [ α, β ] ∈ I . For example, if α = ( i, j, k ) ∈ W and β = ( s, p, q ) ∈ W , we have θ α β = ϑ ijk spq .

By flattening U into a vector, we can think of θ as a matrix with indices α and β . Eq. 4 can be rewritten in the flattened form:

In the flattened representation, we can write the action of σ ∈ S on vec ( U ) by a permutation matrix P σ . Equivariance requires that P σ θ vec ( U ) = θP σ vec ( U ) for any σ ∈ S . Since the input U can be anything, we get the following constraint on θ :

We can re-express the action of S on U (Eq. 1) as an action on the index space U . For α ∈ U , we write σ ( α ) which is defined:

for any σ ∈ S . We extend this definition into an action of S on I :

We can now rewrite the equivariance constraint in Eq. 8 using indices α, β :

By relabeling α ← σ -1 ( α ) , we can rewrite this condition θ α β = θ σ ( α ) σ ( β ) . Hence, equivariance is equivalent to requiring that θ share parameters within orbits under the action of S on its indices α, β [54, Prop 3.1]. We characterize these orbits in Sec. B.2.1, yielding Table 7 which gives the parameter sharing as a system of constraints on the indices of θ . These constraints characterize the space of S -equivariant T ( · , θ ) . In Sec. B.2.2, we conclude the proof by showing that under parameter sharing, T ( · , θ ) is equivalent to Eq. 3.

B.2.1 Parameter orbits and sharing

We now partition the parameters of θ 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 θ is I = U × U . There are four subsets of I :

4. I V V := V × V : Contains [ α, β ] = [( i, j ) , ( s, p )] , indexing parameters ψ ij sp .

Consider the block of indices I WW = W × W , containing [ α, β ] = [( i, j, k ) , ( s, p, q )] indexing parameters ϑ α β . Since the S -action never changes the layer indices ( i, s ) , we can independently consider orbits within sub-blocks of indices I WW i,s = { [( i, j, k ) , ( s, p, q )] | ∀ j, k, p, q } . The number of orbits within each sub-block I WW i,s depends on the relationship between the layer indices i and s : they are either the same layer ( s = i ), they are adjacent ( s = i -1 or s = i +1 ), or they are non-adjacent ( s / ∈ { i -1 , i, i +1 } ). We now analyze the orbits of sub-blocks for a few cases.

̸

If s = i -1 , then choose any two indices [ α (1) , β (1) ] , [ α (2) , β (2) ] ∈ I WW i,s where the first satisfies p = k and the second satisfies p = k . Then the orbits of each index are:

̸

̸

We see that these two orbits actually partition the entire sub-block of indices I WW i,s , with each orbit characterized by whether or not p = k . We introduce the parameters a i,i -1 ϑ (for the first orbit) and b i,i -1 ϑ (for the second orbit). Under equivariant parameter sharing, all parameters of ϑ corresponding I WW i,s are equal to either a i,i -1 ϑ or b i,i -1 ϑ , depending on whether p = k or p = k .

̸

If s = i +1 , we instead choose any two indices where the first satisfies j = q and the second satisfies j = q . Then the sub-block of indices I WW i,s is again partitioned into two orbits:

̸

depending on the condition j = q . We name two parameters a i,i +1 ϑ and c i,i +1 ϑ for this sub-block, with one for each orbit.

We can repeat this process for sub-blocks of I WW where i = s and s / ∈ { i -1 , i, i +1 } , as well as for the other three blocks of I . 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, s ) 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 U . Hence the number of distinct parameters after parameter sharing simply grows with the number of layer pairs, i.e. O ( L 2 ) .

B.2.2 Equivalence to equivariant NF-Layer definition

All that remains is to show that the map T ( · ; θ ) : U → U with S -equivariant parameter sharing (Table 7) is equivalent to the NF-Layer H we defined in Eq. 3.

Consider a single term from Eq. 4 where s = i -1 . Substituting using the constraints of Table 7, we simplify:

̸

We can then reparameterize b i,i -1 ← b i,i -1 -a i,i -1 , resulting in two terms that appear in Eq. 3. We can simplify every term of Eq. 4 in a similar manner using the parameter sharing of Table 7, reducing the general layer to the S -equivariant NF-Layer.

B.3 Proof of Proposition 2

Analogous to Sec. B.2, we begin with general parameterized linear maps T ( · ; θ ) : U → R d and characterize the parameter sharing implied by S -invariance. We then show that T ( · , θ ) under these constraints is equivalent to f ◦ P , where f : R 2 L → R d is a parameterized linear map (i.e., a fully connected layer) and P ( U ) = ( W (1) ⋆,⋆ , · · · , W ( L ) ⋆,⋆ , v (1) ⋆ , · · · , v ( L ) ⋆ ) . We will assume d = 1 for the rest of this section, but it is straightforward to generalize our arguments to arbitrary d .

In general, we can parameterize linear maps U → R by θ ∈ R dim( U ) , where:

where α is an index for either weights or biases as in Sec. B.2. The invariance constraint T ( U ) = T ( σU ) implies:

where σ ( α ) is defined by Eq. 9.

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, j, k ) ∈ W form an orbit, and we introduce a shared parameter a i . Similarly, for each i all ( i, j ) ∈ V form an orbit, and we introduce a shared parameter b i for each. We can then rewrite the layer:

This is simply a linear combination of the output of P ( U ) parameterized by { ( a i , b i ) } , and can be written as the composition f ◦ P where f is a fully connected layer.

C NF-Layers for the HNP setting

C.1 Equivariant NF-Layer

Because an expression for the ˜ S -equivariant NF-Layer analogous to Eq. 3 would be unwieldy, we instead define the layer in terms of its parameter sharing (Tables 8-11) on θ .

We can derive HNP-equivariant parameter sharing of θ using a similar strategy to Sec. B.2.1: we partition the index spaces I WW , I WV , I V W , I V V into orbits under the action of ˜ S , and share parameters within each corresponding orbit of ϑ, ϕ, φ, ψ . The resulting parameter sharing is different from the NP-setting because while the action of S on U could permute the rows and columns of every weight and bias, the action of ˜ S on U does not affect the columns of W (1) or the rows of W ( L ) , v ( L ) , which correspond to input and output dimensions (respectively).

The orbits are again analyzed within sub-blocks defined by the values of the layer indices ( i, s ) . As with the NP setting, there are broadly four types of sub-blocks based on whether i = s , s = i -1 , s = i + 1 , or s / ∈ { i -1 , i, i + 1 } . However, there are now additional considerations based on whether i or s is an input or output layer. For example, consider the sub-block of I WW where i = s = 1 , which we denote I WW 1 , 1 . The action on the indices in this sub-block can be written σ [ α, β ] = [(1 , σ 1 ( j ) , k ) , (1 , σ 1 ( p ) , q )] . Importantly, the column indices k, q are never permuted since they correspond to the input layer. We see that I WW 1 , 1 contains two orbits for each k ∈ J 1 ..n 0 K and q ∈ J 1 ..n 0 K , with the two orbits characterized by whether or not j = p . Hence we have 2 n 2 0 orbits and Table 8 introduces 2 n 2 0 parameters { a 1 , 1 ,k,q ϑ , b 1 , 1 ,k,q ϑ ∣ ∣ ∣ k, q ∈ J 1 ..n 0 K } for this sub-block of parameters.

Now consider another sub-block of I WW where 1 < i = s < L . Now the action of ˜ S on indices in this sub-block can be written σ [ α, β ] = [( i, σ i ( j ) , σ i -1 ( k )) , ( i, σ i ( p ) , σ i -1 ( q ))] . Then we have a total of two orbits characterized by whether or not k = p , rather than 2 n 2 0 orbits for the i = 1 case. Tables 8-11 present the complete parameter sharing for each of ϑ, ϕ, φ, ψ , resulting from analyzing every possible orbit within any sub-block of I WW , I WV , I V W , I V V .

C.2 Invariant NF-Layer

While the NP-invariant NF-Layer sums over the rows and columns of every weight and bias, under HNPassumptions there is no need to sum over the columns of W (1) (inputs) or the rows of W ( L ) , v ( L ) (outputs). So the HNP invariant NF-Layer ˜ P : U → R 2 L + n 0 +2 n L is defined:

where W (1) ⋆, : 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

̸

̸

̸

̸

̸

̸

ϕ

ijk

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

̸

into an invariant NF-Layer with mean pooling. The output of the invariant NF-Layer is flattened and projected to R 1 , 000 . The resulting vector is then passed through an MLP with two hidden layers, each with 1,000 units and ReLU activations. The output is linearly projected to a scalar and passed through a sigmoid function. Since the output of the model can 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 on τ on the validation set, which takes 1 hour on a Titan RTX GPU.

D.2 Predicting 'winning ticket' masks from initialization

Concretely, the encoder learns the posterior distribution q θ ( Z | U 0 , M ) where Z ∈ R dim( U ) × C is the latent variable for the winning tickets and C is the number of latent channels. The decoder learns p θ ( M | U 0 , Z ) , and both encoder and decoder are implemented using our equivariant NF-Layers. For the prior p ( Z ) we choose the isometric Gaussian distribution, and train using the evidence lower bound (ELBO):

The initialization and sparsity mask are concatenated so the input to the encoder q θ is ( U, M ) ∈ R dim( U ) × 2 . After the bottleneck, we concatenate the latent variables and the original mask along the channels, i.e. the decoder input is ( U 0 , Z ) ∈ R dim( U ) × ( C +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 2 fully-connected layers trained on CIFAR-10. In each dataset, we include 400 pairs for training and hold out 50 for evaluation. The hyperparameter details are in Table 13. The encoder and decoder models contain 4 equivariant NF-Layers with 64 hidden channels within each layer. The latent variable is 5 dimensions. Training takes 5H on a Titan RTX GPU.

D.3 Classifying INRs

We use SIREN [58] for our INRs of CIFAR-10, FashionMNIST, and MNIST. For the SIREN models, we used a 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 of 5 × 10 -5 . Datasets were split into 45,000 training images, 5,000 validation images, and 10,000 (MNIST, CIFAR-10) or 20,000 (FashionMNIST) test images. We trained 10 copies (MNIST, FashionMNIST) or 20 copies (CIFAR-10) of SIRENs on each training image with different initializations, and a 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 [12], and we 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 as an optimizer and we train for 1,000 steps.

We also trained neural functionals with three equivariant NF-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 IO-encoding, we used sinusoidal position encoding with a maximum frequency of 10 and 6 frequency bands (dimension 13). The training hyperparameters are shown in Table 14, and training took ∼ 4 H on a Titan RTX GPU.

Table 14: Hyperparameters for classifying INRs on MNIST and CIFAR-10 using neural functionals.

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

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 i be the INR weights for the i th image and SIREN ( x, y ; U ) be the output of the INR parameterized by U at coordinates ( x, y ) . We edit the INR

weights U ′ i = U i + γ · NFN ( U i ) , and γ is a learned scalar initialized to 0 . 01 . Letting f i ( x, y ) be the pixel values of the ground truth edited image (obtained from image-space processing), the objective is to minimize mean squared error:

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 3 equivariant NF-Layers with 128 channels, one invariant NF-Layer (mean pooling) followed by 4 linear layers with 1,000 hidden neurons. Every layer uses ReLU activation. The training hyperparameters can be found in Table 15, and training takes ∼ 1 hour on a Titan RTX GPU.

E Additional experiments and analysis

E.1 Analyzing train vs test accuracies in INR classification

Table 16 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 is that in addition to better test performance, NFN typically achieve much higher training accuracies compared to the MLPs. MLPAug, the MLP variant trained with permutation augmentations, especially struggles to achieve low training error. This indicates that training architectures to respect neuron permutation symmetries is extremely challenging, perhaps because these symmetry groups are extremely large [1]. Therefore, encoding equivariance or invariance via architecture design, as we do with NFNs, is much more effective.

E.2 Non-weight-space 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 and ScanNet-10. The results, shown in Table 16, indicate that applying state of the art architectures to pixels or 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 17 and 18, show that IO-encoding adds a very small (sometimes negligible) boost to NFNNP performance, though it never hurts. Since NFNNP often performs as well as or better than NFNHNP, 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 NFNPT on lottery ticket mask prediction by plotting the mask score predicted for a given initialization value at each layer. To make the visualization clear we train

Predicted Mask Score

Layer 0 Weights

Layer 1 Weights

Layer 2 Weights

Layer 0 Biases

Layer 1 Biases

Layer 2 Biases

-22.0

36

22.5

8-

23.0

40

23.5

24.0

42

24.5

44

-25.0

0.03

-0.02

0.01

0.00

0.01

0.02

0.03

-0.03

-0.02

-0.01

0.00

0.01

0.02

00

-5

-36

-10

-37

15

-38

6-

20

-40

25

-41

OE-

-0.075-0.050-0.025

0.000

0.025

0.050

0.075

-0.075-0.050

-0.025

0.000

0.025

0.050

0.075

10

13

12

11

0

10

-5

9

8-

-10

-0.075-0.050-0.025

0.000

0.025

0.050

0.075

-0.075

0.050

-0.025

0.000

0.025

0.050

0.075

Input Weight Value

E.5 Predicting MLP generalization from weights

In addition to predicting generalization on the Small CNN Zoo benchmark (Section 3.1), we 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-selected checkpoints from each run to construct a dataset of 20,000 (weight, test accuracy) pairs. Runs are partitioned according to a 90% / 10% split for training and testing.

We evaluate NFNHNP and NFNNP on this task and compare them to the STATNN baseline [64] which predicts test accuracy from hand-crafted features extracted from the weights. Table 19 shows that NFNNP and STATNN are broadly comparable, while NFNHNP consistently outperform other methods across both datasets in two measures of correlation: Kendall's tau and R 2 . These results confirm that processing the raw weights with permutation equivariant neural functionals can lead to greater predictive power when assessing generalization from weights.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\nThe 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.\n\nWEAKNESSES:\nI 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.\n\nQUESTIONS:\nThe paper is very clearly written and I don't have any remaining questions.\n\nLIMITATIONS:\nSome limitations and potential future work are discussed in the conclusion.\nIt should be added that one could consider more general permutation group actions.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-22"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n* 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\nWEAKNESSES:\n* 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\nQUESTIONS:\nI 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?\n\nLIMITATIONS:\nAs 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 2, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-06-23"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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\nSTRENGTHS:\n**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.\n\nWEAKNESSES:\n**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\nQUESTIONS:\n**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\nLIMITATIONS:\n- 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.", "TITLE": "Official Review", "RECOMMENDATION": 7, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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\nSTRENGTHS:\n- 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\nWEAKNESSES:\n- 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\nQUESTIONS:\n- 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\nLIMITATIONS:\nA limitation section is included.", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 5, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS:\nOriginality: 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\nWEAKNESSES:\n1. 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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nThere 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 4, "DATE": "2023-07-06"}