SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Allan Zhou $^{1}$  Kaien Yang $^{1}$  Kaylee Burns $^{1}$  Adriano Cardace $^{2}$  Yiding Jiang $^{3}$
Samuel Sokota<sup>3</sup> J. Zico Kolter<sup>3</sup> Chelsea Finn<sup>1</sup>
$^{1}$ Stanford University  $^{2}$ University of Bologna  $^{3}$ Carnegie Mellon University
ayz@cs.stanford.edu
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<sup>1</sup>.
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
https://github.com/AllanYangZhou/nfn
Figure 1: The internal operation of our permutation equivariant neural functionals (NFNs). The NFN processes the input weights through a series of equivariant NF-Layers, with each one producing weight-space features with varying numbers of channels. In this example, a neuron permutation symmetry simultaneously permutes the rows of  $W^{(2)}$  and the columns of  $W^{(3)}$ . This permutation propagates through the NFN in an equivariant manner.
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).
Table 1: Permutation symmetries of  $L$ -layer feedforward networks with  $n_0, \ldots, n_L$  neurons at each layer. All feedforward networks are invariant under hidden neuron permutations (HNP), while NP assumes that input and output neurons can also be permuted. We show the corresponding equivariant NF-Layers which process weight-space features from  $\mathcal{U}$ , with  $c_i$  input channels and  $c_o$  output channels.
Group	Abbrev	Permutable layers	Equivariant NF-Layer
Signature	Parameter count
S = ∏L i=0niSn_i	NP	All layers	H:Uci → Uco	O(ci coL^2)
S̅ = ∏L-1i=1niSn_i	HNP	Hidden layers	H̅:Uci → Uco	O (ci co(L + n_0 + n_L)^2)
—	—	None	T:Uci → Uco	ci co dim(U)^2
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.
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)} \in \mathbb{R}^{n_i \times n_{i-1}} \mid i \in [[1..L]]\}$  and biases  $v = \{v^{(i)} \in \mathbb{R}^{n_i} \mid i \in [[1..L]]\}$ . We denote the combined collection  $U := (W, v)$  belonging to weight-space,  $\mathcal{U} := \mathcal{W} \times \mathcal{V}$ .
Since the neurons in a hidden layer  $i \in \{1, \dots, 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  $\tilde{S} \coloneqq S_{n_1} \times \dots \times 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 \coloneqq S_0 \times \dots \times 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  $\sigma = (\sigma_0, \dots, \sigma_L) \in S$ . The action of the neuron permutation group is to permute the rows of each weight matrix  $W^{(i)}$  by  $\sigma_i$ , and the columns by  $\sigma_{i-1}$ . Each bias vector  $v^{(i)}$  is also permuted by  $\sigma_i$ . So the action is  $\sigma U := (\sigma W, \sigma v)$ , where:
$$
[ \sigma W ] _ {j k} ^ {i} = W _ {\sigma_ {i} ^ {- 1} (j), \sigma_ {i - 1} ^ {- 1} (k)} ^ {(i)}, \quad [ \sigma v ] _ {j} ^ {i} = v _ {\sigma_ {i} ^ {- 1} (j)} ^ {(i)}. \tag {1}
$$
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\geq 1$  feature channels, belonging to  $\mathcal{U}^c = \bigoplus_{i = 1}^c\mathcal{U}$ , the direct sum of  $c$  copies of  $\mathcal{U}$ . Concretely, each  $U\in \mathcal{U}^c$  consists of weights  $W = \left\{W^{(i)}\in \mathbb{R}^{n_i\times n_{i - 1}\times c}\mid i\in [[1..L]]\right\}$  and biases  $v = \left\{v^{(i)}\in \mathbb{R}^{n_i\times c}\mid i\in [[1..L]]\right\}$ , with the channels in the final dimension. The action defined in Eq. 1 extends to the multiple channel case if we define  $W_{jk}^{(i)}\coloneqq W_{j,k,:}^{(i)}\in \mathbb{R}^c$  and  $v_{j}^{(i)}\coloneqq v_{j,:}^{(i)}\in \mathbb{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:\mathcal{U}^{c_i}\to \mathcal{U}^{c_o}$  as  $S$ -equivariant if  $\sigma f(U) = f(\sigma U)$  for all  $\sigma \in S$  and  $U\in \mathcal{U}^{c_i}$ , where the action of  $S$  on the input and output spaces is defined by Eq. 1. Similarly, a function  $f:\mathcal{U}^c\to \mathbb{R}$  is  $S$ -invariant if  $f(\sigma U) = f(U)$  for all  $\sigma$  and  $U$ .
If  $f, g$  are equivariant, then their composition  $f \circ g$  is also equivariant; if  $g$  is equivariant and  $f$  is invariant, then  $f \circ 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.
Figure 2: A permutation equivariant NF-Layer takes in weight-space features as input (bottom) and outputs transformed features (top), while respecting the neuron permutation symmetries of feedforward networks. This illustrates the computation of a single output element  $H(W)_{jk}^{i}$ , defined in Eq. 2. Each output is a weighted combination of rows or column sums of the input weights, which preserves permutation symmetry. The first term contributes a weighted combination of row-and-column sums from every input weight, though this is omitted for visual clarity.
$$
H (W) _ {j k} ^ {(i)} = \left(\sum_ {s} a ^ {i, s} W _ {* *} ^ {(s)}\right) + c ^ {i, i} W _ {j, *} ^ {(i)} + c ^ {i, i + 1} W _ {* j} ^ {(i + 1)} + b ^ {i, i} W _ {* k} ^ {(i)} + b ^ {i, i - 1} W _ {k, *} ^ {(i - 1)} + d ^ {i} W _ {j k} ^ {(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(\cdot; \theta): \operatorname{vec}(U) \mapsto \theta \operatorname{vec}(U)$ , where  $\operatorname{vec}(U) \in \mathbb{R}^{\dim(U)}$  is  $U$  flattened as a vector and  $\theta \in \mathbb{R}^{\dim(\mathcal{U}) \times \dim(\mathcal{U})}$  is a matrix of parameters. We show in Appendix B.2 that any  $S$ -equivariant  $T(\cdot; \theta)$  must satisfy a system of constraints on  $\theta$  known as equivariant parameter sharing. We derive this parameter sharing by partitioning the entries of  $\theta$  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  $\theta \mathrm{vec}(U)$  to the NF-Layer we now present. For simplicity we ignore  $\mathcal{V}$  and assume here that  $\mathcal{U} = \mathcal{W}$  and defer the full form to Eq. 3 in the appendix. Then  $H: \mathcal{W}^{c_i} \to \mathcal{W}^{c_o}$  maps input  $(W^{(1)}, \dots, W^{(L)})$  to  $(H(W)^{(1)}, \dots, 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)} \in \mathbb{R}^{n_i \times n_{i-1} \times c_i}$ , the corresponding output is  $H(W)^{(i)} \in \mathbb{R}^{n_i \times n_{i-1} \times c_o}$  with entries computed:
$$
H (W) _ {j k} ^ {(i)} = \left(\sum_ {s} a ^ {i, s} W _ {\star , \star} ^ {(s)}\right) + b ^ {i, i} W _ {\star , k} ^ {(i)} + b ^ {i, i - 1} W _ {k, \star} ^ {(i - 1)} + c ^ {i, i} W _ {j, \star} ^ {(i)} + c ^ {i, i + 1} W _ {\star , j} ^ {(i + 1)} + d ^ {i} W _ {j k} ^ {(i)}. \tag {2}
$$
Note that the terms involving  $W^{(i-1)}$  or  $W^{(i+1)}$  should be omitted for  $i = 0$  and  $i = L$ , respectively, and  $\star$  denotes summation or averaging over either the rows or columns. Recall that in the multi-channel case, each  $W_{jk}^{(i)}$  is a vector in  $\mathbb{R}^{c_i}$  so each parameter is a  $c_o \times 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: \mathcal{U}^{c_i} \to \mathcal{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: \mathcal{U}^{c_i} \to \mathcal{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  $[\sigma H(W)]_{jk}^{(i)} = H(\sigma W)_{jk}^{(i)}$  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(\cdot, \theta)$  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  $\tilde{S}$  is a subgroup of  $S$ ,  $H$  is also  $\tilde{S}$ -equivariant. However, it does not express every possible linear  $\tilde{S}$ -equivariant function. We derive the full  $\tilde{S}$ -equivariant NF-Layer  $\tilde{H}: \mathcal{U} \to \mathcal{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(\cdot;\theta): \mathcal{U}^{c_i} \to \mathcal{U}^{c_o}$  has  $c_i c_o \dim(\mathcal{U})^2$  parameters, the equivariant NF-Layers have significantly fewer free parameters due to parameter sharing. The  $S$ -equivariant layer  $H$  has  $O\left(c_i c_o L^2\right)$ , while the  $\tilde{S}$ -equivariant layer  $\tilde{H}$  has  $O\left(c_i c_o (L + n_0 + n_L)^2\right)$  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^{\text{th}}$  layer. Each bias  $v^{(i)} \in \mathbb{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)} \in \mathbb{R}^{n_i \times n_{i-1} \times w} \mid i \in [1..L]\}$ , where  $n_i \times n_{i-1}$  are the output and input channel dimensions and  $w$  is the filter width. We let  $W_{jk}^{(i)} := W_{j,k,:}^{(i)} \in \mathbb{R}^w$  denote the  $k^{\text{th}}$  filter in the  $j^{\text{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 \in \mathcal{U}^c$  where  $\mathcal{U}$  is the weight-space of a 1D CNN. Then we combine the filter and channel dimensions of the weights, with  $W^{(i)} \in \mathbb{R}^{n_i \times n_{i-1} \times (cw)}$ . This allows us to use the multi-channel NF-Layer  $H: \mathcal{U}^{wc_i} \to \mathcal{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  $\ell$  is FC and layer  $\ell - 1$  is convolutional. If global spatial pooling removes all spatial dimensions from the output of  $\ell - 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  $\tilde{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  $\tilde{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 layer  $P: \mathcal{U} \to \mathbb{R}^{2L}$  by simply summing or averaging the weight matrices and bias vectors across any axis that has permutation symmetry, i.e.,  $P(U) = \left(W_{\star, \star}^{(1)}, \dots, W_{\star, \star}^{(L)}, v_{\star}^{(1)}, \dots, v_{\star}^{(L)}\right)$ . 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  $\mathcal{U}$ :
Proposition 2. Any  $S$ -invariant linear function  $\mathcal{U} \to \mathbb{R}^d$  can be expressed in the form  $f \circ P$ , for some choice of linear  $f: \mathbb{R}^{2L} \to \mathbb{R}^d$ . Proof: See section B.3.
2We define the analogous  $\tilde{S}$ -invariant layer  $\tilde{P}$  in Eq. 20 of the appendix.
Table 2: Test  $\tau$  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.  $\mathrm{NFN_{HNP}}$  outperforms other methods on both datasets. Uncertainties indicate max and min over two runs.
NFNHNP	NFNNP	STATNN
CIFAR-10-GS	0.934 ± 0.001	0.922 ± 0.001	0.915 ± 0.002
SVHN-GS	0.931 ± 0.005	0.856 ± 0.001	0.843 ± 0.000
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  $\mathrm{NFN_{NP}}$  and  $\mathrm{NFN_{HNP}}$ , 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)_{jk}^{i} \coloneqq d^{i}W_{jk}^{(i)}$ . We refer to NFNs that use this pointwise NF-Layer as  $\mathrm{NFN_{PT}}$ .
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  $\mathrm{NFN_{HNP}}$ , 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  $\mathrm{NFN_{HNP}}$  and  $\mathrm{NFN_{NP}}$  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  $\tau$  [31], a popular rank correlation metric with values in  $[-1, 1]$ .