Title: Neural Functional Transformers
Venue: NeurIPS 2023
Accepted: yes
Keywords: equivariance

permutation

implicit neural representation

generalization

transformers

attention
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

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 .

1 Introduction

Deep neural networks have emerged as flexible modeling tools applicable to a 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 [32, 30, 37]. This has fueled interest in techniques that directly operate in weight space to modify or extract information from a 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 a consequence, some existing methods for processing datasets of weights assume a restricted training process that reduces the effective weight space [9, 3, 6].

In contrast, we follow recent work in building permutation equivariant weight-space models called neural functionals , that can process neural network weights 1 without such restrictions [31, 46]. In particular, this work concerns neural functionals that are equivariant to permutations of the weights that correspond to re-arranging neurons. These weight-space permutations, known as neuron permutation symmetries , exactly preserve the network's behavior (see examples in Figure 1). Since equivariance to neuron permutations is typically a 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 [31, 46], suggesting that neural functional architectures can be significantly improved.

While existing neural functionals rely on linear layers (analogous to convolution), some of the most successful architectures in other domains, such as Transformers [41, 8], 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 symmetries 2 (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 functional 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-based objective. INR2ARRAY produces permutation-invariant latents that can be useful for downstream tasks such as INR classification. We also construct NFTs for a 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 parameters. 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 layers, leading to improved performance across weight-space tasks.

2 Preliminaries

Consider a feedforward network having n i neurons at each layer i = 0 , · · · , L , including the input and output layers. The network has weights 3 W = { W ( i ) ∈ R n i × n i -1 | i ∈ J 1 ..L K } 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 c -channel weight-space features W ∈ W c , we have W = { W ( i ) ∈ R n i × n i -1 × c | i ∈ J 1 ..L K } and each W ( i ) jk := W ( i ) j,k, : ∈ R c is a vector rather than a scalar.

Since the neurons in each hidden layer i ∈ { 1 , · · · , L -1 } have no inherent ordering, the feedforward network is invariant to the symmetric group S n i of permutations of the neurons in layer i . We follow Zhou et al. [46] in studying the neuron permutation (NP) group S NP := S 0 ×··· × S n L , 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 σ = ( σ 0 , · · · , σ L ) ∈ S NP. Each σ i permutes the neurons of layer i , which correspond to the output space (rows) of W ( i -1) and the input space (columns) of W ( i ) . Hence the action on weight space features W ∈ W c is denoted σW , where

As illustrated in Figure 1, a neuron permutation must always permute the columns of W ( i ) and the rows of W ( 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) to neuron permutation symmetries [31, 46]. Consider weight-space function classes parameterized by θ ∈ Θ . We say that a function class f : W c × Θ →W c is S NP -equivariant if permuting the input by any σ ∈ S NP always has the same effect as permuting the output by σ :

Similarly, a function class f : W c × Θ → R is S NP -invariant if f ( σW ; θ ) = f ( W ; θ ) for all σ, W, and θ .

W

(

i

+1)

W

(

i

)

σ

i

σ

i

+1

σ

i

σ

i

-

1

· · ·

· · ·

W

(

i

+1)

W

(

i

)

· · ·

· · ·

τ

i

τ

i

˜

τ

Our weight-space layers build off of the dot-product attention mechanism, which takes a query q and a set of N key-value pairs { ( k p , v p ) } N p =1 and computes, in the single-headed case:

If all q, k p , v p are vectors in R d , then the output is also a vector in R d . The definition extends to the situation where q, k, v p are multi-dimensional arrays of equal shape, by taking the dot product between these arrays in the natural way. We can also use multi-headed attention ([41]) without significant modification.

3 Neural Functional Transformers

We now introduce two attention-based layers for processing weight-space features: weight-space self-attention (which is S NPequivariant ) and weight-space cross-attention (which is S NPinvariant ). We then describe how to compose these layers into deep architectures called neural functional transformers (NFTs).

3.1 Equivariant weight-space self-attention

Akey concept motivating the design of our weight-space self-attention layer is the distinction between actual NP symmetries S NP that preserve network behavior, and other permutations of the weights that can generally change network behavior. For example, consider a τ ∈ S dim( W ) that permutes the columns of W ( i ) differently from the rows of W ( i -1) or that moves weights between different weight matrices (see Figure 1). Such permutations of the weight space are called 'false symmetries' [31] since they generally modify network behavior. We are interested in equivariance to the actual weightspace symmetries (neuron permutations) but not to any false symmetries, a property we call minimal equivariance:

̸

Definition 1. A function class f : W c × Θ → W c is minimally S NP-equivariant if it is S NPequivariant (Eq. 2), but not equivariant to any false symmetries. More precisely, for any τ ∈ S dim( W ) such that τ / ∈ S NP, there exist a W ∈ W c and θ ∈ Θ such that f ( τW ; θ ) = τf ( W ; θ ) .

Consider naive self-attention between the dim( W ) weights { W ( i ) jk } , which would be equivariant to any weight-space permutation. This would be S NP-equivariant, 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 by too many symmetries, including those that are not actually relevant to weight-space tasks.

Our weight-space self-attention layer SA ( · ; θ Q , θ K , θ V ) : W c → W c is more sensitive to the distinction between neuron permutations and false symmetries. It operates on weight-space features with c channels and is parameterized by query/key/value projection matrices θ Q , θ K , θ V ∈ R c × c .

W(e+1)

n+1×n

u×1+u

Attn

(n+1 +n1)×n

(2)M

(n +n2)×n-1

n ×n-1

n×n-1

Attn

n1×n2

n-1×n2

W(e-1)

Input

K,Q,V

Output

Each entry of the output feature is computed:

and the notation ( K,V ) ( i ) jk is shorthand for the tuple ( K ( i ) jk , V ( i ) 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 { W ( i ) jk ∈ R c } , but the first two terms give the layer additional structure relevant to neuron permutations in particular, as illustrated by Figure 2. To understand the symmetry properties of the first two terms, consider the top input to ATTN in Figure 2, which is constructed from the rows of W ( ℓ ) and the columns of W ( ell +1) . If the rows of W ( ℓ ) and the columns of W ( ℓ +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 of ATTN is simply permuted, giving equivariance. On the other hand, independently permuting the rows and columns of W ( ℓ ) and W ( ℓ +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.

Since neuron permutations never permute weights between different layers W ( i ) and W ( i ′ ) , we may additionally use learned layer position encodings { ϕ ( i ) ∈ R c | i ∈ J 1 ..L K } to prevent equivariance to such permutations:

We can now compose self-attention with position encoding to obtain SA ◦ LAYERENC.

Theorem 3.1 (Minimal equivariance) . The combined layer SA ◦ LAYERENC is minimally S NPequivariant.

(2)

W

W

(1)

｝

Weight-space Features

W

Q

Output

A

(

Q

,

K

)

K

=

θ

k

W

Project

V

=

θ

v

W

Project

K

V T

A

(

Q

,

K

)

V T

Proof (Sketch). To show equivariance: LAYERENC is clearly equivariant to neuron permutations. We can check SA ( σW ) ( i ) jk = SA ( W ) ( i ) σ -1 i ( j ) ,σ -1 i -1 ( k ) by expanding the left hand side term-by-term using the definition (Eq. 4). To show non-equivariance to false symmetries: we can broadly classify false symmetries into three types and check non-equivariance to each of them. Appendix A.1 provides the full proof.

Practical considerations. The final term of SA (Eq. 6) amounts to self-attention between each weight-space feature W ( i ) jk ∈ R c and requires O ( (dim W ) 2 c ) operations, which is usually intractable.

In practice, we replace this term by a self-attention between the row-column sums W ( i ) ⋆,⋆ = ∑ j,k W ( i ) jk of each weight matrix. This preserves interaction between weight-space features in different weight matrices, while reducing the computational complexity to O ( L 2 c ) . One can verify that this simplified version also maintains minimal S NP-equivariance. Meanwhile, consider the self-attention operation required to compute the first two terms of SA (depicted in Figure 2). If n i = n for all i , then the first two terms require O ( Ln 3 c ) operations, which is feasible for moderate n but can become expensive when processing very wide networks. As a result, weight-space attention layers are usually more computationally expensive than linear weight-space layers.

3.2 Invariant weight-space cross-attention

Stacking weight-space self-attention layers produces equivariant weight-space features, but some situations require producing an S NPinvariant output. Existing approaches achieve invariance by relying on a summation or maximization over the rows and columns of each input W ( i ) [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-attention layer CA : W c → R d is simply a cross attention between a learned query e ∈ R d and the set keys and values produced from weight-space features { W ( i ) jk ∈ R c } :

with θ K , θ V ∈ R d × c being learned projection matrices. By repeating this operation with M different learned embeddings e 1 , · · · , e M , we can easily extend this into an invariant map W c → R M × d . We depict the operation of the multi-embedding case in Figure 3.

3.3 Convolutional weight spaces

Although our description thus far focuses on fully-connected weight-space features W , we can also extend our layers to convolutional weight spaces. Suppose W is the c -channel weight-space feature corresponding to a 1D convolutional network with filter widths k i at each layer i . It contains matrices

W ( i ) ∈ R n i × n i -1 × k i × c . As in Zhou et al. [46], the filter dimension k i can be folded into the channel dimension, creating a feature in R n i × n i -1 × ( k i c ) with k i c channels. The problem is that k i may not be consistent across all i , while dot-product attention operations require that the channel dimensions for each W ( i ) match. We solve this problem by choosing a shared channel size ˜ c and using learned linear projections PROJ i : R n i × n i -1 × ( k i c ) → R n i × n i -1 × ˜ c to guarantee matching channel dimensions. The output of SA ( · ) can then be restored to the original channel dimension by another set of learned linear projections UNPROJ i : R n i × n i -1 × ˜ c → R n i × n i -1 × k i × c .

3.4 Building Neural Functional Transformers

Following Transformer architecture design [41], we form stackable 'blocks' that combine weightspace self attention with LayerNorm, pointwise MLPs, and residual connections. Each block maps W c →W c and preserves S NP-equivariance:

where both the MLP : R c → R c and LayerNorm LN : R c → R c operate 'pointwise,' or independently on each input W ( i ) jk . Since pointwise operations are permutation equivariant, the overall block is an S NP-equivariant map on W c . To build neural functional Transformers (NFTs), we stack multiple blocks into a deep equivariant architecture. For tasks which require S NP-invariance, like classification or learning an 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 a few channels, while our desired hidden channel dimension may be significantly larger (e.g., c = 256 ). We can apply any function g : R c 1 → R c 2 to each W ( i ) jk ∈ R c 1 to project from the input channel dimension to the hidden channel dimension while preserving equivariance. In our experiments, we use random Fourier features [35] to achieve this projection.

4 INR2ARRAY: Learning an invariant latent space for weights

An application of interest for weight-space architectures is to learn a 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 representation 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 to the original INR2VEC formulation:

We implement the encoder with an S NP-invariant neural functional transformer (NFT). This guarantees that the latent representation is invariant to neuron permutations of the weights.

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 a 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 [3], 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 , 1] 2 , though the general case follows naturally. We first split the coordinate grid into M spatial patches P 1 ∪ · · · ∪ P M = [ -1 , 1] 2 . Given a SIREN ( · ; W ) , INR2ARRAY uses an S NP-invariant NFT encoder ENC θ ( · ) to map weights W to a latent array of M vectors, z ∈ R M × d . The decoder DEC θ : R d →W produces a set of weights { ˆ W i ∣ ∣ ∣ i = 1 , · · · , M } , one for each vector z i := z i, : ∈ R d . Each ˆ W i ( z i ) is responsible for parameterizing the SIREN only for the spatial patch P i ⊂ [ -1 , 1] 2 . The objective

MNIST

FashionMNIST

CIFAR-10

3

Original SIREN

1

3

Inr2ArrayNFN

3

Inr2ArrayNFT

is to reconstruct the original INR's content using the decoded weights:

Our INR2ARRAY encoder architecture uses multiple weight-space self-attention layers followed by a weight-space cross-attention layer CA, which is especially well-suited for producing M × d permutation invariant arrays. For the decoder, we use the hypernetwork design of Sitzmann et al. [37], where the hypernetworks are conditioned on the spatial latent z i ∈ R d . In principle, one can also design an invariant encoder using (non-attentive) NFN layers, and we will compare both options in our experiments.

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 to model interesting properties of the trained networks (e.g., generalization) from their weights; and (2) datasets containing the weights of INRs each representing a single data signal, such as an image. Tasks including editing INR weights to modify the represented signal, or predicting target information related to the signal (e.g., image class). For our INR experiments, we construct datasets for MNIST [25], FashionMNIST [42], and CIFAR-10 [23] following the procedure of Zhou et al. [46] exactly.

5.1 INR classification with INR2ARRAY

Classifying the signal represented by an INR directly from its weights is a challenging problem, with current approaches requiring that all INRs share initialization [10, 3, 6]. In the 'vanilla' setting where INRs can be initialized and trained independently, state-of-the-art methods struggle to classify signals from even simple datasets [31, 46]. We train INR2ARRAY to learn a latent representation of INR weights that can be used for more effective classification in this vanilla setting.

We implement both NFT (weight-space self-attention and cross-attention) and NFNNP [46] variants of the INR2ARRAY encoder, and distinguish the two variants by superscripts: INR2ARRAY NFT and INR2ARRAY NFN . For the decoder, we use the hypernetwork design of Sitzmann et al. [37], where the hypernetworks are conditioned on the spatial latent z i ∈ R d . Appendix B.3 describes the implementation and training of each variant in full detail.

We train separate encoders for each INR dataset shown in Table 6 of the appendix, which reports the reconstruction mean square error (Eq 12) on test INR inputs. The NFT encoder consistently achieves lower reconstruction error than the NFN encoder, and INR2ARRAY NFT also produces qualitatively better reconstructions than INR2ARRAY NFN (Figure 4), confirming that NFTs enable higher quality encoding than previous neural functionals.

Once trained, the INR2ARRAY encoder maps INR weights to compact latent arrays that represent the INR's contents. Downstream tasks such as classifying the INR's signal (e.g., image) can now be performed directly on these S NP-invariant latent arrays z ∈ R M × d . Concretely, in K -way

classification the trained encoder ENC : W → R M × d can be composed with any classification head f : R M × d → R K to form an S NP-invariant classifier f ◦ ENC : W → R K . Since the encoder is already trained by the reconstruction objective, we only train f during classification and keep the encoder fixed. We choose to implement f using a Transformer classification head [41], which views the latent array input as a lengthM sequence of d -dimensional vectors.

Existing INR classifiers, which we refer to as DWS [31] and NFN [46], are permutation-invariant architectures that directly map input weights to predicted labels. We also create a modified NFN classifier (Spatial NFN) which produces an M × d array of activations before a convolutional classifier head, to test the impact of spatial latents independent of the INR2ARRAY training process. Appendix B.4 describes the setup for each method in further detail. Finally, we measure the performance of NFT classifiers without INR2ARRAY.

Table 1 shows that INR2ARRAY significantly improves classification test accuracies across the board. In addition, the best performance is achieved by implementing the encoder using our attention-based layers (NFT) compared to linear weight-space layers (NFN). Notably, INR2ARRAY NFT achieves a test accuracy of 64 . 4% on CIFAR-10 (the hardest dataset), an improvement of +17% over the previous best result of 46 . 6% by NFNs. They also achieve an MNIST test accuracy of 98 . 5% , up from 92 . 9% by NFNs and 85 . 7% by DWS. The vanilla NFT performance is comparable to DWS and NFNs (better than the former and worse than the latter), suggesting the importance of INR2ARRAY in addition to the architecture. Finally, the Spatial NFN fails to improve performance over standard NFN classifiers, implying that the particular objective of INR2ARRAY is crucial to make spatial latent spaces useful. The results show that learning a latent array representation of INR weights using INR2ARRAY and NFTs can greatly improve performance on downstream INR tasks like classification.

5.1.1 Term-by-term ablation

Here we investigate the importance of each term in our weight space self-attention layer (Eqs. 4-6). The first two terms of SA, illustrated by Figure 2, amount to self-attention between the rows and columns of adjacent weights, while the third term is a global self-attention between features from all layers. Note that in practice we use a tractable approximation of the third term described in Sec 3.1. For this ablation experiment we either keep only the first two terms ( AdjacentSA ) or only the third term ( GlobalSA ). FullSA denotes the original layer.

Table 2 shows the results of this ablation experiment on

INR2ARRAY NFT 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

Table 2: Test classification accuracy of INR2ARRAY NFT when ablating the terms of SA (Eqs.4-6).

equivariant (but helps propagate information between any two weight matrices). The results emphasize the importance of designing our layer to achieve minimal equivariance rather than naively applying self-attention to the weights.

5.2 Editing INRs

We also evaluate NFTs on editing INR weights to alter their signal, e.g., to modify the represented image. The goal of this task is to edit the weights of a trained SIREN to alter its encoded image, expressed as a difference W ′ ← W + ∆( W ) . Neural functionals are trained to learn ∆( · ) that achieves some desired visual transformation, such as dilating the image. Permutation equivariance is a useful inductive bias here since if ∆( W ) is the desired edit for W , then for any neuron permutation σ the desired edit to σW is σ ∆( W ) . In addition to the MNIST dilation and CIFAR contrast tasks from Zhou et al. [46], we also introduce several new editing tasks: MNIST erosion, MNIST gradient, and FashionMNIST gradient. Gradient tasks roughly anount to edge detection; Figure 5 in the appendix shows sample inputs and targets for each editing task.

We compare NFT editing performance against the two NFN variants [46], with each method using the a similar number of parameters ( ∼ 7 M). For full training details, see Appendix B.2). Table 3 shows the test mean square error (MSE) between the edited INR and the ground truth visual transformation. NFTs consistently outperform NFNs on most INR editing tasks. In addition, Table 5 in the appendix shows that NFTs generally obtain both lower training and test error, indicating that the attention-based architecture enables greater expressivity. Figure 5 shows qualitative samples produced by each editing method on each task.

5.3 Predicting CNN classifier generalization

Whereas the previous experiments have focused on the weight spaces of INRs (which are implemented by small MLPs), we would also like to evaluate how NFTs process other weight spaces such as those belonging to convolutional neural network classifiers. Large-scale empirical studies have produced datasets of trained classifiers under different hyperparameter settings [40, 11], enabling a data-driven approach to modeling generalization from their weights, which could lead to new insights. Prior methods for predicting classifier generalization typically rely on extracting hand-designed features from the weights before using them to predict the test accuracy [18, 43, 40, 19, 28].

We now study how NFTs can model generalization from raw weights using the Small CNN Zoo [40], which contains thousands of CNNs trained on image classification datasets. In addition to comparing against the two neural functional variants NFNNP 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.1 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 handdesigned features baseline STATNN. NFTs slightly outperform NFNNP on each dataset, while NFNHNP achieves the best performance overall. Note that both NFNNP 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 [17]. Permutation symmetries in particular have been studied in the context of neural network loss landscapes [14, 4, 12] and weight-space merging [38, 1]. However, most prior methods for processing weight space objects do not explicitly account for these symmetries [2, 26, 15, 24, 45, 7, 21], although some have tried to encourage permutation equivariance through data augmentation [33, 29]. Another workaround approach is to explicitly restrict the weight space being considered by, for example, fixing the initialization of all networks being processed [6] or meta-learning modulations of a set of base parameters [10, 3]. Instead, we study the problem of encoding permutation symmetries into the neural functional itself, without any restrictions on the weight space being processed.

There are a variety of methods that incorporate symmetries into deep neural network architectures [36, 22, 13]. Examples include (group) convolutions for images [25, 5] and permutation equivariant architectures for general set-structured inputs [34, 44, 16, 39, 27]. Most directly related to our work are that of Navon et al. [31] and Zhou et al. [46], who introduce linear layers equivariant to the neuron permutation symmetries of feedforward networks. We extend their approach by introducing nonlinear equivariant layers based on the attention mechanism, and use them to construct NFTs.

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 on permutationequivariant weight-space models using nonlinear self-attention and cross-attention layers, which can enable greater expressivity compared to existing linear layers.

We empirically evaluate the effectiveness of NFTs on weight-space tasks involving datasets of trained CNN classifiers and implicit neural representations (INRs). Operating on weights alone, NFTs can predict the test accuracy of CNN classifiers, modify the content of INRs, and classify INR signals. We also use NFTs to develop INR2ARRAY, a method for mapping INR weights to compact and S NP-invariant latent representations, which significantly improve performance in downstream tasks such as INR classification.

Some limitations of NFTs include increased computational costs from self-attention relative to linear weight-space layers, and the difficulty of training large NFT architectures stably. Future work may address these limitations, and explore additional applications of NFTs such as for learned optimization or weight-space generative modeling. Overall, our work highlights the potential of attention-based weight-space layers offers a promising direction for the development of more expressive and powerful neural functional architectures.

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 N00014-22-1-2621.

References

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

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

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.03130 , 2023.

J. Brea, B. Simsek, B. Illing, and W. Gerstner. Weight-space symmetry in deep networks gives rise to permutation saddles, connected by equal-loss valleys across the loss landscape. arXiv preprint arXiv:1907.02911 , 2019.

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

L. De Luigi, A. Cardace, R. Spezialetti, P. Z. Ramirez, S. Salti, and L. Di Stefano. Deep learning on implicit neural representations of shapes. arXiv preprint arXiv:2302.05438 , 2023.

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

A. Dosovitskiy, L. Beyer, A. Kolesnikov, D. Weissenborn, X. Zhai, T. Unterthiner, M. Dehghani, M. Minderer, G. Heigold, S. Gelly, et al. An image is worth 16x16 words: Transformers for image recognition at scale. arXiv preprint arXiv:2010.11929 , 2020.

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

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

G. Eilertsen, D. Jönsson, T. Ropinski, J. Unger, and A. Ynnerman. Classifying the classifier: dissecting the weight space of neural networks. arXiv preprint arXiv:2002.05688 , 2020.

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

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

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

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

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.

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

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

Y. Jiang, P. Natekar, M. Sharma, S. K. Aithal, D. Kashyap, N. Subramanyam, C. Lassance, D. M. Roy, G. K. Dziugaite, S. Gunasekar, et al. Methods and analysis of the first competition in predicting generalization of deep learning. In NeurIPS 2020 Competition and Demonstration Track , pages 170-190. PMLR, 2021.

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

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

R. Kondor and S. Trivedi. On the generalization of equivariance and convolution in neural networks to the action of compact groups. In International Conference on Machine Learning , pages 2747-2755. PMLR, 2018.

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

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

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

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

H. Maron, O. Litany, G. Chechik, and E. Fetaya. On learning sets of symmetric elements. In International conference on machine learning , pages 6734-6744. PMLR, 2020.

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

L. Metz, J. Harrison, C. D. Freeman, A. Merchant, L. Beyer, J. Bradbury, N. Agrawal, B. Poole, I. Mordatch, A. Roberts, et al. Velo: Training versatile learned optimizers by scaling up. arXiv preprint arXiv:2211.09760 , 2022.

B. Mildenhall, P. P. Srinivasan, M. Tancik, J. T. Barron, R. Ramamoorthi, and R. Ng. Nerf: representing scenes as neural radiance fields for view synthesis (2020). arXiv preprint arXiv:2003.08934 , 2020.

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

J. J. Park, P. Florence, J. Straub, R. Newcombe, and S. Lovegrove. Deepsdf: Learning continuous signed distance functions for shape representation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition , pages 165-174, 2019.

W. Peebles, I. Radosavovic, T. Brooks, A. A. Efros, and J. Malik. Learning to learn with generative models of neural network checkpoints. arXiv preprint arXiv:2209.12892 , 2022.

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

A. Rahimi and B. Recht. Random features for large-scale kernel machines. Advances in neural information processing systems , 20, 2007.

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

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

N. Tatro, P.-Y. Chen, P. Das, I. Melnyk, P. Sattigeri, and R. Lai. Optimizing mode connectivity via neuron alignment. Advances in Neural Information Processing Systems , 33:15300-15311, 2020.

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

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

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

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

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

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

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

A. Zhou, K. Yang, K. Burns, Y. Jiang, S. Sokota, J. Z. Kolter, and C. Finn. Permutation equivariant neural functionals. arXiv preprint arXiv:2302.14040 , 2023.

A Weight-space self-attention

A.1 Proof of minimal equivariance

We repeat the claim about minimal equivariance for weight-space self-attention and layer position encodings SA ◦ LAYERENC : W c →W c :

Theorem 3.1 (Minimal equivariance) . The combined layer SA ◦ LAYERENC is minimally S NPequivariant.

Proof. We first show S NP-equivariance, before checking that it minimally equivariant.

Equivariance . It is straightforward to check that LAYERENC is equivariant to any permutation that preserves the layer index (which neuron permutations do). At a high level, we check that SA is S NP-equivariant by verifying that SA ( σW ) ( i ) jk = SA ( W ) ( i ) σ -1 i ( j ) ,σ -1 i -1 ( k ) . We proceed term by term.

First, it will be useful to show that ATTN (Eq. 3) is equivariant to permutations π ∈ S d of each key, query, and value input. In particular, let [ πq ][ i ] = q [ π -1 ( i )] permute the entries of the input vectors:

Note that in the first equality we use πq · πk = q · k .

Since the key, query, and value projections are only linear projections along the channel dimension, K ( σW ) = σK ( W ) , where K denotes all keys produced by linearly projection of W by θ K (and similarly for Q,V ). Then consider, for example, how the first term (Eq 4) changes under W ↦→ σW :

Note the distinction between neuron permutations of the weight space [ σQ ] and permutations of individual vectors like σ i -1 Q p, : , and that permutations along the KV set indices (e.g., q in the first set) are ignored since set order is irrelevant. Using the equivariance of ATTN (Eqs. 14-16):

Notice that compared to Eq. 4 the output indices have changed j ↦→ σ -1 i ( j ) and k ↦→ σ -1 i -1 ( k ) , showing that the first term is equivariant. We can repeat this argument in a similar fashion for the second term (Eq. 5), while the third term (Eq. 6) applies ATTN by treating each weight as an individual token, so its equivariance follows immediately from the permutation equivariance of ATTN. This verifies that all terms of SA are S NP-equivariant.

̸

Minimal equivariance . Here our goal is to show non -equivariance to permutations τ ∈ S dim( W ) but τ / ∈ S NP (false symmetries). That is, we should find some input W such that SA ( LAYERENC ( τW )) = τ SA ( LAYERENC ( W )) , for some parameters { θ Q , θ K , θ V } and { ϕ ( i ) } .

It will be useful to define the actions of permutations in S dim( W ) on the indices of the weights, rather than the weights themselves. The index-space for a weight space W is defined as a set of 3-tuples representing layer number, row, and column, respectively:

For α ∈ I , we use the subscripts ℓ, r, c to denote the layer, row, and column indices respectively. That is, α = ( α ℓ , α r , α c ) = ( i, j, k ) .

Notice that the set of arbitrary permutations S dim( W ) is simply the set of all bijections τ : I → I (any index can be arbitrarily permuted to any other index), while S NP is the subgroup of bijections σ that can be written as:

False symmetries can broadly be categorized into three groups. In fact, we can verify that any permutation τ that fails to satisfy one of these three cases must in S NP:

̸

False symmetry that permutes weights across layers. There exists an ( i, j, k ) ∈ I such that τ ( i, j, k ) ℓ = i .

̸

False symmetry where row and column permutations are not independent. Suppose a false symmetry τ does not fall under case (1), i.e., it preserves the layer index. Then it may still fail to have independent permutations of row and columns. In this case, either there is an ( i, j, k ) and q such that τ ( i, j, k ) = ( i, j ′ , k ′ ) , but τ ( i, j, q ) r = j ′ . Or there is an ( i, j, k ) and p such that τ ( i, j, k ) = ( i, j ′ , k ′ ) but τ ( i, p, k ) c = k ′ .

̸

̸

Adjacent permutations decoupled. If a false symmetry τ is not one of the two cases above, there must be an ( i, j, k ) ∈ I and q such that τ ( i, j, k ) = ( i, j ′ , k ′ ) but τ ( i -1 , k, q ) r = k ′ . That is, an instance where one row pair of layer i -1 do not permute the same way as the columns of layer i .

We can now check non-equivariance for τ in each case. To simplify the notation here we will assume c = 1 although the general case is similar:

If τ does not preserve layer index, LAYERENC will be non-equivariant since each layer has a different encoding.

̸

Suppose there is an ( i, j, k ) and q such that τ ( i, j, k ) = ( i, j ′ , k ′ ) , but τ ( i, j, q ) r = j ′ , and let θ Q = θ K = θ V = I . Now consider the input W ∈ W such that all weights are 0 except for W ( i ) jk = W ( i ) jq = 1 . Then we can check that [ τ SA ( W )] i,j ′ ,k ′ = SA ( τW ) i,j ′ ,k ′ , so the layer is not equivariant to τ . We can construct a similar argument for the other case, where there is an ( i, j, k ) and p such that τ ( i, j, k ) = ( i, j ′ , k ′ ) but τ ( i, p, k ) c = k ′ .

̸

̸

̸

In this case there must be an ( i, j, k ) ∈ I and q such that τ ( i, j, k ) = ( i, j ′ , k ′ ) but τ ( i -1 , k, q ) r = k ′ . Again let θ Q = θ K = θ V = I and consider W ∈ W such that all weights are 0 except for W ( i ) jk = W ( i -1) kq = 1 . Then, similar to the above case, we can verify that [ τ SA ( W )] i,j ′ ,k ′ = SA ( τW ) i,j ′ ,k ′ , so the layer is not equivariant to such τ .

Taken together, the three cases above show that although SA ◦ LAYERENC is S NP-equivariant, it is not equivariant to any other permutations τ / ∈ S NP. Hence the layer is minimally equivariant.

A.2 Full description including biases

Let B be the space of biases, and U = W×B the full space of both weights and biases. Then the full weight-space self-attention layer SA : U c →U c is defined, for input U = ( W,b ) :

̸

where Y ( U ) ∈ W c and z ( U ) ∈ B c , with:

We use KV 1 , KV 2 , KV 3 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.

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 is typically a projection from the input channel dimensionality to the hidden channel dimensionality, for which we use random Fourier features.

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 c 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 S NP-invariant NFTs): Whether the invariant pooling is done with summation (HNP/NP [46]) or cross-attention (ours).

CA M : If using cross attention, the number of learned query embeddings e 1 , · · · , e M .

CA dim: Dimension of vectors going into cross attention.

MNIST (Erode)

MNIST (Dilate)

MNIST (Gradient)

Ground truth

M

3

Original INR

Q

3

NFNNP

3

NFT

3

Fashion-MNIST (Gradient)

CIFAR (Contrast)

Ground truth

Original INR

NFNNP

NFT

B.1 Generalization prediction

Each Small CNN Zoo dataset that we 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 on the remaining 12000 . We use the validation data to sweep over Fourier scales of 3 or 10 and try both HNP (summation) or 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 5 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 to 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 ( 4 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 of the best checkpoint can then be used to produce latent array representations of INR inputs.

B.4 INR Classification

After training INR2ARRAY on a given dataset, we select the checkpoint with the smallest reconstruction error on the validation set, and use the encoder to produce a 16 × 256 latent array representation for any input SIREN. We view the latent as a length16 sequence of 256 -dim vectors, and project each vector into R 512 before feeding the array into Transformer encoder with 12 blocks. This produces an output z o ∈ R 16 × 512 ; we take the first output z o 0 ∈ R 512 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 AdamW 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.

Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nDesigning 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\nSTRENGTHS:\n1. 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\nWEAKNESSES:\nMy 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\nQUESTIONS:\n1. 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\nLIMITATIONS:\nAddressed, although more details would be beneficial, see Questions.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 4, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n1. 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\nWEAKNESSES:\nWhile 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).\n\nQUESTIONS:\nPlease see Questions in Weaknesses above.\n\nLIMITATIONS:\nLimitations 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.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 2, "DATE": "2023-07-05"}

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

comments: 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.

STRENGTHS:
1. The overall written is good, and the authors provide code for reproduction.

2. The topic of this paper is cutting-edge and the proposed method is novel.

3. 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?

2. Equation (4) (5) (6) are difficult to read. The authors may represent them via a better way.

3. 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?

4. Do you evaluate your method on 3D dataset?

QUESTIONS:
See "Weaknesses".

LIMITATIONS:
This paper has no potential negative societal impact.
TITLE: Official Review
SOUNDNESS_CORRECTNESS: 3
DATE: 2023-07-08

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS:\n- 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.\n\nWEAKNESSES:\nI 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\nQUESTIONS:\nThe 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\nLIMITATIONS:\nThe authors discuss computational limitations of the vanilla form of their self-attention layer, and discuss computationally feasible alternatives.", "TITLE": "Official Review", "RECOMMENDATION": 6, "REVIEWER_CONFIDENCE": 3, "SOUNDNESS_CORRECTNESS": 3, "DATE": "2023-07-11"}