Equivariant Architectures for Learning in Deep Weight Spaces
Aviv Navon, Aviv Shamsian, Idan Achituve, Ethan Fetaya, Gal Chechik, Haggai Maron
ICML 2023
cs.LG
2023-01-30
DWSNets read raw MLP weights with permutation-equivariant layers, hitting 85.71% MNIST INR accuracy versus 58.98% for alignment, and 71.36% on CIFAR-10 domain adaptation.
Two jobs need a model that reads raw weights. An INR fits a small MLP to an image or a 3D shape, so the sample is that weight vector. The other job is a trained classifier: change its weights for a new distribution, rather than retraining on the new data.
Hidden neurons have no canonical order. One permutation, applied to the rows of one layer and the columns of the next, leaves the function unchanged. Hecht-Nielsen recorded this in 1990. Two parameter vectors can sit far apart and still compute the same map. Flatten them into an ordinary network and those copies look like different samples.
Earlier models fed flattened weights to an MLP or a Transformer and tried to predict hyperparameters or generalization. On INRs, functa learns low-dimensional modulations by meta-learning, so the representation is tied to that training procedure. INR2Vec treats rows from every layer as one global set. Swapping rows across layers generally changes the function. Permutation augmentation and weight alignment push the symmetry into the data pipeline. The layer itself still does not see it.
A DWSNet takes every weight matrix and bias of one MLP. The symmetry group is the product of symmetric groups on the hidden widths. Input and output dimensions stay fixed, because they index data coordinates rather than interchangeable neurons. The first weight matrix is left-multiplied by a permutation, the last is right-multiplied, and each hidden matrix is hit on both sides. Biases follow their layer. The output bias does not move.
Equivariance means the output weights undergo that same reordering. When the task needs a single number, an invariant readout and a small MLP sit on top.
Theorem 5.1 writes every linear equivariant map from this weight space to itself as a block matrix. Per-layer weights and biases form a direct sum, so the map splits into blocks between those pieces: weight to weight, weight to bias, bias to weight, and bias to bias. Each block uses only three operations: sum over an axis, copy a value along an axis, and a dense linear map. Blocks keep their own parameters, so the layer is much smaller than a dense map on flattened weights.
A hidden weight is the sum of four terms. The same-layer block is the equivariant map Hartford et al. gave for a product of two sets: rows are one set, columns another. Neighboring layers use DeepSets, a per-element map plus a sum over the set. Farther hidden layers collapse to a sum scaled by a learned scalar. The first and last layers have one axis the group does not permute, and that axis goes through a dense map. Between biases of two different hidden layers, one scalar remains: sum, then broadcast.
Pointwise nonlinearities sit between these linear layers. Sums may be swapped for max or attention. The experiments keep the linear sum.
Proposition 6.1: with ReLU, a DWSNet can uniformly approximate one forward pass of the input MLP on compact sets, meaning it can compute the network on an input x. Proposition 6.2, under a Lipschitz assumption and mild extra conditions, extends this to maps defined on the functions those MLPs represent. Distinct weights that encode the same function can therefore land on the same output.
Input networks are trained separately, from different seeds. Unless noted, each function is stored once. The INR2Vec column uses that architecture only, without its pretraining.
The label of an image INR is the label of the image it fits.
| Method | MNIST INR | Fashion-MNIST INR |
| MLP | 17.55 ± 0.01 | 19.91 ± 0.47 |
| MLP + permutation aug. | 29.26 ± 0.18 | 22.76 ± 0.13 |
| MLP + alignment | 58.98 ± 0.52 | 47.79 ± 1.03 |
| INR2Vec architecture | 23.69 ± 0.10 | 22.33 ± 0.41 |
| Transformer | 26.57 ± 0.18 | 26.97 ± 0.33 |
| DWSNets | 85.71 ± 0.57 | 67.06 ± 0.29 |
On MNIST the flat MLP reaches 17.55%, permutation augmentation 29.26%, alignment 58.98%, and DWSNets 85.71%. On Fashion-MNIST alignment reaches 47.79% and DWSNets 67.06%.
The frequency-regression figure does not list MSE at each training size. INRs fit sines on [-π, π] with frequency drawn from U(0.5, 10). Figure 3 shows the DWSNet test error below the baselines even when few INRs are available.
A contrastive stage embeds INRs of a sin(bx), then a linear regressor predicts a and b. Both are drawn from U(0, 10). Views add Gaussian noise with standard deviation 0.2 and mask entries with probability 0.5.
| Method | MSE for a, b |
| MLP | 7.39 ± 0.19 |
| MLP + permutation aug. | 5.65 ± 0.01 |
| MLP + alignment | 4.47 ± 0.15 |
| INR2Vec architecture | 3.86 ± 0.32 |
| Transformer | 5.11 ± 0.12 |
| DWSNets | 1.39 ± 0.06 |
Domain adaptation is a harder use of the same idea. Each input network is a CIFAR-10 classifier trained to separate two random classes. The target adds rotation, flips, Gaussian noise, and color jitter. The model emits a residual Δv, the adapted weights are v - Δv, and test time is one forward pass on classifiers and images it has not seen.
| Method | Test accuracy |
| No adaptation | 60.92 ± 0.41 |
| MLP | 64.33 ± 0.36 |
| MLP + permutation aug. | 64.69 ± 0.56 |
| MLP + alignment | 67.66 ± 0.90 |
| INR2Vec architecture | 65.69 ± 0.41 |
| Transformer | 61.37 ± 0.13 |
| DWSNets | 71.36 ± 0.38 |
No adaptation already scores 60.92%. The Transformer stops at 61.37%. Alignment reaches 67.66%, DWSNets 71.36%. On INR classification the same comparison is 85.71% against 58.98%.
Extra INR copies of one Fashion-MNIST image, from new initializations, move accuracy from 67.06% at one copy to 75.12% at ten. In between: 70.22% (2), 70.31% (4), 73.32% (6), 74.87% (8). The one-to-ten gain is about 8 points, and two copies to four copies barely move.
Treat neuron order as a feature and every random initialization looks like a new object. On MNIST, permutation augmentation only goes from 17.55% to 29.26%. Alignment does more, but the matching problem gets harder as networks and datasets grow, so it is a weak preprocessor. An equivariant layer folds that matching into the forward pass. The parameters are scalars and small dense maps.
An INR can be classified or regressed without rendering it back to pixels. A classifier can be adapted by one forward pass that emits a residual. In the adaptation experiment those classifiers are two-class models on CIFAR-10.
The MNIST gap, 85.71% against 58.98% for alignment, is where the symmetry prior earns its keep: the sample is a function. Adaptation, 71.36% against 67.66%, is an early result on editing weights.
The conclusion lists three limits. The layer layout is tied to one MLP depth and one set of widths, so a different architecture needs a new layout. Sharing block parameters across inner layers is left for later and is not tested. Some tasks were hard to train. The paper points to initialization and puts the discussion in Appendix K.5, without naming the failed tasks or giving their numbers in the main text. The implementation is cumbersome.
Scaling symmetry is left out on purpose. In a ReLU net, multiplying one layer by a constant and dividing the next can leave the function unchanged. Proposition 6.2 covers maps on the function space. It does not show that every equivariant map on raw weights can be approximated. Channel permutations in CNNs, and the paired query/key permutation in Transformers, are described and not tested.
INR classification is MNIST and Fashion-MNIST only. Adaptation uses random two-class CIFAR-10 models and one bundled corruption, with no per-corruption breakdown and no large backbone. An appendix experiment on predicting generalization from weights never reports numbers in the main text. INR2Vec is compared without its pretraining, so that column is not a verdict on the original method. Adding INR copies lifts Fashion-MNIST from 67.06% to 75.12%, about 8 points, which means the copy count can move the ranking by itself.