Equivariant Architectures for Learning in Deep Weight Spaces

Aviv NavonAviv ShamsianIdan AchituveEthan FetayaGal ChechikHaggai Maron

article2023ICML100 citations

Introduces a principled neural network architecture that directly processes raw weights and biases of other models by preserving their inherent permutation symmetries, enabling effective learning across tasks like implicit neural representation processing and model domain adaptation.

Bar-Ilan University NVIDIA broadcasting multilayer perceptrons architecture designs

arXiv: 2301.12780

Cover for Equivariant Architectures for Learning in Deep Weight Spaces

Abstract

Designing machine learning architectures for processing neural networks in their raw weight matrix form is a newly introduced research direction. Unfortunately, the unique symmetry structure of deep weight spaces makes this design very challenging. If successful, such architectures would be capable of performing a wide range of intriguing tasks, from adapting a pre-trained network to a new domain to editing objects represented as functions (INRs or NeRFs). As a first step towards this goal, we present here a novel network architecture for learning in deep weight spaces. It takes as input a concatenation of weights and biases of a pre-trained MLP and processes it using a composition of layers that are equivariant to the natural permutation symmetry of the MLP’s weights: Changing the order of neurons in intermediate layers of the MLP does not affect the function it represents. We provide a full characterization of all affine equivariant and invariant layers for these symmetries and show how these layers can be implemented using three basic operations: pooling, broadcasting, and fully connected layers applied to the input in an appropriate manner. We demonstrate the effectiveness of our architecture and its advantages over natural baselines in a variety of learning tasks.

Topics

broadcasting multilayer perceptrons architecture designs novel architecture system architectures layered

1. Introduction
2. Previous Work
3. Preliminaries
4. Permutation Symmetries of Neural Networks
5. A Characterization of Linear Invariant and Equivariant Layers for Weight-Spaces
5.1. Overview and Main Results
5.2. Linear Equivariant Maps for Direct Sums
5.3. Linear Equivariant Layers for Deep Weight-Spaces
5.4. Extension to Other Input Architectures
6. Expressive Power
7. Experiments
7.1. Results
7.2. Analysis of the Results
8. Conclusion and Future Work
9. Acknowledgements
References
A. Related Work
B. Multiple Channels, Invariant Layers and Biases for Equivariant Maps
C. Specification of All Affine Equivariant Layers Between Sub-Representations
D. Linear Maps Between Specific Weight and Bias Spaces
E. More Proofs for Section 5
F. Proofs of Proposition 6.1
G. Proof of Proposition 6.2
H. Alternative Characterization Strategies
I. Computational and Memory Requirements
J. Experimental and Technical Details
K. Additional Experiments
K.1. Predicting the Generalization Error of Neural Networks.
K.2. Dense Representation
K.3. Ablation Study
K.4. The importance of data augmentation and batch normalization
K.5. Challenging cases

Knowls

Knowl 1 — Permutation Symmetry Group and Action on MLP Weight Spaces
definition
Let an $M$ -layer MultiLayer Perceptron (MLP) $f: \mathbb{R}^{d_0} \to \mathbb{R}^{d_M}$ be defined by the recurrence:

$f(x) = x_M, \quad x_m = \sigma(W_m x_{m-1} + b_m) \text{ for } m \in \{1, \dots, M-1\}, \quad x_M = W_M x_{M-1} + b_M, \quad x_0 = x$

where $W_m \in \mathbb{R}^{d_m \times d_{m-1}}$ is the weight matrix and $b_m \in \mathbb{R}^{d_m}$ is the bias vector of the $m$ -th layer, and $\sigma$ is a pointwise activation function. The weight space $\mathcal{V}$ of the MLP is the direct sum:

$\mathcal{V} = \bigoplus_{m=1}^M (\mathcal{W}_m \oplus \mathcal{B}_m)$

where $\mathcal{W}_m = \mathbb{R}^{d_m \times d_{m-1}}$ and $\mathcal{B}_m = \mathbb{R}^{d_m}$ .

The permutation symmetry group $G$ of the weight space corresponds to independent permutations of intermediate neuron activations and is defined as the direct product of symmetric groups:

$G = S_{d_1} \times S_{d_2} \times \dots \times S_{d_{M-1}}$

For a group element $g = (\tau_1, \dots, \tau_{M-1}) \in G$ , where $P_{\tau_m} \in \{0, 1\}^{d_m \times d_m}$ denotes the permutation matrix associated with permutation $\tau_m \in S_{d_m}$ , the representation $\rho(g)$ acts on a concatenated weight vector $v = [W_m, b_m]_{m=1}^M \in \mathcal{V}$ according to:

$\rho(g)v = [W'_m, b'_m]_{m=1}^M$

$W'_1 = P_{\tau_1}^T W_1, \quad b'_1 = P_{\tau_1}^T b_1$

$W'_m = P_{\tau_m}^T W_m P_{\tau_{m-1}}, \quad b'_m = P_{\tau_m}^T b_m \quad \text{for } m \in \{2, \dots, M-1\}$

$W'_M = W_M P_{\tau_{M-1}}, \quad b'_M = b_M$

For any pointwise nonlinearity $\sigma$ , the transformed parameters $\rho(g)v$ represent the exact same mathematical function $f(\cdot)$ as $v$ .
Knowl 2 — Characterization of Linear Equivariant Layers for Deep Weight Spaces
theoretical result
Let $\mathcal{V} = \bigoplus_{m=1}^M (\mathcal{W}_m \oplus \mathcal{B}_m)$ be the weight space of an $M$ -layer MLP with symmetry group $G = S_{d_1} \times \dots \times S_{d_{M-1}}$ . Any linear $G$ -equivariant layer $L: \mathcal{V} \to \mathcal{V}$ decomposes into four direct sum maps: $L_{ww}: \mathcal{W} \to \mathcal{W}$ , $L_{wb}: \mathcal{W} \to \mathcal{B}$ , $L_{bw}: \mathcal{B} \to \mathcal{W}$ , and $L_{bb}: \mathcal{B} \to \mathcal{B}$ , where $\mathcal{W} = \bigoplus_{m=1}^M \mathcal{W}_m$ and $\mathcal{B} = \bigoplus_{m=1}^M \mathcal{B}_m$ .

Each sub-block $T: U \to U'$ between constituent subrepresentations $U, U' \in \{\mathcal{W}_m, \mathcal{B}_\ell\}_{m,\ell=1}^M$ is constructed using combinations of three primitive operations:
1. Pooling ( $\mathrm{POOL}(d)$ ): Summation contracting an unshared set dimension $d$ .
2. Broadcasting ( $\mathrm{BC}(d)$ ): Replicating values to extend a target set dimension $d$ .
3. Dense Linear Transformation ( $\mathrm{LIN}(d, d')$ ): Standard linear mapping for unpermuted free dimensions $d_0, d_M$ .
The sub-blocks are categorized by their index sharing structure:
- Two shared set dimensions (e.g., mapping $\mathcal{W}_m \to \mathcal{W}_m$ for $1 < m < M$ ): Spanned by the 4-parameter two-set linear equivariant layer $L_{\mathrm{Har}}(X)_{ij} = w_1 X_{ij} + w_2 \sum_{i'} X_{i'j} + w_3 \sum_{j'} X_{ij'} + w_4 \sum_{i',j'} X_{i'j'}$ .
- One shared set dimension (e.g., mapping $\mathcal{B}_m \to \mathcal{B}_m$ for $1 \le m < M$ , or adjacent weight-weight transitions $\mathcal{W}_{m \pm 1} \to \mathcal{W}_m$ ): Spanned by the 2-parameter DeepSets layer $L_{\mathrm{DS}}(X)_i = w_1 X_i + w_2 \sum_{j} X_j$ .
- No shared set dimensions (e.g., mapping $\mathcal{W}_j \to \mathcal{W}_i$ where $|i - j| > 1$ and $1 < i, j < M$ ): Spanned by global pooling, multiplication by a scalar parameter, and broadcasting, requiring 1 learnable parameter.
- Free boundary dimensions ( $d_0$ or $d_M$ ): Parameter count equals the product of free dimensions multiplied by $2^k$ , where $k$ is the number of shared set dimensions.
Knowl 3 — Block Decomposition of Linear Equivariant Maps on Direct Sums
theoretical result
Let $(V_m, \rho_m)$ for $m \in \{1, \dots, M\}$ and $(V'_\ell, \rho'_\ell)$ for $\ell \in \{1, \dots, M'\}$ be orthogonal representations of a group $G$ . Let $(V, \rho) = \bigoplus_{m=1}^M V_m$ and $(V', \rho') = \bigoplus_{\ell=1}^{M'} V'_\ell$ be their direct sum representations.

Let $\mathcal{B}_{m\ell}$ be a basis for the linear space of $G$ -equivariant maps from $(V_m, \rho_m)$ to $(V'_\ell, \rho'_\ell)$ . Define $\mathcal{B}^P_{m\ell}$ as the zero-padded embedding of $\mathcal{B}_{m\ell}$ into $\mathbb{R}^{(\sum_\ell d'_\ell) \times (\sum_m d_m)}$ , where each matrix is non-zero only at the sub-block mapping $V_m$ to $V'_\ell$ .

Then $\mathcal{B} = \bigcup_{m=1}^M \bigcup_{\ell=1}^{M'} \mathcal{B}^P_{m\ell}$ is a basis for the space of linear $G$ -equivariant maps from $V$ to $V'$ , and the total dimension of the equivariant map space satisfies:

$\dim(E(\rho, \rho')) = \sum_{m=1}^M \sum_{\ell=1}^{M'} \dim(E(\rho_m, \rho'_\ell)) = \frac{1}{|G|} \sum_{g \in G} \mathrm{tr}(\rho(g)) \cdot \mathrm{tr}(\rho'(g))$

where $\mathrm{tr}(\cdot)$ denotes the matrix trace.
Knowl 4 — Equivariant Constant Biases and Linear Invariant Layers on Weight Spaces
model/method
Let $G = S_{d_1} \times \dots \times S_{d_{M-1}}$ be the permutation symmetry group of the weight space $\mathcal{V} = \bigoplus_{m=1}^M (\mathcal{W}_m \oplus \mathcal{B}_m)$ . Equivariant constant biases and linear invariant maps $L_{\mathrm{inv}}: \mathcal{V} \to \mathbb{R}$ are uniquely defined by the orbits of the group action on the index set of $\mathcal{V}$ :
1. Orbits per sub-representation:
  - Subspace $\mathcal{W}_1 \in \mathbb{R}^{d_1 \times d_0}$ : $d_0$ orbits, indexed by $j \in [d_0]$ with index set $\{(i, j) \mid i \in [d_1]\}$ .
  - Subspace $\mathcal{W}_m \in \mathbb{R}^{d_m \times d_{m-1}}$ ( $1 < m < M$ ): 1 orbit $\{(i, j) \mid i \in [d_m], j \in [d_{m-1}]\}$ .
  - Subspace $\mathcal{W}_M \in \mathbb{R}^{d_M \times d_{M-1}}$ : $d_M$ orbits, indexed by $i \in [d_M]$ with index set $\{(i, j) \mid j \in [d_{M-1}]\}$ .
  - Subspace $\mathcal{B}_m \in \mathbb{R}^{d_m}$ ( $1 \le m < M$ ): 1 orbit $[d_m]$ .
  - Subspace $\mathcal{B}_M \in \mathbb{R}^{d_M}$ : $d_M$ orbits $\{i\}$ for $i \in [d_M]$ .
2. Equivariant Bias Representation: An affine bias term $b \in \mathcal{V}$ satisfying $\rho(g)b = b$ for all $g \in G$ is a linear combination of indicator vectors of these orbits. Specifically, $b$ is constant across all entries for intermediate weight matrices $\mathcal{W}_m$ ( $1 < m < M$ ) and bias vectors $\mathcal{B}_m$ ( $1 \le m < M$ ), column-constant for $\mathcal{W}_1$ , row-constant for $\mathcal{W}_M$ , and unconstrained for $\mathcal{B}_M$ .
3. Linear Invariant Layer: Every linear $G$ -invariant map $L_{\mathrm{inv}}: \mathcal{V} \to \mathbb{R}$ is parameterised as:
  
  $L_{\mathrm{inv}}(v) = \sum_{m=2}^{M-1} \alpha_m \sum_{i,j} (W_m)_{ij} + \sum_{m=1}^{M-1} \beta_m \sum_{i} (b_m)_i + \sum_{j=1}^{d_0} \gamma_j \sum_{i=1}^{d_1} (W_1)_{ij} + \sum_{i=1}^{d_M} \delta_i \sum_{j=1}^{d_{M-1}} (W_M)_{ij} + \sum_{i=1}^{d_M} \mu_i (b_M)_i$
  
  where $\alpha_m, \beta_m, \gamma_j, \delta_i, \mu_i \in \mathbb{R}$ are learnable scalar parameters.
Knowl 5 — Universal Approximation of Functions Defined on MLP Function Spaces
theoretical result
Let $\mathcal{F}_{\mathcal{V}}$ denote the space of functions represented by $M$ -layer MLPs with fixed dimensions $d_0, \dots, d_M$ , ReLU activations, and weight vectors $v \in C_2 \subset \mathcal{V}$ , where $C_2$ is a compact weight set, evaluated on a compact input domain $C_1 \subset \mathbb{R}^{d_0}$ . Assume every function $f_v \in \mathcal{F}_{\mathcal{V}}$ is $L_1$ -Lipschitz with respect to $\|\cdot\|_\infty$ .

Let $g: \mathcal{F}_{\mathcal{V}} \to \mathbb{R}$ be an $L_2$ -Lipschitz function on the function space with respect to the $\|\cdot\|_\infty$ norm (such that $g(f_{v_1}) = g(f_{v_2})$ whenever $f_{v_1} = f_{v_2}$ ). For any finite sample $x^{(1)}, \dots, x^{(N)} \in C_1$ , define the volume function:

$\mathrm{vol}_\sigma(y^{(1)}, \dots, y^{(N)}) = \mathrm{vol}\left(\{v \in C_2 \mid \|(f(x^{(1)}; v), \dots, f(x^{(N)}; v)) - (y^{(1)}, \dots, y^{(N)})\|_\infty \le \sigma\}\right)$

Assume that for all $\sigma > 0$ , $\mathrm{vol}_\sigma$ is continuous in both $(y^{(1)}, \dots, y^{(N)})$ and $\sigma$ , and that there exists $\lambda > 0$ such that $\mathrm{vol}_\sigma > \lambda$ on the image of $C_2$ .

Then, for any $\epsilon > 0$ , there exists an invariant Deep Weight-Space Network (DWSNet) $F$ with ReLU nonlinearities that uniformly approximates $g$ over $C_2$ :

$\max_{v \in C_2} |g(f_v) - F(v)| \le \epsilon$
Knowl 6 — Approximation of Feed-Forward Network Evaluation by DWSNets
theoretical result
Let an MLP architecture with ReLU activation functions be specified by depth $M$ and layer dimensions $d_0, d_1, \dots, d_M$ . Let $K \subset \mathcal{V}$ and $K' \subset \mathbb{R}^{d_0}$ be compact sets of weight parameters and network inputs, respectively.

A Deep Weight-Space Network (DWSNet) with ReLU nonlinearities composed of affine equivariant layers is capable of uniformly approximating the evaluation of the feed-forward forward pass of the input MLP over $(v, x) \in K \times K'$ . That is, for any $\epsilon > 0$ , there exists a DWSNet $F$ such that:

$\sup_{v \in K, x \in K'} \|F([x, v]) - f(x; v)\|_\infty \le \epsilon$

where $f(x; v)$ denotes the output of the MLP parameterized by $v$ on input $x$ .
Knowl 7 — Parameter and Computational Complexity of DWS-Layers
theoretical result
For an $M$ -layer input MLP with uniform layer dimensions $d_m = d$ for all $m \in \{0, \dots, M\}$ , the parameter count and computational complexity of a single DWS-layer compare to a standard fully connected (FC) layer mapping the flattened weight space to itself as follows:
- Parameter Complexity:
  - A DWS-layer requires $\mathcal{O}((M + d)^2)$ parameters. This occurs because each internal block mapping between hidden weight and bias spaces requires $\mathcal{O}(1)$ parameters (at most 4 for weight-to-weight blocks, 2 for bias-to-bias blocks), and only boundary blocks connecting to the input and output dimensions scale with $d$ .
  - An unconstrained fully connected layer mapping the parameter vector of size $\mathcal{O}(M d^2)$ to itself requires $\mathcal{O}((M d^2)^2)$ parameters.
- Time Complexity:
  - Computing a forward pass through a DWS-layer block operates in $\mathcal{O}(d^2)$ time by leveraging linear equivariant pooling and broadcasting primitives.
  - An unconstrained fully connected layer on a $d \times d$ weight matrix representation requires $\mathcal{O}(d^4)$ operations.

Knowl 8 — Image Classification Benchmark on Implicit Neural Representations

data/table

The task evaluates direct classification of images from the learned weights of Implicit Neural Representations (INRs). Each INR is an independent 3-layer SIREN MLP ( $2 \to 32 \to 32 \to 1$ ) trained to reconstruct an individual image from either MNIST or Fashion-MNIST. The models process the raw weights and biases of the INRs directly to predict the 10 image classes. Baseline methods are dimension-matched to have comparable parameter capacity ( $\sim 550\text{K}$ parameters).

Method	MNIST INR Accuracy (%)	Fashion-MNIST INR Accuracy (%)
MLP	$17.55 \pm 0.01$	$19.91 \pm 0.47$
MLP + Perm. aug	$29.26 \pm 0.18$	$22.76 \pm 0.13$
MLP + Alignment	$58.98 \pm 0.52$	$47.79 \pm 1.03$
INR2Vec (Arch.)	$23.69 \pm 0.10$	$22.33 \pm 0.41$
Transformer	$26.57 \pm 0.18$	$26.97 \pm 0.33$
DWSNets (ours)	$\mathbf{85.71 \pm 0.57}$	$\mathbf{67.06 \pm 0.29}$

DWSNets outperforms all baselines by a substantial margin, including the natural baseline of aligning weights across models using the MergeMany permutation alignment algorithm prior to applying an MLP.

Knowl 9 — Network Domain Adaptation via Weight Space Residual Prediction

data/table

The domain adaptation task evaluates equivariant mapping on classifier weight spaces. A DWSNet takes as input the parameter vector $v$ of a 5-layer classifier ( $3072 \to 64 \to 64 \to 64 \to 64 \to 10$ ) trained on clean CIFAR-10 binary classification and outputs a residual weight vector $\Delta v$ such that the adapted classifier $v - \Delta v$ classifies corrupted CIFAR-10 images (corrupted with rotation, flipping, Gaussian noise, and color jittering). The evaluated models have $\sim 4\text{M}$ parameters.

Method	CIFAR-10 $\to$ CIFAR-10-Corrupted Test Accuracy (%)
No adaptation	$60.92 \pm 0.41$
MLP	$64.33 \pm 0.36$
MLP + Perm. aug	$64.69 \pm 0.56$
MLP + Alignment	$67.66 \pm 0.90$
INR2Vec (Arch.)	$65.69 \pm 0.41$
Transformer	$61.37 \pm 0.13$
DWSNets (ours)	$\mathbf{71.36 \pm 0.38}$

DWSNets outperforms the unadapted baseline by $10.44\%$ and outperforms the strongest baseline (MLP + weight alignment) by $3.7\%$ .

Knowl 10 — Self-Supervised Dense Representation Learning on Neural Weight Spaces

data/table

A dataset of 5,000 INRs ( $1 \to 32 \to 32 \to 1$ ) was fitted to 1D sine waves $f(x) = a \sin(bx)$ for $a, b \sim U(0, 10)$ on $[-\pi, \pi]$ . Each model was trained using a SimCLR-style contrastive objective on views generated with Gaussian weight noise ( $\sigma = 0.2$ ) and random parameter masking ( $p = 0.5$ ) to produce a 16-dimensional embedding space. A linear regressor was then trained on top of the fixed embeddings to predict frequency $b$ and amplitude $a$ . All models had $\sim 100\text{K}$ parameters.

Method	Embedding Regression MSE
MLP	$7.39 \pm 0.19$
MLP + Perm. aug	$5.65 \pm 0.01$
MLP + Alignment	$4.47 \pm 0.15$
INR2Vec (Arch.)	$3.86 \pm 0.32$
Transformer	$5.11 \pm 0.12$
DWSNets (ours)	$\mathbf{1.39 \pm 0.06}$

DWSNets achieves an MSE of $1.39$ , significantly outperforming all competing representation learning architectures on neural weight spaces.

Knowl 11 — Limitations of Deep Weight-Space Networks
limitation
Deep Weight-Space Networks (DWSNets) exhibit several structural and practical limitations:
1. Architecture Specificity: The exact equivariant layer connectivity and block dimensions are hardcoded to the specific architecture (depth $M$ and layer dimensions $d_0, \dots, d_M$ ) of the input MLP, preventing a single model from processing networks of variable layer sizes or heterogeneous architectures.
2. Initialization Sensitivity in Pruning: In network pruning experiments on DIV2K INRs, DWSNets tended to drop parameters unevenly, pruning large numbers of weights in some layers while leaving other layers untouched, reflecting challenges in standard weight initialization schemes for DWS architectures.
3. Domain Extrapolation Sensitivity in Higher-Dimensional INRs: When processing INRs trained on 2D images (such as CIFAR-10 INRs), performance was degraded because the INR functions are only constrained on the training bounding box $[0, 1]^2$ . Outside this domain, unconstrained network parameter variations introduce implicit noise that DWSNets learn to exploit in the absence of explicit domain bounds.

Coverage note — Omitted were intermediate ablation results (Table 10), data augmentation/batch normalization study (Table 11), sine wave regression curves (Figure 3), generalization error prediction on Fashion-MNIST classifiers (Figure 7), and brief speculative extensions to CNNs and Transformers (Section 5.4).

References

1.Agarap, A. F. Deep learning using rectified linear units (relu). arXiv preprint arXiv:1803.08375, 2018.
2.Agustsson, E. and Timofte, R. Ntire 2017 challenge on single image super-resolution: Dataset and study. In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR) Workshops, July 2017.
3.Ainsworth, S. K., Hayase, J., and Srinivasa, S. Git re-basin: Merging models modulo permutation symmetries. arXiv preprint arXiv:2209.04836, 2022.
4.Albooyeh, M., Bertolini, D., and Ravanbakhsh, S. Incidence networks for geometric deep learning. arXiv preprint arXiv:1905.11460, 2019.
5.Arjevani, Y. and Field, M. Analytic study of families of spurious minima in two-layer relu neural networks: a tale of symmetry ii. Advances in Neural Information Processing Systems, 34:15162–15174, 2021.
6.Ashmore, S. and Gashler, M. A method for finding similarity between multi-layer perceptrons by forward bipartite alignment. In 2015 International Joint Conference on Neural Networks (IJCNN), pp. 1–7. IEEE, 2015.
7.Azizian, W. and Lelarge, M. Expressive power of invariant and equivariant graph neural networks. In 9th International Conference on Learning Representations, ICLR, 2021.
8.Badrinarayanan, V., Mishra, B., and Cipolla, R. Understanding symmetries in deep networks. arXiv preprint arXiv:1511.01029, 2015.
9.Baker, B., Gupta, O., Raskar, R., and Naik, N. Accelerating neural architecture search using performance prediction. In 6th International Conference on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 - May 3, 2018, Workshop Track Proceedings, 2018.
10.Brea, J., Simsek, B., Illing, B., and Gerstner, W. 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.
11.Bronstein, M. M., Bruna, J., Cohen, T., and Veličković, P. Geometric deep learning: Grids, groups, graphs, geodesics, and gauges. arXiv preprint arXiv:2104.13478, 2021.
12.Bui Thi Mai, P. and Lampert, C. Functional vs. parametric equivalence of relu networks. In 8th International Conference on Learning Representations, 2020.
13.Chang, O., Flokas, L., and Lipson, H. Principled weight initialization for hypernetworks. In International Conference on Learning Representations, 2019.
14.Chen, A. M., Lu, H.-m., and Hecht-Nielsen, R. On the geometry of feedforward neural network error surfaces. Neural computation, 5(6):910–927, 1993.
15.Chen, T., Kornblith, S., Norouzi, M., and Hinton, G. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pp. 1597–1607. PMLR, 2020.
16.Cohen, T. and Welling, M. Group equivariant convolutional networks. In International conference on machine learning, pp. 2990–2999. PMLR, 2016.
17.Cohen, T. S. and Welling, M. Steerable cnns. In 5th International Conference on Learning Representations, ICLR, 2017.
18.Cohen, T. S., Geiger, M., Köhler, J., and Welling, M. Spherical cnns. In 6th International Conference on Learning Representations, ICLR, 2018.
19.Dupont, E., Kim, H., Eslami, S. A., Rezende, D. J., and Rosenbaum, D. From data to functa: Your data point is a function and you can treat it like one. In International Conference on Machine Learning, pp. 5694–5725. PMLR, 2022.
20.Eilertsen, G., Jönsson, D., Ropinski, T., Unger, J., and Ynnerman, A. Classifying the classifier: dissecting the weight space of neural networks. In European Conference on Artificial Intelligence (ECAI 2020), volume 325, pp. 1119–1926, 2020.
21.Elesedy, B. and Zaidi, S. Provably strict generalisation benefit for equivariant models. In International Conference on Machine Learning, pp. 2959–2969. PMLR, 2021.
22.Entezari, R., Sedghi, H., Saukh, O., and Neyshabur, B. The role of permutation invariance in linear mode connectivity of neural networks. In International Conference on Learning Representations, 2021.
23.Esteves, C., Allen-Blanchette, C., Makadia, A., and Daniilidis, K. Learning so (3) equivariant representations with spherical cnns. In Proceedings of the European Conference on Computer Vision (ECCV), pp. 52–68, 2018.
24.Finzi, M., Welling, M., and Wilson, A. G. A practical method for constructing equivariant multilayer perceptrons for arbitrary matrix groups. In International Conference on Machine Learning, pp. 3318–3328. PMLR, 2021.
25.Fulton, W. and Harris, J. Representation theory: a first course, volume 129. Springer Science & Business Media, 2013.
26.Godfrey, C., Brown, D., Emerson, T., and Kvinge, H. On the symmetries of deep learning models and their internal representations. arXiv preprint arXiv:2205.14258, 2022.
27.Hartford, J., Graham, D., Leyton-Brown, K., and Ravanbakhsh, S. Deep models of interactions across sets. In International Conference on Machine Learning, pp. 1909–1918. PMLR, 2018.
28.Hecht-Nielsen, R. On the algebraic structure of feedforward network weight spaces. In Advanced Neural Computers, pp. 129–135. Elsevier, 1990.
29.Hornik, K. Approximation capabilities of multilayer feedforward networks. Neural networks, 4(2):251–257, 1991.
30.Hubara, I., Courbariaux, M., Soudry, D., El-Yaniv, R., and Bengio, Y. Binarized neural networks. Advances in neural information processing systems, 29, 2016.
31.Jaeckle, F. and Kumar, M. P. Generating adversarial examples with graph neural networks. In Uncertainty in Artificial Intelligence, pp. 1556–1564. PMLR, 2021.
32.Keriven, N. and Peyré, G. Universal invariant and equivariant graph neural networks. Advances in Neural Information Processing Systems, 32, 2019.
33.Knyazev, B., Drozdzal, M., Taylor, G. W., and Romero Soriano, A. Parameter prediction for unseen deep architectures. Advances in Neural Information Processing Systems, 34:29433–29448, 2021.
34.Kondor, R. and Trivedi, S. On the generalization of equivariance and convolution in neural networks to the action of compact groups. In International Conference on Machine Learning, pp. 2747–2755. PMLR, 2018.
35.Krizhevsky, A., Hinton, G., et al. Learning multiple layers of features from tiny images. Technical report, University of Toronto, 2009.
36.LeCun, Y., Bottou, L., Bengio, Y., and Haffner, P. Gradient-based learning applied to document recognition. Proceedings of the IEEE, 86(11):2278–2324, 1998.
37.Lee, J., Lee, Y., Kim, J., Kosiorek, A., Choi, S., and Teh, Y. W. Set transformer: A framework for attention-based permutation-invariant neural networks. In International conference on machine learning, pp. 3744–3753. PMLR, 2019.
38.Lim, D., Robinson, J., Zhao, L., Smidt, T., Sra, S., Maron, H., and Jegelka, S. Sign and basis invariant networks for spectral graph representation learning. arXiv preprint arXiv:2202.13013, 2022.
39.Litany, O., Maron, H., Acuna, D., Kautz, J., Chechik, G., and Fidler, S. Federated learning with heterogeneous architectures using graph hypernetworks. arXiv preprint arXiv:2201.08459, 2022.
40.Loshchilov, I. and Hutter, F. Decoupled weight decay regularization. In 7th International Conference on Learning Representations, ICLR, 2019.
41.Lu, J. and Kumar, M. P. Neural network branching for neural network verification. In International Conference on Learning Representations, 2019.
42.Luigi, L. D., Cardace, A., Spezialetti, R., Ramirez, P. Z., Salti, S., and di Stefano, L. Deep learning on implicit neural representations of shapes. In The Eleventh International Conference on Learning Representations, 2023. URL https://openreview.net/forum?id=OoOIW-3uadi.
43.Maron, H., Ben-Hamu, H., Serviansky, H., and Lipman, Y. Provably powerful graph networks. Advances in neural information processing systems, 32, 2019a.
44.Maron, H., Ben-Hamu, H., Shamir, N., and Lipman, Y. Invariant and equivariant graph networks. In 7th International Conference on Learning Representations, ICLR, 2019b.
45.Maron, H., Fetaya, E., Segol, N., and Lipman, Y. On the universality of invariant networks. In International conference on machine learning, pp. 4363–4371. PMLR, 2019c.
46.Maron, H., Litany, O., Chechik, G., and Fetaya, E. On learning sets of symmetric elements. In International Conference on Machine Learning, pp. 6734–6744. PMLR, 2020.
47.Mildenhall, B., Srinivasan, P. P., Tancik, M., Barron, J. T., Ramamoorthi, R., and Ng, R. Nerf: Representing scenes as neural radiance fields for view synthesis. Communications of the ACM, 65(1):99–106, 2021.
48.Morris, C., Ritzert, M., Fey, M., Hamilton, W. L., Lenssen, J. E., Rattan, G., and Grohe, M. Weisfeiler and leman go neural: Higher-order graph neural networks. In Proceedings of the AAAI conference on artificial intelligence, volume 33, pp. 4602–4609, 2019.
49.Morris, C., Lipman, Y., Maron, H., Rieck, B., Kriege, N. M., Grohe, M., Fey, M., and Borgwardt, K. Weisfeiler and leman go machine learning: The story so far. arXiv preprint arXiv:2112.09992, 2021.
50.Neyshabur, B., Salakhutdinov, R. R., and Srebro, N. Path-sgd: Path-normalized optimization in deep neural networks. Advances in neural information processing systems, 28, 2015.
51.Park, J. J., Florence, P., Straub, J., Newcombe, R., and Lovegrove, S. Deepsdf: Learning continuous signed distance functions for shape representation. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 165–174, 2019.
52.Peebles, W., Radosavovic, I., Brooks, T., Efros, A. A., and Malik, J. Learning to learn with generative models of neural network checkpoints. arXiv preprint arXiv:2209.12892, 2022.
53.Peña, F. A. G., Medeiros, H. R., Dubail, T., Aminbeidokhti, M., Granger, E., and Pedersoli, M. Re-basin via implicit sinkhorn differentiation. arXiv preprint arXiv:2212.12042, 2022.
54.Qi, C. R., Su, H., Mo, K., and Guibas, L. J. Pointnet: Deep learning on point sets for 3d classification and segmentation. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 652–660, 2017.
55.Ravanbakhsh, S., Schneider, J., and Poczos, B. Equivariance through parameter-sharing. In International conference on machine learning, pp. 2892–2901. PMLR, 2017.
56.Schürholt, K., Kostadinov, D., and Borth, D. Self-supervised representation learning on neural network weights for model characteristic prediction. Advances in Neural Information Processing Systems, 34:16481–16493, 2021.
57.Schürholt, K., Knyazev, B., Giró-i Nieto, X., and Borth, D. Hyper-representations as generative models: Sampling unseen neural network weights. arXiv preprint arXiv:2209.14733, 2022a.
58.Schürholt, K., Taskiran, D., Knyazev, B., Giró-i Nieto, X., and Borth, D. Model zoos: A dataset of diverse populations of neural network models. arXiv preprint arXiv:2209.14764, 2022b.
59.Simsek, B., Ged, F., Jacot, A., Spadaro, F., Hongler, C., Gerstner, W., and Brea, J. Geometry of the loss landscape in overparameterized neural networks: Symmetries and invariances. In International Conference on Machine Learning, pp. 9722–9732. PMLR, 2021.
60.Singh, S. P. and Jaggi, M. Model fusion via optimal transport. Advances in Neural Information Processing Systems, 33:22045–22055, 2020.
61.Sitzmann, V., Martel, J., Bergman, A., Lindell, D., and Wetzstein, G. Implicit neural representations with periodic activation functions. Advances in Neural Information Processing Systems, 33:7462–7473, 2020.
62.Tancik, M., Srinivasan, P., Mildenhall, B., Fridovich-Keil, S., Raghavan, N., Singhal, U., Ramamoorthi, R., Barron, J., and Ng, R. Fourier features let networks learn high frequency functions in low dimensional domains. Advances in Neural Information Processing Systems, 33: 7537–7547, 2020.
63.Tatro, N., Chen, P.-Y., Das, P., Melnyk, I., Sattigeri, P., and Lai, R. Optimizing mode connectivity via neuron alignment. Advances in Neural Information Processing Systems, 33:15300–15311, 2020.
64.Thomas, N., Smidt, T., Kearnes, S., Yang, L., Li, L., Kohlhoff, K., and Riley, P. Tensor field networks: Rotation-and translation-equivariant neural networks for 3d point clouds. arXiv preprint arXiv:1802.08219, 2018.
65.Unterthiner, T., Keysers, D., Gelly, S., Bousquet, O., and Tolstikhin, I. Predicting neural network accuracy from weights. arXiv preprint arXiv:2002.11448, 2020.
66.Vaswani, A., Shazeer, N., Parmar, N., Uszkoreit, J., Jones, L., Gomez, A. N., Kaiser, Ł., and Polosukhin, I. Attention is all you need. Advances in neural information processing systems, 30, 2017.
67.Velickovic, P., Cucurull, G., Casanova, A., Romero, A., Liò, P., and Bengio, Y. Graph attention networks. In 6th International Conference on Learning Representations, ICLR, 2018.
68.Wang, G., Wang, G., Liang, W., and Lai, J. Understanding weight similarity of neural networks via chain normalization rule and hypothesis-training-testing. arXiv preprint arXiv:2208.04369, 2022.
69.Wang, H., Yurochkin, M., Sun, Y., Papailiopoulos, D., and Khazaeni, Y. Federated learning with matched averaging. In International Conference on Learning Representations, 2019.
70.Wang, R., Albooyeh, M., and Ravanbakhsh, S. Equivariant networks for hierarchical structures. Advances in Neural Information Processing Systems, 33:13806–13817, 2020.
71.Wood, J. and Shawe-Taylor, J. Representation theory and invariant neural networks. Discrete applied mathematics, 69(1-2):33–60, 1996.
72.Xiao, H., Rasul, K., and Vollgraf, R. Fashion-mnist: a novel image dataset for benchmarking machine learning algorithms. arXiv preprint arXiv:1708.07747, 2017.
73.Xu, D., Wang, P., Jiang, Y., Fan, Z., and Wang, Z. Signal processing for implicit neural representations. In Advances in Neural Information Processing Systems, 2022.
74.Xu, K., Hu, W., Leskovec, J., and Jegelka, S. How powerful are graph neural networks? In 7th International Conference on Learning Representations, ICLR, 2019.
75.Yurochkin, M., Agarwal, M., Ghosh, S., Greenewald, K., Hoang, N., and Khazaeni, Y. Bayesian nonparametric federated learning of neural networks. In International Conference on Machine Learning, pp. 7252–7261. PMLR, 2019.
76.Zaheer, M., Kottur, S., Ravanbakhsh, S., Poczos, B., Salakhutdinov, R. R., and Smola, A. J. Deep sets. Advances in neural information processing systems, 30, 2017.

Citation

MLA

Navon, A., et al. “Equivariant Architectures for Learning in Deep Weight Spaces”. International Conference on Machine Learning, vol. 202, 2023, pp. 25790–816, https://proceedings.mlr.press/v202/navon23a.html.

APA

Navon, A., Shamsian, A., Achituve, I., Fetaya, E., Chechik, G., & Maron, H. (2023). Equivariant Architectures for Learning in Deep Weight Spaces. International Conference on Machine Learning, 202, 25790–25816. https://proceedings.mlr.press/v202/navon23a.html

Chicago

Navon, A., A. Shamsian, I. Achituve, E. Fetaya, G. Chechik, and H. Maron. 2023. “Equivariant Architectures for Learning in Deep Weight Spaces”. International Conference on Machine Learning 202: 25790–816. https://proceedings.mlr.press/v202/navon23a.html.

Harvard

Navon, A. et al. (2023) “Equivariant Architectures for Learning in Deep Weight Spaces”, International Conference on Machine Learning. PMLR, pp. 25790–25816. Available at: https://proceedings.mlr.press/v202/navon23a.html.

Vancouver

1. Navon A, Shamsian A, Achituve I, Fetaya E, Chechik G, Maron H (2023) Equivariant Architectures for Learning in Deep Weight Spaces. In: International Conference on Machine Learning. PMLR, pp 25790–25816

BibTeX

@InProceedings{pmlr-v202-navon23a,
  title = 	 {Equivariant Architectures for Learning in Deep Weight Spaces},
  author =       {Navon, Aviv and Shamsian, Aviv and Achituve, Idan and Fetaya, Ethan and Chechik, Gal and Maron, Haggai},
  booktitle = 	 {Proceedings of the 40th International Conference on Machine Learning},
  pages = 	 {25790--25816},
  year = 	 {2023},
  editor = 	 {Krause, Andreas and Brunskill, Emma and Cho, Kyunghyun and Engelhardt, Barbara and Sabato, Sivan and Scarlett, Jonathan},
  volume = 	 {202},
  series = 	 {Proceedings of Machine Learning Research},
  month = 	 {23--29 Jul},
  publisher =    {PMLR},
  pdf = 	 {https://proceedings.mlr.press/v202/navon23a/navon23a.pdf},
  url = 	 {https://proceedings.mlr.press/v202/navon23a.html},
  abstract = 	 {Designing machine learning architectures for processing neural networks in their raw weight matrix form is a newly introduced research direction. Unfortunately, the unique symmetry structure of deep weight spaces makes this design very challenging. If successful, such architectures would be capable of performing a wide range of intriguing tasks, from adapting a pre-trained network to a new domain to editing objects represented as functions (INRs or NeRFs). As a first step towards this goal, we present here a novel network architecture for learning in deep weight spaces. It takes as input a concatenation of weights and biases of a pre-trained MLP and processes it using a composition of layers that are equivariant to the natural permutation symmetry of the MLP’s weights: Changing the order of neurons in intermediate layers of the MLP does not affect the function it represents. We provide a full characterization of all affine equivariant and invariant layers for these symmetries and show how these layers can be implemented using three basic operations: pooling, broadcasting, and fully connected layers applied to the input in an appropriate manner. We demonstrate the effectiveness of our architecture and its advantages over natural baselines in a variety of learning tasks.}
}

Metadata:DOI registry

Access the Paper

This paper is available from its original source. Click below to access the PDF.

Open PDF

License: https://creativecommons.org/licenses/by/4.0/

Equivariant Architectures for Learning in Deep Weight Spaces

Abstract

Topics

Table of Contents

Knowls

Knowl 1 — Permutation Symmetry Group and Action on MLP Weight Spaces

Knowl 2 — Characterization of Linear Equivariant Layers for Deep Weight Spaces

Knowl 3 — Block Decomposition of Linear Equivariant Maps on Direct Sums

Knowl 4 — Equivariant Constant Biases and Linear Invariant Layers on Weight Spaces

Knowl 5 — Universal Approximation of Functions Defined on MLP Function Spaces

Knowl 6 — Approximation of Feed-Forward Network Evaluation by DWSNets

Knowl 7 — Parameter and Computational Complexity of DWS-Layers

Knowl 8 — Image Classification Benchmark on Implicit Neural Representations

Knowl 9 — Network Domain Adaptation via Weight Space Residual Prediction

Knowl 10 — Self-Supervised Dense Representation Learning on Neural Weight Spaces

Knowl 11 — Limitations of Deep Weight-Space Networks

References

Citation

BibTeX

Access the Paper