Research map / Architectures and efficiency / RNNs and normalization
Equivariant Deep Weight Space Alignment
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes Deep-Align, an equivariant DWSNet-based network that learns to predict neural network weight alignments (permutations) in a single forward pass, trained without labeled data, and evaluated on MLP, CNN, INR and federated learning merging.
Problem
Weight alignment (finding permutations between two networks' weights) is NP-hard and existing optimization-based methods are slow or suboptimal, which hampers model merging, loss landscape analysis, and federated learning.
Approach
A learned, permutation-equivariant architecture (Siamese DWSNet encoder, activation-space mapping, generalized outer product, Sinkhorn/Hungarian projection) that respects the H=G×G and transposition symmetries of the alignment map, trained with self-generated labeled pairs plus unsupervised alignment and LMC losses.
Key results
- On CIFAR10 CNN, Deep-Align + Sinkhorn reaches barrier 0.08 vs 0.31 for Sinkhorn and 0.66 for Weight Matching.
- On VGG16 CIFAR10, Deep-Align + Sinkhorn barrier 0.31 vs 0.92 for Sinkhorn; Deep-Align alone 0.67 vs Weight Matching 2.07.
- In federated learning on CIFAR10 with 50 clients, Deep-Align reaches 69.86% vs 67.87% for FedAvg.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Initialize Sinkhorn with Deep-Align alignment | CNN CIFAR10 alignment · CNN CIFAR10 | Barrier: 0.23 ± 0.01 -> 0.08 ± 0.00 | Helped |
| Initialize Sinkhorn with Deep-Align alignment | CNN STL10 alignment · CNN STL10 | Barrier: 0.38 ± 0.01 -> 0.23 ± 0.00 | Helped |
| Initialize Sinkhorn with Deep-Align alignment | VGG11 CIFAR10 alignment · VGG11 CIFAR10 | Barrier: 0.29 ± 0.00 -> 0.09 ± 0.01 | Helped |
| Initialize Sinkhorn with Deep-Align alignment | VGG16 CIFAR10 alignment · VGG16 CIFAR10 | Barrier: 0.67 ± 0.01 -> 0.31 ± 0.01 | Helped |
| Use Deep-Align instead of the strongest listed baseline | Federated CIFAR10 model averaging · CIFAR10 50 clients | accuracy: 68.05 ± 2.35 -> 69.86 ± 1.13 | Helped |
| Use Deep-Align instead of the strongest listed baseline | Federated CIFAR10 model averaging · CIFAR10 100 clients | accuracy: 63.38 ± 2.59 -> 66.31 ± 1.19 | Helped |
| Use Deep-Align instead of the strongest listed baseline | Federated CIFAR10 model averaging · CIFAR10 200 clients | accuracy: 60.22 ± 1.92 -> 65.52 ± 1.78 | Helped |
| Use Deep-Align instead of the strongest listed baseline | Federated STL10 model averaging · STL10 10 clients | accuracy: 48.01 ± 0.72 -> 49.48 ± 2.23 | Helped |
| Use Deep-Align instead of the strongest listed baseline | Federated STL10 model averaging · STL10 25 clients | accuracy: 47.08 ± 0.51 -> 48.22 ± 1.08 | Helped |
| Use Deep-Align instead of the strongest listed baseline | Federated STL10 model averaging · STL10 50 clients | accuracy: 45.33 ± 0.95 -> 48.14 ± 0.88 | Helped |
| Add the alignment loss to training (ablation) | Deep-Align for Sine INRs · Sine INRs | alignment performance | Hurt |
| Add the alignment loss to training (ablation) | Deep-Align for CIFAR10 INRs · CIFAR10 INRs | alignment performance | Hurt |
Surprising findings
- Deep-Align alone is on par with or better than Sinkhorn and Weight Matching at inference cost of a single forward pass, without optimizing the test pair.
Limitations
- Requires pretraining a network on a set of weight-space samples for each architecture.
- Architecture is specific to the input network architecture because of the DWSNet encoder.
- Inference with Deep-Align requires running a DWSNet encoder over the input weights; the inputs-with-varying-depth extension is only preliminary.
Experimental setup
Datasets: MNIST CIFAR10 STL10
Benchmarks: MNIST MLP CIFAR10 MLP CIFAR10 CNN STL10 CNN VGG11 CIFAR10 VGG16 CIFAR10 Sine INRs CIFAR10 INRs federated CIFAR10/STL10
Metrics: Barrier AUC alignment runtime test accuracy
Baselines: Naive averaging Weight Matching Activation Matching Sinkhorn WM + Sinkhorn Sinkhorn-L2
Builds on: DWSNets Activation Matching Weight Matching Sinkhorn permutation learning
Closest papers by meaning
- Learning symmetries via weight-sharing with doubly stochastic tensors (NeurIPS 2024 · similarity 0.88)
- Scale Equivariant Graph Metanetworks (NeurIPS 2024 · similarity 0.86)
- Gradient-Weight Alignment as a Train-Time Proxy for Generalization in Classification Tasks (NeurIPS 2025 · similarity 0.86)
- Improving Equivariant Networks with Probabilistic Symmetry Breaking (ICLR 2025 · similarity 0.85)
- A Tale of Two Symmetries: Exploring the Loss Landscape of Equivariant Models (NeurIPS 2025 · similarity 0.85)
- REVISITING MULTI-PERMUTATION EQUIVARIANCE THROUGH THE LENS OF IRREDUCIBLE REPRESENTATIONS (ICLR 2025 · similarity 0.85)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.