SlowGuess/NeurIPS_2023_OCR huggingface dataset:
is to reconstruct the original INR's content using the decoded weights:
$$
\mathcal {L} (\theta , W) = \sum_ {i = 1} ^ {M} \sum_ {x \in P _ {i}} \left(\text {S I R E N} \left(x; \hat {W} _ {i}\right) - \text {S I R E N} (x; W)\right) ^ {2}, \tag {12}
$$
$$
\hat {W} _ {i} = \operatorname {D e c} _ {\theta} \left(z _ {i}\right), \quad z = \operatorname {E n c} _ {\theta} (W). \tag {13}
$$
Our INR2ARRAY encoder architecture uses multiple weight-space self-attention layers followed by a weight-space cross-attention layer CA, which is especially well-suited for producing  $M \times d$  permutation invariant arrays. For the decoder, we use the hypernetwork design of Sitzmann et al. [37], where the hypernetworks are conditioned on the spatial latent  $z_{i} \in \mathbb{R}^{d}$ . In principle, one can also design an invariant encoder using (non-attentive) NFN layers, and we will compare both options in our experiments.
Our experiments broadly involve two types of weight-space datasets: (1) datasets of neural networks (e.g., CNN image classifiers) trained with varying hyperparameters, where the goal is to model interesting properties of the trained networks (e.g., generalization) from their weights; and (2) datasets containing the weights of INRs each representing a single data signal, such as an image. Tasks including editing INR weights to modify the represented signal, or predicting target information related to the signal (e.g., image class). For our INR experiments, we construct datasets for MNIST [25], FashionMNIST [42], and CIFAR-10 [23] following the procedure of Zhou et al. [46] exactly.
5.1 INR classification with INR2ARRAY
Classifying the signal represented by an INR directly from its weights is a challenging problem, with current approaches requiring that all INRs share initialization [10, 3, 6]. In the "vanilla" setting where INRs can be initialized and trained independently, state-of-the-art methods struggle to classify signals from even simple datasets [31, 46]. We train INR2ARRAY to learn a latent representation of INR weights that can be used for more effective classification in this vanilla setting.
We implement both NFT (weight-space self-attention and cross-attention) and  $\mathrm{NFN_{NP}}$  [46] variants of the INR2ARRAY encoder, and distinguish the two variants by superscripts: INR2ARRAY $^{\mathrm{NFT}}$  and INR2ARRAY $^{\mathrm{NFN}}$ . For the decoder, we use the hypernetwork design of Sitzmann et al. [37], where the hypernetworks are conditioned on the spatial latent  $z_{i} \in \mathbb{R}^{d}$ . Appendix B.3 describes the implementation and training of each variant in full detail.
We train separate encoders for each INR dataset shown in Table 6 of the appendix, which reports the reconstruction mean square error (Eq 12) on test INR inputs. The NFT encoder consistently achieves lower reconstruction error than the NFN encoder, and INR2ARRAY $^{\text{NFT}}$  also produces qualitatively better reconstructions than INR2ARRAY $^{\text{NFN}}$  (Figure 4), confirming that NFTs enable higher quality encoding than previous neural functionals.
Once trained, the INR2ARRAY encoder maps INR weights to compact latent arrays that represent the INR's contents. Downstream tasks such as classifying the INR's signal (e.g., image) can now be performed directly on these  $S_{\mathrm{NP}}$ -invariant latent arrays  $z \in \mathbb{R}^{M \times d}$ . Concretely, in  $K$ -way
Table 1: Test accuracy (\%) for weight-space INR classification in the MNIST, FashionMNIST, and CIFAR-10 datasets. INR2ARRAY significantly improves over current state-of-the-art results in this setting. For the NFN baseline [46] we report the higher performance out of their NP and HNP variants. Note that the MNIST and FashionMNIST results reported for DWS in Navon et al. [31] are on their own independently constructed INR datasets, while we use the datasets from Zhou et al. [46] for all methods for consistent comparison. Uncertainties indicate standard error over three runs.
MNIST	FashionMNIST	CIFAR-10
INR2VEC [6]	19.1 ± 0.18	23.3 ± 0.36	16.7 ± 0.24
DWS [31]	74.4 ± 0.14	64.8 ± 0.69	41.5 ± 0.43
NFN [46]	92.9 ± 0.38	75.6 ± 1.07	46.6 ± 0.13
Spatial NFN	92.9 ± 0.46	70.8 ± 0.53	45.6 ± 0.11
NFT	89.9 ± 1.03	72.7 ± 0.05	44.8 ± 0.32
INR2ARRAYaNFN (Ours)	94.6 ± 0.00	76.7 ± 0.00	45.4 ± 0.00
INR2ARRAYaNFT (Ours)	98.5 ± 0.00	79.3 ± 0.00	63.4 ± 0.00
classification the trained encoder  $\mathrm{ENC}:\mathcal{W}\to \mathbb{R}^{M\times d}$  can be composed with any classification head  $f:\mathbb{R}^{M\times d}\to \mathbb{R}^K$  to form an  $\mathcal{S}_{\mathrm{NP}}$ -invariant classifier  $f\circ \mathrm{ENC}:\mathcal{W}\to \mathbb{R}^K$ . Since the encoder is already trained by the reconstruction objective, we only train  $f$  during classification and keep the encoder fixed. We choose to implement  $f$  using a Transformer classification head [41], which views the latent array input as a length- $M$  sequence of  $d$ -dimensional vectors.
Existing INR classifiers, which we refer to as DWS [31] and NFN [46], are permutation-invariant architectures that directly map input weights to predicted labels. We also create a modified NFN classifier (Spatial NFN) which produces an  $M \times d$  array of activations before a convolutional classifier head, to test the impact of spatial latents independent of the INR2ARRAY training process. Appendix B.4 describes the setup for each method in further detail. Finally, we measure the performance of NFT classifiers without INR2ARRAY.
Table 1 shows that INR2ARRAY significantly improves classification test accuracies across the board. In addition, the best performance is achieved by implementing the encoder using our attention-based layers (NFT) compared to linear weight-space layers (NFN). Notably, INR2ARRAY achieves a test accuracy of  $64.4\%$  on CIFAR-10 (the hardest dataset), an improvement of  $+17\%$  over the previous best result of  $46.6\%$  by NFNs. They also achieve an MNIST test accuracy of  $98.5\%$ , up from  $92.9\%$  by NFNs and  $85.7\%$  by DWS. The vanilla NFT performance is comparable to DWS and NFNs (better than the former and worse than the latter), suggesting the importance of INR2ARRAY in addition to the architecture. Finally, the Spatial NFN fails to improve performance over standard NFN classifiers, implying that the particular objective of INR2ARRAY is crucial to make spatial latent spaces useful. The results show that learning a latent array representation of INR weights using INR2ARRAY and NFTs can greatly improve performance on downstream INR tasks like classification.
5.1.1 Term-by-term ablation
Here we investigate the importance of each term in our weight space self-attention layer (Eqs. 4-6). The first two terms of SA, illustrated by Figure 2, amount to self-attention between the rows and columns of adjacent weights, while the third term is a global self-attention between features from all layers. Note that in practice we use a tractable approximation of the third term described in Sec 3.1. For this ablation experiment we either keep only the first two terms (AdjacentSA) or only the third term (GlobalSA). FullSA denotes the original layer.
Table 2: Test classification accuracy of INR2ARRAYNFT when ablating the terms of SA (Eqs.4-6).
Model	MNIST
FullSA	98.5
AdjacentSA	97.6
GlobalSA	37.4
Table 2 shows the results of this ablation experiment on
INR2ARRAY $^{\mathrm{NFT}}$  performance in the MNIST classification task. We see that our full layer (FullSA) performs best, abating the final term (AdjacentSA) only slightly degrades performance, and abating the first two terms (GlobalSA) drastically harms performance. This is interesting since the first two terms are necessary to achieve minimal equivariance, while the third term alone is not minimally
Table 3: Test mean squared error to the target visual transformation for five INR editing tasks. NFTs generally achieve lower test error than NFN variants. Uncertainties indicate standard error over three seeds.
NFNHNP [46]	NFNNP [46]	NFT (Ours)
MNIST (erode)	0.0228 ± 0.0003	0.0223 ± 0.0000	0.0194 ± 0.0002
MNIST (dilate)	0.0706 ± 0.0005	0.0693 ± 0.0009	0.0510 ± 0.0004
MNIST (gradient)	0.0607 ± 0.0013	0.0566 ± 0.0000	0.0484 ± 0.0007
FashionMNIST (gradient)	0.0878 ± 0.0002	0.0870 ± 0.0001	0.0800 ± 0.0002
CIFAR (contrast)	0.0204 ± 0.0000	0.0203 ± 0.0000	0.0200 ± 0.0002
Table 4: (Small CNN Zoo [40] benchmark.) Rank correlation  $\tau$  for predicting the generalization of CNN image classifiers with unseen hyperparameters trained CIFAR-10-GS and SVHN-GS (GS=grayscale). The NFT outperforms  $\mathrm{NFN_{NP}}$  and hand-picked features (STATNN), while  $\mathrm{NFN_{HNP}}$  performs best overall. Uncertainties indicate standard error over two runs.
NFNHNP [46]	NFNNP [46]	STATNN [40]	NFT (Ours)
CIFAR-10-GS	0.934 ± 0.001	0.922 ± 0.001	0.915 ± 0.002	0.926 ± 0.001
SVHN-GS	0.931 ± 0.005	0.856 ± 0.001	0.843 ± 0.000	0.858 ± 0.000
equivariant (but helps propagate information between any two weight matrices). The results emphasize the importance of designing our layer to achieve minimal equivariance rather than naively applying self-attention to the weights.
5.2 Editing INRs
We also evaluate NFTs on editing INR weights to alter their signal, e.g., to modify the represented image. The goal of this task is to edit the weights of a trained SIREN to alter its encoded image, expressed as a difference  $W' \gets W + \Delta(W)$ . Neural functionals are trained to learn  $\Delta(\cdot)$  that achieves some desired visual transformation, such as dilating the image. Permutation equivariance is a useful inductive bias here since if  $\Delta(W)$  is the desired edit for  $W$ , then for any neuron permutation  $\sigma$  the desired edit to  $\sigma W$  is  $\sigma \Delta(W)$ . In addition to the MNIST dilation and CIFAR contrast tasks from Zhou et al. [46], we also introduce several new editing tasks: MNIST erosion, MNIST gradient, and FashionMNIST gradient. Gradient tasks roughly amount to edge detection; Figure 5 in the appendix shows sample inputs and targets for each editing task.
We compare NFT editing performance against the two NFN variants [46], with each method using the a similar number of parameters ( $\sim 7\mathrm{M}$ ). For full training details, see Appendix B.2). Table 3 shows the test mean square error (MSE) between the edited INR and the ground truth visual transformation. NFTs consistently outperform NFNs on most INR editing tasks. In addition, Table 5 in the appendix shows that NFTs generally obtain both lower training and test error, indicating that the attention-based architecture enables greater expressivity. Figure 5 shows qualitative samples produced by each editing method on each task.
5.3 Predicting CNN classifier generalization
Whereas the previous experiments have focused on the weight spaces of INRs (which are implemented by small MLPs), we would also like to evaluate how NFTs process other weight spaces such as those belonging to convolutional neural network classifiers. Large-scale empirical studies have produced datasets of trained classifiers under different hyperparameter settings [40, 11], enabling a data-driven approach to modeling generalization from their weights, which could lead to new insights. Prior methods for predicting classifier generalization typically rely on extracting hand-designed features from the weights before using them to predict the test accuracy [18, 43, 40, 19, 28].
We now study how NFTs can model generalization from raw weights using the Small CNN Zoo [40], which contains thousands of CNNs trained on image classification datasets. In addition to comparing against the two neural functional variants  $\mathrm{NFN_{NP}}$  and  $\mathrm{NFN_{HNP}}$  [46], we also show the performance of
a hand-designed features approach called STATNN [40]. Each generalization predictor is trained on thousands of CNN weights produced under varying initialization and optimization hyperparameters, and is then tested on held out weights produced using unseen hyperparameters. Appendix B.1 describes the experimental setup in full detail.
Table 4 shows the test performance of each method using the Kendall rank correlation coefficient  $\tau$  [20]. Across both datasets, neural functionals operating on raw weights outperform the hand-designed features baseline STATNN. NFTs slightly outperform  $\mathrm{NFN_{NP}}$  on each dataset, while  $\mathrm{NFN_{HNP}}$  achieves the best performance overall. Note that both  $\mathrm{NFN_{NP}}$  and NFTs use layers that assume input and output neurons are permutable (NP). Although we use input/output position encoding to remove the stronger NP assumptions, the results suggest that HNP designs may be naturally better suited to this task.
It is well known that the weight spaces of neural networks contain numerous symmetries, i.e., transformations that preserve the network's behavior [17]. Permutation symmetries in particular have been studied in the context of neural network loss landscapes [14, 4, 12] and weight-space merging [38, 1]. However, most prior methods for processing weight space objects do not explicitly account for these symmetries [2, 26, 15, 24, 45, 7, 21], although some have tried to encourage permutation equivariance through data augmentation [33, 29]. Another workaround approach is to explicitly restrict the weight space being considered by, for example, fixing the initialization of all networks being processed [6] or meta-learning modulations of a set of base parameters [10, 3]. Instead, we study the problem of encoding permutation symmetries into the neural functional itself, without any restrictions on the weight space being processed.
There are a variety of methods that incorporate symmetries into deep neural network architectures [36, 22, 13]. Examples include (group) convolutions for images [25, 5] and permutation equivariant architectures for general set-structured inputs [34, 44, 16, 39, 27]. Most directly related to our work are that of Navon et al. [31] and Zhou et al. [46], who introduce linear layers equivariant to the neuron permutation symmetries of feedforward networks. We extend their approach by introducing nonlinear equivariant layers based on the attention mechanism, and use them to construct NFTs.
This work introduces neural functional transformers, a novel class of weight-space models designed with neuron permutation symmetries in mind. Our approach extends recent work on permutation-equivariant weight-space models using nonlinear self-attention and cross-attention layers, which can enable greater expressivity compared to existing linear layers.
We empirically evaluate the effectiveness of NFTs on weight-space tasks involving datasets of trained CNN classifiers and implicit neural representations (INRs). Operating on weights alone, NFTs can predict the test accuracy of CNN classifiers, modify the content of INRs, and classify INR signals. We also use NFTs to develop INR2ARRAY, a method for mapping INR weights to compact and  $S_{\mathrm{NP}}$ -invariant latent representations, which significantly improve performance in downstream tasks such as INR classification.
Some limitations of NFTs include increased computational costs from self-attention relative to linear weight-space layers, and the difficulty of training large NFT architectures stably. Future work may address these limitations, and explore additional applications of NFTs such as for learned optimization or weight-space generative modeling. Overall, our work highlights the potential of attention-based weight-space layers offers a promising direction for the development of more expressive and powerful neural functional architectures.
We thank Adriano Cardace and Yoonho Lee for interesting discussions and suggestions related to the development of this paper. AZ and KB are supported by the NSF Graduate Research Fellowship Program. This project was supported by Juniper Networks and ONR grant N00014-22-1-2621.
[40] T. Unterthiner, D. Keysers, S. Gelly, O. Bousquet, and I. Tolstikhin. Predicting neural network accuracy from weights. arXiv preprint arXiv:2002.11448, 2020.
[41] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[42] H. Xiao, K. Rasul, and R. Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017.
[43] S. Yak, J. Gonzalvo, and H. Mazzawi. Towards task and architecture-independent generalization gap predictors. arXiv preprint arXiv:1906.01550, 2019.
[44] M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. Salakhutdinov, and A. Smola. Deep sets. doi: 10.48550. arXiv preprint ARXIV.1703.06114, 2017.
[45] C. Zhang, M. Ren, and R. Urtasun. Graph hypernetworks for neural architecture search. arXiv preprint arXiv:1810.05749, 2018.
[46] A. Zhou, K. Yang, K. Burns, Y. Jiang, S. Sokota, J. Z. Kolter, and C. Finn. Permutation equivariant neural functionals. arXiv preprint arXiv:2302.14040, 2023.
A Weight-space self-attention
A.1 Proof of minimal equivariance
We repeat the claim about minimal equivariance for weight-space self-attention and layer position encodings SA  $\circ$  LAYERENC:  $\mathcal{W}^c\to \mathcal{W}^c$ :
Proof. We first show  $S_{\mathrm{NP}}$ -equivariance, before checking that it minimally equivariant.
Equivalence. It is straightforward to check that LAYERENC is equivariant to any permutation that preserves the layer index (which neuron permutations do). At a high level, we check that SA is  $S_{\mathrm{NP}}$ -equivariant by verifying that  $\mathrm{SA}(\sigma W)_{jk}^{(i)} = \mathrm{SA}(W)_{\sigma_i^{-1}(j),\sigma_{i - 1}^{-1}(k)}^{(i)}$ . We proceed term by term.
First, it will be useful to show that ATTN (Eq. 3) is equivariant to permutations  $\pi \in S_d$  of each key, query, and value input. In particular, let  $[\pi q][i] = q[\pi^{-1}(i)]$  permute the entries of the input vectors:
$$
\begin{array}{l} \operatorname {A T T N} \left(\pi q, \left\{\pi k _ {p}, \pi v _ {p} \right\} _ {p = 1} ^ {N}\right) [ i ] = \left[ \sum_ {p} \left(\frac {\exp \left(q \cdot k _ {p}\right)}{\sum_ {p ^ {\prime}} \exp \left(q \cdot k _ {p ^ {\prime}}\right)}\right) (\pi v _ {p}) \right] [ i ] (14) \\ = \sum_ {p} \left(\frac {\exp \left(q \cdot k _ {p}\right)}{\sum_ {p ^ {\prime}} \exp \left(q \cdot k _ {p ^ {\prime}}\right)}\right) v _ {p} \left[ \pi^ {- 1} (i) \right] (15) \\ = \operatorname {A T T N} \left(q, \left\{k _ {p}, v _ {p} \right\} _ {p = 1} ^ {N}\right) \left[ \pi^ {- 1} (i) \right] (16) \\ \end{array}
$$
Note that in the first equality we use  $\pi q\cdot \pi k = q\cdot k$
Since the key, query, and value projections are only linear projections along the channel dimension,  $K(\sigma W) = \sigma K(W)$ , where  $K$  denotes all keys produced by linearly projection of  $W$  by  $\theta_K$  (and similarly for  $Q, V$ ). Then consider, for example, how the first term (Eq 4) changes under  $W \mapsto \sigma W$ :
$$
\operatorname {A t t n} \left(\sigma Q _ {j,:} ^ {(i)}, \left\{(\sigma K, \sigma V) _ {:, q} ^ {(i - 1)} \right\} _ {q} \bigcup \left\{(\sigma K, \sigma V) _ {p,:} ^ {(i)} \right\} _ {p}\right) [ k ] = \tag {17}
$$
$$
\operatorname {A t t n} \left(\sigma_ {i - 1} Q _ {\sigma_ {i} ^ {- 1} (j),,:}, \left\{\left(\sigma_ {i - 1} K _ {:, q} ^ {(i - 1)}, \sigma_ {i - 1} V _ {:, q} ^ {(i - 1)}\right) \right\} _ {q} \bigcup \left\{\left(\sigma_ {i - 1} K _ {p,:} ^ {(i)}, \sigma_ {i - 1} V _ {p,:} ^ {(i)}\right) \right\} _ {p}\right) [ k ] \tag {18}
$$
Note the distinction between neuron permutations of the weight space  $[\sigma Q]$  and permutations of individual vectors like  $\sigma_{i-1}Q_{p,:}$ , and that permutations along the KV set indices (e.g.,  $q$  in the first set) are ignored since set order is irrelevant. Using the equivariance of ATTN (Eqs. 14-16):
$$
\operatorname {A T T N} \left(Q _ {\sigma_ {i} ^ {- 1} (j),:,} ^ {(i)} \left\{(K, V) _ {:, q} ^ {(i - 1)} \right\} _ {q} \bigcup \left\{\left(K, V\right) _ {p,:} ^ {(i)} \right\} _ {p}\right) [ \sigma_ {i - 1} ^ {- 1} (k) ] \tag {19}
$$
Notice that compared to Eq. 4 the output indices have changed  $j \mapsto \sigma_i^{-1}(j)$  and  $k \mapsto \sigma_{i-1}^{-1}(k)$ , showing that the first term is equivariant. We can repeat this argument in a similar fashion for the second term (Eq. 5), while the third term (Eq. 6) applies ATTN by treating each weight as an individual token, so its equivariance follows immediately from the permutation equivariance of ATTN. This verifies that all terms of SA are  $S_{\mathrm{NP}}$ -equivariant.
Minimal equivariance. Here our goal is to show non-equivariance to permutations  $\tau \in S_{\dim(\mathcal{W})}$  but  $\tau \notin S_{\mathrm{NP}}$  (false symmetries). That is, we should find some input  $W$  such that  $\mathrm{SA}(\mathrm{LAYERENC}(\tau W)) \neq \tau \mathrm{SA}(\mathrm{LAYERENC}(W))$ , for some parameters  $\{\theta_Q, \theta_K, \theta_V\}$  and  $\{\phi^{(i)}\}$ .
It will be useful to define the actions of permutations in  $S_{\dim(\mathcal{W})}$  on the indices of the weights, rather than the weights themselves. The index-space for a weight space  $\mathcal{W}$  is defined as a set of 3-tuples representing layer number, row, and column, respectively:
$$
\mathbb {I} = \left\{\left(i, j, k\right) \mid i = \llbracket 1.. L \rrbracket , j = \llbracket 1.. n _ {i} \rrbracket ; k = \llbracket 1.. n _ {i - 1} \rrbracket \right\}. \tag {20}
$$
For  $\alpha \in \mathbb{I}$ , we use the subscripts  $\ell, r, c$  to denote the layer, row, and column indices respectively. That is,  $\alpha = (\alpha_{\ell}, \alpha_{r}, \alpha_{c}) = (i, j, k)$ .
Notice that the set of arbitrary permutations  $S_{\dim(\mathcal{W})}$  is simply the set of all bijections  $\tau : \mathbb{I} \to \mathbb{I}$  (any index can be arbitrarily permuted to any other index), while  $S_{\mathrm{NP}}$  is the subgroup of bijections  $\sigma$  that can be written as:
$$
\sigma (i, j, k) = (i, \sigma_ {i} (j), \sigma_ {i - 1} (k)), \quad \forall (i, j, k) \in \mathbb {I}. \tag {21}
$$
False symmetries can broadly be categorized into three groups. In fact, we can verify that any permutation  $\tau$  that fails to satisfy one of these three cases must in  $S_{\mathrm{NP}}$ :
1. False symmetry that permutes weights across layers. There exists an  $(i,j,k) \in \mathbb{I}$  such that  $\tau(i,j,k)_{\ell} \neq i$ .
2. False symmetry where row and column permutations are not independent. Suppose a false symmetry  $\tau$  does not fall under case (1), i.e., it preserves the layer index. Then it may still fail to have independent permutations of row and columns. In this case, either there is an  $(i,j,k)$  and  $q$  such that  $\tau(i,j,k) = (i,j',k')$ , but  $\tau(i,j,q)_r \neq j'$ . Or there is an  $(i,j,k)$  and  $p$  such that  $\tau(i,j,k) = (i,j',k')$  but  $\tau(i,p,k)_c \neq k'$ .
3. Adjacent permutations decoupled. If a false symmetry  $\tau$  is not one of the two cases above, there must be an  $(i,j,k) \in \mathbb{I}$  and  $q$  such that  $\tau(i,j,k) = (i,j',k')$  but  $\tau(i - 1,k,q)_r \neq k'$ . That is, an instance where one row pair of layer  $i - 1$  do not permute the same way as the columns of layer  $i$ .
We can now check non-equivariance for  $\tau$  in each case. To simplify the notation here we will assume  $c = 1$  although the general case is similar: