PaperlayerView exact revision on arXiv

Source-linked AI summary

Equivariant Architectures for Learning in Deep Weight Spaces

Aviv Navon, Aviv Shamsian, Idan Achituve, Ethan Fetaya, Gal Chechik, Haggai Maron

arXiv:2301.12780v2cs.LG

TL;DR

Learning directly from neural weights is challenging because deep weight spaces have a distinctive symmetry structure. The paper introduces DWSNets, which process MLP weights and biases with permutation-equivariant layers and characterize how such layers can be implemented. It reports advantages over natural baselines and proves that DWSNets can approximate forward passes of input networks.

  • Problem

    Learning directly from neural weights remains an open design problem because weight spaces have a distinctive structure and existing approaches make assumptions that limit their application to other network types.

    Source: §1 · p. 1

  • Method

    DWSNets process concatenated MLP weights and biases using layers equivariant to intermediate-neuron permutations, with affine equivariant layers characterized and implemented through pooling, broadcasting, or dense layers.

    Source: §1 · pp. 1–2

  • Results

    The architecture significantly outperforms natural baselines based on data augmentation and weight-space alignment, and DWSNets can approximate a forward pass of an input network.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The results support symmetry-based architectures as a principled approach to processing MLP weight spaces across tasks including INR classification and network adaptation.

    Source: §1 · p. 2

    The equivalent layer structure is tailored to a specific MLP architecture; training some tasks was difficult, and implementation is somewhat complicated.

    Source: §8 · p. 9

Abstract

from arXiv · show

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.

1. Introduction

The paper addresses how to learn directly from neural-network weights, whose structure is not captured by standard tensor- or sequence-based architectures. It develops permutation-equivariant networks for MLP weight spaces, characterizes their affine layers, and reports advantages over natural baselines.

1. Introduction

  • The paper targets direct weight-space learning for tasks including adapting trained networks to new data distributions and analyzing datasets of INR or NeRF representations.Prior approaches leave unclear how they extend to other neural networks because they assume particular input dimensions or training procedures.Source: §1 · p. 1
  • DWSNets process concatenated MLP weights and biases with layers equivariant to permutations of intermediate neurons.The paper characterizes affine equivariant layers and implements them through pooling, broadcasting, or dense linear layers, reducing parameters relative to fully connected networks.Source: §1 · pp. 1–2
  • The authors prove DWSNets can approximate a forward pass of an input network and certain functions defined on the functions represented by input MLPs.Convolutional networks and transformers are discussed as possible extensions, not as the focus of this work.Source: §1 · p. 2
  • The architecture significantly outperforms natural baselines based on data augmentation and weight-space alignment across the demonstrated tasks.The tasks include processing INRs and standard neural networks.Source: §1 · p. 2
  • The paper contributes a symmetry-based design approach, a characterization of affine equivariant layers, expressive-power analysis, and application results.The applications range from INR classification to adapting networks to new domains.Source: §1 · p. 2

2. Previous Work

Earlier work applies generic networks to neural weights or develops methods specialized to particular settings, including INRs. The paper contrasts its approach with these methods, noting that it can work on any MLP and is agnostic to training procedure.

2. Previous Work

  • Prior work used standard networks on flattened weights or weight statistics, learned low-dimensional INR modulations, used self-supervision, processed spatial derivatives, generated networks conditionally, or applied set architectures to INRs.The authors state that their method can work on any MLP and is agnostic to how it was trained, unlike the modulation-based approach.Source: §2 · p. 2

3. Preliminaries

The preliminaries define MLPs as sequential fully connected networks and describe equivariant networks as compositions of affine equivariant layers and pointwise nonlinearities. For invariant tasks, an invariant layer and an MLP can follow the equivariant network.

3. Preliminaries

  • An MLP is a sequence of fully connected layers, with weights, biases, and pointwise nonlinearities specifying its forward computation.The parameters are represented as a concatenation of the layer weight matrices and bias vectors.Source: §3 · p. 3
  • Equivariant networks compose affine equivariant layers with pointwise nonlinearities.Each affine layer has a linear equivariant map and a constant equivariant bias.Source: §3 · p. 3
  • For invariant tasks, the equivariant network is followed by a linear invariant map and an MLP..Source: §3 · p. 3

4. Permutation Symmetries of Neural Networks

The paper models MLP weight space as a direct sum of the layer-specific weight-matrix and bias-vector spaces. Its symmetry group independently permutes each intermediate layer, preserving the represented function; other symmetries are outside the paper’s scope.

4. Permutation Symmetries of Neural Networks

  • Intermediate-neuron permutations rearrange rows and columns of successive MLP weight matrices without changing the represented function.The figure illustrates this symmetry for a three-layer MLP; the invariance holds for any pointwise nonlinearity.Source: Figure 1 · p. 3, §4 · p. 3
  • The MLP weight space is a direct sum of spaces for each layer’s weight matrix and bias vector.The summands are individually invariant under the defined action.Source: §4 · pp. 3–4
  • The symmetry group is the product of symmetric groups over intermediate layer dimensions, allowing each intermediate layer to be permuted independently.The group acts on the concatenated weights and biases by the corresponding row and column permutations.Source: §4 · p. 3
  • The architecture considers permutation symmetries but leaves scaling transformations and other weight-space symmetries for future work.The paper notes that other weight-space symmetries exist beyond those considered.Source: §4 · p. 4

5. A Characterization of Linear Invariant and Equivariant Layers for Weight-Spaces

The section characterizes equivariant layers on weight spaces by decomposing weights and biases into sub-representations and specifying equivariant maps between them. The resulting blocks are implementable with pooling, broadcasting, and fully connected layers, with extensions discussed for CNNs and Transformers.

5.2. Linear Equivariant Maps for Direct Sums

  • Decomposing V into direct sums reduces equivariant-map characterization to maps between constituent sub-representations, whose blocks can be implemented independently.The general result states that each block contains a basis element for an equivariant map between its corresponding representations.Source: §5.1 · p. 4, §5.2 · p. 5

5.1. Overview and Main Results

  • Every equivariant layer on the weight space V has a block-matrix form over weight and bias sub-representations, with each block implementable using pooling, broadcasting, or fully connected layers.The characterization covers linear equivariant maps; the paper calls these DWS-layers and builds DWSNets by interleaving them with pointwise nonlinearities.Source: §5.1 · p. 4

5.3. Linear Equivariant Layers for Deep Weight-Spaces

  • Pooling, broadcasting, and dense linear maps construct the equivariant blocks, and parameter counting establishes that the proposed layers span the full space of such maps.The construction uses DeepSets or product-of-sets equivariant layers for shared set indices, and pooling or broadcasting to contract or extend set dimensions.Source: §5.3 · p. 6, Bias-to-bias layers · pp. 6–7

5.4. Extension to Other Input Architectures

  • The characterization is discussed for CNNs and Transformers as well as MLPs, with examples of channel-permutation symmetries that preserve the represented function.For CNNs, simultaneously permuting adjacent layers’ channel dimensions preserves the function; the text also describes a symmetry of Transformer attention weights.Source: §5.4 · p. 7

6. Expressive Power

The section examines whether equivariant DWSNets retain expressive power despite their restricted hypothesis class. It proves approximation results for feed-forward evaluation and for certain functions on the function space represented by input MLPs.

6. Expressive Power

  • DWSNets with ReLU nonlinearities uniformly approximate a feed-forward procedure on an input MLP and an input x, when their weight vector and x lie in compact sets.This is the stated result for an MLP architecture with ReLU nonlinearities.Source: §6 · p. 7
  • Under additional mild assumptions, DWSNets with ReLU nonlinearities can uniformly approximate L-Lipschitz functions on the function space represented by MLPs with weights in a compact subset.The Lipschitz condition is with respect to the infinity norm on functions, and the input domain is compact.Source: sets · p. 7
  • This result concerns functions on the input MLPs’ function space, rather than every G-equivariant function on the weight space V.The approximated functions give the same result for all weights representing the same function.Source: sets · pp. 7–8

7. Experiments

The experiments evaluate DWSNets on tasks involving INRs as data representations and networks as input-output mappings. Across several tasks, DWSNets outperform other methods, usually by a large margin, and scale better with data than weight alignment.

7. Experiments

  • The evaluation covers INR-based data tasks and adaptation of networks to new domains, comparing architectures on shared data, loss, and training procedures.Baselines include vectorized-weight MLPs, MLPs with random permutation augmentation, weight alignment, INR2Vec without its pretraining, and a transformer.Source: §7 · p. 8

7.1. Results

  • On sine-wave frequency regression, DWSNets significantly outperform baseline methods even with few training INRs.The task predicts the frequency of a test INR fitted to a sine wave, with training-set size varied to assess generalization.Source: §7.1 · p. 8
  • For MNIST and Fashion-MNIST image classification from INR weights, DWSNets outperform all baselines by a large margin.The task is to recognize the class of the image represented by each INR.Source: §7.1 · p. 8
  • Dense representations of sine-wave INRs are evaluated qualitatively with t-SNE and quantitatively by linear-regression MSE for predicting amplitude and frequency.The INRs represent a sin(bx) with a,b sampled from U(0,10); the training procedure uses noisy and randomly masked views.Source: §7.1 · p. 8, Table 2 · p. 8
  • The domain-adaptation task predicts residual weight updates so an input classifier performs well on corrupted CIFAR-10; Table 3 reports test accuracy.Training uses CIFAR-10 classifiers trained to distinguish random class pairs, while testing targets unseen classifiers and images.Source: §7.1 · p. 9, Table 3 · p. 9
  • Using multiple randomly initialized INRs per Fashion-MNIST image improves generalization by approximately 8%.The experiment varies the number of INR copies from 1 to 10 and highlights the need to match copy counts across model evaluations.Source: §7.1 · p. 9, Table 4 · p. 9

