Neural Functional Transformers
Authors: Allan Zhou
Kaien Yang
Yiding Jiang
Kaylee Burns
Winnie Xu
Samuel Sokota
J Zico Kolter
Chelsea Finn
Venue: NeurIPS 2023 poster (2023)
Keywords: equivariance
permutation
implicit neural representation
generalization
transformers
attention
TL;DR: Neural Functional Transformers (NFTs) are deep permutation-equivariant weight-space models based on attention mechanisms, enabling greater expressivity and improved performance in tasks such as learning latent representation of weights.

Abstract: The recent success of neural networks as implicit representation of data has driven growing interest in neural functionals: models that can process other neural networks as input by operating directly over their weight spaces. Nevertheless, constructing expressive and efficient neural functional architectures that can handle high-dimensional weight-space objects remains challenging. This paper uses the attention mechanism to define a novel set of permutation equivariant weight-space layers and composes them into deep equivariant models called neural functional Transformers (NFTs). NFTs respect weight-space permutation symmetries while incorporating the advantages of attention, which have exhibited remarkable success across multiple domains. In experiments processing the weights of feedforward MLPs and CNNs, we find that NFTs match or exceed the performance of prior weight-space methods. We also leverage NFTs to develop Inr2Array, a novel method for computing permutation invariant latent representations from the weights of implicit neural representations (INRs). Our proposed method improves INR classification accuracy by up to $+17\\%$ over existing methods. We provide an implementation of our layers at https://github.com/AllanYangZhou/nfn.

Neural Functional Transformers
Allan Zhou'
Kaien Yang',
Yiding Jiang? Kaylee Burns'
Winnie Xu
Samuel Sokota²
J. Zico Kolter²
Chelsea Finn'
Stanford University
Carnegie Mellon University
syz8cs.stanfors
Abstract
The recent success of neural networks as implicit representation of data has driven growing interest in neural functionals: models that can process other neural net- works as input by operating directly over their weight spaces. Nevertheless, con- structing expressive and efficient neural functional architectures that can handle nigh-dimensional weight-space objects remains challenging. This paper uses the attention mechanism define novel set of permutation equivariant weight-space layers and composes them into deep equivariant models called neural functional Transformers (NFTs). NFTss respect weight-space permutation symmetries while incorporating the advantages of attention, which have exhibited remarkable success across multiple domains. In experiments processing the weights of feedforward MLPs and CNNs, we find that NFTs match or exceed the performance of prior weight-space methods. We also leverage NFTs develop INR2ARRAY, novel method for computing permutation invariant latent representations from the weights of implicit neural representations (INRs). Our proposed method improves INR classification accuracy by up to +17% over existing methods. We provide an implementation of our layers at
com www.ooooollllaii cot..///o//aaaaaa///oo//iitii
Introduction
Deep neural networks have emerged as flexible modeling tools applicable to wide variety of different fields, ranging from natural language processing to vision to the natural sciences. The sub-field of implicit neural representations (INRs) has achieved significant success in using neural networks to represent data such as 3D surfaces or scenes 332 30 37). This has fueled interest in techniques that directly operate in weight space to modify or extract information from neural network through its weights or gradients. However, developing models that can process weight-space objects is challenging due to their high dimensional nature. As consequence, some existing methods for processing datasets of weights assume a restricted training process that reduces the effective weight space 6.
In contrast, we follow recent work in building permutation equivariant weight-space models called neural functionals, that can process neural network weight!] without such restrictions [3 46).
In particular, this work concerns neural functionals that are equivariant to permutations of the weights that correspond re-arranging neurons. These weight-space permutations, known neuron permutation symmetries, exactly preserve the network's behavior (see examples in Figure T]. Since equivariance neuron permutations is typically useful inductive bias for weight-space tasks, neural functionals achieve superior generalization compared to non-equivariant architectures. Despite this, their performance on weight-space tasks remains significantly worse than convolutional networks' performance on analogous image-space tasks [1 46). suggesting that neural functional architectures can be significantly improved.
For clarity, use "weights" (and "weight space") to describe the weights (and space they belong to) for the network being processed by a neural functional. We use "parameters" for the weights of the neural functional.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

