# Improved Generalization of Weight Space Networks via Augmentations

arXiv: https://arxiv.org/abs/2402.04081 · HTML: https://mlsanity.com/paper/2402.04081/

## Key facts

- Cited by 24 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 36 references that are in the map.
- First posted on arXiv 6 Feb 2024.

## Builds on (most cited first, 28 of 28)

- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Rethinking the Inception Architecture for Computer Vision](https://mlsanity.com/paper/1512.00567/) (2015)
- [mixup: Beyond Empirical Risk Minimization](https://mlsanity.com/paper/1710.09412/) (2017)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [UMAP: Uniform Manifold Approximation and Projection for Dimension Reduction](https://mlsanity.com/paper/1802.03426/) (2018)
- [DeepSDF: Learning Continuous Signed Distance Functions for Shape Representation](https://mlsanity.com/paper/1901.05103/) (2019)
- [Implicit Neural Representations with Periodic Activation Functions](https://mlsanity.com/paper/2006.09661/) (2020)
- [Dynamic Graph CNN for Learning on Point Clouds](https://mlsanity.com/paper/1801.07829/) (2018)
- [Occupancy Networks: Learning 3D Reconstruction in Function Space](https://mlsanity.com/paper/1812.03828/) (2018)
- [Decoupled Weight Decay Regularization](https://mlsanity.com/paper/1711.05101/) (2017)
- [Principal Neighbourhood Aggregation for Graph Nets](https://mlsanity.com/paper/2004.05718/) (2020)
- [Git Re-Basin: Merging Models modulo Permutation Symmetries](https://mlsanity.com/paper/2209.04836/) (2022)
- [The Role of Permutation Invariance in Linear Mode Connectivity of Neural Networks](https://mlsanity.com/paper/2110.06296/) (2021)
- [Point Convolutional Neural Networks by Extension Operators](https://mlsanity.com/paper/1803.10091/) (2018)
- [G-Mixup: Graph Data Augmentation for Graph Classification](https://mlsanity.com/paper/2202.07179/) (2022)
- [Predicting Neural Network Accuracy from Weights](https://mlsanity.com/paper/2002.11448/) (2020)
- [Equivariant Architectures for Learning in Deep Weight Spaces](https://mlsanity.com/paper/2301.12780/) (2023)
- [PointMixup: Augmentation for Point Clouds](https://mlsanity.com/paper/2008.06374/) (2020)
- [Self-Supervised Learning for Domain Adaptation on Point-Clouds](https://mlsanity.com/paper/2003.12641/) (2020)
- [Permutation Equivariant Neural Functionals](https://mlsanity.com/paper/2302.14040/) (2023)
- [Classifying the classifier: dissecting the weight space of neural networks](https://mlsanity.com/paper/2002.05688/) (2020)
- [Re-basin via implicit Sinkhorn differentiation](https://mlsanity.com/paper/2212.12042/) (2022)
- [Graph Metanetworks for Processing Diverse Neural Architectures](https://mlsanity.com/paper/2312.04501/) (2023)
- [Neural Functional Transformers](https://mlsanity.com/paper/2305.13546/) (2023)
- [Relational Attention: Generalizing Transformers for Graph-Structured Tasks](https://mlsanity.com/paper/2210.05062/) (2022)
- [Equivariant Deep Weight Space Alignment](https://mlsanity.com/paper/2310.13397/) (2023)
- [Graph Mixup with Soft Alignments](https://mlsanity.com/paper/2306.06788/) (2023)
- [Set-based Neural Network Encoding Without Weight Tying](https://mlsanity.com/paper/2305.16625/) (2023)

## Built on by (most cited first, 24 of 24)

- [Symmetry in Neural Network Parameter Spaces](https://mlsanity.com/paper/2506.13018/) (2025)
- [We Should Chart an Atlas of All the World's Models](https://mlsanity.com/paper/2503.10633/) (2025)
- [Learning on LoRAs: GL-Equivariant Processing of Low-Rank Weight Spaces for Large Finetuned Models](https://mlsanity.com/paper/2410.04207/) (2024)
- [The Empirical Impact of Neural Parameter Symmetries, or Lack Thereof](https://mlsanity.com/paper/2405.20231/) (2024)
- [Deep Linear Probe Generators for Weight Space Learning](https://mlsanity.com/paper/2410.10811/) (2024)
- [End-to-End Implicit Neural Representations for Classification](https://mlsanity.com/paper/2503.18123/) (2025)
- [Connecting NeRFs, Images, and Text](https://mlsanity.com/paper/2404.07993/) (2024)
- [Structure Is Not Enough: Leveraging Behavior for Neural Network Weight Reconstruction](https://mlsanity.com/paper/2503.17138/) (2025)
- [Position: Weight Space Should Be a First-Class Generative AI Modality](https://mlsanity.com/paper/2605.18632/) (2026)
- [DeepWeightFlow: Re-Basined Flow Matching for Generating Neural Network Weights](https://mlsanity.com/paper/2601.05052/) (2026)
- [Foldable SuperNets: Scalable Merging of Transformers with Different Initializations and Tasks](https://mlsanity.com/paper/2410.01483/) (2024)
- [Learning Model Representations Using Publicly Available Model Hubs](https://mlsanity.com/paper/2510.02096/) (2025)
- [ARC: Anchored Representation Clouds for High-Resolution INR Classification](https://mlsanity.com/paper/2503.15156/) (2025)
- [SINR: Sparsity Driven Compressed Implicit Neural Representations](https://mlsanity.com/paper/2503.19576/) (2025)
- [A Model Zoo on Phase Transitions in Neural Networks](https://mlsanity.com/paper/2504.18072/) (2025)
- [The Impact of Model Zoo Size and Composition on Weight Space Learning](https://mlsanity.com/paper/2504.10141/) (2025)
- [Discovering Hidden Gems in Model Repositories](https://mlsanity.com/paper/2601.22157/) (2026)
- [Dynamic Neural Graph Encoding of Inference Processes in Deep Weight Space](https://mlsanity.com/paper/2607.02166/) (2026)
- [On the Internal Representations of Graph Metanetworks](https://mlsanity.com/paper/2503.09120/) (2025)
- [Metanetworks as Regulatory Operators: Learning to Edit for Requirement Compliance](https://mlsanity.com/paper/2512.15469/) (2025)
- [Is Class Signal Clustered or Routed in Task-Induced Implicit Neural Representation Weight Spaces?](https://mlsanity.com/paper/2605.08281/) (2026)
- [Beyond Structural Symmetries: Linear Mode Connectivity via Neuron Identifiability](https://mlsanity.com/paper/2606.04754/) (2026)
- [Permutation-Equivariant Flow Matching for Alignment-Free Neural Weight Generation](https://mlsanity.com/paper/2609.32833/) (2026)
- [CrossGMN: Graph Metanetworks for Cross-Architecture Weight-Space Transformations](https://mlsanity.com/paper/2610.01649/) (2026)

## Abstract (arXiv)

Learning in deep weight spaces (DWS), where neural networks process the weights of other neural networks, is an emerging research direction, with applications to 2D and 3D neural fields (INRs, NeRFs), as well as making inferences about other types of neural networks. Unfortunately, weight space models tend to suffer from substantial overfitting. We empirically analyze the reasons for this overfitting and find that a key reason is the lack of diversity in DWS datasets. While a given object can be represented by many different weight configurations, typical INR training sets fail to capture variability across INRs that represent the same object. To address this, we explore strategies for data augmentation in weight spaces and propose a MixUp method adapted for weight spaces. We demonstrate the effectiveness of these methods in two setups. In classification, they improve performance similarly to having up to 10 times more data. In self-supervised contrastive learning, they yield substantial 5-10% gains in downstream classification.