SlowGuess/NeurIPS_2023_OCR huggingface dataset:
1. If  $\tau$  does not preserve layer index, LAYERENC will be non-equivariant since each layer has a different encoding.
2. Suppose there is an  $(i,j,k)$  and  $q$  such that  $\tau(i,j,k) = (i,j',k')$ , but  $\tau(i,j,q)_r \neq j'$ , and let  $\theta_Q = \theta_K = \theta_V = I$ . Now consider the input  $W \in \mathcal{W}$  such that all weights are 0 except for  $W_{jk}^{(i)} = W_{jq}^{(i)} = 1$ . Then we can check that  $[\tau \mathrm{SA}(W)]_{i,j',k'} \neq \mathrm{SA}(\tau W)_{i,j',k'}$ , so the layer is not equivariant to  $\tau$ . We can construct a similar argument for the other case, where there is an  $(i,j,k)$  and  $p$  such that  $\tau(i,j,k) = (i,j',k')$  but  $\tau(i,p,k)_c \neq k'$ .
3. In this case there must be an  $(i,j,k) \in \mathbb{I}$  and  $q$  such that  $\tau(i,j,k) = (i,j',k')$  but  $\tau(i - 1,k,q)_r \neq k'$ . Again let  $\theta_Q = \theta_K = \theta_V = I$  and consider  $W \in \mathcal{W}$  such that all weights are 0 except for  $W_{jk}^{(i)} = W_{kq}^{(i-1)} = 1$ . Then, similar to the above case, we can verify that  $[\tau \mathrm{SA}(W)]_{i,j',k'} \neq \mathrm{SA}(\tau W)_{i,j',k'}$ , so the layer is not equivariant to such  $\tau$ .
Taken together, the three cases above show that although SA  $\circ$  LAYERENC is  $S_{\mathrm{NP}}$ -equivariant, it is not equivariant to any other permutations  $\tau \notin S_{\mathrm{NP}}$ . Hence the layer is minimally equivariant.
A.2 Full description including biases
Let  $\mathcal{B}$  be the space of biases, and  $\mathcal{U} = \mathcal{W} \times \mathcal{B}$  the full space of both weights and biases. Then the full weight-space self-attention layer  $\mathrm{SA}: \mathcal{U}^c \to \mathcal{U}^c$  is defined, for input  $U = (W, b)$ :
$$
\operatorname {S A} (U; \theta_ {Q}, \theta_ {K}, \theta_ {V}) = (Y (U), z (U)), \tag {22}
$$
where  $Y(U)\in \mathcal{W}^c$  and  $z(U)\in \mathcal{B}^c$ , with:
$$
Y (W, b) _ {j k} ^ {(i)} = \operatorname {A T T N} \left(Q _ {j,:} ^ {(i)}, K V _ {1}\right) _ {k} + \operatorname {A T T N} \left(Q _ {:, k} ^ {(i)}, K V _ {2}\right) _ {j} + \operatorname {A T T N} \left(Q _ {j k} ^ {(i)}, K V _ {3}\right), \tag {23}
$$
$$
z (W, b) _ {j} ^ {(i)} = \operatorname {A T T N} \left(q ^ {(i)}, K V _ {2}\right) _ {j} + \operatorname {A T T N} \left(q _ {j} ^ {(i)}, K V _ {3}\right), \text {w h e r e} \tag {24}
$$
$$
K V _ {1} = \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}} \bigcup \left\{\left(k, v\right) ^ {(i - 1)} \right\} \tag {25}
$$
$$
K V _ {2} = \left\{ \right.\left. (K, V) _ {:; q} ^ {(i)} \right\} _ {q = 1} ^ {n _ {i - 1}} \bigcup \left\{\left. (K, V) _ {p,:} ^ {(i + 1)} \right\} _ {p = 1} ^ {n _ {i + 1}} \bigcup \left\{\left. (k, v) ^ {(i)} \right\}\right.\right. \tag {26}
$$
$$
K V _ {3} = \left\{\left. (K, V) _ {p q} ^ {(s)} \mid \forall s, p, q \right\} \bigcup \left\{\left. (k, v) _ {p} ^ {(s)} \mid \forall s, p \right\}, \text {a n d} \right. \right. \tag {27}
$$
$$
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)}, \tag {28}
$$
$$
q _ {j} ^ {(i)} := \theta_ {Q} b _ {j} ^ {(i)}, \quad k _ {j} ^ {(i)} := \theta_ {K} b _ {j} ^ {(i)}, \quad v _ {j} ^ {(i)} := \theta_ {V} b _ {j} ^ {(i)}. \tag {29}
$$
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.
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.
Hyperparameters	Predicting gen.	Editing INRs	INR2ARRAY
# blocks	4	4	6
# channels	256	256	256
MLP hidden dim	1024	512	1024
Fourier scale	3	3	3
Fourier size	128	128	128
# attn heads	4	8	4
dropout p	0.1	0.1	0
invariant layer	HNP	-	CA
CA M	-	-	16
CA dim	-	-	256
Optimizer	Adam	Adam	AdamW
Learning rate	0.001	0.001	0.0001
Weight decay coeff	0	0	0.01
LR warmup steps	10K	10K	10K
Total params	46M	7M	22M
Training steps	75K	50K	200K
Training time	5H	2H	13H
Here we summarize the hyperparameters (and other training information) for the NFTs in each experiment in this paper. NFTs are described by the following hyperparameters:
- # blocks: Number of equivariant weight-space blocks  $\text{BLOCK}(\cdot)$
- # 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_{\mathrm{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, \dots, e_M$ .
- CA dim: Dimension of vectors going into cross attention.
Figure 5: Samples of target image transformations, as well as original and edited INRs for each editing task in Table 3. Despite the fact that the NFT quantitatively outperforms the  $\mathrm{NFN_{HNP}}$  and  $\mathrm{NFN_{NP}}$ , compelling qualitative differences in the edited images can be difficult to identify due to small image size.
B.1 Generalization prediction
Each Small CNN Zoo dataset that 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.
Table 5: Mean squared error to ground truth visual transformation for five INR editing tasks. NFTs generally achieve lower both lower training and test error than NFN variants without using more parameters, suggesting greater expressivity from attention. Uncertainties indicate standard error over three seeds.
\( \mathrm{NFN}_{\mathrm{HNP}}\left[46\right] \)	\( \mathrm{NFN}_{\mathrm{NP}}\left[46\right] \)	NFT (Ours)
MNIST(erode)	Train	0.0217 ± 0.0004	0.0212 ± 0.0002	0.0195 ± 0.0004
Test	0.0225 ± 0.0001	0.0223 ± 0.0000	0.0194 ± 0.0002
MNIST(dilate)	Train	0.0575 ± 0.0010	0.0671 ± 0.0017	0.0473 ± 0.0006
Test	0.0656 ± 0.0000	0.0693 ± 0.0009	0.0510 ± 0.0004
MNIST(gradient)	Train	0.0536 ± 0.0020	0.0547 ± 0.0009	0.0474 ± 0.0005
Test	0.0552 ± 0.0003	0.0566 ± 0.0000	0.0484 ± 0.0007
FashionMNIST(gradient)	Train	0.0851 ± 0.0012	0.0853 ± 0.0008	0.0795 ± 0.0009
Test	0.0859 ± 0.0001	0.0870 ± 0.0001	0.0800 ± 0.0002
CIFAR (contrast)	Train	0.0191 ± 0.0002	0.0185 ± 0.0002	0.0192 ± 0.0005
Test	0.0204 ± 0.0000	0.0203 ± 0.0000	0.0200 ± 0.0002
Table 6: Mean squared error (MSE) of INR2ARRAY reconstructions on test INR weights. Using an NFT encoder in INR2ARRAY achieves better reconstruction error, produces higher quality samples, and leads to better downstream performance.
INR2ARRAYNFN	INR2ARRAYNFT
MNIST	0.046 ± 0.001	0.027 ± 0.003
FashionMNIST	0.085 ± 0.006	0.070 ± 0.006
CIFAR-10	0.117 ± 0.008	0.036 ± 0.006
B.4 INR Classification
After training INR2ARRAY on a given dataset, we select the checkpoint with the smallest reconstruction error on the validation set, and use the encoder to produce a  $16 \times 256$  latent array representation for any input SIREN. We view the latent as a length-16 sequence of 256-dim vectors, and project each vector into  $\mathbb{R}^{512}$  before feeding the array into Transformer encoder with 12 blocks. This produces an output  $z^o \in \mathbb{R}^{16 \times 512}$ ; we take the first output  $z_0^o \in \mathbb{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.
Table 7: Classification train and test accuracies  $(\%)$  on datasets of image INRs (MNIST, FashionMNIST, and CIFAR-10). INR2ARRAY outperforms prior state-of-the-art results in this setting (DWS and NFN). For the NFN baseline [46] we report the higher performance out of their NP and HNP variants. Uncertainties indicate standard error over three runs.
MNIST	FashionMNIST	CIFAR-10
Train	Test	Train	Test	Train	Test
DWS [31]	-	85.7 ± 0.57	-	65.5 ± 0.48	-	-
NFN [46]	96.2 ± 0.24	92.9 ± 0.38	84.1 ± 1.02	75.6 ± 1.07	61.0 ± 6.60	46.6 ± 0.13
Spatial NFN (Ours)	95.3 ± 0.59	92.9 ± 0.46	76.4 ± 0.87	70.8 ± 0.53	52.3 ± 0.30	45.6 ± 0.11
INR2ARRAYNFN (Ours)	100 ± 0.00	94.6 ± 0.00	92.8 ± 0.00	76.7 ± 0.00	98.5 ± 0.00	45.4 ± 0.00
INR2ARRAYNFT (Ours)	100 ± 0.00	98.5 ± 0.00	97.6 ± 0.00	79.3 ± 0.00	100 ± 0.00	63.4 ± 0.00