While existing neural functionals rely on linear layers (analogous to convolution), some of the most successful architectures in other domains, such as Transformers 4 rely on nonlinear attention mechanisms. Motivated by this fact, our work develops novel equivariant weight-space layers based on attention. While naive self-attention between input weights is permutation equivariant, it cannot distinguish true neuron permutation symmetries (which preserve network behavior) from "false" permutation symmetrie (which do not preserve network behavior). In contrast, we prove that our weight-space self-attention layer is minimally equivariant: it is equivariant to only neuron permutations, which are the actual weight-space symmetries.
When composed into deep architectures, our weight-space attention layers give rise to neural func- tional transformers (NFTs). As an immediate application. we use NFTs to develop INR2ARRAY, a method for mapping INR weights into compact latent representations that are trained through a reconstruction-basee -based objective. INRRARRAY produces permutation-invarian latents that can be useful for downstream tasks such as INR classification. We also construct NFT's for variety of other tasks such as weight-space editing or generalization prediction.
NFTs are competitive with or outperform existing methods on each task without using more param- eters. Notably, NFTs and INR2ARRAY improve downstream INR classification accuracy over the best existing methods on multiple datasets, by up to +17%. Overall, our experiments show that attention- based neural functionals can be more expressive and powerful than existing architectures that rely on linear weight-space space layers, leading to improved performance across weight-space tasks.
2 Preliminaries
Consider a feedforward network having neurons at each layer i = 0, 0,... L. including the input and output layers. The network has weights W = W(') E R E [1..L]] } belonging to weight space, W. More generally, we are interested in multi-channel weight space features that can arise from stacking multiple weight space objects such as weights and gradients. For e-
channel weight-space features W E W we have W = { ((() E Rm;x", 1xc i [I..L]] } and
each w:( W j.k.: E IR is vector rather than scalar.
Since the neurons in each hidden layer i {1..... L L-1) 1} have no inherent ordering, the feedforward network is invariant to the symmetric group Sn, of permutations of the neurons in layer i. We follow Zhou et al. [46] in studying the neuron permutation (NP) group SNP == So x ... X Suy= which further assumes permutation symmetry of the input and output layers. Although this simplifying assumption is usually too strong, it can be effectively corrected in practice using position encodings at the input and output layers [46].
Consider a neuron permutation
(ao...., .1)) SNP- Each a; permutes the neurons of layer i, which correspond to the output space (rows) of Wi-1) and the input space (columns) of Wi). Hence the action on weight space features W E W denoted oW, where
[ow]s0
(i) jk o7 (1)
(0
(oo: oL) E SXP-
As illustrated in Figure] 1 a neuron permutation must always permute the columns of W() and the rows of (i-1) simultaneously in order to preserve network behavior.
This work is focused on principled architecture design for neural functionals, i.e., architectures that process weight space features [46]. Recent work has shown the benefit of incorporating neuron permutation symmetries into neural functional architectures by enforcing equivariance (or invariance) neuron permutation symmetries [31 46). Consider weight-space function classes parameterized by 0 E O. We say that a function class f We x θ - W is Sxp-equivariant if permuting the input by any E Sxp always has the same effect as permuting the output by a:
a s(W:0) = ((oW; 0),
Vo SNp, W E W 5 e.
(2)
Similarly, function class f W x θ - R is Syp-invariant if f(o = f(W:0) 8) for all o, W, and 0.
2 Morcover, naive self-attention hetween input weights is computationally intractable.
For simplicity, we omit the biases from discussion here, the Appendix for full detail.

---

W
0441
Figure I: Left: An illustration of how neuron permutations = (ao, "0a,""" ar) act on weight matrices W = (W((), W())). Each o, simultaneously permutes the rows and columns of adjacent weight matrices W W(+1), Right: Two examples of false symmetries, i.e., permutations that don't t preserve network behavior: (I) when т; permutes rows and columns of adjacent matrices permute differently and (2) when 7 permutes weights across layers.
Our weight-space layers build off of the dot-product attention mechanism, which takes a query y and a set N key-value pairs (kips Up) 1==1 and computes, in the single-heade headed case:
exp(a kip) ATTN (a. (kip up) ) Σer up Ep exp(y kip') P
If all q. kp, Up are vectors in Rd, then the output is also a vector in Rd The definition extends to the situation where k, Up are multi-dimensiona arrays of equal shape, by taking the dot product between these arrays the natural way. We can also use multi-headed attention ([41) without significant modification.
3 Neural Functional Transformers
We now introduce two attention-basee layers for processing weight-space features: weight-space self-attention (which is Sxp-equivariant) and weight-space cross-attention (which is Sxp-invariant). We then describe how to compose these layers into deep architectures called neural functional transformers (NFTs).
3.1 Equivariant weight-space self-attention
A key concept motivating the design of our weight-space self self-attention layer is the distinction between actual NP symmetries SNP that preserve network behavior, and other permutations of the weights that can generally change network behavior. For example, consider a Saim(W) that permutes the columns of Wi( differently from the rows of ((i11) or that moves weights between different weight matrices (see Figure []. Such permutations of the weight space are called "false symmetries" [31] since they generally modify network behavior. We are interested in equivariance the actual weight- space symmetries (neuron permutations) but not to any false symmetries, a property we call minimal cquivariance:
Definition 1. A function class f W x θ W is minimally Sxp-equivariant equivariant (. but not equivariant to any false symmetries. More precisely, for any T E &Saim(W) such that t & Syr. there exist a W E W and 0 E θ such that f(rW; 0) + F ([(W:
SNP-
Consider naive self-attention between the dim(W) weights (w(2) }. which would be equivariant any weight-space permutation. This would be Sxp-cquivariant, but would also be equivariant to any false permutation symmetries, meaning that it is not minimally equivariant. In other words, naive self-attention is overly constrained too many symmetries, including those that are not actually relevant to weight-space tasks.
Our weight-space self-attention layer SA(:; 0o, 0x, 0v) We W is more sensitive to the distinction between neuron permutations and false symmetries. operates on weight-space features with channels and is parameterized by query/key/value projection matrices lq 0x. By E Rexe

---

Attn
Weir
Attn
Input
K,Q,V
Output
Figure 2: An illustration of Eqs of our weight-space self-attention layer, the single channel case. The inputs to ATTN (Eq 3 are matrices with rows and columns corresponding to the sequence and feature dimensions, respectively. The final term (Eq 6, not shown here) simply computes self attention between all weights.
Each entry of the output feature is computed:
SA (W:0q. 2k. 0v 0v)
ATTN
Jil (K, V)
{
U{ (K, (K, 12 }
(4)
4=1
p=1
k
+ ATTN (1) (K, (K, V)
(K, (K, V) (6+1) n++1
(5)
ATTN (o()) [ (K, V)(+)
Vss.p. }): where
(6)
12 Sk 0 0
K = 0x w-()
v(21 Ol ((() jk
(7)
and the notation (K, V) shorthand for the tuple (K
v(() jk jk
Appendix A.2 defines the more
general version of this layer, which handles inputs with both weights and biases.
The final term (Eq. 6 is simply naive self-attention between the inputs { (V()) ER²), but the first two terms give the layer additional structure relevant to neuron permutations in particular, illustrated by Figure 2 To understand the symmetry properties of the first two terms, consider the top input to ATTN in Figure 2 which constructed from the rows (() and the columns of (ell++))) the rows of (()) and the columns of (+++1) are permuted simultaneously. the pairwise dot products between any two vectors in the input is preserved and the attention matrix is unchanged. Then the output ATTN is simply permuted, giving equivariance. On the other hand, independently permuting the rows and columns of (ff) and ((++1) will in general change the attention matrix, breaking equivariance. Since such independent permutations are false symmetries, this behavior is exactly what we need to achieve minimal equivariance.
jk
Since neuron permutations never permute weights between different layers Wi)) and W('') may additionally use learned layer position encodings ( Re i [1..L] prevent equivariance to such permutations:
LAYERENC (W: (1())) )
W(() (0)) ib
(1)
(
We can now compose self-attention with position encoding obtain SA LAYERENC. Theorem 3.1 (Minimal equivariance). The combined layer SA LAYERENC is minimally SNP- equivariant.

---

W
Project
Project
0g
W
0,
W
1)(()
Flatten
Output
A(Q.K)
Q
weight-space Features
A(Q. K)v*
An illustration of weight-space cross-attention layer, which pools weight-space features into a set of SNe-invariant vectors. The layer uses set of learned queries Q E R M xe to attend over keys and values produced from the weight space features.
Figure 3:
Proof (Sketch). To show equivariance: LAVERENC clearly equivariant to neuron permutations. We can check SA(oW) jk SA(W) o, "0,, by expanding the left hand side term-by. ferm-by-term using
the definition (Eq. 4 To show non-equivari to false symmetries: we can broadly classify false symmetries into three types and check non-equivariance to each of them. Appendix A. provides the full proof.
Practical considerations. The final term of SA (Eq. (6) amounts to self-attention between each weight-space feature w E I and requires O ((dim 3) operations, which is usually intractable.
In practice, we replace this term by self-attention between the row-column sums W E,, w(( of each weight matrix. This preserves interaction between weight-space features in different weight matrices, while reducing the computational complexity version also maintains minimal Sxp-equivariance. Meanwhile, consider the self-attention operation required to compute the first two terms of SA (depicted in Figure Z}. nı for i, then the first two terms require (L))cc) operations, which is feasible for moderate n but can become expensive when processing very wide networks. As result, weight-space attention layers are usually more computationally expensive than linear weight-space layers.
(I One can verify that this simplified
3.2 Invariant weight-space cross-attention
Stacking weight-space cef-attention layers produces equivariant weight-space features, but some situations require producing an Syp-invariant output. Existing approaches achieve invariance by
relying on a summation or maximization over the rows and columns of each input (ii) [31 46).
It is natural to extend these existing invariant layers by using attention over the input's entries, which has similar permutation invariant properties as summation and maximization. Our weight-space cross-atee cross-attention layer CA We Rd is simply cross attention between leareed query Rd and the set keys and values produced from weight-space features { (V(() E R:
CA (W; 0x 0v) ATTN (e, (OxWWa)) By W(()) vss.. ))
(9)
with 0K. dy E dx being learned projection matrices. By repeating this operation with M different learned embeddings e.. EMM depict the operation of the multi-embedding case in Figure
casily extend this into invariant map W + R Mxd We
3.3 Convolutional weight spaces
Although our description thus far focuses on fully-connected weight-space features W, we can also extend our layers convolutional weight spaces. Suppose W the e-channel weight-space feature corresponding ID convolutional network with filter widths k; each layer i. It contains matrices
5

---

W(O) E RM 1x4xxe, As in Zhou et al. [46] the filter dimension K can be folded into the channel dimension, creating a feature in R=,**; (x(e,e) with k,c channels. The problem is that k; may not be consistent across all i, while dot-product attention operations require that the channel dimensions for each w(( match. We solve this problem by choosing shared channel size č and using learned linear projections PROJ; I X (aexxnnnnxx_.. 1x((;c) + R" x 1xx to guarantee matching channel dimensions. The output of (-) can then be restored to the original channel dimension by another set of learned linear projections UNPROJ; [P4xx--1xE
3.4 Building Neural Functional Transformers
Following Transformer architecture design we form stackable "blocks" that combine weight- space self attention with LayerNorm, pointwise MLPs, and residual connections. Each block maps W + W and preserves -queeaiivreaance:
BLOCK(W) Z - MLP(LN(2))
(10)
Z = W + SA(LN(W)),
(11)
where both the MLP: : RC - RC and LayerNorm IN IR - RR operate "pointwise," or indepen- dently each imput w Since pointwise operations permutation equivariant, the overall block is an Sxp-equivariant map on W. To build neural functional Transformers (NFTs), we stack multiple blocks into a deep equivariant architecture. For tasks which require SNp Svp-invariance, like classification or learning invariant latent space of weights (Section 4). we can apply weight-space cross-attention on top of the features produced by the final BLOCK (.). The invariant activations produced by CA(-) can then be fed into an MLP that produces the final output.
Additionally, the NFT's input is often one or few channels, while our desired hidden channel dimension may be significantly larger (e.g., = 256). We can apply any function g RR + R to jk _ project from the input channel dimension to the hidden channel dimension while preserving equivariance. In our experiments, use random Fourier features 35 to achieve this projection.
each w
4 INRZARRAY: Learning an invariant latent space for weights
An application of interest for weight-space architectures is learn compact latent representation of weights, which can be useful for downstream tasks such as classifying the signal represented by an implicit neural representation (INR). The recently introduced INR2VEC (6] can learn a rep- resentation of INR weights by using a reconstruction objective, but requires every INR share the same initialization. In this section, we leverage neural functionals to extend the INR2VEC framework to the independent initialization setting, which is significantly more challenging but removes any restrictions on the INR training process.
INR2VEC uses an encoder-decoder setup to map INR weights into useful latent representations. We make two key changes the original INR2 VC formulation:
We implement the encoder with an Sxp-invariant neural functional transformer (NFT). This guarantees that the latent representation is invariant to neuron permutations of the weights. 2. We allow the latent representation to be a spatially meaningful array of vectors. In particular, each vector in the latent array is responsible for encoding single spatial patch of the INR.
The first modification is a useful inductive bias because the signal represented by an INR should be invariant to neuron permutations. The second modification is inspired by Spatial Functa 33 which found that giving INR latent spaces spatial structure was helpful in their meta-learning setting. We call our representation learning approach INR2ARRAY due to its modified latent space structure.
The following explanation focuses on 2D INRs that represent images on the coordinate grid [-1, 1j2 though the general case follows naturally. We first split the coordinate grid into M spatial patches P. U... J PM = 1|1, Given SIREN(-; W). INR2ARRAY uses an Sxp-invariant NFT encoder E-Co(-) to map weights W to latent array of M vectors, Mxd. The decoder DEC, .dd W
produces set of weights W; i = 1.... M }. one for each vector == Zi.: E Rdd, Each W,();)
is responsible for parameterizing the SIREN only for the spatial patch P. [-1, 1 The objective

---

MNIST
RashionMNIST
CIFAR-10
Original SIREN
5
1
/
3
In12AraayON
/
3
In12AraayFi
S
/
3
Figure 4: INR2ARRAY reconstructions for random samples from each INR dataset. The NFT encoder produces higher qualitatively better reconstructed SIRENs on each dataset, especially CIFAR-10. Blocking artifacts are due to the spatial latent structure, and suggest further room for improvement.
is to reconstruct the original INR's content using the decoded weights:
M [) = Σ 1=1 FE
SIREN
(2; (I:W::) SIREN (z; (c):
(12)
W. = (2))
2 ENCo(W).
(13)
Our INR2ARRAY encoder architecture multiple weight-space self-attention layers followed by weight-space cross-attention layer CA, which is especially well-suited for producing M x d permutation invariant arrays. For the decoder, we use the hypernetwork design of Sitzmann al. (37. where the hypernetworks are conditioned on the spatial latent 2; E Rd. In principle, one can also design an invariant encoder using (non-attentive) NEN layers, and we will compare both options in our experiments.
5 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 model inter- esting 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 image. Tasks in- cluding editing INR weights to modify the represented signal, or predicting target information related the signal (e.g., image class). For our INR experiments, we construct datasets for MNIST (25).
MNSHIonMNIST [2). 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 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-ar methods struggle 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 NFN 46 variants of the INR2ARRAY encoder, and distinguish the two variants by superscripts: INR2ARRAY NET and A where the hypernetworks are conditioned on the spatial latent z; E IRd. Appendix B.3 describes the implementation and training of each variant in full detail.
For the decoder, we use the hypernetwork design of Sitzmann et al. [7).
We train separate encoders for each INR dataset shown in Table of the appendix, which reports the reconstruction mean square error (Eq [2) on test INR inputs. The NFT encoder consistently achieves lower reconstruction better reconstructions than 1 NEN (Figure 4 confirming that NFTs enable higher quality encoding than previous neural functionals.
than the NFN encoder, and NNR2ARRAY also produces qualitatively
Once trained, the INR2ARRAY encoder maps INR weights 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 these Sye-invariant latent arrays z Xd Concretely, in K-way
7

---

Table 1: Test accuracy (%) for seight-space INR classification in the MNIST, FashionMNIST, and CIFAR-10 datasets. INR2ARRAY significantly improves over current state-of-the-ar 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 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-I0
INC2 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
Spatial NFN
92.9 + 0.46
70.8 0.53
45.6
NFT
89.9 + 1.03
72.7 0.05
4..2
TIR2ARRAY'S NFN NFN (Ours)
94.6 + 0.00
76.7 + 0.00
45.4 0.00
INR2ARRAY NET (Ours)
98.5 + 0.00
79.3 + 0.00
63.4 0.00
classification the trained encoder ENC: W + R Mxd can be composed with any classification head f: : IR M xd + R to form an Snv-invariant classifier f ENC : W + RK 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 [11. 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 eermutation-ivvaiinn architectures that directly map input weights predicted labels. We also create a modified NFN classifier (Spatial NFN) which produces an M x d array of activations before a convolutional classifier head, test the impact of spatial latents independent of the INR2ARRAY training process. Appendix B B.4 describes the setup for each method in further detail. Finally, we measure the performance of NFT classifiers without INR2ARRAY.
Table D 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, INR2 ARAY achieves 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 MNIST test accuracy of 98.5%, up from 92.9% by NFNs and 95%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 INR ARRAA AY in addition to the architecture. Finally, the Spatial NFN fails to improve performance over standard NFN classifiers, implying that the particular objective of INR2A 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. 66.) The first two terms SA, illustrated by Figure 2 amount to self self-attention between the rows and columns of adjacent weights, while the third term is a global self- ooff-attention between features from all layers. Note that in practice we use a tractable imatoni 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 INR2ARRAY?S NF- when ablating the terms of SA (
Model
MNIST
FullSA
98.5
AdjacentSA
97.6
GlobalSA
37.4
Table 2 shows the results of this ablation experiment on INR2ARRAY performance in the MNIST classification task. We see that our full layer (FullSA) performs best, ablating the final term (AdjacentSA) only slightly degrades performance, and ablating 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
8

---

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.
NFN 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 т for predicting the generaliza- tion of CNN image classifiers with unseen hyperparameters trained CIFAR- and SVHN-GS (GS=grayscalee). The NFT outperforms NFNNP and hand-picked features (STATNN), while NEN performs best overall. Uncertainties indicate standard error over two runs.
HFN 446
NFNNr 46
STATNN 40
NET (Ours)
CIFAR-IO-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 em- phasize 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 a difference W W + A(W). Neural functionals are trained learn A--) that achieves some desired visual transformation, such as dilating the image. Permutation equivariance is useful inductive bias here since if (WW) is the desired edit for W. then for any neuron permutation the desired edit to oW is (WW). 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 MNSS gradient. Gradient tasks roughly anount to edge detection; Figure in the appendix shows sample inputs and targets for each editing task.
We compare NFT editing performance against the two NEN variants [46,, with each method using the similar number of parameters (~ 7M). For full training details, 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 most INR editing tasks. In addition, Table in the appendix shows that NFTs generally obtain both lower training and test error, indicating that the attention- 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 those belonging to convolutional neural network classifiers. Large-scale empirical studies have produced datasets of trained classifiers under different hyperparameter settings [40] [T enabling 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
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. addition to comparing against the two neural functional variants NENNP and NFNHNP [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. describes the experimental setup in full detail.
Table 4 shows the test performance of each method using the Kendall rank correlation coefficient 20 Across both datasets, neural functionals operating on raw weights outperform the hand- designed features baseline STATNN. NFTs slightly outperform NFNNp on each dataset, while NENNNP achieves the best performance overall. Note that both NFNsP 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.
6 Related Work
It is well known that the weight spaces of neural networks contain numerous symmetries, i.e., transformations that preserve the network's behavior 177. Permutation symmetries particular have been studied in the context of neural network loss landscapes [14] 12 and weight-space merging [38] []. However, most prior methods for processing weight space objects do not explicitly 211, although some have tried to encourage permutation equivariance through data augmentation 33 299] 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 [[]] 3). Instead, study the problem of encoding permutation symmetries into the neural functional itself, without any restrictions on the weight space being processed.
account for these symmetries 12 26 15
24
45
7
There are a variety of methods that incorporate symmetries into deep neural network architectures B6 22
Examples include (group) convolutions for images (25 and permutation equivariant
architectures for general set-structured inputs [34 44 [6, 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.
7 Conclusion
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 permutation- cquivariant weight-space models using nonlinear self-attention and cross-attention layers, which can enable greater expressivity compared existing linear layers.
We empirically evaluate the effectiveness NFTs on weight- 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, method for mapping INR weights to compact and Sxp-invariant latent representations, which significantly improve performance in downstream tasks such 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 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.
Acknowledgments and Disclosure of Funding
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 000014-22-1-26211
10

---

References
[1] K. Ainsworth, J. Hayase, and S. Srinivasa. Git re-basin: Merging models modulo permutation symmetries. arXiv preprint wrXiv:2209.48833 2022.
[2] 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.
[3] M. Bauer, E. Dupont, A. Brock, D. Rosenbaum, J. Schwarz, and H. Kim. Spatial functa: Scaling functa to imagenet classification and generation. arXiv preprint arXiv:2302.03 2023.
[4] 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 190.021907.02911,220 2019.
[5] T. Cohen and M. Welling. Group equivariant convolutional networks. In International confer- ence on machine learning, pages 2990-2900200 2999.. PMLR, 2016.
[6] L.. De Luigi, A. Cardace, R. Spezialetti, 7. Ramirez, Salti, and Di Stefano. Deep learning on implicit neural representations of shapes. arXiv preprint arXiv:2302.054553 2023.
[7] L. Deutsch, E. Nijkamp, and Y. Yang. A generative model for sampling high-performance and diverse weights for neural networks. arXiv preprint arXiv: 028998 2019.
[8] A. Dosovitskiy, L.. Beyer, A. Kolesnikov, D. Weissenbomn X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010. 1999,, 2020.
[9] E. Dupont, Y. W. Teh, and A. Doucet. Generative models as distributions of functions. arXiv preprint
2021.
[10] E. Dupont, H. Kim, S. Eslami, D. Rezende, and D.. Rosenbaum. From data functa: Your data point is function and you should treat it like one. arXiv preprint arXiv:2201.12200 2022.
[11] G. Eilertsen, D. Jönsson, T. Ropinski, J. Unger, and A. Ynnerman. Classifying the classifier: dissecting the weight space of neural networks. arXiv preprint arXiv:2002.05668. 2020.
[12] R. Entezari, Sedghi, Saukh, and B. Neyshabur. The role of permutation invariance in linear mode connectivity of neural networks. arXiv preprint arXiv:2 10.0666,, 2021.
[13] M. Finzi, M. Welling, and A. G. Wilson. A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups International Conference on Machine Learning, pages 3318-3328. PMLR, 2021.
[14] T. Garipov, P. Izmailov, D. Podoprikhin, D. P. Vetrov, and A. G. Wilson. Loss surfaces, mode connectivity, and fast ensembling DNNs. Advances neural information processing systems, 31, 2018.
[15] D. Ha, A. Dai, and Q. Le. Hypernetorrss. arXiv preprint arXiv:1609.09100 2016.
[16] 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.
[17] R. Hecht-Nielsen. On the algebraic structure of feedforward network weight spaces. In Advanced Neural Computers, pages 129-135. Elsevier, 1990.
[18] Y. Jiang, D. Krishnan, H. Mobahi, and Bengio. Predicting the generalization gap in deep networks with margin distributions. In International Conference on Learning Representations, 2019. URL httphhts:///pereee////e
[19] Y. Jiang, P. Natekar, M. Sharma, S. K. Aithal, D. Kashyap, Subramanyam, C. Lassance, D. M. Roy, G. K. Dziugaite, S. Gunasekar, et al. Methods and analysis of the first competition in predicting generalization deep learning. In NewrlPS 2020 Competition and Demonstration Track, pages 170-190. PMLR, 2021.

---

[20] M. G. Kendall. A new measure of rank correlation. Biometrika, 10(1/2) 1938.
[21] B. Knyazev, M. Drozdzal, G. W. Taylor, and A. Romero Soriano. Parameter prediction for unseen deep architectures. Advances Neural Information Processing Systems, 34:29433- 29448, 2021.
[22] R. Kondor and 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.
[23] A. Krizhevsky, G. Hinton, et al. Learning multiple layers of features from tiny images. 2009.
[24] D. Krueger, C.-W. Huang, R. Islam, R. Tumer, A. Lacoste, and A. Courville. Bayesian hypernetworks. arXiv preprint 171.:4171 2017.
[25] 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.
[26] K. Li and Malik. Learning optimize. arXiv preprint wrXiv:1606.01885 2016.
[27] H. Maron, O. Litany, G. Chechik, and E. Fetaya. On learning sets of symmetric elements. In International conference machine learning, pages 6734-6744, PMLR, 2020.
[28] C. H. Martin and M. W. Mahoney, Implicit reffraegularizationn in deep neural networks: Evidence from random matrix theory and implications for learning. The Journal of Machine Learning Research, 22(1):747--7511, 2021.
[29] L. Metz, J. Harrison, C. D. Freeman, A. Merchant, L. Beyer, Bradbury, N. Agrawal, B. Poole, I. Mordatch, A. Roberts, et al. Velo: Training versatile learned optimizers by scaling up. arXiv 2022.
preprint arXiv:2211.09770
[30] 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.
[31] 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. 2023.
[32] 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.
[33] 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.
[34] C. Qi, H. Su, K. Mo, and L. Guibas. Pointnet: deep learning on point sets for 3d classification and segmentation. evpr (2017). arXiv preprint arXiv: 1rXiv:1612.00599 2016.
[35] A. Rahimi and B. Recht. Random features for large-scale kernel machines. Advances in neural information processing systems, 20, 2007.
[36] S. Ravanbakhsh, J. Schneider, and B. Poezos. Equivariance through parameter-sharing. In International conference on machine learning, pages 2892-2901. PMLR, 2017.
[37] V. Sitzmann, J. Martel, A. Bergman, D. Lindell, and G. Wetzstein. Implicit neural representations with periodic activation functions. Advances Neural Information Processing Systems, 33: 7462-7473, 2020.
[38] N. Tatro, P.-Y. Chen, P. Das, I. Melnyk, P. Sattigeri, and R. Lai. Optimizing mode connectivity via neuron alignment. Advances in Neural Information Processing Systems, 33:1 53:15300-15311, 2020.
[39] F. H. Thiede, T. Hy, and Kondor. The general theory of permutation equivarant neural networks and higher order graph variational encoders, arXiv preprint arXiv:2004.03999 2020.
12

---

[40] T. Unterthiner, D. Keysers, S. Gelly, O. Bousquet, and Tolstikhin. Predicting neural network accuracy from weights. arXiv preprint 1rXiv:2002.11448, 2020.
[41] A. Vaswani, Shazeer, N. Parmar, 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. Pashion-mnis novel image dataset for benchmarking machine learning algorithms, 2017.
[43] S. Yak, J. Gonzalvo, and H. Mazzawi. Towards task and architeeturr--iddeeeneent generalization gap predictors, arXiv preprint wrXiv:1906.015500 2019.
[44] M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. Salakhutdinov, and A. Smola. Deep sets. doi: 10.48550. arXiv preprint ARXIV//703.06/ 14, 2017.
[45] C. Zhang, M. Ren, and R. Urtasun. Graph hypernetorrks for neural architecture search. arXiv preprint arXiv:1810.057493 2018.
[46] A. Zhou, K. Yang, K. Burns, Y. Jiang, S. Sokota, Z. Kolter, and C. Finn. Permutation equivariant neural functionals, arXiv preprint arXiv:2302. 14040, 2023.
13

---

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 LAYERENC: W W
Theorem 3.1 (Minimal equivariance). The combined layer SA LAYERENC is minimally Sup- equivariant.
Proof. We first show Sup-equivariance, before checking that it minimally equivariant.
Equivariance. is straightforward to check that LAvERENC equivariant to any permutation that preserves the layer index (which neuron permutations do). At high level, we check that SA is Sxp-equivariant by verifying that SA(oW)
SA(W) (1) We proceed term by term.
First, it will be useful to show that ATTN (Eq. query, and value input. In particular, let [Tq] q( (i)) permute the entries of the input vectors:
is equivariant permutations E Sa of each key,
exply kp)
ATTN (mg, (nkpp Tpp )=1)) [i]
Σ
(TTU)
[1]
(14)
Ep expla kz))
expla kip)
(15)
En exply kip))
ATTN(4, (Kips up) (p) (1)]
(16)
Note that in the first equality we use æq- Tk q
Since the key, query, and value projections are only linear projections along the channel dimension, K(oW)= = a K (W), where K denotes all keys produced by linearly projection of W by 0x (and similarly for Q, V). Then consider, for example, how the first term (Eq 4 changes under W + oW:
Attn (o0%% ((oK,or) U(iok,or)5)) [K] = (17)
Attn (0-1Q07'0).0 ((o- oz
, ( } U of o. v(0)
[K]
(18)
Note the distinction between neuron permutations of the weight space [aQ] and permutations of individual vectors like F=-1Qp:. and that permutations along the KV set indices (e.g., in the first set) are ignored since set order irrelevant. Using the equivariance of ATTN (Eqs. 16]:
ATTN ( U((K
((K,
(o (A)1
(19)
Notice that compared to Eq. 4 the output indices have changed j + o. ' and k + - (l)) showing that the first term is cquivariant. We can repeat this argument 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 equivariance follows immediately from the permutation equivariance of ATTN. This verifies that all terms of SA are Sxr-equivariant.
Minimal equivariance. Here our goal is to show non-equivariance permutations Sai((W but 7 & SNP (false symmetries). That is, we should find some input W such that SA(LAYEREN((WW) + +SA(LAYERENC(MW) for some parameters {0g, 0k, Ov and ()()) }.
will be useful define the actions of permutations in Sai((W) on the indices of the weights, rather than the weights themselves. The index space for a weight space W is defined as set of 3-tuples representing layer number, row, and column, respectively:
= { (i.j.)) [1.4], [1..m ::k k [L.m4-1]] }.
(20)
14

---

For I, we use the subscripts É, r, to denote the layer, row, and column indices respectively. That is, a (ay, ac) kk).
Notice that the set of arbitrary permutations Sam((w) simply the set of all bijections r : I + I (any index be arbitrarily permuted to any other index), while SNp is the subgroup of bijections that can be written as:
((.j,k) = (atk))=(;,0000) (1,0(()),; V(t,,,k) E
(21)
False symmetries can broadly be categorized into three groups. In fact, verify that any permutation 7 that fails to satisfy one of these three cases must SNP:
1. False symmetry that permutes weights across layers. There exists an (i.jk) E II such that ri, (.,jk)), k)2 *
2. False symmetry where row and column permutations are not independent. Suppose false symmetry T does not fall under case (1), i.c., 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 such that rii.j, k) = (i, J', k'), ). but J q), + jr Or there an (i, j, k) and such that 7(i, j. k) (i.3", k') but F(2, p)k # k'.
k)) a false symmetry T is not one of the two cases above,
3. Adjacent permutations decoupled.
J', k') but 7(i - 1, 1, q)r # k' That is, an instance where one row pair of layer i do not permute the same way as the columns of layer
there must be an (i,j. k) II and g such that (i..j,
We can now check non-equivariance for each case. To simplify the notation here we will assume I although the general case is similar:
1. If 7 does not preserve layer index, LAYERENC will be non-equivariant since each layer has a different encoding.
2. Suppose there is an j. k)) and q such that 7(i, j. k-) (i,J', ".J'',' K"), but (().).))- * j', and let fq = 0x = By = I. Now consider the input W E W such that all weights are except for w
jg 1. Then we can check that +.SA(W). * SA(TW))s's w so the layer is not equivariant т. We can construct similar argument for the other case, where there is an (i. j, k) and such that j(i,j, k) = (6,j', k') but 7 (i. p. k))
jk (
K.
3. this case there must be an (i, k) Tl and q such that (i,, j. k) = (i. k) but 7(i
1, k, 9)r * k'. Again let 0q 0g By = I and consider W E W such that all weights
are except for ((() = W (1-1) 1. Then, similar to the above case, we can verify that [TSA(W))ejrA) o BA(TWW).''ke so the layer is not equivariant to such T.
kg
Taken together, the three cases above show that although SA LAVERENC exp-equivariant, is not equivariant to any other permutations 7 & Sxp. Hence the layer is minimally equivariant.
A.2 Full description including biases
Let be the space of biases, and 24 = W the full space of both weights and biases. Then the full weight-space self-attention layer SA 2 Lt is defined, for input U (W, b):
SA(U; 0x, 0v) -
(22)
15

---

where Y(U) E W and z(U) E B" with:
Y(W, 61.66) ATTN KV,, + ATTN (0....V;, + ATTN (05x-VVs)
(23)
=W, 1),
ATTN
AATN ( KV3))
where
(24)
KV; (K (K.
{
(K.
U{ (ke, e)(6-1))
(25)
KV2 = { (K
U (K, V)
11+1)
p=1 (k: (H)
(26)
KV. = { V) Vs,p, 4 U (k, 0)
[vs.p and
(27)
20 8 W K 0x
w ==
V O.
(28)
jk
jk
4j We use KV,, 2V2 VV3 to denote the three sets of key-value pairs involved in the layer computation and distinguish the linear projections of weights and biases using uppercase (Q, K, V) and lowercase (q, k. v), respectively.
(8) 00b,)
Ovb,"
(29)
B Experiment details
Recall from Section 3.4 that NFTs consist of "blocks" that include weight-space self attention, layer normalization, and feedforward MLPs. Additionally, the first layer typically a projection from the input channel dimensionality to the hidden channel dimensionality, for which we use random Fourier features.
Hyperparameters blocks
Predicting gen.
Editing INRs
INRZARRAY
4
6
# channels
256
256
256
MLP hidden dim
1024
512
1024
Fourier scale
3
Fourier size
128
128
128
attn heads
8
dropout p
0.1
0.1
0
invariant layer
HNP
CA
CA M
16
CA dim Optimizer
256
Adam
Adam
AdamW
Learning rate
0.001
0.001
0.0001
Weight decay coeff
0
0.01
LR warmup steps
10K
10K
10K
Total params
46M
7M
22M
Training steps
75K
50K
200K
Training time
5H
2H
13H
Here we summarize the hyperparameters (and other training information) for the NFTs in each experiment in this paper. NFTs are described by the following hyperparameters:
# blocks: Number of equivariant weight-space blocks BLOCK(-)
# channels: channel dimension used by the blocks.
MLP hidden dim: the hidden layer size of the feedforward MLPs within each block
Fourier scale: standard deviation of random vectors used to compute random Fourier features
Fourier size: number of random fourier features
# attn heads: Number of attention heads used by weight-space self-attention.
dropout p: Dropout probability for attention matrix and in feedforward networks.
invariant layer (for Sxp-invariant NFTs): Whether the invariant pooling is done with sum-
mation (HNP/NP (46)) cross-altention (ours).
CA M: If using cross attention, the number of learned query embeddings e1, M
CA dim: Dimension vectors going into cross attention.
16

---

MNIST (Erode)
MNIST (Dilate)
MNIST (Gradient)
Ground truth
3
8
4
S
4
9
6
8
6A
Original INR
3
8
4
S
4
9
6
8
2
1
MLPPEN
3
8
4
S
4
9
6
8
2
0
6
NENsr
3
8
4
S
4
9
8
2
0
NFT
3
8
4
S
4
9
8
Fashion-MNIST (Gradient)
CIFAR (Contrast)
Ground truth
Toa
Original INR
MLPNIP
NENN
Figure 5: Samples of target image transformations, as well as original and edited INRs for each editing task in T Despite the fact that the NFT quantitatively outperforms the NFN and NFNNp, compelling qualitative differences in the edited images can be difficult to identify due to small image size.
B.1
Generalization prediction
Each Small CNN Zoo dataset that consider (CIFAR-10-GS and SVHN-GS) contains the weights of 30000 CNNs trained with varying hyperparameters. The weights are split into 15000 for training and validation and the rest for testing. We use 20% of the non-test data for validation, and train the remaining 12000. We use the validation data to sweep over Fourier scales of 3 or 10 and try both HNP (summation) CA (cross attention) layers to achieve invariance.
B.2
Style editing
We implement four visual transformations: erode reduces the size of objects in an image by removing pixels on the boundary; dilate adds pixels to object boundaries and can be thought of as the opposite of erode; morphological gradient performs edge detection by computing the difference between dilation and erosion; and contrast magnifies color differences in an image, making objects more distinguishable. Figure shows visual samples of the input and target for each task.
B.3 INR2ARRAY
We split the SIRENs of each dataset into training, validation, and testing sets, and train the encoder and decoder using only the training set. We use the validation error sweep over the following hyperparameters (for both the NFT and NFN variants): # blocks (4 vs 6), # channels (256 vs 512). MLP hidden dim (512 vs 1024), Fourier scale (3 vs 10), attn heads vs 8), and dropout (0.1 vs 0). After sweeping these hyperparameters, we use the best hyperparameter configuration to train the encoder and decoder and do early stopping with the validation error. The encoder portion the best checkpoint can then be used to produce latent array representations of INR inputs.
17

---

Table 5: Mean squared error to ground truth visual transformation for five INR editing tasks. NFTs generally achieve lower both lower training and test error than NFN variants without using more parameters, suggesting greater expressivity from attention. Uncertainties indicate standard error over three seeds.
HNN
46
NFN 46
NFT (Ours)
Train
0.0217 = 0.0004
0.0212 + 0.0002
0.0195 + 0.0004
MNIST (erode)
Test
0.0225 = 0.0001
0.0223 + 0.0000
0.0194 = 0.0002
Train
0.0575 = 0.0010
0.0671 0.0017
0.0473 0.0006
MNIST (dilate)
Test
0.0656 = 0.0000
0.0693 + 0.0009
0.0510 = 0.0004
Train
0.0536 = 0.0020
0.0547 = 0.0009
0.0474 = 0.0005
MNIST (gradient)
Test
0.0552 = 0.0003
0.0566 + 0.0000
0.0484 = 0.0007
Train
0.0851 + 0.0012
0.0853 + 0.0008
0.0795 + 0.0009
FashionMNIST (gradient)
Test
0.0859 0.0001
0.0870 + 0.0001
0.0800 0.0002
Train
0.0191 + 0.0002
0.0185 0.0002
0.0192 0.0005
CIFAR (contrast)
Test 0.0204 + 0.0000
0.0203 + 0.0000
0.0200 = 0.0002
Table 6: Mean squared error (MSE) of INR2ARRAY reconstructions test INR weights. Using an NFT encoder in INR2ARRAY achieves better reconstruction error, produces higher quality samples, and leads better downstream performance.
INR2ARRAY NFN
INR2ARRAY
MNIST
0.046 + 0.001
0.027 0.003
MNSSI
0.085 + 0.006
0.070 + 0.006
CIFAR-10
0.117 + 0.008
0.036 + 0.006
B.4 INR Classification
After training INR2ARRAY on given dataset, we select the checkpoint with the smallest reconstruc- tion еrror the validation set, and use the encoder produce a 16 x 256 latent array representation for any input SIREN. We view the latent as a length-16 sequence of 256-dim vectors, and project each vector into R before feeding the array into Transformer encoder with 12 blocks. This produces an output 2 E IR 16x512, we take the first output =8 E R512 and feed it into a 2-layer MLP with 512 hidden units that produces the classification logits. We train with cross-entropy loss for 100 epochs, using the Adam W optimizer and mixup augmentation. Using the validation accuracy, we sweep over weight decay (0.1 vs 0.01) and mixup alpha (0 vs 0.2 vs 1) and perform early stopping.
Table 7: Classification train and test accuracies (%) on datasets of image INRs (MNIST, FashionM- NIST, and CIFAR-10), 10). INR2ARRAY outperforms prior state-of-the-ard results in this setting (DWS and NFN), For the NFN baseline [46) we report the higher performance out of their NP and HNP variants. Uncertainties indicate standard error over three runs.
MNIST
MNSHI
CIFAR-10
Train
Test
Train
Test
Train
Test
DWS
R 46
85.7 1 0.57
65.5 1 0.48
NEN
96.2 0.24
92.9 + 0.38
84.1 + 1.02
75.6 + 1.07
61.0 + 6.60
46.6 + 0.13
S Spatial NEN (Ours)
95.3 + 0.59
92.9 -: 0.46
76.4
70.8
+: 0.53
52.3 +
0.30 0.30
15.6
INR2ARRAY NPN (Ours)
100 + 0.00
94.6 - 0.00
92.8 + 0.00
76.7
+ 0.00
98.5 4.00
45.4 0.00
INR2ARRAY NPT (Ours)
100 0.00
98.5 0.00
97.6 0.00
79.3 0.00
100
6.04
18


Reviews:
[{"review_id": "zokddScBhv", "reviewer": "Reviewer_BGFR", "summary": "The authors consider the principled design of neural functionals, i.e. neural\nnetworks that take other neural networks as inputs. Example tasks include\npredicting the generalisation error of a trained network for a given set of\nhyperparameters, or trying to reconstruct the images or 3D object encoded in an\nimplicit neural representation (INR, e.g. Mildenhall et al. arXiv:2003.08934).\n\nThe authors propose a transformer-style architecture that consists of a stack of\ntransformer blocks, each one consisting of a self-attention mechanism followed\nby an two-layer fully-connected neural net. The key challenge is to design a\nself-attention mechanism that respects the various symmetries of that preserve\nthe network's behaviour. Naïve self-attention is equivariant to weight\npermutations that change network outputs. Previous work (Navon et al. [31] and\nZhou et al. [46]) designed linear layers that are equivariant to neuron\npermutation symmetries. \n\nThe key contribution of this paper is to propose a particular form of the\nself-attention mechanism for neural functionals. It really is a combination of\nthree attention mechanisms (eq. 4-6), which has however to be approximated in a\nrather heavy way: from self-attention between weight-space features to\nself-attention between row-column *sums* of the weight matrices. The proposed\nmechanism works for both fully-connected layers and convolutional layers.\n\nGiven that weight space tasks remain computationally challenging, the authors\nalso propose an embedding method that maps implicit neural representations into\narrays, an extension of the inr2vec method of De Luigi et\nal. [arXiv:2302.05438]. Its performance compares favourably to other methods\nwhen tested on weight-space INR classification of standard image data sets,\nwhere we seek to reconstruct the image encoded in the weights of an INR.", "questions": "The authors introduce a neural functional architecture that is transformer-based\nand can deal with fully-connected and convolutional networks. An obvious\nquestion is whether this approach can be extended to transformers? I would\nappreciate some discussion of this issue from the authors.\n", "limitations": "The authors discuss computational limitations of the vanilla form of their self-attention layer, and discuss computationally feasible alternatives.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- The authors propose a well-designed architecture for a neural functionals.\n- The authors show that the self-attention has several desirable properties,\n  even in its computationally simplified form.\n- The proposed method works well on some tasks chosen by the authors.\n\nAs I have no practical experience with neural functionals, I may have missed\nsome relevant related work or may misjudge the novelty of the proposed solution,\nhence my low confidence.", "weaknesses": "I believe the article could benefit from some discussion of key concepts regarding neural functionals, for\nexample the type of tasks in standard benchmarks, to make\nthe article easier to read on its own.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your feedback and thoughtful review.\n\n> I believe the article could benefit from some discussion of key concepts regarding neural functionals, for example the type of tasks in standard benchmarks, to make the article easier to read on its own.\n\nWe will expand the Experiments (Section 5) to elaborate on the types of weight-space tasks one would solve when using neural functionals. We would also be happy to elaborate, both on openreview and in the manuscript, on any other specific topics that were insufficiently explained.\n\n> An obvious question is whether this approach can be extended to transformers?\n\nThis is a good idea--Transformer weight spaces also contain permutation symmetries, for example because they contain feedforward MLP blocks. So it should be feasible to extend NFTs (and neural functionals in general) to create minimally equivariant layers for processing Transformer weights.\n\nThis extension is not trivial or immediately obvious, since Transformers’ permutation symmetries are more complex than those of feedforward MLPs/CNNs–for example, because they contain residual connections that tie together neurons from different layers and prevent them from being independently permuted. We believe that extending NFTs and neural functionals to non-feedforward architectures is a very useful and important future direction for making this framework more broadly applicable."}, {"review_id": "ptv9baQ5pc", "reviewer": "Reviewer_e4bE", "summary": "This paper aims to construct neural functional architectures that can handle high-dimensional weight-space objects. Specifically, this paper designs a novel class of weight-space models with neuron permutation symmetries using nonlinear self-attention and cross-attention layers. The proposed method is called neural functional transformers (NFTs). NFTs enable greater expressivity and higher classification accuracy on implicit neural representations (INRs) over existing methods.", "questions": "See \"Weaknesses\".", "limitations": "This paper has no potential negative societal impact.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. The overall written is good, and the authors provide code for reproduction.\n\n2. The topic of this paper is cutting-edge and the proposed method is novel.\n\n3. The proposed method is simple and its performance for INR classification is good. ", "weaknesses": "1. It seems straightforward to make self-attention permutation equivariant. What is the actual contribution of section 3.1? Do you compare the performance of original attention and your modified attention?\n\n2. Equation (4) (5) (6) are difficult to read. The authors may represent them via a better way.\n\n3. In table 3, the performance is worse than NFN which is inconsistent with table 1 and table 2. Can you explain the reason of this?\n\n4. Do you evaluate your method on 3D dataset?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your suggestions and insightful questions.\n\n> It seems straightforward to make self-attention permutation equivariant. What is the actual contribution of section 3.1? Do you compare the performance of original attention and your modified attention?\n\nThe core of our contribution is that naively applying self-attention is permutation equivariant, but not _minimally_ so. If we treat each $W^{(i)}\\_{jk} \\in \\mathbb{R}^c$ as a distinct input vector, then self attention on these $\\dim(\\mathcal{W})$ inputs would be equivariant to _any_ permutation of the weights, including permutations that are not true symmetries of the weight space (see examples in Figure 1). This means the layer would be overly constrained by symmetries that are not relevant to weight space tasks. On the other hand, the layer we define (Eqs 4-7) is minimally equivariant, which is the core theoretical contribution of Section 3.1.\n\nTo confirm the importance of minimal equivariance, we have added ablation experiments where we remove terms from the definition of our equivariant self-attention layer (see top-level reply).\n\nIn these ablations _GlobalSA_ is basically a naive application of self attention between features of every weight matrix, and it performs very poorly. This shows that the particular definition of our layer, especially the first two terms which give minimal equivariance, is important for good performance in practice.\n\n> Equation (4) (5) (6) are difficult to read. The authors may represent them via a better way.\n\nWe have created an illustration of Eqs 4-6, please see our attached PDF response. This should help illuminate the shapes of the inputs, outputs, and intermediate computations. We will add this figure to our paper.\n\n> In table 3, the performance is worse than NFN which is inconsistent with table 1 and table 2. Can you explain the reason of this?\n\nThis is an interesting question, and could be due to a number of factors that separate the small CNN zoo from, e.g., the INR classification tasks. The NFN performance on the small CNN zoo benchmark (table 3) is already quite high (state of the art on the benchmark). Additionally, the size of the training dataset (12,000 training samples) is significantly smaller than for the INR datasets (45,000 training samples, plus data augmentation). So it may be that NFTs only benefit at greater scales and data quantities. We will add this to our paper's discussion of the experimental results.\n\n> Do you evaluate your method on 3D dataset?\n\nAlthough Inr2Array could be extended to learn a 3D spatial latent encoding each 3D INR, our experiments focused on 2D INR benchmarks. From private correspondence with the authors of [1], it turns out that NFNs already perform quite well on 3D shape INR classification benchmarks, e.g. where each INR encodes a signed distance field for a ShapeNet-10 shape. This is in contrast with 2D INRs where, for example, NFNs classify CIFAR INRs with _much_ lower accuracy than standard CNN baselines do on vanilla CIFAR.\n\nSo we focused on 2D INRs since they are, surprisingly, more challenging in the weight space. This aligns with findings from Functa and Spatial Functa [2,3], which also found that classifying 2D INRs was actually more challenging than classifying 3D INRs.\n\n[1] Zhou et al. Permutation Equivariant Neural Functionals.\n\n[2] Dupont et al. From data to functa: Your data point is a function and you can treat it like one.\n\n[3] Bauer et al. Spatial Functa: Scaling Functa to ImageNet Classification and Generation."}, {"review_id": "bndzigShKL", "reviewer": "Reviewer_zT4L", "summary": "This paper introduces a new architecture called Neural Functional Transformer (NFT), that uses an attention-based layer that is equivariant to neuron permutations. NFT is an architecture that is used for tasks where the model is taking an input another neural network's weights (in their case, either an MLP or CNN), which may be e.g. an INR of a some signal or the weights of another trained model. They benchmark their architecture on two INR tasks (image classification, image editing) and predicting classifier generalization from weights.", "questions": "- I think there's a typo in the caption of Figure 1 in the last sentence.\n- The weight spaces handled by the proposed layer are those of CNNs and MLPs. Do you think there are ways to extend this to other commonly used architectures like transformers?\n- Neuron permutations come from the fact that MLPs have layers composed like $W_1\\sigma(W_2x)$ where $\\sigma$ is a point-wise nonlinearity. From this, given a permutation matrix $P$, we get $W_1P^T\\sigma(PW_2x) = W_1\\sigma(W_2x)$ because $\\sigma(PW) = P\\sigma(W)$ for any matrix $W$. Is this property satisfied for matrices $P$ that are not permutation matrices? If so, we would also want equivariance to those transformations of the weight space. \n- I found Section 3.3 slightly unclear to read, particularly lines 155-158.\n- Where is the layer encoding in section 3.4?\n- In Section 5.2, do you train two INRs with the same initialization on the the original and edited images and then take their weight difference as the target for the NFT?\n- Typo in \"amount\" in line 266.\n- In section 5.3, define the abbreviation HNP for clarify (I think only NP was defined in the beginning of the paper).\n", "limitations": "Adequately addressed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "The paper presents the problem of building architectures that operate on other network's weights, which is significant and has useful applications to e.g. working with INRs. The architecture they propose based on attention is original, interesting and clearly presented. The writing of the paper is clear throughout. Their experiments are relevant and explained well. ", "weaknesses": "Although there are no major weaknesses in this work, I found the empirical performance of NFT slightly underwhelming. \n- For CIFAR-10 classification, why isn't there a comparison to Spatial Functa (or usual functa) (which was the inspiration for using spatially structured latent representations)? As shown in Spatial Functa's table 2, they achieve 90% top-1 accuracy, which is significantly higher than 63% achieved by NFT. Both approaches are for image classification using INRs for the images, hence a comparison seems appropriate. Regarding using the same initialization for different INRs, rather than being a weakness, wouldn't that be the go-to approach anyway because it cuts the INR training time by using techniques like MAML to pick the initialization? \n- What about direct image classification with NFTs instead of first encoding the INR as an array and then performing classification with that array instead of the original image array?\n- For generalization prediction, why isn't there a method corresponding to NFT_HNP in Table 3?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thorough analysis and review. We have updated our paper to correct the minor typos and missing definitions you found.\n\n> What about direct image classification with NFTs instead of first encoding the INR as an array and then performing classification with that array instead of the original image array?\n\nIf we understand the question correctly, is the question about why classify INRs at all instead of directly performing the classification in image space?\n\nWe share the motivation of [1,2,3] in operating on INR weights since INRs are continuous representations of signals that decouple the memory cost of the representation from the actual spatial resolution. For example, high resolution gridded representations of 3D data can be very expensive and unwieldy to work with.\n\nAlso, there would be no need to use NFTs (or any neural functional architecture) in that case since they are designed specifically for handling weight inputs, not images.\n\nPlease let us know if your question had a different intended meaning.\n\n> For generalization prediction, why isn't there a method corresponding to NFT_HNP in Table 3?\n\nIn this paper, we only define equivariant and invariant attention layers with respect to the Neuron Permutation (NP) group, and not the Hidden Neuron Permutation (HNP). The NP group is simpler to analyze and design layers for, and the finding in [1] is that the resulting layers can be more computationally efficient.\n\nExtending NFTs to HNP would require non-trivial effort and implementation complexity, but could be an interesting direction for future work, especially given that the results of Table 3 do suggest a hypothetical NFT_HNP could have benefits.\n\n> The weight spaces handled by the proposed layer are those of CNNs and MLPs. Do you think there are ways to extend this to other commonly used architectures like transformers?\n\nYes, we believe that this is another interesting direction for future work. Transformer weight spaces also contain permutation symmetries, not least because they contain feedforward MLP blocks. So it should be feasible to extend NFTs (and neural functionals in general) to create minimally equivariant layers for processing Transformer weights. This extension is not trivial or immediately obvious, since Transformers’ permutation symmetries are more complex than those of feedforward MLPs/CNNs–for example, because they contain residual connections that tie together neurons from different layers and prevent them from being independently permuted. So a first step would be to characterize this more complex symmetry structure, then derive equivariant/invariant layers for it.\n\n> Is this property satisfied for any matrices that are not permutation matrices? If so, we would also want equivariance to those transformations\n\nYes, for activation functions like ReLU we also get a scaling symmetry, so permutation symmetries would not be the only kind. For a positive scalar $c$ we have $\\frac{1}{c}W_2 \\sigma(cW_1 x) = W_2 \\sigma (W_1 x)$. \n\nFor NFTs we focus on permutation symmetries because one is typically dealing with weights produced by an optimization process like SGD, where these scaling symmetries typically do not appear. As [2] found: “SGD’s implicit regularization balances weight norms and, therefore, scale invariance does not seem to play an important role in understanding symmetries of solutions found by SGD.” So there is some evidence from prior literature that accounting for permutation symmetries is sufficient in practice.\n\n> I found Section 3.3 slightly unclear to read, particularly lines 155-158.\n\nWe will revise this section to make it clearer. Here we are assuming that the input to the NFT are convolutional weight filters, instead of MLP weight matrices. For a 1D convolution, the filter has shape $\\mathbb{R}^{n_i \\times n_{i-1} \\times k_i}$, where $n_i,n_{i-1}$ are the input and output channels _of the convolution_. Meanwhile, $c$ is the number of channels _of the NFT layer which is operating on the convolution filters_, hence the shape given on L158.\n\nWe welcome any additional feedback if this is still unclear.\n\n> Where is the layer encoding in section 3.4?\n\nIf the question is referring to the random fourier features (RFF): it is applied right before the first NFT layer, in order to guarantee a consistent channel dimension throughout the NFT.\n\nSuppose the NFT input has a single channel, so each $W^{(i)}\\_{jk} \\in \\mathbb{R}$ is a scalar. Then the RFF independently maps each scalar weight into a vector of length $c$, i.e. it is a function $\\mathbb{R} \\rightarrow \\mathbb{R}^c$.\n\n> In Section 5.2, do you train two INRs with the same initialization on the original and edited images and then take their weight difference as the target for the NFT?\n\nIn Section 5.2, the training loss is computed directly in image space, so we never have to train an INR on the edited image at all. To elaborate, the NFT produces edited weights for the INR, and the INR produces a predicted image. The loss is computed between the edited image (target) and the INR’s predicted image. The loss is then backpropagated _through_ the INR back to the NFT that produced the INR’s weights.  We will add this clarification to 5.2.\n\n\n[1] Zhou et al. Permutation Equivariant Neural Functionals.\n\n[2] Navon et al., Equivariant architectures for learning in deep weight spaces, ICML 2023.\n\n[3] Luigi et al., Deep Learning on Implicit Neural Representations of Shapes, ICLR 2023.\n\n[4] Entezari et al. The Role of Permutation Invariance in Linear Mode Connectivity of Neural Networks."}, {"review_id": "hbhQOv2kJm", "reviewer": "Reviewer_5PwT", "summary": "The paper studies an interesting problem of encoding weights of neural network using another network (Neural Functional Transformer or NFT). NFTs are built based on previous works, mainly DWS [31] and NFN [46], proposing permutation equivariant linear models to encode neural network weights. NFT's main contribution compared to these is a more expressive model based on Transformers that is permutation equivariant only to those weight permutations that result in the same function. While DWS already proposed such a minimal permutation equivariance, this paper achieves better performance due to leveraging Transformers and other tricks such as spatial INRs (INR2ARRAY).\nOverall, NFTs achieve better performance when encoding the weights of INRs as well as competitive performance in other tasks.", "questions": "Please see Questions in Weaknesses above.", "limitations": "Limitations are briefly discussed, which is appreciated. The authors could empirically compare time/memory complexity of their method to NFN, DWS and naive self-attention, which could strengthen the paper.", "rating": 6, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The topic of encoding neural network weights is very interesting with a lot of potential.\n2. The idea of combining Transformers with minimal equivariance seems as a powerful and promising way to approach this problem.\n3. The improvements in most experiments are significant. \n4. The paper is written and presented relatively well.\n", "weaknesses": "While I really like the problem the paper tackles and using Transformers with minimal equivariance, unfortunately there are many weak points described below. I will be willing to raise my score if these are addressed. \n\n1. Intuitive explanation of the difference between Eq. 4 and 5 is missing (why the union of sets is used and how is it implemented, by concatenation?) Perhaps, some illustration would help. Moreover, it would be very useful to see ablation results with each of the terms (4,5,6) ablated in the model. \n\n2. The paper mentioned that the weight space can include different features such as weights and gradients, but the details of what features are actually used in this work are missing. It's a bit confusing to mention gradients if they are never used, because it creates mental links with completely different literature. Why self-attention between different weight space features (dimensionalities from 1 to c) are needed? \n\n3. L186 mentions the limitation of Inr2Vec \"requires every INR share the same initialization\", but in the Experiments this advantage of the proposed method is not analyzed w.r.t. Inr2Vec, so it's unclear if this advantage is really useful. Also, Inr2Vec is not compared in Table 1 for some reason.\n\n4. In Table 1, why there are no results for INR2ARRAY with DWS and no DWS results for CIFAR-10? Since, DWS has source code available I believe this should be straightforward to evaluate and would strengthen the paper if the proposed NFT still dominates. Why INR2ARRAY only improves when combined with NFT and not with NFN? \n\n5. Section 5.1 is confusing, because it mixes methods, experiments and related work together.\n\n6. In Table 3, why NFN can be combined with HNP, while NFT cannot? Using HNP seems important in this experiment, so combining it with NFT could improve the paper.\n\n7. Editing INRs is an interesting experiment, but it is not very clear what would be the realistic use case. The experiment in DWS \"Adapting a network to a new domain\" seems more useful in this sense. \n\n8. Overall, the experiments seem quite esoteric with unclear real world use cases. For example, the motivation for encoding the weights of INR is not very convincing. This is a general problem with the topic of encoding of neural network weights as researchers seem to have not found really useful and realistic use cases of these models. The experiment \"Predicting CNN classifier generalization\" tackles a more concrete and useful application, but these also seem tiny networks with a few thousand parameters (Small CNN Zoo in [40]). So overall, this paper would be stronger if some other more practical applications are considered and/or the model is scaled to encoding more realistic networks (like ResNets in the \"NeRN - Learning Neural Representations for Neural Networks. ICLR 2023.\" paper).\n\n9. Some relevant papers/methods are not mentioned/compared in this paper. These include \"Self-supervised representation learning on neural network weights for model characteristic prediction. NeurIPs 2021.\", which discussed and compared in DWS, and \"NeRN - Learning Neural Representations for Neural Networks. ICLR 2023.\"\n\n**Minor**\n\nThe paper needs some polishing to improve the reading experience. Some abbreviations like NFN and SIRENS were used before they are introduced/cited, e.g. L201: \" Given a SIREN ...\". There are some occasional typos (e.g. L266 \"anount\") and the hyperrefs are broken (always jump to the first page).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your review--we appreciate the detailed and thoughtful questions. We will work to disentangle Section 5.1 and improve the clarity of the writing. Here we aim to answer them individually:\n\n> Intuitive explanation of the difference between Eq. 4 and 5 is missing (why the union of sets is used and how is it implemented, by concatenation?) Perhaps, some illustration would help. \n\nWe provide (in our PDF response) a diagram illustrating Equations 4 and 5 visually. In Eq 4 the rows of $W^{(i)}$ attend to themselves and the columns of $W^{(i-1)}$. In Eq 5 the columns of $W^{(i)}$ attend to themselves and the rows of $W^{(i+1)}$. But notice that computing Eq 5 for output $i-1$ and Eq 4 for output $i$ amount to a single self attention, as shown in the diagram. Set union (concatenation, as you inferred) is necessary to perform this attention over the rows and columns of different matrices altogether.\n\n> Moreover, it would be very useful to see ablation results with each of the terms (4,5,6) ablated in the model.\n\nWe performed this ablation and report the results in our top level reply. We found that the first two terms, which give minimal equivariance, are the important ones, while the last term, which is basically a self attention between features of each weight matrix, is less important.\n\n> Why self-attention between different weight space features (dimensionalities from 1 to c) are needed?\n\nEven if your input is just a single set of weights (c=1), the hidden features of a neural functional such as an NFT always results in more channels (c>1) in the hidden layers. Just as convolutional networks might take in an image with 1 or 3 channels, but 500 channels in its hidden layers. Hence NFT layers need to handle features with arbitrary c.\n\n> L186 mentions the limitation of Inr2Vec \"requires every INR share the same initialization\", but in the Experiments this advantage of the proposed method is not analyzed w.r.t. Inr2Vec, so it's unclear if this advantage is really useful. Also, Inr2Vec is not compared in Table 1 for some reason.\n\nWe will provide an empirical comparison with Inr2Vec on the benchmarks of Table 1. See the top-level reply for initial Inr2Vec results on MNIST, which confirm that Inr2Vec performs poorly on these tasks (19% test accuracy on MNIST). We will also add results for FashionMNIST and CIFAR once they are completed.\n\n> In Table 1, why there are no results for INR2ARRAY with DWS and no DWS results for CIFAR-10? Since, DWS has source code available I believe this should be straightforward to evaluate and would strengthen the paper if the proposed NFT still dominates. \n\nWe have added DWS results for CIFAR-10, see the top-level reply for details. The relative performance of DWS to the other methods in CIFAR-10 is consistent with the other datasets (MNIST, FMNIST).\n\nWe did not have the bandwidth to implement DWS with Inr2Array in the short response period. We do note that DWS performance tends to be lower than that of NFNs on each benchmark independent of Inr2Array, so it's unclear this would add significant new information over our existing Inr2Array+NFN results.\n\n> Why INR2ARRAY only improves when combined with NFT and not with NFN?\n\nWe believe this is because NFT is a more expressive and powerful architecture than NFNs, which was our goal in developing the nonlinear layers that underly NFTs. This makes it better at extracting useful representations from weights for downstream tasks.\n\n> In Table 3, why NFN can be combined with HNP, while NFT cannot? \n\nNFN_HNP and NFN_NP and two different variants whose layers are designed under different symmetry assumptions. In [1], designing layers under NP assumptions lead to layers that are simpler and more computationally efficient, so we designed our NFT layers under NP assumptions for simplicity and tractability purposes.\n\nAn HNP variant of the NFT layers would be helpful in some settings and is an interesting direction for future work--though it is likely nontrivial to implement and train scalably in practice.\n\n> Overall, the experiments seem quite esoteric with unclear real world use cases.\n\nWe agree that applications such as predicting classifier generalization would be generally useful, though currently they are at small scale. We would like to emphasize that the development of architectures that can process raw weights is very recent, and we believe that our experiments (including the INR tasks) represent promising early steps towards larger scale applications. The scale of our benchmarks in line with that of similar papers published at very recent conferences [2,3], and some of our results (e.g., INR classification) significantly improve performance and results over those works.\n\n> Some relevant papers/methods are not mentioned/compared in this paper.\n\nThank you for the references, we will cite and discuss those works in our revised paper.\n\n[1] Zhou et al. Permutation Equivariant Neural Functionals.\n\n[2] Navon et al., Equivariant architectures for learning in deep weight spaces, ICML 2023.\n\n[3] Luigi et al., Deep Learning on Implicit Neural Representations of Shapes, ICLR 2023."}, {"review_id": "bkNjA0IFwf", "reviewer": "Reviewer_73nt", "summary": "Designing architectures for processing the weights of neural networks is an exciting and relatively new area with great implications. This paper proposes a novel, attention-based architecture for processing NN weights. The proposed architecture is equivariant to the permutation symmetries of neural weight spaces. Furthermore, the authors propose INR2Array, a method for learning an invariant latent space for INR weights.\n", "questions": "1. Can the proposed Arch. be easily extended to the HNP design and setup? From Table 3 it appears this could be very beneficial for some tasks. \n1. The authors stated the difficulty of training large NFT architectures stably as a limitation of the method (Section 7). Since learning over weight spaces is a relatively novel task, it would be very beneficial if the authors could elaborate on any limitations/findings/failure cases from their experiments in the Appendix.\n", "limitations": "Addressed, although more details would be beneficial, see Questions.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "1. The paper is generally well-structured and well-written.\n1. The authors propose a novel architecture for processing the weights of neural networks. The new attention-based architecture achieves comparable or superior results compared to baseline methods.\n1. This is a natural and necessary extension of previous work on linear equivariant weight space layers.\n", "weaknesses": "My main concerns are the limited novelty w.r.t prior works and the missing/unfair evaluation results, which makes it difficult to assess the performance of the proposed approach w.r.t to prior work, and the effect of the different part of the proposed method (e.g., INR2Array^NFT vs. NFT).\n\n1. While attention-based equivariant layers are a natural extension, the technical novelty of the paper is somewhat limited with respect to prior works ([1], [2], [3]).\n1. Missing results for NFT without INR2Array for the INR classification task in Table 1. This is a necessary comparison for evaluating the effect of the main contributions of the paper, attention-based base layers for weight spaces, and INR2Array.\n1. The evaluation against DWS in Table 1 is flawed and unfair: The results for DWS are taken from [1], which uses a single INR copy for each image. The results for other baselines and the proposed method use 10 INR copies. As discussed in Section 7 of [1], additional INR copies can significantly improve the performance. Also, a more direct evaluation should use the same INR dataset, as choices of the INR arch., hyperparameters, and training procedure can affect the results of methods trained using this INR dataset. Using different INR datasets can potentially greatly vary the results. Finally, the augmentation scheme used in this paper is different than the one used in [1].\n1. Line 120: Need to provide an explanation/intuition for the provided claim.\n1. The equivariance proof provided for Theorem 3.1 is insufficient. While I believe the results are valid, the claim (“To show this, we can expand the left hand side term-by-term using the definition of SA.”) should be extended, and more details are needed. Also, it would be beneficial to provide a concrete example (using a specific layer) to show equivariance.\n1. As stated by the authors, the proposed layers are computationally expensive. Please provide runtime (or complexity) comparisons w.r.t linear equivariant weight-space layers.\n1. Additional comments on experiments: \n    - Please report the number of parameters for each method (In Table 1).\n    - What is the NFN (and INR2Array^NFN) variant used in Table 1 (NP or HNP)?\n    - It would be beneficial to include ablation on design choices, specifically the use of Random Fourier features and the use of Layer norm.\n    - Please provide more details on the use of mixup augmentations: are you performing the mixup in weight space?\n    - Also, it would be beneficial to compare the results of NFT with/without mixup augmentation.\n\nMinor: \n1. In Sec. 2, it would be clearer if the authors explicitly state that elements from the NP group, generally do not preserve the function defined by the set of weights.\n\n[1] Navon et al., Equivariant architectures for learning in deep weight spaces, ICML 2023.\n\n[2] Zhou et al., Permutation Equivariant Neural Functionals, 2023.\n\n[3] Luigi et al., Deep Learning on Implicit Neural Representations of Shapes, ICLR 2023.\n\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your careful review of our work. We will try to address your questions and feedback individually here:\n\n> the technical novelty of the paper is somewhat limited with respect to prior works ([1], [2], [3])\n\nIn addition to extending the equivariant framework of [1,2], we highlight our theoretical analysis of minimal equivariance in our self attention layer (Theorem 3.1), which is notable relative to [1,2] since those focus on theoretical properties of linear layers, while our layer is nonlinear. Moreover, in Inr2Array we add a spatial latent structure which is crucial to achieve good performance on the 2D image INR tasks, so our contribution in Inr2Array is not merely combining equivariant layers with Inr2Vec.\n\n> Missing results for NFT without INR2Array for the INR classification task in Table 1. \n\nWe will run this and report the results when they are available.\n\n> The results for DWS are taken from [1] [...] using different INR datasets can potentially greatly vary the results.\n\nWe've put the corrected comparison in the top-level reply. The relative ordering of methods by performance remains broadly unchanged. We will also add the corrected results to the paper.\n\n> Line 120: Need to provide an explanation/intuition for the provided claim [“Self attention between these n_i + n_{i-2} vectors is only equivariant if if the rows and columns of W^{(i-1)} and W^{(i)} permute simultaneously”]\n\nWe will add the following explanation:\n> Permuting the rows of $W^{(i-1)}$ and columns of $W^{(i)}$ corresponds to permuting the feature dimension of the inputs to Attn (Eq 3). If the features are permuted simultaneously, the pairwise dot products between any two vectors are preserved and the attention matrix is unchanged, producing an output whose features are simply permuted (equivariance). However, independent permutations of the Attn inputs will in general change the dot products and hence the attention matrix, breaking equivariance. For a more detailed explanation, see the proof of non-equivariance to false symmetries in Appendix A.1, in particular cases 2 and 3.\n\nThe new diagram of SA that we have provided should also aid in illustrating this argument visually.\n\n> The claim [equivariance of SA] should be extended, and more details are needed\n\nSee the expanded proof of this claim in the top-level reply, which we will add to the Appendix.\n\n> runtime (or complexity) comparisons w.r.t linear equivariant weight-space layers.\n\nUsing the same (A40) GPU, INR2Array with NFT takes 2.79x times as long per training step compared to with NFNs (NP). We will add full runtime comparisons to the paper.\n\n> Please report the number of parameters for each method (In Table 1).\n\nWe will add the number of parameters for each baseline method in Table 1. The number of parameters for the NFT method in each experiment is already reported in Appendix B. For quick reference:\n\n| Method            | # params                                    |\n|-------------------|---------------------------------------------|\n| DWS (32 chan)     | 0.5M (F/MNIST), 1M (CIFAR)                                      |\n| DWS (512 chan)    | 71M (F/MNIST), 134M (CIFAR)                 |\n| NFN_NP            | 45M                                         |\n| INR2Array (NFN)   | 31M                                         |\n| INR2Array (NFT)   | 22M                                         |\n\n\n> What is the NFN (and INR2Array^NFN) variant used in Table 1 (NP or HNP)?\n\nThat used NP. We are running an HNP variant for comparison and will include its results in the revised paper\n\n> [Various questions about mixup and its effect]\n\nMixup (with $\\alpha=1$) is applied to the latent embeddings of INRs produced by Inr2Array–not the raw weights--for downstream INR classification.\n\nAs suggested, we also tried removing mixup. The results suggest that mixup provides a very small benefit in most cases, and a modest benefit on CIFAR-10:\n\n| Dataset      | INR2Array^NFT | Without mixup |\n|--------------|---------------|---------------|\n| **MNIST**       | 98.5          | 98.3          |\n| **FashionMNIST** | 79.3          | 78.8          |\n| **CIFAR-10**     | 63.4          | 60.2          |\n\nWe will include these details in the revised paper.\n\n> It would be clearer if the authors explicitly state that elements from the NP group,  generally do not preserve the function defined by the set of weights.\n\nWe will modify Section 2 to make this point about Neuron Permutations clear.\n\n> [can] the proposed Arch. be easily extended to the HNP design?\n\nYes, it is likely that an HNP extension of the NFT layers would be helpful in some settings, and should be feasible. Though it may perhaps be complicated to implement and computationally expensive in some settings.\n\n> It would be beneficial to include ablation on design choices [...] Authors could elaborate on any limitations/findings/failure cases\n\nAlthough we have not had the bandwidth to run these ablations yet, we will aim to do so and include the results in the revised paper, in addition to our findings on how to train large NFT architectures stably."}]

Author rebuttal:
We thank the reviewers for their thoughtful feedback and questions. Here are some changes and new experiments we have done that will help strengthen this paper:

* Created a diagram illustrating the equivariant self attention (SA) layer (see PDF response)
* Performed ablation experiments on the individual terms in SA to understand the importance of each
* Added empirical Inr2Vec comparison
* Corrected the DWS comparison in Table 1 and added DWS results for CIFAR-10
* Ablated the effect of mixup on the Inr2Array latents in classification tasks (see reply to 73nt)
* Expanded the proof of Theorem 3.1.

## Ablating terms in SA

As suggested, we ran an ablation on the terms in our equivariant self-attention layer. As illustrated by the diagram in our response PDF, the first two terms collectively represent attention between the rows and columns of _adjacent_ weight matrices. The final term is a _global_ attention between features from every layer. Note that in practice we use a tractable variant of the third term described at L133.

For the ablation we either keep only the first two terms (**AdjacentSA**) or only the third term (**GlobalSA**). “FullSA” denotes the original layer.

| Model                      | MNIST test accuracy |
|-----------------------------|----------------------|
| Inr2Array (NFT, FullSA)     | 98.5                 |
| Inr2Array (NFT, AdjacentSA) | 97.6                 |
| INR2Array (NFT, GlobalSA)   | 37.4                 |
| Inr2Vec [1]                 | 19.4                 |

We see that ablating the final term (AdjacentSA) is roughly comparable with the original method, while ablating the first two terms (GlobalSA) drastically harms performance. This is interesting since the first two terms provide minimal equivariance, while the third term alone is not minimally equivariant (but helps propagate information between any two weight matrices). The results emphasize the importance of minimal equivariance.

We will run this ablation on the other benchmarks and will include the results in the revised paper.

## Inr2Vec comparison

We are also running Inr2Vec [1] on the classification tasks (Table 1), as suggested. See the table above for the initial results on MNIST, which empirically confirm that our method outperforms Inr2Vec by a wide margin, and also confirms the limitations of Inr2Vec on settings without shared initialization. 

We are running this comparison for the other datasets (Fashion MNIST, CIFAR) and will include the results in the revised paper.

## Corrected DWS comparison (+ CIFAR10)
As pointed out by 73nt, the results for DWS [2] in Table 1 are those reported by [2] itself, but INR datasets were constructed differently than the ones we use, which were produced by [3]. We have corrected this issue by running DWS ourselves on the latter datasets, and present the results below. This also adds DWS results for CIFAR-10, which were missing before. With the corrected numbers, the relative performance of different methods in Table 1 remains unchanged.

| Test accuracy  | MNIST | FashionMNIST | CIFAR |
|----------------|-------|--------------|-------|
| DWS (32 channel)| 74.7 | 67.5         | 42.3  |
| DWS (512 channel)| 61.6 | 62.0        | 42.9  |
| Inr2Array^NFT (ours) | 98.5 | 79.3        | 63.4  |

**Baseline details**: We trained the DWSNets of two different widths (width 32 and 512), with the same data augmentation scheme as for the other methods. For each dataset and channel size we swept learning rates in $[1e-3, 5e-3, 1e-4, 5e-4]$, following the protocol in [2].

## Expanded proof of Thm 3.1
Reviewer 73nt pointed out that the proof of equivariance of SA could be made more detailed. We provide that here.

**Lemma**: attention (Eq 3) is equivariant to permutations $\pi \in S_d$ to each
key/value/query input. In particular, let $[\pi q][i] = q[\pi^{-1}(i)]$ permute the entries
of the input vectors. Then:

$$Attn\left(\pi q, \left\\{\pi k\_p, \pi v\_p\right\\}\_{p=1}^N\right)[i]
= \left[\sum_p \left(\frac{\exp(q \cdot k_p)}{\sum\_{p'} \exp(q \cdot k\_{p'})}\right) (\pi v_p)\right][i]
= \sum_p \left(\frac{\exp(q \cdot k_p)}{\sum\_{p'} \exp(q \cdot k\_{p'})}\right) v_p[\sigma^{-1}(i)]
= Attn\left(q, \left\\{k\_p, v\_p\right\\}_{p=1}^N\right)[\pi^{-1}(i)]$$

Note that in the first equality we use $\pi q \cdot \pi k = q \cdot k$.

**Equivariance of Eqs 4-7**: As the KQV 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) change under $W \mapsto \sigma W$:

$$Attn\left(\[\sigma Q\]^{(i)}\_{j,:}, \left\\{ (\sigma K, \sigma V)^{(i-1)}\_{:,q} \right\\}\_q \bigcup \left\\{ (\sigma K, \sigma V)^{(i)}\_{p,:}\right\\}\_p\right)\[k\] = Attn\left( \sigma\_{i-1} Q\_{\sigma\_i^{-1}(j),:}, \left\\{ (\sigma\_{i-1} K^{(i-1)}\_{:,q}, \sigma\_{i-1} V^{(i-1)}\_{:,q}) \right\\}\_q \bigcup \left\\{ (\sigma\_{i-1} K^{(i)}\_{p,:}, \sigma\_{i-1} V^{(i)}\_{p,:}) \right\\}\_p \right)\[k\]$$

Note that permutations along the KV set indices are ignored since set order is irrelevant. Applying the lemma, we have:

$$Attn\left(Q^{(i)}\_{\sigma\_i^{-1}(j),:}, \left\\{ (K, V)^{(i-1)}\_{:,q} \right\\}\_q \bigcup \left\\{ (K,V)^{(i)}\_{p,:}\right\\}\_p\right)\[\sigma\_{i-1}^{-1}(k)\]$$

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)$.

We can similarly show equivariance for the second term (Eq 5), and the third term (Eq 6) is equivariant to any permutation since it is treating each weight as a token. We omit the full details due to character limit, but will add them to the paper.

## References

[1] De Luigi et al. Deep Learning on Implicit Neural Representations of Shapes.

[2] Navon et al. Equivariant Architectures for Learning in Deep Weight Spaces.

[3] Zhou et al. Permutation equivariant neural functionals.

Meta-review: The paper suggests a novel transformer-like architecture that can equivariantly process the weights of other neural networks, building on recent works that suggested linear layers for that setup. In addition, the authors propose INR2array - a novel approach for computing invariant representations for weights of implicit neural representations.  The reviewers found this extension to be natural and well-designed and appreciated the experimental results. On the other hand, reviewers were concerned by some cases in which the architecture performs worse than previous linear architectures, by the lack of experiments on 3D data, lack of clarity, and some missing experiments and baselines and unfair evaluation The authors posted a rebuttal that addressed some of these concerns. 

All reviewers finally converged to an accept decision.
Decision: Accept (poster)