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

TL;DR

DeepWeightFlow introduces a flow-based model that efficiently generates high-quality, diverse neural network weights across architectures and sizes without fine-tuning, surpassing previous methods in speed and scalability.

Contribution

The paper presents DeepWeightFlow, a novel flow matching approach for direct weight space generation, incorporating neural network canonicalization to handle symmetries and improve efficiency.

Findings

Generates diverse neural network weights without fine-tuning.

Scales to large models like ResNet and ViT efficiently.

Enables rapid generation of hundreds of networks for transfer learning.

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.