DeepWeightFlow: Re-Basined Flow Matching for Generating Neural Network Weights

TL;DR

DeepWeightFlow introduces Flow Matching directly in neural network weight space to generate complete, high-quality weight sets across architectures without task-conditioned training. It leverages canonicalization (Git Re-Basin for MLP/ResNet and TransFusion for ViT) to resolve permutation symmetries, enabling efficient, diverse weight generation that does not require post-generation fine-tuning. The approach scales to large weights using Incremental and Dual PCA and achieves strong transferability across datasets, outperforming several diffusion-based baselines in speed and sometimes in performance. This work opens a path toward rapid, scalable generation of diverse neural networks for vision and NLP tasks, with practical implications for model distribution, transfer learning, and on-device adaptation.

Abstract

Building efficient and effective generative models for neural network weights has been a research focus of significant interest that faces challenges posed by the high-dimensional weight spaces of modern neural networks and their symmetries. Several prior generative models are limited to generating partial neural network weights, particularly for larger models, such as ResNet and ViT. Those that do generate complete weights struggle with generation speed or require finetuning of the generated models. In this work, we present DeepWeightFlow, a Flow Matching model that operates directly in weight space to generate diverse and high-accuracy neural network weights for a variety of architectures, neural network sizes, and data modalities. The neural networks generated by DeepWeightFlow do not require fine-tuning to perform well and can scale to large networks. We apply Git Re-Basin and TransFusion for neural network canonicalization in the context of generative weight models to account for the impact of neural network permutation symmetries and to improve generation efficiency for larger model sizes. The generated networks excel at transfer learning, and ensembles of hundreds of neural networks can be generated in minutes, far exceeding the efficiency of diffusion-based methods. DeepWeightFlow models pave the way for more efficient and scalable generation of diverse sets of neural networks.

Figures (2)

• Figure 1: Schematic depiction of DeepWeightFlow. a) We construct a training dataset of weights by fully training neural networks with weights $W_1,\dots,W_L$ on a given target task. b) Optionally, we use canonicalization, i.e., choosing a canonical representative $\tilde{W}_i$ from the same orbit as $W_i$, to break the permutation symmetry in parameter space. c) We train a flow model $p_{\hat{\theta}}$ for efficient generation of high-performance weights $(W_1,\dots,W_L)\sim p_{\hat{\theta}}$ for the target task.
• Figure 2: Maximum IoU vs test set accuracy for MNIST classifying MLPs. Lower maximum IoU implies greater diversity in the neural network weights. The left panels are generated and original neural networks (from the DeepWeightFlow training set) with different scales of Gaussian noise added to the original neural networks. The middle panels show that the generated neural networks and the original neural networks with noise added, which overlap in the left panels, are concretely different. The right panels contain the original and generated neural networks with different source distributions. All panels include 500 generated neural networks.