7.2. Analysis of the Results

  • Across several learning tasks, DWSNets outperform all other methods, usually by a large margin, and scale significantly better with data than network alignment.Alignment is challenging and becomes more difficult with large input networks or datasets.Source: §7.2 · p. 9

8. Conclusion and Future Work

The conclusion presents an equivariant approach to learning directly over neural weight spaces, while noting architectural and training limitations. The related-work discussion situates the approach among methods for processing neural networks, and the appendix characterizes invariant maps, multi-channel layers, and equivariant biases.

8. Conclusion and Future Work

  • The method’s layer structure is tailored to a specific MLP architecture, and some tasks were difficult to train, presumably because suitable initialization was hard.The implementation is also described as somewhat complicated.Source: §8 · p. 9

A. Related Work

  • Prior approaches process network weights with generic networks, focus on INRs, or impose assumptions that limit their scope; this work introduces inductive biases for weight-space learning.The related-work discussion notes that some INR methods require function evaluation or aligned meta-learned representations.Source: A. Related Work · p. 14

B. Multiple Channels, Invariant Layers and Biases for Equivariant Maps

  • Multiple input and output feature channels form direct sums of representation spaces, with equivariant maps assembled from maps between their component representations.The block-matrix view separates maps among weight and bias spaces; repeated colors across matrices need not denote related functions.Source: B. Multiple Channels, Invariant Layers… · p. 15, Figure 5 · p. 16
  • Invariant maps on weight spaces combine orbit-based sums of inner weights and biases, weighted sums of input-weight columns and output-weight rows, and an inner product with the output bias.The orbit characterization also determines the form of invariant layers and equivariant biases.Source: B. Multiple Channels, Invariant Layers… · p. 16

C. Specification of All Affine Equivariant Layers Between Sub-Representations

The appendix specifies equivariant layer implementations between weight and bias subspaces, then explains how their dimensions are counted. Shared set dimensions and free dimensions affect the number of equivariant parameters, while unshared set dimensions do not.

C. Specification of All Affine Equivariant Layers Between Sub-Representations

  • Tables 5–8 specify the implementation and dimensionality of equivariant maps for weight-to-weight, bias-to-bias, weight-to-bias, and bias-to-weight blocks.The block matrices organize maps among specific weight and bias spaces; repeated asterisked entries indicate the same layer type at a different position.Source: C. Specification of All Affine… · p. 16, Table 5 · p. 16, Table 6 · p. 17

D. Linear Maps Between Specific Weight and Bias Spaces

  • The proof checks that listed layer types are equivariant and linearly independent, then matches their parameter count to the dimension of the corresponding equivariant-map space.The calculation is organized by the types of maps between weight and bias subspaces.Source: D. Linear Maps Between Specific Weight… · p. 17
  • Shared set dimensions contribute a factor of bell(2) = 2, free dimensions contribute their dimensionality, and unshared set dimensions contribute bell(1) = 1.The appendix illustrates the dimension count with the two-shared-set-indices case, whose layer is four-dimensional.Source: D. Linear Maps Between Specific Weight… · pp. 17–18

E. More Proofs for Section 5

The proof establishes a basis for equivariant maps by showing that the proposed basis elements are equivariant and linearly independent, and that their count matches the dimension of the map space.

E. More Proofs for Section 5

  • The proposed block-supported basis elements are linear, linearly independent, and equivariant because each block map is equivariant.Independence follows because the basis elements occupy non-overlapping blocks.Source: E. More Proofs for Section 5 · p. 19
  • The basis is complete when its size matches the dimension of the equivariant-map space, which the proof computes using a group-average trace formula.The proof uses traces of the input and output representations to calculate that dimension.Source: E. More Proofs for Section 5 · pp. 19–20

F. Proofs of Proposition 6.1

The proof constructs equivariant layers that approximate a forward pass of an input MLP on a compact domain. It simulates propagation layer by layer, approximating the Hadamard product with an MLP.

F. Proofs of Proposition 6.1

  • On a compact domain, the proof assumes ReLU activations for both networks for simplicity, though it says this assumption is not necessary in general.The approximation target is the input MLP’s output, with the DWSNet also taking the input x and the MLP weights v.Source: F. Proofs of Proposition 6.1 · p. 20
  • The construction simulates each MLP layer using equivariant operations including summation, broadcasting, and pointwise nonlinearities, while an MLP approximates the Hadamard product.It stores the running approximation in an extra channel and propagates it through the weight and bias spaces.Source: F. Proofs of Proposition 6.1 · p. 20
  • The composition approximates the input MLP’s forward pass uniformly, with the Hadamard-product approximation justified on compact domains.The proof states that uniform approximation is preserved under composition of continuous functions on compact domains.Source: F. Proofs of Proposition 6.1 · pp. 20–21

