SlowGuess/NeurIPS_2023_OCR huggingface dataset:
In Table 2, we show the results on two challenging subsets of Small CNN Zoo corresponding to CNNs trained on CIFAR-10-GS and SVHN-GS (GS stands for grayscale). We see that  $\mathrm{NFN_{HNP}}$  consistently performs the best on both datasets by a significant margin, showing that having access to the full weights can increase predictive power over hand-designed features as in STATNN. Because the input and output dimensionalities are small on these datasets,  $\mathrm{NFN_{HNP}}$  only uses moderately more ( $\sim 1.4\times$ ) parameters than  $\mathrm{NFN_{NP}}$  with equivalent depth and channel dimensions, while having significantly better performance.
3.2 Classifying implicit neural representations of images and 3D shapes
Given the rise of implicit neural representations (INRs) that encode data such as images and 3Dscenes [61, 43, 7, 49, 58, 45, 14, 15], it is natural to wonder how to extract information about the
Table 3: Classification test accuracies (\%) for implicit neural representations of MNIST, FashionMNIST, and CIFAR-10. Equivariant architectures such as NFNs and DWSNets [47] outperform the non-equivariant MLP baselines, even when the MLP has permutation augmentations. Our  $\mathrm{NFN_{NP}}$  variant consistently outperforms all other methods across each dataset. Uncertainties indicate standard error over three runs.
NFHNP	NFNNP	DWSNets	MLP	MLPAug
MNIST	92.5 ± 0.071	92.9 ± 0.218	74.4 ± 0.143	14.5 ± 0.035	21.0 ± 0.172
FashionMNIST	72.7 ± 1.53	75.6 ± 1.07	64.8 ± 0.685	12.5 ± 0.111	15.9 ± 0.181
CIFAR-10	44.1 ± 0.471	46.6 ± 0.072	41.5 ± 0.431	16.9 ± 0.250	18.9 ± 0.432
Table 4: Classification test accuracies  $(\%)$  for datasets of implicit neural representations (INRs) of either ShapeNet-10 [6] or ScanNet-10 [10]. Our equivariant NFNs outperform the MLP baselines and recent non-equivariant methods such as inr2vec [12]. Uncertainties indicate standard error over three runs.
NFHNP	NFNNP	MLP	MLPAug	inr2vec[12]
ShapeNet-10	86.9 ± 0.860	88.7 ± 0.461	25.4 ± 0.121	33.8 ± 0.126	39.1 ± 0.385
ScanNet-10	64.1 ± 0.572	65.9 ± 1.10	32.9 ± 0.351	45.5 ± 0.126	38.2 ± 0.409
original data directly from INR weights. Compared to discrete signal representations (pixels, voxels, etc.), the advantage of directly processing INRs is that one can easily be agnostic to varying signal sizes and resolutions.
In this task, our goal is to classify the contents of INRs given only the weights as input. We consider datasets of SIRENs [58] that encode images (MNIST [38], FashionMNIST [66], and CIFAR [35]) and 3D shapes (ShapeNet-10 and ScanNet-10 [53]). For image datasets each SIREN network represents the mapping from pixel coordinate to RGB (or grayscale) value for a single image, while for 3D shapes each network is a signed (or unsigned) distance function encoding a single shape. Each dataset of SIREN weights is split into training, validation, and testing sets. We construct and train invariant neural functionals to classify the INRs, and compare their performance against the MLP and  $\mathrm{MLP_{Aug}}$  baselines, which are three-layer MLPs with ReLU activations and 1,000 hidden units per layer.
As Navon et al. [47] test the performance of DWSNets on their independently constructed 2D-image INR datasets, we also present the results of training DWSNets on our own 2D-image INR datasets. Table 3 show that  $\mathrm{NFN_{NP}}$  consistently achieves the highest test accuracies of any method on the 2D-image tasks. More broadly, equivariant architectures significantly outperform the non-equivariant MLP approaches, even with permutation data augmentations.
For the 3D-shape datasets we also report the performance of inr2vec [12], a recent non-equivariant method for classifying 3D shapes from INR weights. Note that inr2vec's original setting assumes that all INRs in a dataset are trained from the same shared initialization, whereas our problem setting makes no such assumption and allows INRs to be trained from random and independent initializations. As expected, Table 4 shows that  $\mathrm{NFN_{HNP}}$  and  $\mathrm{NFN_{NP}}$  achieve significantly higher test accuracies than inr2vec, as well as the non-equivariant MLP baselines.
In addition to superior generalization, Appendix Table 16 shows that the NFNs are also better at fitting the training data compared to non-equivariant architectures. The MLPs achieve low train accuracy, even with an equal number of parameters as the NFNs. Interestingly,  $\mathrm{NFN_{NP}}$  matches or exceeds  $\mathrm{NFN_{HNP}}$  performance on both CIFAR-10 and the 3D-shape datasets while using fewer parameters (e.g.,  $35\%$  as many parameters on CIFAR-10). Finally, we note that no weight space methods (including NFN) match the performance of near state-of-the-art methods on discrete data representations such as 2D image arrays and point clouds (see Section E.2 for details). All weight space methods still lack the geometric inductive biases that, e.g., convolutional networks have in image tasks.
Table 5: Test accuracy (\%) of training with winning tickets (95% sparsity masks) produced either by running IMP or predicted by an NFN. We also show the performance of Random ticket (random mask of equivalent sparsity level), and Dense training (no sparsity). We show results for MLPs (trained on MNIST) and CNNs (trained on CIFAR-10). Uncertainties show standard error over initializations.
Dense	IMP	Random	NFNNP	NFNPT
CIFAR-10	63.1 ± 0.06	44.0 ± 0.06	21.1 ± 0.26	41.4 ± 0.08	42.6 ± 0.07
MNIST	97.8 ± 0.0	96.2 ± 0.04	89.6 ± 0.36	94.8 ± 0.01	95.0 ± 0.01
Figure 3: In weight-space style editing, an NFN directly edits the weights of an INR to alter the content it encodes. In this example, the NFN edits the weights to dilate the encoded image.
Method	Contrast (CIFAR-10)	Dilate (MNIST)
MLP	0.031	0.306
MLPAug	0.029	0.307
NFNPT	0.029	0.197
NFNHNP	0.021	0.070
NFNNP	0.020	0.068
Table 6: Test mean squared error (lower is better) between weight-space editing methods and ground-truth image-space transformations.
3.3 Predicting "winning ticket" masks from initialization
The Lottery Ticket Hypothesis [19, 20, LTH] conjectures the existence of winning tickets, or sparse initializations that train to the same final performance as dense networks, and showed their existence in some settings through iterative magnitude pruning (IMP). IMP retroactively finds a winning ticket by pruning trained models by magnitude; however, finding the winning ticket from only the initialization without training remains challenging.
We demonstrate that permutation equivariant neural functionals are a promising approach for finding winning tickets at initialization by learning over datasets of initializations and their winning tickets. Let  $U_0 \in \mathcal{U}$  be an initialization and let the sparsity mask  $M \in \{0,1\}^{\dim(\mathcal{U})}$  be a winning ticket for the initialization, with zeros indicating that the corresponding entries of  $U_0$  should be pruned. The goal is to predict a winning ticket  $\hat{M}$  given a held out initialization  $U_0$ , such that the MLP initialized with  $U_0$  and sparsity pattern  $\hat{M}$  will achieve a high test accuracy after training.
We construct a conditional variational autoencoder [32, 59, cVAE] that learns a generative model of the winning tickets conditioned on initialization and train on datasets of (initialization, ticket) pairs found by one step of IMP with a sparsity level of  $P_{m} = 0.95$  for both MLPs trained on MNIST and CNNs trained on CIFAR-10. Table 5 compares the performance of tickets predicted by equivariant neural functionals against IMP tickets and random tickets. We generate random tickets by randomly sampling sparsity mask entries from Bernoulli  $(1 - P_{m})$ . In this setting,  $\mathrm{NFN}_{\mathrm{HNP}}$  is prohibitively parameter inefficient, but  $\mathrm{NFN}_{\mathrm{NP}}$  is able to recover test accuracies that are close to that of IMP pruned networks in CIFAR-10 and MNIST, respectively. Somewhat surprisingly,  $\mathrm{NFN}_{\mathrm{PT}}$  performs just as well as the other NFNs, indicating that one can approach IMP performance in these settings without considering interactions between weights or layers. Appendix E.4 further analyzes how  $\mathrm{NFN}_{\mathrm{PT}}$  learns to prune.
3.4 Weight-space style editing
Another potentially useful application of neural functionals is to edit (i.e., transform) the weights of a given INR to alter the content that it encodes. In particular, the goal of this task is to edit the weights of a trained SIREN to alter its encoded image (Figure 3). We evaluate two editing tasks:
Figure 4: Random qualitative samples of INR editing behavior on the Dilate (MNIST) and Contrast (CIFAR-10) editing tasks. The first row shows the image produced by the original INR, while the rows below show the result of editing the INR weights with an NFN. The difference between MLP neural functionals and equivariant neural functionals is especially pronounced on the more challenging Dilate tasks, which require modifying the geometry of the image. In the Contrast tasks, the MLP baseline produces dimmer images compared to the ground truth, which is especially evident in the second and third columns.
(1) making MNIST digits thicker via image dilation (Dilate), and (2) increasing image contrast on CIFAR-10 (Contrast). Both of these tasks require neural functionals to process the relationships between different pixels to successfully solve the task.
To produce training data for this task, we use standard image processing libraries [28, OpenCV] to dilate or increase the contrast of the MNIST and CIFAR-10 images, respectively. The training objective is to minimize the mean squared error between the image generated by the NFN-edited INR and the image produced by image processing. We construct equivariant neural functionals to edit the INR weights, and compare them against MLP-based neural functionals with and without permutation augmentation.
Table 6 shows that permutation equivariant neural functionals  $(\mathrm{NFN}_{\mathrm{HNP}}$  and  $\mathrm{NFN_{NP}}$  ) achieve significantly better test MSE when editing held out INRs compared to other methods, on both the Dilate (MNIST) and Contrast (CIFAR-10) tasks. In other words, they produce results that are closest to the "ground truth" image-space processing operations for each task. The pointwise ablation  $\mathrm{NFN_{PT}}$  performs significantly worse, indicating that accounting for interactions between weights and layers is important to accomplishing these tasks. Figure 4 shows random qualitative samples of editing by different methods below the original (pre-edit) INR. We observe that NFNs are more effective than  $\mathrm{MLP_{Aug}}$  at dilating MNIST digits and increasing the contrast in CIFAR-10 images.
The permutation symmetries of neurons have been a topic of interest in the context of loss landscapes and model merging [21, 4, 62, 17, 1]. Other works have analyzed the degree of learned permutation symmetry in networks that process weights [64] and studied ways of accounting for symmetries when measuring or encouraging diversity in the weight-space [13]. However, these symmetries have not been a key consideration in architecture design for processing weight-space objects [2, 39, 23, 36, 69, 13, 33]. Instead, existing approaches try to encourage permutation equivariance through data augmentation [51, 44]. In contrast, this work directly encodes the equivariance of the weight-space into our architecture design, which can result in much higher data and computational efficiency, as evidenced by the success of convolutional neural networks [37].
Our work follows a long line of literature that incorporates structure and symmetry into neural network architectures [37, 8, 54, 34, 9, 18]. This includes works that design permutation equivariant layers, originally for processing sets and graphs [52, 68, 25, 41]. More generally, equivariant layers have been developed for processing arbitrary rank- $k$  tensors (and correspondingly, hypergraphs) under higher-order permutation actions [40, 63, 48]. In this context, we can view feedforward networks as graphs that factor into a special layered structure where each (hidden) layer of neurons is an independently permutable set of nodes, and the weights are adjacency matrices specifying the edge weights between nodes of adjacent layers. Then our equivariant NF-Layers give the maximal set
of linear functions on these special graph structures. As discussed in Section 1, Navon et al. [47] recently developed an equivariant weight-space layer that is equivalent to our NF-Layer in the HNP setting. Our work introduces the NP setting to improve parameter efficiency and scalability over the HNP setting, and extends beyond the fully connected case to handle convolutional weight-space inputs.
This paper proposes a novel symmetry-inspired framework for the design of neural functional networks (NFNs), which process weight-space features such as weights, gradients, and sparsity masks. Our framework focuses on the permutation symmetries that arise in weight-spaces due to the particular structure of neural networks. We introduce two equivariant NF-Layers as building blocks for NFNs, which differ in their underlying symmetry assumptions and parameter efficiency, then use them to construct a variety of permutation equivariant neural functionals. Experimental results across diverse settings demonstrate that permutation equivariant neural functionals outperform prior methods and are effective for solving weight-space tasks.
Limitations and future work. Although we believe this framework is a step toward the principled design of effective neural functionals, there remain multiple directions for improvement. One such direction would concern extending the NF-Layers beyond feedforward weight spaces, in order to process the weights of more complex architectures such as ResNets [26] and Transformers [65]. Another useful direction would involve reducing the activation sizes produced by NF-Layers, in order to scale neural functionals to process the weights of very large networks. Finally, improvements to NF-Layers could account for the other symmetries of neural network weight spaces, such as scaling symmetries in ReLU networks [22].
We thank Andy Jiang and Will Dorrell for helpful early discussions about theoretical concepts, and Yoonho Lee and Eric Mitchell for reviewing an early draft of this paper. AZ and KB are supported by the NSF Graduate Research Fellowship Program. This work was also supported by Apple.
[1] S. K. Ainsworth, J. Hayase, and S. Srinivasa. Git re-basin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209.04836, 2022.
[2] M. Andrychowicz, M. Denil, S. Gomez, M. W. Hoffman, D. Pfau, T. Schaul, B. Shillingford, and N. De Freitas. Learning to learn by gradient descent by gradient descent. Advances in neural information processing systems, 29, 2016.
[3] S. Bengio, Y. Bengio, J. Cloutier, and J. Gescei. On the optimization of a synaptic learning rule. In Optimality in Biological and Artificial Networks?, pages 281-303. Routledge, 2013.
[4] J. Brea, B. Simsek, B. Illing, and W. Gerstner. Weight-space symmetry in deep networks gives rise to permutation saddles, connected by equal-loss valleys across the loss landscape. arXiv preprint arXiv:1907.02911, 2019.
[5] M. M. Bronstein, J. Bruna, T. Cohen, and P. Velicković. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478, 2021.
[6] A. X. Chang, T. Funkhouser, L. Guibas, P. Hanrahan, Q. Huang, Z. Li, S. Savarese, M. Savva, S. Song, H. Su, J. Xiao, L. Yi, and F. Yu. ShapeNet: An Information-Rich 3D Model Repository. Technical Report arXiv:1512.03012 [cs.GR], Stanford University — Princeton University — Toyota Technological Institute at Chicago, 2015.
[7] Z. Chen and H. Zhang. Learning implicit fields for generative shape modeling. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 5939-5948, 2019.
[66] H. Xiao, K. Rasul, and R. Vollgraf. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms, 2017.
[67] S. Yak, J. Gonzalvo, and H. Mazzawi. Towards task and architecture-independent generalization gap predictors. arXiv preprint arXiv:1906.01550, 2019.
[68] M. Zaheer, S. Kottur, S. Ravanbakhsh, B. Poczos, R. Salakhutdinov, and A. Smola. Deep sets. doi: 10.48550. arXiv preprint ARXIV.1703.06114, 2017.
[69] C. Zhang, M. Ren, and R. Urtasun. Graph hypernetworks for neural architecture search. arXiv preprint arXiv:1810.05749, 2018.
A Equivariant NF-Layer pseudocode 15
B NF-Layers for the NP Setting 16
B.1 Full definition 16
B.2 Proof of Proposition 1 17
B.2.1 Parameter orbits and sharing 18
B.2.2 Equivalence to equivariant NF-Layer definition 19
B.3 Proof of Proposition 2 19
C NF-Layers for the HNP setting 20
C.1 Equivariant NF-Layer 20
C.2 Invariant NF-Layer 20
D Additional experimental details 20
D.1 Predicting generalization 20
D.2 Predicting "winning ticket" masks from initialization 22
D.3 Classifying INRs 23
D.4 Weight-space style editing 23
E Additional experiments and analysis 24
E.1 Analyzing train vs test accuracies in INR classification 24
E.2 Non-weight-space baselines for image classification 24
E.3 Ablating positional encoding 24
E.4 Interpreting learned lottery ticket masks 24
E.5 Predicting MLP generalization from weights 27
A Equivariant NF-Layer pseudocode
Here we present a multi-channel implementation of the  $S$ -equivariant NF-Layer presented in Eq. 2 (which ignores biases), using PyTorch [50] and Einops-like [55] pseudocode. That is, it implements a linear layer  $H: \mathcal{W}^{c_i} \to \mathcal{W}^{c_o}$ , where  $c_i$  and  $c_o$  are the number of input and output channels.
Note that our actual implementation differs from this pseudocode in a few ways: (1) it supports the full weight-space  $\mathcal{U}$  which includes biases, (2) it supports convolution weights as well as fully connected weights, and (3) it initializes parameters based on the fan-in of the NF-Layer, instead of from  $\mathcal{N}(0,1)$ .
1 class NLayer(nnModule):  
2 def __init__(self, L, co, ci):  
3 super().__init_(   )  
4 # initialize weights. co=output channels, ci=input channels.  
5 self.A = nn_PARAMETER(torch.random(L, L, co, ci))  
6 self.B = nn_PARAMETER(torch.random(L, co, ci))  
7 self.B Prev = nn_PARAMETER(torch.random(L, co, ci))  
8 self.C = nn_PARAMETER(torch.random(L, co, ci))  
9 self.C_next = nn_PARAMETER(torch.random(L, co, ci))
def forward(self, W): # Input W is a list of L weight-space tensors with shapes: # [(B, ci, n1, n0), ..., ((B, ci, nL, nL-1))] # We return a list of L tensors with shapes: # [(B, co, n1, n0), ..., ((B, co, nL, nL-1))] # where ci, co are the input and output channels. # Compute [W(1), ..., W(L)]. Each has shape (B, ci, ni-1). row_means = [w.mean(-2) for w in W] # Compute [W(1), ..., W(L)]. Each has shape (B, ci, ni). col_means = [w.mean(-1) for w in W] # Compute [W(1), ..., W(L)] with shape (B, ci, L). rowcol_means = torch.stack([w.mean(dim=(-2, -1)) for w in W], -1) out_W = [] for i, (w, a, b, b Prev, c, c_next, d) in enumerate(zip(W, self.A, self.B, self.B Prev, self.C, self.C_next, self.D)): #Calculates  $\sum_{s}a^{i,s}W_{\star,\star}^{(s)}$ . h1 = einsum(a, rowcol_means, 'co ci s, B ci s -> B co () ()') #Calculates  $b^{i,i}W_{\star,k}^{(i)}$ . h2 = einsum(b, row_means[i], 'co ci, b ci n_im1 -> B co () n_im1') if i > 0: #Calculates  $b^{i,i-1}W_{k,\star}^{(i-1)}$ . h2 += einsum(b Prev, col_means[i-1], 'co ci, B ci n_im1 -> B co () n_im1') #Calculates  $c^{i,i}W_{j,\star}^{(i)}$ . h3 = einsum(c, col_means[i], 'co ci, B ci n_i -> b co n_i ()') if i < self.L - 1: #Calculates  $c^{i,i+1}W_{\star,j}^{(i+1)}$ . h3 += einsum(c_next, row_col[i+1], 'co ci, B ci n_i -> B co n_i ()') #Calculates  $d^{i}W_{jk}^{(i)}$ . h4 = einsum(d, w, 'co ci, B ci n_i n_im1 -> B co n_i n_im1') #Calculates H(W(i)) with shape (B, co, ni, ni-1). out_W.append(h1 + h2 + h3 + h4) return out_W
B NF-Layers for the NP Setting
This section presents further details related to the  $S$ -equivariant layer  $H: \mathcal{U} \to \mathcal{U}$  and the  $S$ -invariant layer  $P$ . First, Sec. B.1 gives the full definition of  $H$  as a function of both weight-space and bias-space features. Sec. B.2 provides a proof that  $H(\cdot)$  expresses any  $S$ -equivariant linear function on  $\mathcal{U}$  (Proposition 1). Sec. B.3 proves that combining our invariant layer  $P$  with a fully connected layer can express any  $S$ -invariant function from  $\mathcal{U}$  to  $\mathbb{R}^d$  (Proposition 2).
B.1 Full definition
We present the full definition of the  $S$ -equivariant NF-Layer  $H: \mathcal{U} \to \mathcal{U}$ , which completes Eq. 2 by also including the biases. The layer processes a set of weights  $U = (W, v)$  and outputs arrays  $(Y, z)$
10   
11   
12   
13   
14   
15   
16   
17   
18   
19   
20   
21   
22   
23   
24   
25   
26   
27   
28   
29   
30   
31   
32   
33   
34   
35   
36   
37   
38   
39   
40   
41   
42   
43   
44   
45
Table 7:  $S$ -equivariant parameter sharing for linear maps  $\operatorname{vec}(U) \mapsto \theta \operatorname{vec}(U)$ . Parameter sharing is a system of constraints on the entries of  $(\vartheta, \phi, \varphi, \psi) = \theta$ . Each table is organized by the layer indices  $(i, s)$ . For example, the first table says that for any  $(i, s)$  where  $s = i - 1$ , we constrain  $\vartheta_{spq}^{ijk} = \vartheta_{s,p',q'}^{i,j',k'} = a_{\vartheta}^{i,i - 1}$  for any  $j, k, p, q$  and  $j', k', p', q'$  where  $k \neq p$  and  $k' \neq p'$ . After parameter sharing, we observe that there are only a constant number of free parameters for each  $(i, s)$  pair, adding up to  $O(L^2)$  parameters total.
s=i-1	s=i	s=i+1	other s
θijkspq	{a0i,i-1θj}k≠p	{a0i, j≠p,k≠q	{b0i, j≠p,k≠q	{a0i, j≠q}c0i,k=q	{a0i, j≠q}c0i,k=q	{a0i,s}j=q	a0i,sθj
b0i, i-1θj=k=p	d0i,j=p,k=q
s=i-1	s=i	other s
φijk/sp	{aφi,i-1k≠p	{aφi,ij≠p	{aφj≠p	aφi,sj}
bφi,i-1k=p	bφi,k=j=p	aφj}
s = i	s = i + 1	other s
φijspq	{a^i}_{\varphi}	j ≠ p	{a^i}_{\varphi}	j ≠ q	a^i}_{\varphi}
b^i}_{\varphi}	j = p	b^i}_{\varphi}	j = q
s = i	other s
ψijsp	{aiψi, j ≠ p}	ai,sψ
{bjψ, j = p}
with the same dimensions. The learnable parameters are in blue:
$$
H (U) := (Y, z) \in \mathcal {W} \times \mathcal {V} \tag {3}
$$
$$
\begin{array}{l} Y _ {j k} ^ {(i)} := \sum_ {s} \left(a _ {\vartheta} ^ {i, s} W _ {\star , \star} ^ {(s)} + a _ {\phi} ^ {i, s} v _ {\star} ^ {(s)}\right) + b _ {\vartheta} ^ {i, i} W _ {\star , k} ^ {(i)} + b _ {\vartheta} ^ {i, i - 1} W _ {k, \star} ^ {(i - 1)} \\ + b _ {\phi} ^ {i} v _ {j} ^ {(i)} + c _ {\vartheta} ^ {i, i} W _ {j, \star} ^ {(i)} + c _ {\vartheta} ^ {i, i + 1} W _ {\star , j} ^ {(i + 1)} + c _ {\phi} ^ {i} v _ {k} ^ {(i - 1)} + d _ {\vartheta} ^ {i} W _ {j k} ^ {(i)} \\ \end{array}
$$
$$
z _ {j} ^ {(i)} := \sum_ {s} \left(a _ {\varphi} ^ {i, s} W _ {\star , \star} ^ {(s)} + a _ {\psi} ^ {i, s} v _ {\star} ^ {(s)}\right) + b _ {\varphi} ^ {i, i} W _ {j, \star} ^ {(i)} + b _ {\varphi} ^ {i, i + 1} W _ {\star , j} ^ {(i + 1)} + b _ {\psi} ^ {i} v _ {j} ^ {(i)},
$$
where  $\star$  denotes summation or averaging over a dimension.
Here we show that the layer defined in Eq. 3 is not only  $S$ -equivariant, it can express every linear equivariant function  $\mathcal{U} \to \mathcal{U}$ . Our strategy is to first consider general linear layers  $T(\cdot; \theta): \mathcal{U} \to \mathcal{U}$  parameterized by  $\theta \in \Theta$ , and then constrain to  $\theta$  such that  $T(\cdot; \theta)$  is  $S$ -equivariant. Applying this constraint will directly yield Eq. 3.
For our purposes, it is sometimes convenient to distinguish layer, row, and column indices of entries in  $U$ , so we split the parameters  $\theta = (\vartheta, \phi, \varphi, \psi)$  and write  $T: \mathcal{U} \to \mathcal{U}$  in the form:
$$
T (U; \theta) := (Y (U), z (U)) \quad \in \mathcal {W} \times \mathcal {V} = \mathcal {U} \tag {4}
$$
$$
Y (U) _ {j k} ^ {(i)} := \sum_ {s = 1} ^ {L} \sum_ {p = 1} ^ {n _ {s}} \sum_ {q = 1} ^ {n _ {s - 1}} v _ {s p q} ^ {i j k} W _ {p q} ^ {(s)} + \sum_ {s = 1} ^ {L} \sum_ {p = 1} ^ {n _ {s}} \phi_ {s p} ^ {i j k} v _ {p} ^ {(s)} \tag {5}
$$