SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Allan Zhou $^{1}$   
Kaien Yang $^{1}$   
Yiding Jiang $^{2}$   
Kaylee Burns $^{1}$   
Winnie Xu $^{1}$   
Samuel Sokota $^{2}$   
J. Zico Kolter $^{2}$   
Chelsea Finn $^{1}$ $^{1}$ Stanford University  
 $^{2}$ Carnegie Mellon University  
ayz@cs.stanford.edu
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.
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<sup>1</sup> 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.
For clarity, we 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.
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<sup>2</sup> (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.
Consider a feedforward network having  $n_i$  neurons at each layer  $i = 0, \dots, L$ , including the input and output layers. The network has weights  $W = \{W^{(i)} \in \mathbb{R}^{n_i \times n_{i-1}} \mid i \in [1..L]\}$  belonging to weight space,  $\mathcal{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 \in \mathcal{W}^c$ , we have  $W = \{W^{(i)} \in \mathbb{R}^{n_i \times n_{i-1} \times c} \mid i \in [1..L]\}$  and each  $W_{jk}^{(i)} := W_{j,k,:}^{(i)} \in \mathbb{R}^c$  is a vector rather than a scalar.
Since the neurons in each hidden layer  $i \in \{1, \dots, 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_{\mathrm{NP}} := S_0 \times \dots \times 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  $\sigma = (\sigma_0, \dots, \sigma_L) \in S_{\mathrm{NP}}$ . Each  $\sigma_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 \in \mathcal{W}^c$  is denoted  $\sigma W$ , where
$$
\left[ \sigma W \right] _ {j k} ^ {(i)} = W _ {\sigma_ {i} ^ {- 1} (j), \sigma_ {i - 1} ^ {- 1} (k)} ^ {(i)}, \quad \sigma = \left(\sigma_ {0}, \dots , \sigma_ {L}\right) \in \mathcal {S} _ {\mathrm {N P}}. \tag {1}
$$
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  $\theta \in \Theta$ . We say that a function class  $f: \mathcal{W}^c \times \Theta \to \mathcal{W}^c$  is  $S_{\mathrm{NP}}$ -equivariant if permuting the input by any  $\sigma \in S_{\mathrm{NP}}$  always has the same effect as permuting the output by  $\sigma$ :
$$
\sigma f (W; \theta) = f (\sigma W; \theta), \quad \forall \sigma \in \mathcal {S} _ {\mathrm {N P}}, W \in \mathcal {W} ^ {c}, \theta \in \Theta . \tag {2}
$$
Similarly, a function class  $f: \mathcal{W}^c \times \Theta \to \mathbb{R}$  is  $S_{\mathrm{NP}}$ -invariant if  $f(\sigma W; \theta) = f(W; \theta)$  for all  $\sigma, W,$  and  $\theta$ .
2Moreover, naive self-attention between input weights is computationally intractable.
For simplicity, we omit the biases from our discussion here, see the Appendix for full detail.
Figure 1: Left: An illustration of how neuron permutations  $\sigma = (\sigma_0, \dots, \sigma_L)$  act on weight matrices  $W = (W^{(0)}, \dots, W^{(L)})$ . Each  $\sigma_i$  simultaneously permutes the rows and columns of adjacent weight matrices  $W^{(i)}, W^{(i+1)}$ . Right: Two examples of false symmetries, i.e., permutations that don't preserve network behavior: (1) when  $\tau_i$  permutes the rows and columns of adjacent matrices permute differently and (2) when  $\tilde{\tau}$  permutes weights across layers.
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)\}_{p=1}^N$  and computes, in the single-headed case:
$$
\operatorname {A T T N} \left(q, \left\{\left(k _ {p}, v _ {p}\right) \right\} _ {p = 1} ^ {N}\right) = \sum_ {p} v _ {p} \left(\frac {\exp \left(q \cdot k _ {p}\right)}{\sum_ {p ^ {\prime}} \exp \left(q \cdot k _ {p ^ {\prime}}\right)}\right). \tag {3}
$$
If all  $q, k_p, v_p$  are vectors in  $\mathbb{R}^d$ , then the output is also a vector in  $\mathbb{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_{\mathrm{NP}}$ -equivariant) and weight-space cross-attention (which is  $S_{\mathrm{NP}}$ -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-attention layer is the distinction between actual NP symmetries  $S_{\mathrm{NP}}$  that preserve network behavior, and other permutations of the weights that can generally change network behavior. For example, consider a  $\tau \in S_{\dim(\mathcal{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 weight-space symmetries (neuron permutations) but not to any false symmetries, a property we call minimal equivariance:
Definition 1. A function class  $f: \mathcal{W}^c \times \Theta \to \mathcal{W}^c$  is minimally  $S_{NP}$ -equivariant if it is  $S_{NP}$ -equivariant (Eq. 2), but not equivariant to any false symmetries. More precisely, for any  $\tau \in S_{\dim(\mathcal{W})}$  such that  $\tau \notin S_{NP}$ , there exist a  $W \in \mathcal{W}^c$  and  $\theta \in \Theta$  such that  $f(\tau W; \theta) \neq \tau f(W; \theta)$ .
Consider naive self-attention between the  $\dim (\mathcal{W})$  weights  $\{W_{jk}^{(i)}\}$ , which would be equivariant to any weight-space permutation. This would be  $S_{\mathrm{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  $\mathrm{SA}(\cdot;\theta_Q,\theta_K,\theta_V): \mathcal{W}^c \to \mathcal{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  $\theta_Q, \theta_K, \theta_V \in \mathbb{R}^{c \times c}$ .
Figure 2: An illustration of Eqs 4-5 of our weight-space self-attention layer, in 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:
$$
\begin{array}{l} \operatorname {S A} \left(W; \theta_ {Q}, \theta_ {K}, \theta_ {V}\right) _ {j k} ^ {(i)} = \operatorname {A T T N} \left(Q _ {j,:} ^ {(i)}, \left\{\left(K, V\right) _ {:, q} ^ {(i - 1)} \right\} _ {q = 1} ^ {n _ {i - 2}} \bigcup \left\{\left(K, V\right) _ {p,:} ^ {(i)} \right\} _ {p = 1} ^ {n _ {i}}\right) _ {k} (4) \\ + \operatorname {A T T N} \left( \right.Q _ {:, k} ^ {(i)}, \left\{ \right.\left. (K, V) _ {:, q} ^ {(i)} \right\} _ {q = 1} ^ {n _ {i - 1}} \bigcup \left\{ \right.\left. (K, V) _ {p,:} ^ {(i + 1)} \right\} _ {p = 1} ^ {n _ {i + 1}}\left. \right) _ {j} \left. \right. (5) \\ + \operatorname {A T T N} \left(Q _ {j k} ^ {(i)}, \left\{\left(K, V\right) _ {p q} ^ {(s)} \mid \forall s, p, q \right\}\right), \text {w h e r e} (6) \\ Q _ {j, k} ^ {(i)} := \theta_ {Q} W _ {j k} ^ {(i)}, \quad K _ {j, k} ^ {(i)} := \theta_ {K} W _ {j k} ^ {(i)}, \quad V _ {j, k} ^ {(i)} := \theta_ {V} W _ {j k} ^ {(i)}, (7) \\ \end{array}
$$
and the notation  $(K,V)_{jk}^{(i)}$  is shorthand for the tuple  $\left(K_{jk}^{(i)},V_{jk}^{(i)}\right)$ . 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_{jk}^{(i)}\in \mathbb{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^{(\ell)}$  and the columns of  $W^{(ell + 1)}$ . If the rows of  $W^{(\ell)}$  and the columns of  $W^{(\ell +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^{(\ell)}$  and  $W^{(\ell +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  $\{\phi^{(i)} \in \mathbb{R}^c \mid i \in [1..L]\}$  to prevent equivariance to such permutations:
$$
\text {L A Y E R E N C} \left(W; \{\phi^ {(i)} \} _ {i = 1} ^ {L}\right) _ {j k} ^ {(i)} = W _ {j k} ^ {(i)} + \phi^ {(i)}. \tag {8}
$$
We can now compose self-attention with position encoding to obtain SA  $\circ$  LAYERENC.
Theorem 3.1 (Minimal equivariance). The combined layer SA  $\circ$  LAYERENC is minimally  $S_{NP}$ -equivariant.
Figure 3: An illustration of our weight-space cross-attention layer, which pools weight-space features into a set of  $S_{\mathrm{NP}}$ -invariant vectors. The layer uses a set of learned queries  $Q \in \mathbb{R}^{M \times c}$  to attend over keys and values produced from the weight space features.
Proof (Sketch). To show equivariance: LAYERENC is clearly equivariant to neuron permutations. We can check  $\mathrm{SA}(\sigma W)_{jk}^{(i)} = \mathrm{SA}(W)_{\sigma_i^{-1}(j),\sigma_{i - 1}^{-1}(k)}^{(i)}$  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_{jk}^{(i)} \in \mathbb{R}^c$  and requires  $O\left((\dim \mathcal{W})^2 c\right)$  operations, which is usually intractable. In practice, we replace this term by a self-attention between the row-column sums  $W_{\star,\star}^{(i)} = \sum_{j,k} W_{jk}^{(i)}$  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_{\mathrm{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_{\mathrm{NP}}$ -invariant 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  $\mathbf{CA}:\mathcal{W}^c\to \mathbb{R}^d$  is simply a cross attention between a learned query  $e\in \mathbb{R}^d$  and the set keys and values produced from weight-space features  $\{W_{jk}^{(i)}\in \mathbb{R}^c\}$ :
$$
\operatorname {C A} (W; e, \theta_ {K}, \theta_ {V}) = \operatorname {A T T N} \left(e, \left\{\left(\theta_ {K} W _ {p q} ^ {(s)}, \theta_ {V} W _ {p q} ^ {(s)}\right) \mid \forall s, p, q \right\}\right), \tag {9}
$$
with  $\theta_K, \theta_V \in \mathbb{R}^{d \times c}$  being learned projection matrices. By repeating this operation with  $M$  different learned embeddings  $e_1, \dots, e_M$ , we can easily extend this into an invariant map  $\mathcal{W}^c \to \mathbb{R}^{M \times 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)}\in \mathbb{R}^{n_i\times n_{i - 1}\times k_i\times c}$ . As in Zhou et al. [46], the filter dimension  $k_{i}$  can be folded into the channel dimension, creating a feature in  $\mathbb{R}^{n_i\times n_{i - 1}\times (k_ic)}$  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  $\tilde{c}$  and using learned linear projections  $\mathrm{PROJ}_i:\mathbb{R}^{n_i\times n_{i - 1}\times (k_ic)}\to \mathbb{R}^{n_i\times n_{i - 1}\times \tilde{c}}$  to guarantee matching channel dimensions. The output of SA(\cdot) can then be restored to the original channel dimension by another set of learned linear projections UNPROJ:  $\mathbb{R}^{n_i\times n_{i - 1}\times \tilde{c}}\rightarrow \mathbb{R}^{n_i\times n_{i - 1}\times k_i\times c}$ .
3.4 Building Neural Functional Transformers
Following Transformer architecture design [41], we form stackable "blocks" that combine weight-space self attention with LayerNorm, pointwise MLPs, and residual connections. Each block maps  $\mathcal{W}^c\rightarrow \mathcal{W}^c$  and preserves  $S_{\mathrm{NP}}$ -equivariance:
$$
\operatorname {B L O C K} (W) = Z + \operatorname {M L P} (\operatorname {L N} (Z)) \tag {10}
$$
$$
Z = W + \operatorname {S A} (\operatorname {L N} (W)), \tag {11}
$$
where both the MLP:  $\mathbb{R}^c\to \mathbb{R}^c$  and LayerNorm LN:  $\mathbb{R}^c\rightarrow \mathbb{R}^c$  operate "pointwise," or independently on each input  $W_{jk}^{(i)}$ . Since pointwise operations are permutation equivariant, the overall block is an  $S_{\mathrm{NP}}$ -equivariant map on  $\mathcal{W}^c$ . To build neural functional Transformers (NFTs), we stack multiple blocks into a deep equivariant architecture. For tasks which require  $S_{\mathrm{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( $\cdot$ ). The invariant activations produced by CA( $\cdot$ ) 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: \mathbb{R}^{c_1} \to \mathbb{R}^{c_2}$  to each  $W_{jk}^{(i)} \in \mathbb{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:
1. We implement the encoder with an  $S_{\mathrm{NP}}$ -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 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} \cup \dots \cup P_{M} = [-1,1]^{2}$ . Given a SIREN  $(\cdot; W)$ , INR2ARRAY uses an  $S_{\mathrm{NP}}$ -invariant NFT encoder  $\mathrm{ENC}_{\theta}(\cdot)$  to map weights  $W$  to a latent array of  $M$  vectors,  $z \in \mathbb{R}^{M \times d}$ . The decoder  $\mathrm{DEC}_{\theta}: \mathbb{R}^d \to \mathcal{W}$  produces a set of weights  $\left\{\hat{W}_i \mid i = 1, \dots, M\right\}$ , one for each vector  $z_i := z_{i,:} \in \mathbb{R}^d$ . Each  $\hat{W}_i(z_i)$  is responsible for parameterizing the SIREN only for the spatial patch  $P_i \subset [-1,1]^2$ . The objective
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.