G. Proof of Proposition 6.2

The section establishes an approximation result for DWSNets under stated assumptions, then discusses alternative equivariant-layer characterization and computational costs relative to fully connected layers.

G. Proof of Proposition 6.2

  • Under compact-domain, Lipschitz, and volume-regularity assumptions, a ReLU DWSNet approximates any L2-Lipschitz function on the represented MLP function space to arbitrary accuracy.The proof approximates evaluations on a finite input net, smooths the function of those evaluations, and then approximates that continuous map with an MLP.Source: G. Proof of Proposition 6.2 · p. 21
  • The proof constructs a DWSNet to approximate the input MLP’s outputs on fixed inputs, then composes it with an MLP approximation of the target function.The inputs are broadcast into the weight representation; finite-net evaluations support the continuous approximation step.Source: G. Proof of Proposition 6.2 · pp. 21–22

H. Alternative Characterization Strategies

  • The authors characterize equivariant maps directly on the weight-and-bias decomposition, noting that irreducible-representation methods offer simpler blocks but may require translation back for implementation.They cite Schur’s Lemma as the basis for the alternative characterization.Source: H. Alternative Characterization… · p. 22

I. Computational and Memory Requirements

  • With equal feature dimensions d, a DWS-layer uses O((M + d)^2) parameters, compared with O((Md^2)^2) for a fully connected layer.The comparison assumes a network with M input-MLP layers and does not account for parallel computations.Source: §I · p. 22
  • Under the stated block-operation accounting, DWS-layers have asymptotically lower feedforward time complexity than fully connected layers.For the cited matrix-to-matrix block, the comparison is O(d^2) operations for the Hartford layer versus O(d^4) for a fully connected layer.Source: §I · p. 23

J. Experimental and Technical Details

The appendix details data preparation, augmentation, alignment, and training choices, then reports supplementary results for generalization prediction and dense representations.

J. Experimental and Technical Details

  • Input networks are trained independently, normalized entrywise using dataset means and standard deviations, and split into training, validation, and test sets.For INRs, the appendix describes image-space transformations such as rotation, translation, and scaling, alongside Gaussian and dropout noise.Source: J. Experimental and Technical Details · p. 23
  • The model-alignment baseline uses MergeMany on a 1,000-model subsample at error 1e-3 after the default-error MNIST run exceeded 24 hours without finishing one iteration.Experiments generally use four hidden equivariant layers, max pooling, matched-capacity baselines, ReLU, batch normalization, AdamW, and three random seeds.Source: J. Experimental and Technical Details · p. 24

K.1. Predicting the Generalization Error of Neural Networks.

  • DWSNets achieved the lowest error and significantly outperformed most baselines when predicting Fashion-MNIST classifiers’ held-out test error from their weights.Figure 7 reports 100×MSE averaged over three random seeds.Source: §K.1 · p. 25, Figure 7 · p. 25

K.2. Dense Representation

  • DWSNets produced a 2D embedding grouping sine-wave models with similar frequencies and amplitudes, while other baselines lacked the same apparent structure.The authors describe gradual change across embedding regions as well.Source: §K.2 · p. 25

K.3. Ablation Study

The ablations identify the most contributing architecture blocks and report benefits from data augmentation and batch normalization; the challenging cases describe limitations encountered in pruning and INR classification.

K.3. Ablation Study

  • On MNIST INR classification, the weight-to-weight block contributes most, adding other blocks improves performance, and the diagonal variant performs similarly to the full model.Table 10 presents the block ablation.Source: Table 10 · p. 25, §K.3 · p. 26

K.4. The importance of data augmentation and batch normalization

  • Data augmentation and batch normalization are reported as consistently beneficial across experiments and improve DWSNet performance on MNIST INR classification.Table 11 reports the techniques’ effect on that task.Source: §K.4 · p. 26, Table 11 · p. 26

K.5. Challenging cases

  • In the pruning experiment, DWSNets tended to prune many parameters from some layers while leaving other layers untouched.The authors identify proper initialization as a possible way to address this issue.Source: §K.5 · pp. 26–27
  • CIFAR10 classification from INRs remained unsatisfactory despite outperforming baselines; the authors suggest out-of-domain implicit signals may contribute.The INR is informative on [0, 1]^2, and the proposed possible remedy is encouraging constant output outside that domain.Source: §K.5 · p. 27

Related papers

Loading 2301.12780v2…